workload changes and fp8_kv.py creation

[kwang3170]

Changed workload.py and added fp8_kv.py. fixed branching issue for original fp8 branch

[github-actions]

Code Review by Gemini

--- a/scripts/fp8_kv.py
+++ b/scripts/fp8_kv.py
@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
 """Serial fp8-KV A/B for vllm-decode — one engine per process, full GPU budget.
 
     python scripts/fp8_ab.py            # runs both legs, prints the delta
     python scripts/fp8_ab.py bf16 OUT   # (internal) single leg, used by the driver
 
 Each leg runs in its **own process** and builds exactly **one** engine, so only one
 engine is ever resident on the GPU. That removes the coexisting-engines constraint
 that forced ``GITM_VLLM_GPU_MEM=0.45`` — the old flow measured the bf16 baseline,
 then called ``gitm_restart_fn`` to build the fp8 candidate alongside the still-live
 baseline, so both had to fit at once (0.45 + 0.45 < 1.0). Serially, each engine gets
 the whole budget: **0.75 by default here**, and 0.9+ works too.
 
 Run as a FILE, never ``python -c`` or a stdin heredoc — a leg may build its engine
 under the ``spawn`` start method, which re-imports ``__main__`` in the child.
 
 Env (all optional; anything you export wins):
 
     GITM_VLLM_GPU_MEM     GPU memory fraction per engine (default 0.75 here)
     GITM_VLLM_MODEL       model (default: ungated Llama-3-8B mirror)
     GITM_VLLM_PROMPTS     number of prompts (default 512)
     GITM_VLLM_MAX_TOKENS  output tokens per prompt (default 2048)
     GITM_VLLM_PROFILE     "torch" -> also capture a PyTorch/kineto trace per leg
 
 To sweep the memory budget, just run it once per level:
 
     for m in 0.45 0.75 0.95; do GITM_VLLM_GPU_MEM=$m python scripts/fp8_ab.py; done
 
 See docs/vllm_decode_runbook.md for the full procedure.
 """
 
 from __future__ import annotations
 
 import os
 import subprocess
 import sys
 import tempfile
 from pathlib import Path
+import shutil
 
 #: leg name -> vLLM ``kv_cache_dtype``. "auto" is bf16 for this model.
 LEGS = {"bf16": "auto", "fp8": "fp8"}
 
 #: A serial A/B keeps only one engine resident, so the 0.45 cap the coexisting
 #: two-engine design needed no longer applies. Export GITM_VLLM_GPU_MEM to override.
 DEFAULT_GPU_MEM = "0.75"
 os.environ.setdefault("GITM_VLLM_GPU_MEM", DEFAULT_GPU_MEM)
 
 PROFILE_TORCH = os.environ.get("GITM_VLLM_PROFILE", "").lower() == "torch"
 PROFILE_DIR = os.environ.get("VLLM_TORCH_PROFILER_DIR", "/root/.cache/gitm/torchprof")
 
 
 def _run_leg(leg: str) -> float:
     """Build ONE engine in this leg's KV dtype and return decode throughput (tok/s)."""
     # Imported here so the driver process never pulls in torch/vLLM.
     from gitm.scheduler.loop import LoopConfig
     from gitm.workloads import get_factory, set_decode_run_defaults
 
     os.environ["GITM_VLLM_KV_DTYPE"] = LEGS[leg]
     if PROFILE_TORCH:
         # Per-leg trace dir, created before the engine (hence the child) is built.
         d = Path(PROFILE_DIR) / leg
         d.mkdir(parents=True, exist_ok=True)
         os.environ["VLLM_TORCH_PROFILER_DIR"] = str(d)
 
     env = set_decode_run_defaults()
     if PROFILE_TORCH:
         # torch/kineto and the CUPTI injection tracer both want the single CUPTI
         # subscriber slot; disable injection so the profiler wins.
         os.environ.pop("CUDA_INJECTION64_PATH", None)
 
     print(
         f"[{leg}] kv_cache_dtype={LEGS[leg]} "
         f"gpu_mem={os.environ['GITM_VLLM_GPU_MEM']} "
         f"model={env['GITM_VLLM_MODEL']} "
         f"load={env['GITM_VLLM_PROMPTS']}x{env['GITM_VLLM_MAX_TOKENS']}"
     )
 
     run = get_factory("vllm-decode")(LoopConfig(workload="vllm-decode"))
     eng = run.engine
     if PROFILE_TORCH:
         eng.start_profile()
     tps = eng.gitm_throughput_fn(eng)
     if PROFILE_TORCH:
         eng.stop_profile()
     return tps
 
 
 def main(argv: list[str]) -> None:
     # --- child mode: run one leg, hand the number back via a file ---
     if len(argv) == 2 and argv[0] in LEGS:
         leg, out_path = argv
         tps = _run_leg(leg)
         Path(out_path).write_text(repr(tps))
         print(f"[{leg}] {tps:,.0f} tok/s")
         return
 
     # --- driver: one process per leg, so only one engine is ever resident ---
     gpu_mem = os.environ["GITM_VLLM_GPU_MEM"]
     tmp = Path(tempfile.mkdtemp(prefix="gitm-fp8ab-"))
     results: dict[str, float] = {}
+    try:
     for leg in LEGS:
         out = tmp / f"{leg}.txt"
         print(f"\n{'=' * 72}")
         print(f"=== leg: {leg}  (kv_cache_dtype={LEGS[leg]}, gpu_mem={gpu_mem})")
         print(f"{'=' * 72}", flush=True)
-        rc = subprocess.run(
-            [sys.executable, os.path.abspath(__file__), leg, str(out)]
-        ).returncode
-        if rc != 0 or not out.exists():
-            raise SystemExit(f"leg {leg!r} failed (exit {rc}) — see output above")
+        subprocess.run(
+            [sys.executable, os.path.abspath(__file__), leg, str(out)], check=True
+        )
+        if not out.exists():
+            raise SystemExit(f"leg {leg!r} failed — output file not found")
         results[leg] = float(out.read_text())
+    finally:
+        shutil.rmtree(tmp)
 
     base, cand = results["bf16"], results["fp8"]
     print(f"\n{'=' * 72}")
     print(f"gpu_memory_utilization = {gpu_mem}   (serial A/B — one engine at a time)")
     print(f"baseline bf16 KV: {base:,.0f} tok/s")
     print(f"fp8 KV:           {cand:,.0f} tok/s   ({(cand / base - 1) * 100:+.1f}%)")
     if PROFILE_TORCH:
         print(f"\ntraces under {PROFILE_DIR}/ (bf16/ and fp8/) -> https://ui.perfetto.dev")
 
 
 if __name__ == "__main__":
     main(sys.argv[1:])

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

Filename/docstring mismatch — The file is fp8_kv.py but the module docstring says fp8_ab.py in three places (the python invocation examples). This will confuse users running the script directly.

Child mode argv parsing is fragile — In main, the child mode check is len(argv) == 2 and argv[0] in LEGS. argv is sys.argv[1:], so the driver spawns with [leg, out_path] → len == 2 ✓. But if a user accidentally passes any extra flag (e.g. --help), it silently falls through to driver mode and spawns infinite children. A dedicated --child sentinel would be safer.

float(out.read_text()) on repr(tps) — repr() of a float is fine for round-tripping standard values, but if tps is inf or nan (e.g. a degenerate run), float("inf") works but the delta arithmetic (cand / base - 1) will produce misleading output silently. Worth a guard.

set_decode_run_defaults return value assumed to contain keys — The env dict is used immediately for printing (env['GITM_VLLM_MODEL'] etc.). If that function doesn't populate those keys, you get a KeyError with no useful diagnostic. No fallback or .get() is used.

---

🔒 Security

subprocess.run without env= inherits the full parent environment — This is intentional for config forwarding, but GITM_VLLM_KV_DTYPE is set on the parent's os.environ before the subprocess is spawned, which means the child correctly inherits it. However, any secrets already in the environment (HF tokens, cloud credentials) are also inherited. This is probably acceptable in a benchmarking context but worth noting in the runbook.

tempfile.mkdtemp result is never cleaned up — tmp is created but there's no finally / atexit cleanup. On repeated runs this accumulates directories under /tmp. Minor, but worth a shutil.rmtree(tmp) in a finally block.

---

⚡ Performance

CUDA_INJECTION64_PATH is popped after os.environ["GITM_VLLM_KV_DTYPE"] is set but before the engine is built — The comment says "disable injection so the profiler wins," which is correct. However, since the child process inherits the parent's environment (which still has CUDA_INJECTION64_PATH), popping it in _run_leg inside the child is the right place. Confirm the parent never runs _run_leg directly in non-profiling mode to avoid accidentally stripping the injector for the driver process.

---

📊 Reproducibility

DEFAULT_GPU_MEM is applied at import time via os.environ.setdefault — This side-effects the environment on import, which means simply import fp8_kv from another module changes GITM_VLLM_GPU_MEM if it wasn't already set. Move this into main() or guard it with if __name__ == "__main__".

# Current — runs at import time, surprising for a library consumer
os.environ.setdefault("GITM_VLLM_GPU_MEM", DEFAULT_GPU_MEM)

No random seed setting — If set_decode_run_defaults or the vLLM engine uses any sampling with temperature > 0, outputs vary across legs. For a throughput-only benchmark this is likely fine, but documenting the assumption (or asserting temperature=0 / greedy) would strengthen reproducibility claims.

---

💡 Suggestions

Missing newline at EOF — The diff shows \ No newline at end of file. Small, but it breaks cat-based tooling and triggers linter warnings.

No timeout on subprocess.run — A hung child (e.g. waiting on a CUDA OOM that doesn't raise) will block the driver indefinitely. Add timeout= with a reasonable bound (e.g. 30 min) and surface the TimeoutExpired as a clear error.

rc = subprocess.run(
    [sys.executable, os.path.abspath(__file__), leg, str(out)],
    timeout=1800,
).returncode

LEGS is a plain dict — iteration order is CPython 3.7+ insertion order, so bf16 always runs first. This is correct and intentional but worth a one-line comment, since the A/B interpretation depends on order being stable.

The workloads.py change accepts any string for kv_cache_dtype — No validation against vLLM's accepted values ("auto", "fp8", "fp8_e5m2", "fp8_e4m3"). A typo silently passes through to vLLM, which will raise an opaque internal error. A small allowlist check with a clear message would help:

_VALID_KV_DTYPES = {"auto", "fp8", "fp8_e5m2", "fp8_e4m3"}
if _kv_dtype and _kv_dtype not in _VALID_KV_DTYPES:
    raise ValueError(f"GITM_VLLM_KV_DTYPE={_kv_dtype!r} not in {_VALID_KV_DTYPES}")