diff --git a/gitm/workloads.py b/gitm/workloads.py index b9924b2..1c59b0d 100644 --- a/gitm/workloads.py +++ b/gitm/workloads.py @@ -555,6 +555,10 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: GITM_VLLM_MODEL HF model id (default facebook/opt-125m) GITM_VLLM_PROMPTS number of prompts to decode (default 64) GITM_VLLM_MAX_TOKENS tokens to generate per prompt (default 128) + GITM_VLLM_KV_DTYPE kv_cache_dtype for the engine (e.g. "fp8"); unset + -> vLLM's "auto". Lets a serial A/B build each leg + directly in its dtype, one engine per process, so + each can use the full GITM_VLLM_GPU_MEM budget. GITM_VLLM_SYNTHETIC "1" -> CPU-only decode stand-in instead of vLLM (exercises the wire/registry path with no GPU or vLLM; produces no GPU kernels) @@ -586,6 +590,14 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: _base_kwargs: dict[str, Any] = {} if _gpu_mem is not None: _base_kwargs["gpu_memory_utilization"] = float(_gpu_mem) + # Optional KV-cache dtype (GITM_VLLM_KV_DTYPE, e.g. "fp8"). Lets the baseline + # engine be built directly in the target dtype, so a serial A/B can run each + # leg in its own process at the full GPU budget — instead of the restart path, + # which builds the candidate alongside the live baseline and therefore caps + # both at ~0.45. Unset (the default) leaves vLLM's own "auto". + _kv_dtype = os.environ.get("GITM_VLLM_KV_DTYPE") + if _kv_dtype: + _base_kwargs["kv_cache_dtype"] = _kv_dtype # Disable CUDA graphs so CUPTI captures decode kernels on platforms where # graph-replayed kernels aren't attributed (and speed up every engine build). # Inherited by restart candidates (kwargs = dict(_base_kwargs)) for a fair A/B. diff --git a/scripts/fp8_kv.py b/scripts/fp8_kv.py new file mode 100644 index 0000000..fc18005 --- /dev/null +++ b/scripts/fp8_kv.py @@ -0,0 +1,122 @@ +"""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 + +#: 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] = {} + 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") + results[leg] = float(out.read_text()) + + 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:]) \ No newline at end of file