diff --git a/.gitignore b/.gitignore index 6d57d4e..27a6a12 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,5 @@ htmlcov/ .claude venv takehome +.venv311/ +.python-version diff --git a/benchmarks/hft/HFT_HEADROOM_TOOLS.md b/benchmarks/hft/HFT_HEADROOM_TOOLS.md new file mode 100644 index 0000000..aa6651a --- /dev/null +++ b/benchmarks/hft/HFT_HEADROOM_TOOLS.md @@ -0,0 +1,81 @@ +# HFT Task A (headroom report) + Task B (apply + prove) — tools and plan + +Two scripts + one library that deliver Task A now (no pod) and set up Task B +(needs pod). Drop the three files in benchmarks/hft/. + +## Files +- hft_library.yaml — 3 candidate interventions in the real + InterventionSpec schema (review: null pending Adit) +- headroom_kernel_rank.py — Task A: ranks them by the real predict_delta() +- (Task B uses gitm/optimizer/apply.py — already in the repo) + +## Task A — run the headroom report (NO POD NEEDED) + + cd benchmarks/hft + python3 headroom_kernel_rank.py + +Output: ranked list of 3 candidate interventions, each with coverage %, +predicted fractional delta, and predicted Δ events/sec. Current ranking: + + 1 hft_parquet_h2d_overlap 89.9% cover ~22.5% +4.65M eps + 2 hft_multistream_symbol_shards 42.1% cover ~6.3% +1.31M eps + 3 hft_groupby_scan_fuse 10.1% cover ~0.8% +0.17M eps + +This satisfies Task A's "done when": >=3 candidates, each with a predicted +Δ events/sec tied to a residual (the covered kernels). + +## What's measured vs estimated (say this to Adit) + +MEASURED (real): + - kernel times / coverage % — from the seed-42 nsys profile + - the predict_delta() function — the same one the optimizer loop uses + +ESTIMATED (literature, cited per spec, NOT measured on our workload): + - expected_delta_mean (0.25 / 0.15 / 0.08) — from cuDF / NVIDIA docs + - therefore the predicted-Δ magnitudes are estimates + +So: the RANKING is defensible (driven by real coverage). The NUMBERS are +estimates until Task B measures them. The "#1 gets us to 25.4M / over target" +line depends on the 0.25 estimate holding; at the low end (0.10) it's ~22.5M, +still short. Present it as "predicted to potentially cross target, pending +measurement," not a promise. + +## Task B — apply + prove the #1 intervention (NEEDS POD) + +Goal: apply hft_parquet_h2d_overlap to the harness, prove it speeds things up +WITHOUT changing the derived metrics (VWAP/microprice), via a checksum gate. + +The pieces that already exist in the repo: + - gitm/optimizer/apply.py — apply_intervention(spec, applicator, + min_keep_delta): snapshot -> apply -> measure + -> keep|rollback. Use this as the gate. + - The Applicator seam — you implement snapshot/apply/measure/restore + for the harness (e.g. toggle prefetch depth). + +What you build for Task B: + 1. A checksum: run harness, hash the derived-metric output (e.g. sha256 of + the sorted VWAP/microprice arrays). This is the correctness gate. + 2. An HarnessApplicator implementing the Applicator protocol: + - snapshot(): record current reader settings + baseline checksum + - apply(spec): set reader_prefetch_depth=2 (overlap decode/H2D) + - measure(spec): re-run harness, return (eps_after - eps_before)/eps_before + - restore(): put settings back + 3. Wrap it: apply_intervention(spec, HarnessApplicator(), min_keep_delta=0.0) + BUT also assert checksum_after == checksum_before, else force rollback. + 4. Provenance report: claim (headroom #1) -> evidence (nsys coverage) -> + intervention (prefetch_depth=2) -> delta (measured eps before/after) + + checksum match. If checksum mismatches, report NO speedup (correctness + beats speed). + +Task B's "done when": provenance report shows a measured speedup with +checksum-identical output, and reports no win if the checksum mismatches. + +HONEST SCOPE: Task B is real engineering — implementing the overlap in the +harness (cuDF chunked/prefetched Parquet read), wiring the Applicator, and +measuring on the pod. It is NOT a copy-paste job. Do it once pod access is +back, and expect to iterate on the actual cuDF read path. + +## Validation done locally +headroom_kernel_rank.py was run against the real seed-42 kernel sums and +produces the ranked table above using the repo's real predict_delta(). The +ranking logic is verified; the magnitudes carry the estimate caveat above. diff --git a/benchmarks/hft/HFT_TASK_TOOLS.md b/benchmarks/hft/HFT_TASK_TOOLS.md new file mode 100644 index 0000000..6e01373 --- /dev/null +++ b/benchmarks/hft/HFT_TASK_TOOLS.md @@ -0,0 +1,66 @@ +# HFT Task 1 + Task 3 — tools and run order + +Three scripts that finish the predicted-graph integration (Task 1) and the +Granger causality analysis (Task 3). All three need the pod (the data and +nsys live there). Run them in this order. + +## What's measured vs estimated (read this first) + +The predicted-graph residuals depend on byte counts. Right now hft_graph.py +uses ESTIMATES: + - compression_ratio = 3.0 (a guess; zstd-1 on int columns is ~2.5-4x) + - per-stage bytes_moved (assumed multiples of the uncompressed size) + +The observed kernel times ARE real (from nsys). So the residual MAGNITUDES +(the "592x" numbers) are only as good as the byte estimates. Step 1 below +replaces the compression-ratio guess with a measured value. The Granger +RANKING (step 3) is more robust to the byte estimates than the magnitudes, +because it depends on relative timing across shards, not absolute residual +size — but at 20 samples it's still suggestive, not conclusive. + +## Step 1 — measure the real compression ratio (no nsys, just file reads) + + python3 measure_compression.py /workspace/hft_numba_seed42 + +Copy the printed "compression_ratio (parquet meta)" value and replace +compression_ratio=3.0 in HFTDatasetSpec (benchmarks or gitm/planner/hft_graph.py). +Re-run `python3 -m gitm.planner.hft_residual_demo` to get corrected magnitudes. + +## Step 2 — profile 20 shards separately (needs nsys + GPU) + + bash profile_per_shard.sh /workspace/hft_numba_seed42 /workspace/shard_profiles 20 + +Profiles 20 single-shard harness runs (5M events each) -> 20 nsys profiles -> +20 samples per kernel op. This is the part that takes a while (20 nsys runs). + +WHY 20 AND NOT 5: grangercausalitytests at maxlag=2 needs >= 8 observations +or statsmodels raises "Maximum allowable lag is 0". Even at 8-20 the p-values +are noisy. 20 is a practical floor; more shards = more reliable. + +## Step 3 — run Granger causality on the residual series + + python3 parse_shard_profiles.py /workspace/shard_profiles + +Parses the 20 CSV summaries, builds a per-op residual time series against the +predicted graph, and runs gitm.optimizer.attribution.attribute(). Output is a +ranked table of (cause -> effect, p_value). A low p for +zstd_decompress -> parquet_decode would confirm decompression is upstream of +the decode slowdown, not just correlated with it. + +## Validated locally + +The parse + Granger chain was tested with 20 synthetic shards carrying a known +zstd->parquet lagged dependence. attribute() correctly ranked +zstd_decompress -> parquet_decode at p=0.0000 and the unrelated merge_sort +pairs as non-significant (p>0.18). So the code path works; on real data the +result depends on what the actual per-shard timings show. + +## Honest limitations to mention to Adit + +1. compression_ratio is estimated until step 1 is run on real files. +2. Granger at n=20 is suggestive, not conclusive — more shards would help. +3. roofline still has no decompression-bound category; the residuals are vs + the memory-bandwidth floor, which these ops were never going to hit. The + "right" fix is a new prediction category calibrated against nvCOMP zstd + throughput numbers — that's a proposal for Adit, not something to merge + into the shared roofline.py unilaterally. diff --git a/benchmarks/hft/datasets.md b/benchmarks/hft/datasets.md index d992906..7d2ae04 100644 --- a/benchmarks/hft/datasets.md +++ b/benchmarks/hft/datasets.md @@ -44,3 +44,32 @@ make manifest # -> manifest.yaml (sha256 + byte count per Parquet file) make verify # re-hash $GITM_SCRATCH/staging/hft and confirm byte-identical ``` Manifest sha256: + +## How It Was Built + +### Generator +Written from scratch in C++ (`benchmarks/hft/generator/main.cpp`). +Built using CMake against pyarrow's bundled Arrow/Parquet C++ libraries. + +### Arrival Model +Hawkes process with parameters: +- mu=100 (baseline rate) +- alpha=0.6 (excitation per event) +- beta=0.8 (decay rate) +Inter-arrival times sampled as dt = -log(uniform) / lambda (inverse transform sampling). + +### Price Model +Integer random walk in tick space. Mid price moves +1 or -1 on each trade event (type=2). + +### Output +Written in 5M event chunks to avoid OOM on large datasets. +Same seed always produces identical bytes (mt19937_64 RNG). + +### Manifest +Generated by `benchmarks/hft/gen_manifest.py`. +Contains sha256 and byte count per file for reproducibility verification. + +### Profiling +Benchmarked using `benchmarks/hft/harness.py` (cuDF/CuPy GPU pipeline). +Profiled with NVIDIA Nsight Systems (nsys). +Results documented in `benchmarks/hft/results.md`. diff --git a/benchmarks/hft/gen_manifest.py b/benchmarks/hft/gen_manifest.py index ea10eeb..53a0165 100644 --- a/benchmarks/hft/gen_manifest.py +++ b/benchmarks/hft/gen_manifest.py @@ -9,7 +9,7 @@ files = [] for seed in seeds: - p = data_root / f"datasets/hft/hft_1b_seed{seed}/part0.parquet" + p = data_root / f"datasets/hft/hft_1b_seed{seed}/part-00000.parquet" sha256 = hashlib.sha256(p.read_bytes()).hexdigest() size = p.stat().st_size files.append({ diff --git a/benchmarks/hft/generator_comparison.md b/benchmarks/hft/generator_comparison.md new file mode 100644 index 0000000..54df63d --- /dev/null +++ b/benchmarks/hft/generator_comparison.md @@ -0,0 +1,112 @@ +# HFT Generator Comparison Report + +## Overview +I tested and built four different approaches to generate 1B HFT events, progressively improving speed while maintaining data realism. The final version is 7.4x faster than C++ while keeping exact Hawkes timestamps. + +## What Makes the Data Realistic + +The realism comes from three things, not from Numba or GPU: + +1. Hawkes process for timestamps - in real markets, one big trade triggers more trades. Events cluster in bursts. Hawkes models this: every event jumps the arrival rate up by alpha (0.6), then it decays back to baseline at rate beta (0.8). Simple uniform random gaps would be unrealistic. + +2. Realistic event mix - real order books have ~55% new orders, ~35% cancellations, ~10% trades. We use exactly this distribution instead of uniform random. + +3. Price random walk - price moves up or down by 1 tick every time a trade happens, simulating real market microstructure. + +Numba and GPU do NOT change the realism. They just make the same math run faster. + +## Why Numba + +The Hawkes process has a sequential dependency - each timestamp depends on the previous lambda value: + + lambda_i = MU + (lambda_{i-1} - MU) * exp(-BETA * dt_i) + ALPHA + +This cannot be parallelized directly. In plain Python this loop runs at 0.7M/sec. Numba JIT compiles the exact same loop to native CPU machine code using LLVM - running at 34.8M/sec. 53x faster, same exact output. + +Attempted fully vectorized GPU Hawkes using parallel prefix scan but thinning approximation produced non-monotonic timestamps. Numba is the correct solution - exact, fast, stable. + +## Generators Tested + +### 1. C++ Generator (benchmarks/hft/generator/main.cpp) +Sequential loop per event. Hawkes process for timestamps. Writes to Parquet in 5M chunks to avoid OOM. Single threaded. + +- Time for 1B events: ~10 minutes +- Events/sec: ~1.67M/sec +- Hawkes: exact +- Event mix: uniform random + +### 2. GPU Single Thread (CuPy) +Generates all 6 fields for 5M events at once on GPU. Transfers to CPU, writes to Parquet. Sequential - waits for write before next chunk. No Hawkes. + +- Time for 1B events: ~4m 25s +- Events/sec: ~3.8M/sec +- Hawkes: simple uniform gaps +- Event mix: uniform random + +### 3. GPU Parallel simple +GPU generates chunk N while 3 CPU threads write simultaneously using ThreadPoolExecutor. GPU never waits for disk. Still no Hawkes. + +- Time for 1B events: 1m 18s +- Events/sec: 12.9M/sec +- Hawkes: simple uniform gaps +- Event mix: uniform random + +### 4. Numba + GPU Parallel (final version) +Best of all approaches combined: +- Numba JIT compiles Hawkes loop to native CPU - 53x faster than plain Python, exact output +- CuPy generates symbol, side, price, size, type on GPU simultaneously +- Realistic event mix: 55% add, 35% cancel, 10% trade +- 3 parallel write threads - disk never idle + +- Time for 1B events: ~1m 22s +- Events/sec: ~12.4M/sec +- Hawkes: exact +- Event mix: realistic 55/35/10 + +## Results + +| Generator | Time for 1B | Events/sec | Speedup vs C++ | Exact Hawkes | Realistic Mix | +|---|---|---|---|---|---| +| Adit Python (numpy) | ~30 min | 0.56M/sec | 0.3x | No | No | +| C++ CPU | ~10 min | 1.67M/sec | 1x | Yes | No | +| GPU single thread | ~4m 25s | 3.8M/sec | 2.3x | No | No | +| GPU parallel simple | 1m 18s | 12.9M/sec | 7.7x | No | No | +| Numba + GPU parallel | 1m 22s | 12.4M/sec | 7.4x | Yes | Yes | + +## Why GPU Parallel is Fast + +Three things happen simultaneously: +- GPU generates fields for chunk N using thousands of GPU cores +- Write thread 1 writes chunk N-2 to disk +- Write thread 2 writes chunk N-1 to disk + +Nothing waits. Like a pit crew where refueling, tire change, and wing adjustment all happen at once. + +## Bottleneck at Each Stage + +| Stage | C++ | GPU single | GPU parallel | Numba + GPU | +|---|---|---|---|---| +| Timestamp generation | slow CPU loop | uniform only | uniform only | Numba JIT fast | +| Field generation | slow CPU loop | GPU fast | GPU fast | GPU fast | +| Parquet write | sequential | sequential | 3x parallel | 3x parallel | +| Bottleneck | generation | write | disk bandwidth | disk bandwidth | + +## RNG Speed Test (5M numbers) +- numpy CPU: 0.204s = 24M/sec +- CuPy GPU: 0.000s = 10B+/sec +- Speedup: 423x + +## Three Seeds Validated + +| Seed | Rows | Files | Events/sec | +|---|---|---|---| +| 42 | 1,000,000,000 | 200 | 12.4M/sec | +| 43 | 1,000,000,000 | 200 | 12.4M/sec | +| 44 | 1,000,000,000 | 200 | 12.4M/sec | + +Seed variance: 0% - identical throughput across all seeds. + +## Conclusion +Final generator (Numba + GPU + parallel writes) is 7.4x faster than C++ with exact Hawkes timestamps and realistic event mix. The key insight is that realism and speed are not tradeoffs - Numba makes exact Hawkes fast enough to match GPU field generation speed. + +Next step: explore Arrow IPC format instead of Parquet to reduce write overhead further. diff --git a/benchmarks/hft/gpu_generator.py b/benchmarks/hft/gpu_generator.py new file mode 100644 index 0000000..c8e6d09 --- /dev/null +++ b/benchmarks/hft/gpu_generator.py @@ -0,0 +1,102 @@ +import argparse +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import cupy as cp +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from numba import njit + +MU = 100.0 +ALPHA = 0.6 +BETA = 0.8 + +@njit(cache=True) +def hawkes_numba(u, n, t0, lam0): + ts = np.empty(n, dtype=np.int64) + t = t0 + lam = lam0 + for i in range(n): + dt = -np.log(u[i]) / lam + t += dt + lam = MU + (lam - MU) * np.exp(-BETA * dt) + ALPHA + ts[i] = int(t * 1e9) + return ts, t, lam + +def warmup(): + u = np.random.default_rng(0).random(100) + hawkes_numba(u, 100, 0.0, MU) + cp.random.default_rng(0).integers(0, 100, size=10_000) + cp.cuda.Stream.null.synchronize() + +def generate_chunk(rng, rng_np, n, t_hawkes, lambda_hawkes): + # Hawkes timestamps via Numba (fast CPU) + u = rng_np.random(n) + ts_ns, t_hawkes, lambda_hawkes = hawkes_numba(u, n, t_hawkes, lambda_hawkes) + + # All other fields on GPU + symbol_id = cp.asnumpy(rng.integers(0, 512, size=n, dtype=cp.int32)) + side = cp.asnumpy(rng.integers(0, 2, size=n, dtype=cp.int8)) + price = cp.asnumpy(rng.integers(9000, 11000, size=n, dtype=cp.int64)) + size_arr = cp.asnumpy(rng.integers(1, 1000, size=n, dtype=cp.int32)) + roll = cp.asnumpy(rng.random(n, dtype=cp.float32)) + etype = np.where(roll < 0.55, 0, np.where(roll < 0.90, 1, 2)).astype(np.int8) + + return { + 'ts_ns': ts_ns, 'symbol_id': symbol_id, 'side': side, + 'price': price, 'size': size_arr, 'type': etype, + }, t_hawkes, lambda_hawkes + +def write_chunk(data, path): + table = pa.table({ + 'ts_ns': pa.array(data['ts_ns'], type=pa.int64()), + 'symbol_id': pa.array(data['symbol_id'], type=pa.int32()), + 'side': pa.array(data['side'], type=pa.int8()), + 'price': pa.array(data['price'], type=pa.int64()), + 'size': pa.array(data['size'], type=pa.int32()), + 'type': pa.array(data['type'], type=pa.int8()), + }) + pq.write_table(table, path, compression='zstd', compression_level=1) + +def generate(n_events, seed, out_dir, chunk_size=5_000_000, n_writers=3): + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + print("Warming up...") + warmup() + rng = cp.random.default_rng(seed) + rng_np = np.random.default_rng(seed) + n_chunks = n_events // chunk_size + t_hawkes = 0.0 + lambda_hawkes = MU + t0 = time.perf_counter() + + with ThreadPoolExecutor(max_workers=n_writers) as executor: + futures = [] + for i in range(n_chunks): + data, t_hawkes, lambda_hawkes = generate_chunk( + rng, rng_np, chunk_size, t_hawkes, lambda_hawkes) + path = out_dir / f"part-{i:05d}.parquet" + futures.append(executor.submit(write_chunk, data, path)) + if len(futures) > n_writers: + futures.pop(0).result() + if i % 10 == 0: + elapsed = time.perf_counter() - t0 + print(f"Progress: {(i+1)*chunk_size/1e6:.0f}M/{n_events/1e6:.0f}M, {elapsed:.1f}s", end='\r') + for f in futures: + f.result() + + elapsed = time.perf_counter() - t0 + print(f"\nDone: {n_events/1e6:.0f}M events in {elapsed:.1f}s = {n_events/elapsed/1e6:.1f}M events/sec") + print("FINISHED") + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--events", type=int, default=100_000_000) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--out", type=str, default="/workspace/test_gpu_output") + p.add_argument("--chunk-size", type=int, default=5_000_000) + p.add_argument("--writers", type=int, default=3) + args = p.parse_args() + generate(args.events, args.seed, args.out, args.chunk_size, args.writers) diff --git a/benchmarks/hft/headroom_kernel_rank.py b/benchmarks/hft/headroom_kernel_rank.py new file mode 100644 index 0000000..53e5403 --- /dev/null +++ b/benchmarks/hft/headroom_kernel_rank.py @@ -0,0 +1,136 @@ +"""Headroom report for the HFT execution path (Task A). + +observe -> decide: take the captured kernel times from the seed-42 nsys +profile, treat each candidate intervention in hft_library.yaml as a +counterfactual, and rank them by gitm.optimizer.replay.predict_delta() -- +the SAME predictor the optimizer loop uses. Output is a ranked list of >=3 +candidate interventions, each with a predicted delta (fractional + events/sec) +tied to the kernels (residual) it covers. + +This is the "observe -> decide" half. The "prove" half (apply the #1 lever, +checksum-gate it, measure real before/after) is Task B and needs the pod. + +HONESTY: predict_delta = coverage x expected_delta_mean. coverage is real +(measured kernel times). expected_delta_mean is a literature estimate from each +spec's cited source, NOT measured on our workload (see hft_library.yaml header). +So the ranking is defensible but the magnitudes are estimates until Task B. + +Run: + python3 headroom_kernel_rank.py + python3 headroom_kernel_rank.py --baseline-eps 20700000 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import yaml + +from gitm.kernels.spec import InterventionSpec +from gitm.optimizer.replay import _applies, predict_delta +from gitm.tracer.schema import KernelEvent, Trace + +# Real per-kernel times (ns) from the seed-42, 25M-event nsys profile. +# Source: benchmarks/hft/results.md "Top GPU Kernels". These are the aggregate +# kernel sums; the per-shard trace (Task 3 tooling) refines them. +OBSERVED_KERNELS_NS = { + "zstd::decompression_kernel": int(252.076091 * 1e6), + "cudf::io::parquet::detail::decode_page_data_generic": int(139.100450 * 1e6), + "cub::detail::merge_sort::DeviceMergeSortMergeKernel": int(26.555531 * 1e6), + "cub::detail::merge_sort::DeviceMergeSortBlockSortKernel": int(17.371249 * 1e6), +} + +# Default baseline throughput for the events/sec delta (seed-42 sequential 25M). +DEFAULT_BASELINE_EPS = 20_700_000 + + +def build_trace_from_observed() -> Trace: + """Construct a Trace from the observed kernel sums (serialized on stream 0).""" + events: list[KernelEvent] = [] + t = 0 + for name, dur in OBSERVED_KERNELS_NS.items(): + events.append( + KernelEvent( + kind="kernel", start_ns=t, end_ns=t + dur, + stream_id=0, device_id=0, name=name, + ) + ) + t += dur + return Trace( + workload_id="hft-lob-replay", fingerprint="seed42-25M", run_id="headroom", + device_count=1, vendor="nvidia", captured_at_ns=0, duration_ns=t, events=events, + ) + + +def covered_kernels(spec: InterventionSpec, trace: Trace) -> list[str]: + return [k.name for k in trace.kernels() if _applies(spec, k.name)] + + +def main() -> None: + ap = argparse.ArgumentParser(description="HFT headroom / intervention ranking") + ap.add_argument( + "--library", type=Path, + default=Path(__file__).parent / "hft_library.yaml", + help="HFT intervention library YAML", + ) + ap.add_argument( + "--baseline-eps", type=float, default=DEFAULT_BASELINE_EPS, + help="baseline events/sec to scale the predicted delta into events/sec", + ) + args = ap.parse_args() + + trace = build_trace_from_observed() + total_ms = trace.duration_ns / 1e6 + + with open(args.library) as fh: + raw = yaml.safe_load(fh) or {} + specs = [InterventionSpec.model_validate(e) for e in raw.get("interventions", [])] + + rows = [] + for spec in specs: + delta = predict_delta(trace, spec) # fractional wall-clock improvement + cov = covered_kernels(spec, trace) + cov_ns = sum( + (k.end_ns - k.start_ns) for k in trace.kernels() if k.name in cov + ) + rows.append({ + "name": spec.name, + "delta_frac": delta, + "delta_eps": delta * args.baseline_eps, + "coverage_pct": 100 * cov_ns / trace.duration_ns, + "lo": spec.expected_delta_lo, + "hi": spec.expected_delta_hi, + "n_kernels": len(cov), + }) + + rows.sort(key=lambda r: r["delta_frac"], reverse=True) + + print(f"HFT headroom report (trace total = {total_ms:.1f} ms, " + f"baseline = {args.baseline_eps/1e6:.1f}M events/sec)") + print("=" * 78) + print(f"{'rank':<5}{'intervention':<32}{'cover%':>7}{'pred Δ':>9}" + f"{'Δ events/sec':>14}") + print("-" * 78) + for i, r in enumerate(rows, 1): + print(f"{i:<5}{r['name']:<32}{r['coverage_pct']:>6.1f}%" + f"{r['delta_frac']*100:>8.1f}%{r['delta_eps']:>14,.0f}") + print("-" * 78) + + top = rows[0] + print(f"\n#1: {top['name']}") + print(f" covers {top['coverage_pct']:.1f}% of GPU kernel time " + f"({top['n_kernels']} kernel type(s))") + print(f" predicted speedup: {top['delta_frac']*100:.1f}% " + f"(source range {top['lo']*100:.0f}-{top['hi']*100:.0f}%)") + print(f" predicted gain: +{top['delta_eps']:,.0f} events/sec " + f"-> {(args.baseline_eps + top['delta_eps'])/1e6:.1f}M events/sec") + print() + print("NOTE: predicted deltas use literature expected_delta_mean values") + print("(cited per spec), scaled by MEASURED kernel coverage. Real deltas") + print("come from Task B: apply #1 via gitm/optimizer/apply.py behind a") + print("derived-metric checksum gate, measure before/after on the pod.") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/hft/hft_library.yaml b/benchmarks/hft/hft_library.yaml new file mode 100644 index 0000000..b645af3 --- /dev/null +++ b/benchmarks/hft/hft_library.yaml @@ -0,0 +1,98 @@ +# Candidate intervention library for the HFT LOB-replay benchmark. +# +# Same schema as gitm/kernels/library.yaml (the vLLM-decode library), but the +# levers here target the data-loading path the nsys profile flagged: zstd +# decompression (47.5%), Parquet page decode (26.2%), and the top-of-book merge +# sort (8.3%). See benchmarks/hft/results.md for the stall profile. +# +# IMPORTANT: every expected_delta_* below is a PRE-REVIEW ESTIMATE from the +# cited source, NOT measured on our workload. review: null on all entries until +# Adit signs off. The real numbers come from the rollback-gated live apply +# (Task B) via gitm/optimizer/apply.py + a checksum correctness gate. +# +# applies_to_kernels uses substring match against the real kernel names seen in +# the seed-42 nsys profile: +# zstd::decompression_kernel +# cudf::io::parquet::detail::...decode_page_data_generic... +# cub::detail::merge_sort::DeviceMergeSort{Block,Merge}Kernel + +interventions: + # --- Candidate 1: overlap Parquet decode / H2D with compute ----------------- + - name: hft_parquet_h2d_overlap + summary: >- + Pipeline Parquet decode + host-to-device copy of shard N+1 against GPU + compute of shard N, so decompression/decode stops serializing in front of + the replay kernels. Targets the 73% data-stall directly. + knob: reader_prefetch_depth + value: 2 + applies_to_kernels: [decode_page_data, zstd::decompression] + expected_delta_mean: 0.25 + expected_delta_lo: 0.10 + expected_delta_hi: 0.40 + source: https://docs.rapids.ai/api/cudf/stable/user_guide/io/io/ + applicability: + workloads: [hft-lob-replay] + requires_hardware: [A100, H100] + other: >- + Estimate based on cuDF chunked-Parquet + prefetch guidance; H2D was + 91.2% of memory time in our trace, so overlap headroom is large but + the realized fraction depends on PCIe vs compute balance. + safety: + tier: low_risk + requires_rollback_window_s: 60 + notes: >- + Read-path only; derived-metric checksum (VWAP/microprice) must be + identical before keep. + review: null + + # --- Candidate 2: multi-stream across symbol shards ------------------------- + - name: hft_multistream_symbol_shards + summary: >- + Launch the per-symbol-shard decode + groupby on separate CUDA streams so + independent shards run concurrently instead of serializing. Targets the + serialized-launch fraction in the trace. + knob: num_cuda_streams + value: 4 + applies_to_kernels: [decode_page_data, DeviceMergeSort, DeviceReduce] + expected_delta_mean: 0.15 + expected_delta_lo: 0.05 + expected_delta_hi: 0.30 + source: https://developer.nvidia.com/blog/gpu-pro-tip-cuda-7-streams-simplify-concurrency/ + applicability: + workloads: [hft-lob-replay] + requires_hardware: [A100, H100] + other: >- + Benefit scales with how independent the symbol shards are; if one symbol + dominates volume the concurrency win shrinks. Estimate, not measured. + safety: + tier: moderate + requires_rollback_window_s: 120 + notes: >- + Multi-stream can change reduction order; checksum gate must allow a + documented float tolerance or use integer-stable accumulation. + review: null + + # --- Candidate 3: fused groupby-scan to cut sync ---------------------------- + - name: hft_groupby_scan_fuse + summary: >- + Fuse the top-of-book groupby with the VWAP cumulative scan to remove the + device sync between them, cutting the sync stall between merge_sort and + the reduction kernels. + knob: fuse_groupby_scan + value: true + applies_to_kernels: [DeviceMergeSort, DeviceScan, DeviceReduce] + expected_delta_mean: 0.08 + expected_delta_lo: 0.02 + expected_delta_hi: 0.15 + source: https://docs.rapids.ai/api/cudf/stable/user_guide/guide-to-udfs/ + applicability: + workloads: [hft-lob-replay] + requires_hardware: [A100, H100] + other: >- + Smallest of the three — sync was only ~7-8% of the trace — but also the + lowest risk. Estimate from cuDF fusion guidance, not measured. + safety: + tier: low_risk + requires_rollback_window_s: 60 + notes: Pure reordering; checksum must be bit-identical. + review: null diff --git a/benchmarks/hft/manifest.yaml b/benchmarks/hft/manifest.yaml index 9eb9331..47138e4 100644 --- a/benchmarks/hft/manifest.yaml +++ b/benchmarks/hft/manifest.yaml @@ -1,16 +1,16 @@ datasets: - bytes: 7763865521 - path: /home/ash/gitm-data/datasets/hft/hft_1b_seed42/part0.parquet + path: /home/ash/gitm-data/datasets/hft/hft_1b_seed42/part-00000.parquet rows: 1000000000 seed: 42 sha256: 69857a2c41130aed20d12e34353259c95b49ee7c1134393cca858330873ba7bf - bytes: 7767347789 - path: /home/ash/gitm-data/datasets/hft/hft_1b_seed43/part0.parquet + path: /home/ash/gitm-data/datasets/hft/hft_1b_seed43/part-00000.parquet rows: 1000000000 seed: 43 sha256: 1ff57b1b485d08d636b6d7bc6ce785eb8d822f82ed7e1a04699381099fca72df - bytes: 7771192742 - path: /home/ash/gitm-data/datasets/hft/hft_1b_seed44/part0.parquet + path: /home/ash/gitm-data/datasets/hft/hft_1b_seed44/part-00000.parquet rows: 1000000000 seed: 44 sha256: 4391dd934ce09f5b8f475cb16ecb43efec0233d98d91747cdd67a3c14fb718ef diff --git a/benchmarks/hft/measure_compression.py b/benchmarks/hft/measure_compression.py new file mode 100644 index 0000000..696f413 --- /dev/null +++ b/benchmarks/hft/measure_compression.py @@ -0,0 +1,82 @@ +"""Measure the REAL compression ratio + per-shard byte counts. + +Replaces the estimated ``compression_ratio = 3.0`` in hft_graph.py with a +number measured from the actual Parquet shards. Run this on the pod where +/workspace/hft_numba_seed42/ lives, then paste the printed compression_ratio +back into HFTDatasetSpec. + +What it measures, per shard: + - on-disk compressed bytes (os.path.getsize) + - uncompressed bytes (num_rows * 26, the fixed-width row size) + - parquet metadata's own uncompressed estimate (total_uncompressed_size) + +The (compressed, uncompressed) pair gives the real ratio. We print both the +26-bytes/row figure and Parquet's own accounting because they can differ +(Parquet stores some per-page overhead + dictionary pages that the naive +26*rows ignores). + +Run: + python3 measure_compression.py /workspace/hft_numba_seed42 +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pyarrow.parquet as pq + +BYTES_PER_ROW = 26 # ts_ns i64 + symbol_id i32 + side i8 + price i64 + size i32 + type i8 + + +def main(stage: str) -> None: + paths = sorted(Path(stage).glob("part-*.parquet")) + if not paths: + print(f"no parquet shards in {stage}") + return + + total_disk = 0 + total_rows = 0 + total_pq_uncompressed = 0 + + # Sample first 5 shards (matches the 25M-event warm window) + report all-200 totals. + print(f"{'shard':<22} {'rows':>12} {'disk_MB':>10} {'pq_uncomp_MB':>14} {'ratio':>7}") + print("-" * 70) + for i, p in enumerate(paths): + md = pq.ParquetFile(p).metadata + disk = p.stat().st_size + rows = md.num_rows + pq_uncomp = sum( + md.row_group(rg).total_byte_size for rg in range(md.num_row_groups) + ) + total_disk += disk + total_rows += rows + total_pq_uncompressed += pq_uncomp + if i < 5: + ratio = pq_uncomp / disk if disk else 0 + print( + f"{p.name:<22} {rows:>12,} {disk / 1e6:>10.2f} " + f"{pq_uncomp / 1e6:>14.2f} {ratio:>7.2f}" + ) + + naive_uncomp = total_rows * BYTES_PER_ROW + ratio_naive = naive_uncomp / total_disk if total_disk else 0 + ratio_pq = total_pq_uncompressed / total_disk if total_disk else 0 + + print("-" * 70) + print(f"shards: {len(paths)}") + print(f"total rows: {total_rows:,}") + print(f"total on-disk: {total_disk / 1e9:.3f} GB") + print(f"naive uncompressed: {naive_uncomp / 1e9:.3f} GB (26 B/row)") + print(f"parquet uncompressed: {total_pq_uncompressed / 1e9:.3f} GB (metadata)") + print() + print(f"compression_ratio (naive 26B/row): {ratio_naive:.3f}") + print(f"compression_ratio (parquet meta): {ratio_pq:.3f}") + print() + print("ACTION: replace compression_ratio=3.0 in HFTDatasetSpec") + print(f" with the parquet-meta ratio above ({ratio_pq:.3f}).") + + +if __name__ == "__main__": + stage = sys.argv[1] if len(sys.argv) > 1 else "/workspace/hft_numba_seed42" + main(stage) diff --git a/benchmarks/hft/parse_shard_profiles.py b/benchmarks/hft/parse_shard_profiles.py new file mode 100644 index 0000000..afdce77 --- /dev/null +++ b/benchmarks/hft/parse_shard_profiles.py @@ -0,0 +1,129 @@ +"""Parse 5 per-shard nsys profiles -> residual series -> Granger (Task 3). + +Reads the cuda_gpu_kern_sum CSVs produced by profile_per_shard.sh, maps the +three HFT kernels (zstd / parquet decode / merge sort) to predicted-graph +ops, computes per-shard r_kt residuals against hft_graph, then feeds the +series into gitm.optimizer.attribution.attribute() for the real Granger +F-test. + +With 5 shards we get 5 samples/op, which clears attribute()'s max_lag+2=4 +minimum. The output ranks which stage Granger-causes the others' slowdown. + +Run on the pod after profile_per_shard.sh: + python3 parse_shard_profiles.py /workspace/shard_profiles + +HONESTY NOTE: residuals are computed against hft_graph's predicted times, +which still depend on HFTDatasetSpec's byte estimates. Run +measure_compression.py first and patch the real ratio in, or the residual +*magnitudes* will be off (the Granger *ranking* is more robust to this since +it's about relative timing across shards, not absolute residual size). +""" + +from __future__ import annotations + +import csv +import sys +from pathlib import Path + +# Map substrings in the nsys kernel "Name" column to our predicted-graph ops. +KERNEL_OP_MAP = { + "zstd": "zstd_decompress", + "decode_page_data": "parquet_decode", + "DeviceMergeSortMergeKernel": "merge_sort", + "DeviceMergeSortBlockSortKernel": "merge_sort", +} + + +def kernel_to_op(name: str) -> str | None: + for needle, op in KERNEL_OP_MAP.items(): + if needle in name: + return op + return None + + +def parse_one_csv(csv_path: Path) -> dict[str, float]: + """Sum kernel time (ns) per op for a single shard's profile.""" + per_op: dict[str, float] = {} + with open(csv_path, newline="") as f: + reader = csv.DictReader(f) + # nsys cuda_gpu_kern_sum columns: "Time (%)","Total Time (ns)",...,"Name" + time_col = next( + (c for c in (reader.fieldnames or []) if "Total Time" in c), None + ) + name_col = "Name" if "Name" in (reader.fieldnames or []) else None + if not time_col or not name_col: + return per_op + for row in reader: + op = kernel_to_op(row[name_col]) + if op is None: + continue + try: + t_ns = float(row[time_col].replace(",", "")) + except (ValueError, KeyError): + continue + per_op[op] = per_op.get(op, 0.0) + t_ns + return per_op + + +def main(profile_dir: str) -> None: + d = Path(profile_dir) + csvs = sorted(d.glob("shard_*_cuda_gpu_kern_sum.csv")) or sorted( + d.glob("shard_*.csv") + ) + if not csvs: + print(f"no per-shard CSVs in {d}. Run profile_per_shard.sh first.") + return + + # Build observed kernel-time series (ms) per op, one entry per shard. + obs_series: dict[str, list[float]] = {} + for csv_path in csvs: + per_op = parse_one_csv(csv_path) + for op, t_ns in per_op.items(): + obs_series.setdefault(op, []).append(t_ns / 1e6) # ns -> ms + + print(f"parsed {len(csvs)} shard profiles") + for op, series in obs_series.items(): + print(f" {op:<18} {len(series)} samples {[round(x, 1) for x in series]}") + + # Build residuals vs predicted graph, scaled to 5M-event shards. + try: + from gitm.optimizer.attribution import attribute + from gitm.optimizer.monitor import KernelResidual, Residuals + from gitm.planner.hft_graph import HFTDatasetSpec, predict_hft_graph + except Exception as e: # noqa: BLE001 + print(f"\ncould not import gitm modules ({e}).") + print("Run this from the repo root with the env that has pydantic/statsmodels.") + return + + # Predicted graph sized for ONE shard (5M events), not 25M. + shard_spec = HFTDatasetSpec(name="hft_shard_5M", n_events=5_000_000) + g = predict_hft_graph(dataset=shard_spec) + pred_ms = {n.op: n.prediction.t_pred_s * 1e3 for n in g.nodes} + + res = Residuals() + n_shards = min((len(v) for v in obs_series.values()), default=0) + for shard_i in range(n_shards): + for op, series in obs_series.items(): + t_obs = series[shard_i] + t_pred = pred_ms.get(op, 1e-9) + r_kt = (t_obs - t_pred) / max(t_pred, 1e-12) + res.per_kernel.append(KernelResidual(op=op, layer=None, r_kt=r_kt, r_mt=None)) + + print(f"\nbuilt {len(res.per_kernel)} residual points " + f"({n_shards} shards x {len(obs_series)} ops)") + + ranked = attribute(res, g, max_lag=2) + if not ranked.hypotheses: + print("\nattribute() returned no hypotheses.") + print("Likely <4 samples/op (need more shards) or statsmodels missing.") + return + + print("\nGranger causality ranking (lower p = stronger causal signal):") + print(f"{'cause':<18} -> {'effect':<18} {'p_value':>10} {'direction'}") + print("-" * 64) + for h in ranked.top(10): + print(f"{h.cause_op:<18} -> {h.effect_op:<18} {h.p_value:>10.4f} {h.direction}") + + +if __name__ == "__main__": + main(sys.argv[1] if len(sys.argv) > 1 else "/workspace/shard_profiles") diff --git a/benchmarks/hft/profile_per_shard.sh b/benchmarks/hft/profile_per_shard.sh new file mode 100644 index 0000000..a07e319 --- /dev/null +++ b/benchmarks/hft/profile_per_shard.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Per-shard profiling for Granger causality (Task 3). +# +# attribute() runs grangercausalitytests at maxlag=2, which needs at LEAST +# 8 observations (samples) per op -- not the naive max_lag+2=4. And even at +# 8-20 samples Granger p-values are noisy: treat the ranking as suggestive, +# not conclusive. One nsys run on the whole 25M window gives 1 sample/op, +# so we profile many shards separately, one profile per shard = 1 sample/op +# each. Default N=20 shards (=100M events profiled) for a usable series. +# +# After this completes, parse_shard_profiles.py turns the 5 .sqlite files +# into a residual time series and runs attribute(). +# +# Run on the pod: +# bash profile_per_shard.sh /workspace/hft_numba_seed42 +# +# NOTE: this profiles 5 single-shard runs, NOT the 25M concatenated run. +# Each run loads exactly one part-0000N.parquet (5M events). That's the +# point -- we want per-shard kernel times as independent samples. + +set -euo pipefail + +STAGE="${1:-/workspace/hft_numba_seed42}" +OUT="${2:-/workspace/shard_profiles}" +HARNESS="benchmarks/hft/harness.py" + +mkdir -p "$OUT" + +# Make a temp staging dir holding one shard at a time, because the harness +# globs part-*.parquet from a directory. We symlink a single shard in, profile, +# then swap. seed is parsed from the stage dir name (…seed42 -> 42). +SEED="$(basename "$STAGE" | grep -oE '[0-9]+$')" + +NSHARDS="${3:-20}" +for ((i=0; i=25M): below target | Variance spec (<2%): slightly above spec + +### Event count comparison (sequential) +| Max Events | Mean Events/sec | Variance | +|------------|-----------------|----------| +| 10M | 8.9M | 10% | +| 15M | 11.6M | 22% | +| 25M | 20.7M | 3.4% | + +25M is clearly the best-balanced configuration of those tested. + +### Parallel runs (3 seeds simultaneously, staggered start) +| Config | Mean Events/sec | Variance | Notes | +|---|---|---|---| +| sleep=3s stagger, 15M events | 12.6M | 1.5% | passes variance spec, below throughput target | +| sleep=2s stagger, 25M events | 18.9M | 6.3% | unstable | +| no stagger, 25M events | crashes (Exit 1) | - | cuDF memory-pool conflict per process | + +Parallel harness execution on a single A100 is unreliable: each Python +process creates its own cuDF memory pool, so 2-3 simultaneous processes at +25M events frequently OOM. A short stagger (2-3s) avoids crashes but adds +variance. Sequential execution is the recommended approach for this harness. + +## Original Baseline (single-file format, original pod, C++ generator) +| Seed | Run 1 | Run 2 | Run 3 | Mean | +|------|-------|-------|-------|------| +| 42 | 29.5M | 31.1M | 29.1M | 29.9M | +| 43 | 31.1M | 29.5M | 30.0M | 30.2M | +| 44 | 30.9M | 29.8M | 30.8M | 30.5M | + +Mean: 30.2M events/sec | Variance: ~2% +Target (>=25M): PASS | Variance spec (<2%): PASS + +## Key Finding - Multi-shard vs Single File Format + +Original data: 1 file per seed (1B rows). Current data: 200 files per seed +(5M rows each, matching generate.py's shard size). + +| Format | Events/sec | +|---|---| +| Single file (1 x 1B rows) | ~30M | +| 200 shards (200 x 5M rows) | ~20M | + +Patched harness.load_events() to read only the files needed for max_events +(5 files for 25M events) instead of all 200 - throughput barely changed +(20.6M vs 20.7M), so file count alone isn't the main cause. The remaining +gap is likely pod/hardware variability (different RunPod instance, cuDF +26.6.0 vs 26.04) rather than the shard format itself. + +## Stall Profile (nsys, seed 42, 25M events) + +| | CPU | Data-stall | Sync | GPU active | +|---|---|---|---|---| +| Expected | <5% | 10-25% | 5-15% | 60-80% | +| Measured (original pod) | ~5% | ~73% | ~7% | ~15% | +| Measured (new pod) | ~5% | ~73% | ~8% | ~15% | + +Stall profile is consistent across pods and formats. + +### Top GPU Kernels +| Kernel | Original pod | New pod | +|---|---|---| +| zstd decompression | 49.3% | 47.5% | +| Parquet page decode | 24.1% | 26.2% | +| Merge sort (top-of-book) | 7.1% | 8.3% | + +Host-to-Device memory transfers: 91.2% of total memory time (7.98 GB, original pod). + +## Saturation Rule +GPU active (~15%) is well below the 85% threshold. No swap to 500M events required. + +## Task 1 - Predicted Graph Integration (W2) + +Built gitm/planner/hft_graph.py - a predicted execution graph for the HFT +pipeline (zstd_decompress, parquet_decode, merge_sort, vwap_reduce), using +the same roofline model as gitm/planner/graph.py (LLM decode predictions). + +### Residuals (predicted vs observed, seed 42, 25M events) + +| Op | Predicted (ms) | Observed (ms) | r_kt | Bound | +|---|---|---|---|---| +| zstd_decompress | 0.425 | 252.076 | 592x | memory | +| parquet_decode | 0.638 | 139.100 | 217x | memory | +| merge_sort | 0.638 | 43.927 | 68x | memory | +| vwap_reduce | 0.159 | not separately profiled | - | memory | + +### Key Finding + +All profiled ops run 68-592x slower than the memory-bandwidth-bound floor +predicts. roofline()'s compute-vs-memory model has no category for ops +bound by sequential/branchy algorithms (zstd entropy decoding, Parquet +page decode) rather than HBM bandwidth or FLOPs. This is a gap in the +planner's prediction model, not specific to HFT - it would affect any +data-loading stage prediction. + +### Task 3 - Granger Causality Next Step + +gitm.optimizer.attribution.attribute() requires >=4 samples per op +(max_lag=2 default). Current data has 1 sample per op (aggregate kernel +sum from one nsys run on seed 42). Next step: profile each of the 5 +shards in the 25M-event run separately to get 5 samples/op, enabling +Granger analysis of which stage causally drives the others' slowdown. + +Code: gitm/planner/hft_graph.py, gitm/planner/hft_residual_demo.py +Run: python3 -m gitm.planner.hft_residual_demo + +## Profiling Artifacts +- nsys profiles generated during testing (not committed to git - see + .gitignore for *.nsys-rep / *.sqlite) + +## Next Steps +1. Investigate remaining ~10M events/sec gap between pods (cuDF version, + pod hardware) since file-count fix alone didn't close it +2. Profile per-shard kernel times to enable Granger causality (Task 3) +3. Extend hft_graph.py with a non-bandwidth prediction category for + decode/decompress-bound ops (Task 1 follow-up) diff --git a/benchmarks/hft/spec.md b/benchmarks/hft/spec.md index 1471542..2923693 100644 --- a/benchmarks/hft/spec.md +++ b/benchmarks/hft/spec.md @@ -1,6 +1,6 @@ # HFT benchmark — spec -> Owner: Ash — baseline + profiling + spec doc. +> Owner: Ash baseline + profiling + spec doc. ## 1. Input definition Synthetic limit-order-book stream at 1×10⁹ events per seed, Parquet @@ -32,3 +32,60 @@ reductions; CPU is low because Arrow handles ingest off the hot path. **Saturation rule:** if measured GPU active > 85 %, flag for review same day — fall back to a 500 M-event shard and document. + +## 5. Profiling Methodology + +Profiled using NVIDIA Nsight Systems (nsys) on RunPod A100 80GB PCIe. + +Command: +nsys profile --trace cuda,nvtx,osrt --output /root/data/hft_baseline_1 \ + python3 benchmarks/hft/harness.py --seed 42 --stage /root/data --max-events 25000000 + +Results saved to: +- /root/data/hft_baseline_1.nsys-rep +- /root/data/hft_baseline_1.sqlite + +View with: +nsys stats /root/data/hft_baseline_1.nsys-rep + +## 6. Reproduction Steps + +**1. Clone repo:** +```bash +git clone git@github.com:GitM-Labs/runtime.git +cd runtime && git checkout hft_datageneration +pip install -e ".[dev]" +``` + +**2. Build the generator:** +```bash +cd benchmarks/hft/generator && mkdir -p build && cd build +export ARROW_LIB=$(python3 -c "import pyarrow; print(pyarrow.get_library_dirs()[0])") +ln -sf $ARROW_LIB/libparquet.so.2400 $ARROW_LIB/libparquet.so +cmake .. && make -j$(nproc) +``` + +**3. Generate datasets:** +```bash +export LD_LIBRARY_PATH=$ARROW_LIB:$LD_LIBRARY_PATH +./hft_gen 1000000000 42 /root/data/hft_1b_seed42/part-00000.parquet +./hft_gen 1000000000 43 /root/data/hft_1b_seed43/part-00000.parquet +./hft_gen 1000000000 44 /root/data/hft_1b_seed44/part-00000.parquet +``` + +**4. Run baseline (3x per seed, take means):** +```bash +export GITM_BENCH_STAGE="/root/data" +python3 benchmarks/hft/harness.py --seed 42 --stage /root/data --max-events 25000000 +``` + +**5. Run nsys profile:** +```bash +nsys profile --trace cuda,nvtx,osrt --output /root/data/hft_baseline_1 \ + python3 benchmarks/hft/harness.py --seed 42 --stage /root/data --max-events 25000000 +``` + +**6. View results:** +```bash +nsys stats /root/data/hft_baseline_1.nsys-rep +``` diff --git a/benchmarks/hft/test_gpu_generator.py b/benchmarks/hft/test_gpu_generator.py new file mode 100644 index 0000000..ab7e668 --- /dev/null +++ b/benchmarks/hft/test_gpu_generator.py @@ -0,0 +1,41 @@ +import time + +import cupy as cp +import pyarrow as pa +import pyarrow.parquet as pq + +n = 5_000_000 +seed = 42 + +rng = cp.random.default_rng(seed) +symbol_id = rng.integers(0, 512, size=n, dtype=cp.int32) +side = rng.integers(0, 2, size=n, dtype=cp.int8) +price = rng.integers(9000, 11000, size=n, dtype=cp.int64) +size = rng.integers(1, 1000, size=n, dtype=cp.int32) +etype = rng.integers(0, 3, size=n, dtype=cp.int8) +gaps = rng.integers(1, 2000, size=n, dtype=cp.int64) +ts_ns = cp.cumsum(gaps) + +table = pa.table({ + 'ts_ns': pa.array(cp.asnumpy(ts_ns)), + 'symbol_id': pa.array(cp.asnumpy(symbol_id)), + 'side': pa.array(cp.asnumpy(side)), + 'price': pa.array(cp.asnumpy(price)), + 'size': pa.array(cp.asnumpy(size)), + 'type': pa.array(cp.asnumpy(etype)), +}) + +# Test no compression +t0 = time.perf_counter() +pq.write_table(table, '/workspace/test_nocomp.parquet', compression='none') +print(f"No compression: {time.perf_counter()-t0:.3f}s") + +# Test zstd +t0 = time.perf_counter() +pq.write_table(table, '/workspace/test_zstd.parquet', compression='zstd', compression_level=1) +print(f"zstd-1: {time.perf_counter()-t0:.3f}s") + +# Test snappy +t0 = time.perf_counter() +pq.write_table(table, '/workspace/test_snappy.parquet', compression='snappy') +print(f"snappy: {time.perf_counter()-t0:.3f}s") diff --git a/benchmarks/hft/test_hawkes_gpu.py b/benchmarks/hft/test_hawkes_gpu.py new file mode 100644 index 0000000..e978b0c --- /dev/null +++ b/benchmarks/hft/test_hawkes_gpu.py @@ -0,0 +1,57 @@ +import time + +import numpy as np +from numba import njit + +MU = 100.0 +ALPHA = 0.6 +BETA = 0.8 + +@njit(cache=True) +def hawkes_numba(u, n, t0, lam0): + """JIT-compiled Hawkes — exact, fast, sequential.""" + ts = np.empty(n, dtype=np.int64) + t = t0 + lam = lam0 + for i in range(n): + dt = -np.log(u[i]) / lam + t += dt + lam = MU + (lam - MU) * np.exp(-BETA * dt) + ALPHA + ts[i] = int(t * 1e9) + return ts, t, lam + +def hawkes_cpu_plain(n, t0, lam0, seed): + rng = np.random.default_rng(seed) + ts = np.empty(n, dtype=np.int64) + t = t0 + lam = lam0 + for i in range(n): + u = rng.random() + dt = -np.log(u) / lam + t += dt + lam = MU + (lam - MU) * np.exp(-BETA * dt) + ALPHA + ts[i] = int(t * 1e9) + return ts + +n = 1_000_000 +rng_np = np.random.default_rng(42) +u = rng_np.random(n) + +# Warmup numba JIT +print("Warming up Numba JIT...") +_ = hawkes_numba(u[:100], 100, 0.0, MU) + +# Numba +t0 = time.perf_counter() +ts_numba, _, _ = hawkes_numba(u, n, 0.0, MU) +numba_time = time.perf_counter() - t0 +print(f"Numba JIT: {numba_time:.3f}s = {n/numba_time/1e6:.1f}M/sec") + +# Plain CPU +t0 = time.perf_counter() +ts_cpu = hawkes_cpu_plain(n, 0.0, MU, 42) +cpu_time = time.perf_counter() - t0 +print(f"Plain CPU: {cpu_time:.3f}s = {n/cpu_time/1e6:.1f}M/sec") + +print(f"Speedup: {cpu_time/numba_time:.1f}x") +print(f"Monotonic: {bool(np.all(np.diff(ts_numba) > 0))}") diff --git a/benchmarks/hft/test_rng_speed.py b/benchmarks/hft/test_rng_speed.py new file mode 100644 index 0000000..4bc5758 --- /dev/null +++ b/benchmarks/hft/test_rng_speed.py @@ -0,0 +1,33 @@ +import time + +import cupy as cp +import numpy as np + +# Warmup CuPy first +print("Warming up CuPy...") +_warm = cp.random.default_rng(0).integers(0, 100, size=100_000, dtype=cp.int32) +cp.cuda.Stream.null.synchronize() +print("Done\n") + +for n in [5_000_000, 50_000_000, 100_000_000, 500_000_000]: + # CPU numpy - 3 runs, take mean + cpu_times = [] + for _ in range(3): + rng_cpu = np.random.default_rng(42) + t0 = time.perf_counter() + _ = rng_cpu.integers(1, 1000, size=n, dtype=np.int32) + cpu_times.append(time.perf_counter() - t0) + cpu_time = sum(cpu_times) / 3 + + # GPU cupy - 3 runs, take mean + gpu_times = [] + for _ in range(3): + rng_gpu = cp.random.default_rng(42) + cp.cuda.Stream.null.synchronize() + t0 = time.perf_counter() + _ = rng_gpu.integers(1, 1000, size=n, dtype=cp.int32) + cp.cuda.Stream.null.synchronize() + gpu_times.append(time.perf_counter() - t0) + gpu_time = sum(gpu_times) / 3 + + print(f"n={n/1e6:.0f}M: numpy={cpu_time:.3f}s cupy={gpu_time:.3f}s speedup={cpu_time/gpu_time:.1f}x") diff --git a/gitm/planner/hft_graph.py b/gitm/planner/hft_graph.py new file mode 100644 index 0000000..18f2728 --- /dev/null +++ b/gitm/planner/hft_graph.py @@ -0,0 +1,134 @@ +"""Predicted execution graph for the HFT LOB-replay benchmark. + +Mirrors gitm.planner.graph (which predicts LLM decode steps) but for the +HFT data-loading + LOB-replay pipeline: zstd decompression, Parquet page +decode, merge sort (top-of-book ordering by symbol_id/ts_ns), and the +VWAP/microprice reductions. + +All four ops here have flops=0, so roofline() always returns +bound="memory" and t_pred_s = bytes_moved / peak_mem_bw_bytes_per_s. + +Numbers are derived from benchmarks/hft/results.md (seed 42, 25M-event +warm window, 200-shard zstd-1 Parquet, schema: ts_ns i64, symbol_id i32, +side i8, price i64, size i32, type i8 = 26 bytes/row uncompressed). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from gitm.planner.graph import Graph, PredictedNode +from gitm.planner.roofline import BatchConfig, HardwareSpec, ModelSpec, roofline + + +@dataclass(frozen=True) +class HFTDatasetSpec: + """Shape of one HFT harness run, used to size the predicted graph. + + compression_ratio is an ESTIMATE (zstd-1 on mixed int columns is + typically 2.5-4x). TODO: replace with the real ratio once file sizes + can be checked on the pod (compressed_size / uncompressed_size). + """ + + name: str = "hft_seed42_25M" + n_events: int = 25_000_000 + bytes_per_event_uncompressed: int = 26 # i64+i32+i8+i64+i32+i8 + compression_ratio: float = 3.0 # ESTIMATE — verify against real files + + @property + def uncompressed_bytes(self) -> int: + return self.n_events * self.bytes_per_event_uncompressed + + @property + def compressed_bytes(self) -> int: + return int(self.uncompressed_bytes / self.compression_ratio) + + +def predict_hft_graph( + dataset: HFTDatasetSpec | None = None, + hw: HardwareSpec | None = None, +) -> Graph: + """Emit a predicted execution graph for the HFT LOB-replay harness. + + Pipeline stages (ordered as they appear in the nsys kernel summary): + 1. zstd_decompress - read compressed column chunks, write decompressed + 2. parquet_decode - decode columnar pages into cuDF columns + 3. merge_sort - sort by (symbol_id, ts_ns) for top-of-book replay + 4. vwap_reduce - cummax/cumsum reductions for microprice + VWAP + + Each PredictedNode.prediction.t_pred_s is the *memory-bandwidth-bound* + floor for that op — the fastest it could possibly run on this GPU if it + were purely bandwidth limited. Comparing this floor to the observed + kernel time (via gitm.optimizer.monitor.residuals) tells us whether the + op is actually bandwidth-bound or bound by something roofline doesn't + model (e.g. sequential entropy decoding, branchy page-decode logic). + """ + dataset = dataset or HFTDatasetSpec() + hw = hw or HardwareSpec() + + # model/batch are unused for HFT but required by the Graph dataclass — + # pass defaults so this stays type-compatible with the LLM planner. + g = Graph(model=ModelSpec(), hw=hw, batch=BatchConfig()) + + comp = dataset.compressed_bytes + uncomp = dataset.uncompressed_bytes + + # 1. zstd decompression: read compressed bytes in, write uncompressed out + g.nodes.append( + PredictedNode( + "zstd_decompress", + layer=None, + prediction=roofline( + "zstd_decompress", flops=0, bytes_moved=comp + uncomp, hw=hw + ), + expected_stream_id=0, + ) + ) + + # 2. Parquet page decode: ~1x uncompressed in, ~1x out -> 2x uncompressed + g.nodes.append( + PredictedNode( + "parquet_decode", + layer=None, + prediction=roofline( + "parquet_decode", flops=0, bytes_moved=2 * uncomp, hw=hw + ), + expected_stream_id=0, + ) + ) + + # 3. Merge sort by (symbol_id, ts_ns): full read + write of reordered rows + g.nodes.append( + PredictedNode( + "merge_sort", + layer=None, + prediction=roofline( + "merge_sort", flops=0, bytes_moved=2 * uncomp, hw=hw + ), + expected_stream_id=0, + ) + ) + + # 4. VWAP / microprice reductions: reads price+size, writes derived cols + g.nodes.append( + PredictedNode( + "vwap_reduce", + layer=None, + prediction=roofline( + "vwap_reduce", flops=0, bytes_moved=int(0.5 * uncomp), hw=hw + ), + expected_stream_id=0, + ) + ) + + return g + + +if __name__ == "__main__": + g = predict_hft_graph() + print(f"Dataset: {HFTDatasetSpec().name}") + print(f"{'op':<18} {'t_pred (ms)':>12} {'bound':>8} {'bytes_moved':>14}") + for n in g.nodes: + p = n.prediction + print(f"{n.op:<18} {p.t_pred_s * 1e3:12.4f} {p.bound:>8} {p.bytes:14,.0f}") + print(f"\ntotal predicted: {g.total_pred_s * 1e3:.4f} ms") diff --git a/gitm/planner/hft_residual_demo.py b/gitm/planner/hft_residual_demo.py new file mode 100644 index 0000000..a1115a1 --- /dev/null +++ b/gitm/planner/hft_residual_demo.py @@ -0,0 +1,78 @@ +"""Predicted-vs-observed residuals for the HFT pipeline (Task 1 demo). + +This is a stand-in for the full nsys-sqlite -> gitm.tracer.schema.Trace +converter (the larger W2 piece). It hardcodes the observed kernel times +from benchmarks/hft/results.md (seed 42, 25M-event run, new pod) and +computes the same residual gitm.optimizer.monitor.residuals() would: + + r_kt = (t_obs - t_pred) / t_pred + +A large positive r_kt means the kernel ran far slower than the +memory-bandwidth-bound floor predicts -- i.e. the op is NOT +bandwidth-limited, and roofline's compute/memory model doesn't capture +what's actually limiting it (e.g. sequential entropy decoding for zstd, +branchy per-row page decode for Parquet). + +Run: + python3 -m gitm.planner.hft_residual_demo +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from gitm.planner.hft_graph import predict_hft_graph + + +@dataclass +class ObservedKernel: + op: str + t_obs_ms: float + source: str + + +# From benchmarks/hft/results.md, "Top GPU Kernels (nsys, seed 42)" / Run 2 +# (new pod, 200-shard format, 25M-event warm window). +OBSERVED = [ + ObservedKernel("zstd_decompress", 252.076091, "zstd::decompression_kernel"), + ObservedKernel("parquet_decode", 139.100450, "decode_page_data_generic"), + # DeviceMergeSortMergeKernel (26.555531ms) + DeviceMergeSortBlockSortKernel + # (17.371249ms) -- both are part of the top-of-book sort. + ObservedKernel("merge_sort", 26.555531 + 17.371249, "DeviceMergeSort*"), +] + + +def main() -> None: + g = predict_hft_graph() + pred_by_op = {n.op: n.prediction for n in g.nodes} + + print(f"{'op':<18} {'pred (ms)':>10} {'obs (ms)':>10} {'r_kt':>10} bound") + print("-" * 60) + + for ok in OBSERVED: + pred = pred_by_op[ok.op] + t_pred_ms = pred.t_pred_s * 1e3 + r_kt = (ok.t_obs_ms - t_pred_ms) / t_pred_ms + print( + f"{ok.op:<18} {t_pred_ms:10.4f} {ok.t_obs_ms:10.3f} " + f"{r_kt:10.1f} {pred.bound}" + ) + + vwap = pred_by_op["vwap_reduce"] + print(f"{'vwap_reduce':<18} {vwap.t_pred_s * 1e3:10.4f} {'?':>10} {'?':>10} {vwap.bound}") + + print() + print("All three profiled ops are ~70-600x slower than the memory-bandwidth") + print("floor predicts. roofline()'s compute-vs-memory model doesn't have a") + print("category for these: zstd entropy decode and Parquet page decode are") + print("bound by sequential/branchy per-element work, not by HBM bandwidth.") + print() + print("Granger causality (gitm.optimizer.attribution.attribute) needs >= 4") + print("samples per op (max_lag=2). This demo has 1 sample per op (the") + print("seed-42 aggregate kernel sum). Next step: profile each of the 5") + print("shards in the 25M run separately -> 5 samples/op, enough to run") + print("attribute() and rank which stage Granger-causes the others.") + + +if __name__ == "__main__": + main() diff --git a/gitm/tracer/vllm_sched.py b/gitm/tracer/vllm_sched.py new file mode 100644 index 0000000..a10a876 --- /dev/null +++ b/gitm/tracer/vllm_sched.py @@ -0,0 +1,142 @@ +"""vLLM scheduler-stats adapter-the engine level signal alongside CUPTI. + +The CUPTI /Kineto trace tells us what the GPU did, it can't tell us why the +batch was the size it was. vLLM scheduler told that:how many sequences are +running vs waiting, KV-cache occupancy, preemptions. Sampling it alongside +the decode traces lets attribution seperate "GPU-bound" from "scheduler-starved". +(e.g. A half-full batch because requests are queued behind KV-cache pressures). + +This adapter duck-types the engine so it works across vLLM versions and in tests: +it reads a ``Stats``-like object if the engine exposes one, otherwise it +introspect ``engine scheduler``(running/waiting/swapped+KV usage). Sampling +is cheap(reads counter), so the loop can poll it at decode-step cadence. +""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class SchedulerSample(BaseModel): + """One point in time read of the vLLM scheduler state.""" + + model_config = ConfigDict(extra="forbid") + + ts_ns: int + num_running: int = 0 + num_waiting: int = 0 + num_swapped: int = 0 + num_preemptions: int = 0 + gpu_cache_usage: float = 0.0 # fraction [0, 1] of KV-cache blocks in use + batch_size: int = 0 # sequences in the current decode batch + + @property + def queue_depth(self) -> int: + """Sequences not currently being decoded(waiting + swapped).""" + return self.num_waiting + self.num_swapped + + @property + def batch_occupancy(self) -> float: + """Running fraction of all live sequences-low = scheduler-starved""" + live = self.num_running + self.queue_depth + return self.num_running / live if live else 0.0 + + +def _first_attr(obj: Any, *names: str, default: Any = None) -> Any: + for n in names: + if hasattr(obj, n): + v = getattr(obj, n) + return v() if callable(v) else v + return default + + +def _len_or_int(v: Any) -> int: + if v is None: + return 0 + if isinstance(v, int): + return v + try: + return len(v) + except TypeError: + return 0 + + +def sample_scheduler(engine: Any, *, ts_ns: int) -> SchedulerSample: + """Read a :class:SchedulerSample from a vLLM engine (or a stand-in). + + Accepts, in order of preference: + - an object with ``get_scheduler_stats()`` / ``get_stats()`` returning a + mapping or stats object, + - an engine exposing ``.scheduler`` with running/waiting/swapped, + - a plain mapping of the field names. + """ + if isinstance(engine, dict): + src: Any = engine + else: + stats = _first_attr(engine, "get_scheduler_stats", "get_stats") + src = stats if stats is not None else engine + + # Mapping form (already-collected stats or a test dict). + if isinstance(src, dict): + running = _len_or_int(src.get("num_running", src.get("running"))) + waiting = _len_or_int(src.get("num_waiting", src.get("waiting"))) + swapped = _len_or_int(src.get("num_swapped", src.get("swapped"))) + preempt = int(src.get("num_preemptions", 0) or 0) + usage = float(src.get("gpu_cache_usage", src.get("gpu_cache_usage_sys", 0.0)) or 0.0) + batch = _len_or_int(src.get("batch_size", running)) + return SchedulerSample( + ts_ns=ts_ns, + num_running=running, + num_waiting=waiting, + num_swapped=swapped, + num_preemptions=preempt, + gpu_cache_usage=usage, + batch_size=batch or running, + ) + + # Object form: introspect a scheduler. + scheduler = _first_attr(src, "scheduler", default=src) + running = _len_or_int(_first_attr(scheduler, "running", "num_running")) + waiting = _len_or_int(_first_attr(scheduler, "waiting", "num_waiting")) + swapped = _len_or_int(_first_attr(scheduler, "swapped", "num_swapped")) + preempt = int(_first_attr(scheduler, "num_preemptions", default=0) or 0) + usage = float(_first_attr(scheduler, "gpu_cache_usage", default=0.0) or 0.0) + return SchedulerSample( + ts_ns=ts_ns, + num_running=running, + num_waiting=waiting, + num_swapped=swapped, + num_preemptions=preempt, + gpu_cache_usage=usage, + batch_size=running, + ) + + +class SchedulerStatsTracker(BaseModel): + """Accumulates :class:SchedulerSample points over a decode run.""" + + model_config = ConfigDict(extra="forbid") + + samples: list[SchedulerSample] = Field(default_factory=list) + + def record(self, engine: Any, *, ts_ns: int) -> SchedulerSample: + s = sample_scheduler(engine, ts_ns=ts_ns) + self.samples.append(s) + return s + + def summarize(self) -> dict[str, Any]: + """Aggregate the series into report-ready scheduler signals.""" + if not self.samples: + return {"n_samples": 0} + n = len(self.samples) + return { + "n_samples": n, + "mean_batch_occupancy": sum(s.batch_occupancy for s in self.samples) / n, + "mean_queue_depth": sum(s.queue_depth for s in self.samples) / n, + "max_queue_depth": max(s.queue_depth for s in self.samples), + "mean_gpu_cache_usage": sum(s.gpu_cache_usage for s in self.samples) / n, + "total_preemptions": max(s.num_preemptions for s in self.samples) - min(s.num_preemptions for s in self.samples), + "starved_fraction": sum(1 for s in self.samples if s.batch_occupancy < 0.5) / n, + } diff --git a/gitm/tracer/vllm_stats.py b/gitm/tracer/vllm_stats.py index c8b5db4..e8b7bd4 100644 --- a/gitm/tracer/vllm_stats.py +++ b/gitm/tracer/vllm_stats.py @@ -90,13 +90,52 @@ def _len_or_none(x: Any) -> int | None: def _schedulers(engine: Any) -> list[Any]: """Resolve the engine's scheduler(s) as a list (vLLM may keep one per PP stage).""" sched = _first_attr( - engine, "scheduler", "engine.scheduler", "llm_engine.scheduler" + engine, + # vLLM V0. + "scheduler", + "engine.scheduler", + "llm_engine.scheduler", + # vLLM V1 (in-process, VLLM_ENABLE_V1_MULTIPROCESSING=0): the scheduler + # lives inside the EngineCore, not on the LLMEngine. Paths vary across + # v0.2x point releases — GPU-validate the exact one on the target build. + "llm_engine.engine_core.engine_core.scheduler", + "llm_engine.engine_core.scheduler", + "engine_core.engine_core.scheduler", + "engine_core.scheduler", ) if sched is None: return [] return list(sched) if isinstance(sched, list | tuple) else [sched] +def _v1_scheduler_stats(scheduler: Any) -> dict[str, Any]: + """vLLM V1 stats via ``scheduler.make_stats()`` — V1 doesn't keep the V0 + running/waiting deques in the same shape, but exposes a stats object with + ``num_running_reqs`` / ``num_waiting_reqs`` / ``kv_cache_usage``. Best-effort; + returns only the fields it actually found. GPU-validate the attr names on the + target vLLM build (they've shifted across v0.2x). + """ + make = getattr(scheduler, "make_stats", None) + if not callable(make): + return {} + try: + stats = make() + except Exception: + return {} + if stats is None: + return {} + out: dict[str, Any] = {} + for field_name, attr in (("num_running", "num_running_reqs"), + ("num_waiting", "num_waiting_reqs")): + v = getattr(stats, attr, None) + if isinstance(v, int): + out[field_name] = v + ku = getattr(stats, "kv_cache_usage", None) + if isinstance(ku, int | float): + out["gpu_cache_usage"] = float(ku) + return out + + def _max_num_seqs(engine: Any) -> int | None: val = _first_attr( engine, @@ -148,6 +187,14 @@ def read_scheduler_stats(engine: Any, *, t_ns: int = 0) -> SchedulerSample | Non sample.gpu_cache_usage = usage saw_any = True + # vLLM V1: fill running/waiting/cache from the scheduler's stats object + # where the V0 deques weren't exposed (they read empty on V1). + for sch in schedulers: + for field_name, val in _v1_scheduler_stats(sch).items(): + if getattr(sample, field_name) is None: + setattr(sample, field_name, val) + saw_any = True + # Total unfinished — a stable public method on LLMEngine across versions. getter = _first_attr( engine, diff --git a/tests/test_vllm_sched.py b/tests/test_vllm_sched.py new file mode 100644 index 0000000..313bdac --- /dev/null +++ b/tests/test_vllm_sched.py @@ -0,0 +1,40 @@ +"""Scheduler-stats adapter: duck-types engine forms and aggregates a series.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from gitm.tracer.vllm_sched import SchedulerStatsTracker, sample_scheduler + + +def test_sample_from_mapping(): + s = sample_scheduler({"num_running": 6, "num_waiting": 4, "gpu_cache_usage": 0.8}, ts_ns=10) + assert s.num_running == 6 + assert s.queue_depth == 4 + assert s.batch_occupancy == 6 / 10 + assert s.gpu_cache_usage == 0.8 + + +def test_sample_from_scheduler_object(): + engine = SimpleNamespace(scheduler=SimpleNamespace(running=[1, 2, 3], waiting=[4], swapped=[])) + s = sample_scheduler(engine, ts_ns=5) + assert s.num_running == 3 + assert s.num_waiting == 1 + assert s.queue_depth == 1 + + +def test_sample_from_get_stats_method(): + engine = SimpleNamespace(get_scheduler_stats=lambda: {"running": 2, "waiting": 8}) + s = sample_scheduler(engine, ts_ns=1) + assert s.num_running == 2 and s.num_waiting == 8 + assert s.batch_occupancy == 0.2 # starved + + +def test_tracker_summarize(): + tr = SchedulerStatsTracker() + tr.record({"num_running": 1, "num_waiting": 9}, ts_ns=1) # starved + tr.record({"num_running": 9, "num_waiting": 1}, ts_ns=2) # healthy + summary = tr.summarize() + assert summary["n_samples"] == 2 + assert summary["max_queue_depth"] == 9 + assert summary["starved_fraction"] == 0.5 diff --git a/tests/test_vllm_stats_v1.py b/tests/test_vllm_stats_v1.py new file mode 100644 index 0000000..3f818fe --- /dev/null +++ b/tests/test_vllm_stats_v1.py @@ -0,0 +1,50 @@ +"""Read scheduler stats from a vLLM V1-shaped engine. + +V1 keeps the scheduler inside EngineCore and exposes a stats object via +`make_stats()` (num_running_reqs / num_waiting_reqs / kv_cache_usage) instead of +the V0 running/waiting deques. Fake that shape to pin the read path off-GPU; the +exact attr names still need GPU validation on the target vLLM build. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from gitm.tracer.vllm_stats import read_scheduler_stats + + +def _v1_engine(running: int, waiting: int, cache: float, max_seqs: int = 256): + stats = SimpleNamespace( + num_running_reqs=running, num_waiting_reqs=waiting, kv_cache_usage=cache + ) + scheduler = SimpleNamespace(make_stats=lambda: stats) + # llm.llm_engine.engine_core.engine_core.scheduler (in-process V1) + return SimpleNamespace( + llm_engine=SimpleNamespace( + engine_core=SimpleNamespace(engine_core=SimpleNamespace(scheduler=scheduler)) + ), + scheduler_config=SimpleNamespace(max_num_seqs=max_seqs), + ) + + +def test_reads_v1_stats_object(): + s = read_scheduler_stats(_v1_engine(running=3, waiting=12, cache=0.87), t_ns=0) + assert s is not None + assert s.num_running == 3 + assert s.num_waiting == 12 + assert s.gpu_cache_usage == 0.87 + assert s.batch_occupancy == 3 / 256 # occupancy derived from V1 num_running + + +def test_none_engine_still_none(): + assert read_scheduler_stats(None) is None + + +def test_v1_engine_without_stats_is_none(): + # a V1-shaped engine whose scheduler exposes nothing readable -> None, no crash. + empty = SimpleNamespace( + llm_engine=SimpleNamespace( + engine_core=SimpleNamespace(engine_core=SimpleNamespace(scheduler=object())) + ) + ) + assert read_scheduler_stats(empty) is None