Hft datageneration vllm sched

[ashwyan]

No description provided.

[github-actions]

Code Review by Gemini

The changes introduce a comprehensive set of tools and documentation for HFT benchmarks, including data generation, profiling, and intervention analysis. A new module for vLLM scheduler statistics is also added, enhancing observability. The code is generally well-structured and documented.

Here are some minor observations and suggestions:

benchmarks/hft/headroom_kernel_rank.py

The covered_kernels function directly uses an internal function _applies from gitm.optimizer.replay. While this is within the same repository, relying on internal functions can sometimes lead to unexpected breakage if the internal API changes. For a benchmark utility script, this might be acceptable, but it's worth noting.

--- a/benchmarks/hft/headroom_kernel_rank.py
+++ b/benchmarks/hft/headroom_kernel_rank.py
@@ -30,6 +30,7 @@
 
 
 def covered_kernels(spec: InterventionSpec, trace: Trace) -> list[str]:
+    # NOTE: Uses internal function _applies from gitm.optimizer.replay.
     return [k.name for k in trace.kernels() if _applies(spec, k.name)]
 
 

gitm/tracer/vllm_sched.py

In the sample_scheduler function, when constructing the SchedulerSample from a dictionary source, the batch_size assignment uses a fallback logic (batch or running). This is a reasonable heuristic, but adding a comment would clarify the intent, especially for cases where batch_size might be explicitly 0 but num_running is not.

--- a/gitm/tracer/vllm_sched.py
+++ b/gitm/tracer/vllm_sched.py
@@ -69,7 +69,9 @@
             num_swapped=swapped,
             num_preemptions=preempt,
             gpu_cache_usage=usage,
-            batch_size=batch or running,
+            # If 'batch_size' is not provided or is 0, default to 'num_running'.
+            # This handles cases where the source doesn't explicitly report batch_size.
+            batch_size=batch or running, 
         )
 
     # Object form: introspect a scheduler.

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

vllm_sched.py — spurious imports that will crash at import time

from numba.core.types import none          # unused; pulls in all of Numba
from pydantic.fields import PropertyT      # not a public symbol in Pydantic v2; ImportError

These will cause an ImportError in any environment without Numba installed alongside vLLM, and PropertyT was removed from the public API in Pydantic v2.

vllm_sched.py — total_preemptions is wrong

"total_preemptions": max(s.num_preemptions for s in self.samples),

num_preemptions is a monotonically-increasing cumulative counter (vLLM increments it over the session). max() gives you the final cumulative value, not the count within the tracked window. You want max(...) - min(...) or last - first.

gpu_generator.py — silently drops the tail when n_events % chunk_size != 0

n_chunks = n_events // chunk_size

If --events 100_000_000 and --chunk-size 5_000_000 this is fine, but any non-multiple silently truncates. No warning is printed and the final count printed still says the full requested number.

gpu_generator.py — both RNGs seeded identically, producing correlated outputs

rng    = cp.random.default_rng(seed)    # GPU fields
rng_np = np.random.default_rng(seed)   # Hawkes timestamps

Both are seeded with the same value, so the NumPy uniform draw u and the CuPy integer draws start from equivalent RNG states. For a benchmark this may be acceptable, but the comparison table claims "exact Hawkes" realism — correlated timestamps and prices are a subtle data-quality issue.

parse_shard_profiles.py — comment/code mismatch
The module docstring says "5 per-shard nsys profiles" and "5 samples / op", but the script and HFT_TASK_TOOLS.md both target 20 shards. The n_shards variable is derived from whatever CSVs are present, so it works at runtime, but the stale "5" in the docstring will mislead anyone reading it.

profile_per_shard.sh — nsys stats output flag misuse

nsys stats --report cuda_gpu_kern_sum --format csv \
  --output "$OUT/shard_$i" "$OUT/shard_$i.nsys-rep"

--output for nsys stats specifies a prefix, so it writes shard_0_cuda_gpu_kern_sum.csv. parse_shard_profiles.py globs for shard_*_cuda_gpu_kern_sum.csv — this matches. However, the || true suppresses all errors including "profile not found", making failures invisible.

test_vllm_stats_v1.py — batch_occupancy assertion is wrong

assert s.batch_occupancy == 3 / 256   # occupancy derived from V1 num_running

batch_occupancy in vllm_stats.py is computed as running / (running + waiting + swapped) = 3 / 15, not 3 / 256. This test will fail unless read_scheduler_stats uses max_num_seqs as the denominator, which is non-obvious and should be documented.

---

🔒 Security

profile_per_shard.sh — potential path injection via basename

SEED="$(basename "$STAGE" | grep -oE '[0-9]+$')"

If $STAGE is attacker-controlled or a mis-typed path without a trailing number, SEED becomes empty. The subsequent mkdir -p "$TMP" and ln -sf then use hft_numba_seed (no number) as a directory name — harmless here but the empty-seed case is never checked. A simple [[ -z "$SEED" ]] && { echo "could not parse seed from $STAGE"; exit 1; } guard is warranted.

---

⚡ Performance

gpu_generator.py — all GPU fields transferred separately (6 round-trips)

symbol_id = cp.asnumpy(rng.integers(...))
side      = cp.asnumpy(rng.integers(...))
price     = cp.asnumpy(rng.integers(...))
...

Six separate cp.asnumpy calls issue six H2D→D2H transfers synchronously. A single cp.cuda.Stream.null.synchronize() followed by building all arrays into one structured NumPy array (or stacking into a 2-D CuPy array then calling cp.asnumpy once) would halve PCIe traffic and reduce kernel launch overhead.

headroom_kernel_rank.py — covered_kernels scans the trace twice

cov = covered_kernels(spec, trace)          # O(n) scan
cov_ns = sum(... for k in trace.kernels()   # second O(n) scan
             if k.name in cov)

With only 4 kernels this is trivial, but the pattern is worth fixing for when the trace grows: fold both into one pass.

---

📊 Reproducibility

gpu_generator.py — GPU RNG state after warmup() is discarded

def warmup():
    ...
    cp.random.default_rng(0).integers(0, 100, size=10_000)  # throwaway rng

warmup() creates its own RNG seeded with 0 and discards it, so it doesn't corrupt generate()'s RNG — fine. But generate() creates rng = cp.random.default_rng(seed) after warmup(), meaning the CUDA RNG warm-up kernel leaves the device in a slightly different internal state that can affect subsequent cuRAND state initialization on some driver versions. Explicitly calling cp.cuda.Stream.null.synchronize() between warmup() and constructing the production RNG is cheap insurance.

gen_manifest.py / manifest.yaml — only part-00000 is hashed
The manifest pins a single shard's SHA-256 per seed, but the generator produces 200 shards. Any corruption in shards 1–199 goes undetected by make verify. Consider hashing all shards or at minimum documenting that the manifest is a spot-check, not a full integrity guarantee.

---

💡 Suggestions

vllm_sched.py — missing __all__; internal helpers are exposed
_first_attr, _len_or_int are prefixed with _ but there is no __all__, so from gitm.tracer.vllm_sched import * would pull them in. Since vllm_stats.py defines the same _first_attr and _len_or_none, these utilities should be factored into a shared private module to avoid drift.

hft_library.yaml — hft_multistream_symbol_shards lists DeviceReduce as a target kernel, but OBSERVED_KERNELS_NS in headroom_kernel_rank.py has no DeviceReduce entry
The coverage computation for candidate 2 will be lower than intended because the DeviceReduce kernels weren't captured. Either add them to OBSERVED_KERNELS_NS with their measured times, or note in the YAML that this kernel was below the profiling threshold.

profile_per_shard.sh — the comment still says "5 .sqlite files"

# After this completes, parse_shard_profiles.py turns the 5 .sqlite files

Should say 20 (or $NSHARDS).

measure_compression.py — only prints first 5 shards but totals all 200
The summary is correct, but a user who spots the "5 shards" in the output header and doesn't read the footer might stop early. A note like (showing first 5 of N shards; totals include all N) would prevent confusion.

test_gpu_generator.py / test_hawkes_gpu.py / test_rng_speed.py — these are ad-hoc scripts, not pytest tests
They write to /workspace/ and print to stdout but have no assertions. They should either be moved to scripts/ or converted to proper tests with pytest and temporary directories, so they don't get accidentally picked up by CI.

[github-actions]

Code Review by Gemini

--- a/benchmarks/hft/gpu_generator.py
+++ b/benchmarks/hft/gpu_generator.py
@@ -36,7 +36,9 @@
     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)
+    # Use np.select for clearer multi-condition logic, especially if more event types are added.
+    conditions = [roll < 0.55, roll < 0.90]
+    choices = [0, 1]
+    etype = np.select(conditions, choices, default=2).astype(np.int8)
 
     return {
         'ts_ns': ts_ns, 'symbol_id': symbol_id, 'side': side,

Reasoning for change: The original np.where nested structure is functional but can become less readable with more conditions. np.select provides a more explicit and scalable way to handle multiple conditional assignments, improving code clarity.

---

--- a/benchmarks/hft/headroom_kernel_rank.py
+++ b/benchmarks/hft/headroom_kernel_rank.py
@@ -60,7 +60,7 @@
 
 
 def covered_kernels(spec: InterventionSpec, trace: Trace) -> list[str]:
-    return [k.name for k in trace.kernels() if _applies(spec, k.name)]
+    return [k.name for k in trace.kernels() if _applies_to_kernel(spec, k.name)]
 
 
 def main() -> None:

Reasoning for change: The function _applies is an internal function (indicated by the leading underscore) from gitm.optimizer.replay. Relying on internal functions can lead to brittle code if the internal implementation changes. It's generally better to use public APIs. If this functionality is intended for external use by tools like headroom_kernel_rank.py, it should be exposed as a public function (e.g., applies_to_kernel) in gitm.kernels.spec or a similar public module. Renaming it to _applies_to_kernel here and assuming it's a local helper for this script, or making it a public function in gitm.kernels.spec would be better. For now, I've renamed it to _applies_to_kernel to indicate it's a helper function within this context, assuming _applies from gitm.optimizer.replay is not meant for direct external consumption.

---

--- a/gitm/tracer/vllm_sched.py
+++ b/gitm/tracer/vllm_sched.py
@@ -10,7 +10,6 @@
 from __future__ import annotations
 
 from typing import Any
-
 from numba.core.types import none
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic.fields import PropertyT

Reasoning for change: The import from numba.core.types import none is unused in this file. Removing unused imports helps keep the code clean and avoids potential confusion.

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

vllm_sched.py — bad imports that will crash at import time

from numba.core.types import none        # unused, wrong package for this module
from pydantic.fields import PropertyT    # internal pydantic API, not public

Both lines will cause ImportError on mismatched versions. Neither symbol is used anywhere in the file.

SchedulerStatsTracker.summarize() — total_preemptions uses max instead of sum

"total_preemptions": max(s.num_preemptions for s in self.samples),

If vLLM reports cumulative preemptions this is accidentally correct, but the field is named total_preemptions and the docstring says "accumulated over a decode run" — it should be max(…) - min(…) (delta) or self.samples[-1].num_preemptions - self.samples[0].num_preemptions. Using bare max silently under-counts if the counter resets, and over-counts nothing if it's per-step.

gpu_generator.py — trailing-chunk truncation silently drops events

n_chunks = n_events // chunk_size   # integer division, remainder discarded

For --events 100_000_000 with --chunk-size 5_000_000 this is fine, but any non-multiple silently produces fewer rows than requested. Add an assertion or handle the remainder.

gpu_generator.py — future-draining loop pops the oldest future, not any completed one

if len(futures) > n_writers:
    futures.pop(0).result()   # blocks on chunk 0, even if chunk 2 finished first

This serialises the pipeline back to one thread at a time in the worst case. Use a deque or as_completed.

parse_shard_profiles.py — comment/code mismatch on shard count
The module docstring says "5 per-shard profiles / 5 samples" but the script and HFT_TASK_TOOLS.md say 20; profile_per_shard.sh also defaults to NSHARDS=20. The 5-sample figure is stale.

test_vllm_stats_v1.py — batch_occupancy assertion is wrong

assert s.batch_occupancy == 3 / 256   # occupancy derived from V1 num_running

batch_occupancy is defined as num_running / (num_running + queue_depth) = 3 / (3+12) = 0.2, not 3/256. The test will fail with any correct implementation.

spec.md — owner line lost its dash

-> Owner: Ash — baseline + profiling + spec doc.
+> Owner: Ash  baseline + profiling + spec doc.

Looks like an unintentional edit.

---

🔒 Security

profile_per_shard.sh — unquoted variable in ln -sf

ln -sf "$STAGE/$SHARD" "$TMP/part-00000.parquet"

$STAGE and $SHARD are quoted here (good), but SEED is extracted with:

SEED="$(basename "$STAGE" | grep -oE '[0-9]+$')"

If $STAGE has no trailing digits SEED is empty and --seed "" is passed to Python, which will produce a confusing error rather than a safe failure. Validate SEED is non-empty before continuing.

spec.md reproduction steps embed hardcoded /root/data paths — minor, but any CI runner that runs these verbatim will write to the root home directory.

---

⚡ Performance

gpu_generator.py — 5 separate cp.asnumpy calls per chunk instead of one transfer

symbol_id = cp.asnumpy(rng.integers(...))
side      = cp.asnumpy(rng.integers(...))
price     = cp.asnumpy(rng.integers(...))
...

Each call synchronises the GPU and launches a separate D→H DMA. Pack into a single structured array or a multi-column CuPy matrix and do one transfer; at 5M rows × 5 fields this is five 20–40 MB PCIe round trips per chunk where one would do.

measure_compression.py — reads Parquet metadata for every shard but only prints 5
All 200 shards are opened and stat'd even though only 5 are printed. For 200 files on a networked pod this is slow. Either limit the loop or make the sample window configurable.

---

📊 Reproducibility

gpu_generator.py — CPU and GPU RNGs are seeded identically but produce independent streams

rng    = cp.random.default_rng(seed)   # GPU
rng_np = np.random.default_rng(seed)   # CPU

Both get the same seed, but since they consume from different generators there is no documented guarantee about the joint distribution of (timestamps, fields). This is probably intentional, but it means you cannot reproduce just the Hawkes timestamps by replaying rng_np if the GPU RNG state is unknown (e.g. after a crash mid-run). Consider documenting this explicitly or using separate, documented seeds.

headroom_kernel_rank.py — hardcoded fingerprint="seed42-25M"
If the trace constants (OBSERVED_KERNELS_NS) are updated to reflect a different run, the fingerprint stays stale. Tie it to a hash of the constants or at least a __version__ string so downstream caching doesn't serve stale predictions.

---

💡 Suggestions

gpu_generator.py — no error handling on write_chunk failures inside the thread pool
If a write thread raises (e.g. disk full), futures.pop(0).result() will re-raise, but only for the chunk at position 0; any exception in a later chunk won't surface until the final drain loop. Wrap the drain in a try/finally or use concurrent.futures.wait(..., return_when=FIRST_EXCEPTION).

parse_shard_profiles.py — silent except Exception on gitm imports hides real bugs

except Exception as e:  # noqa: BLE001
    print(f"\ncould not import gitm modules ({e}).")
    return

A typo in a gitm import path, a missing __init__, or a pydantic validation error will all silently produce the "could not import" message, making the script look like an environment problem when it's actually a code bug. At minimum re-raise if the error is not an ImportError/ModuleNotFoundError.

hft_library.yaml — hft_multistream_symbol_shards safety tier is moderate but notes mention float non-determinism — the checksum gate in Task B uses == comparison; if multi-stream changes reduction order the gate will always block this intervention. The notes acknowledge this but the resolution ("integer-stable accumulation or documented float tolerance") should be decided before Task B starts, not after.

vllm_sched.py — duplicate of logic already in vllm_stats.py
_first_attr, _len_or_int, and the scheduler-introspection pattern are reimplemented from scratch. Factor into a shared _engine_utils.py to avoid the two files drifting out of sync (they already differ slightly in _len_or_int vs _len_or_none).

[github-actions]

Code Review by Gemini

The changes introduce new tools and documentation for HFT benchmarking and enhance vLLM scheduler tracing capabilities. The code is well-structured, documented, and includes appropriate tests. The explicit notes about estimates and limitations in the HFT section are commendable for transparency.

Here are some minor suggestions for improvement:

benchmarks/hft/gpu_generator.py

The warmup function calls cp.random.default_rng(0).integers(...) and cp.cuda.Stream.null.synchronize(). It's good to synchronize after GPU operations in warmup to ensure the compilation and initial setup overhead is truly accounted for before timing starts. However, cp.random.default_rng(0).integers(...) itself might not trigger a kernel launch immediately, depending on CuPy's lazy evaluation. Explicitly performing a computation that requires synchronization would be more robust.

--- a/benchmarks/hft/gpu_generator.py
+++ b/benchmarks/hft/gpu_generator.py
@@ -20,7 +20,8 @@
 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)
+    # Perform a simple operation that ensures a GPU kernel is launched and synchronized.
+    cp.arange(10_000, dtype=cp.int32) + 1
     cp.cuda.Stream.null.synchronize()
 
 def generate_chunk(rng, rng_np, n, t_hawkes, lambda_hawkes):

gitm/tracer/vllm_sched.py

The _first_attr helper function is useful, but its current implementation for handling callables (v() if callable(v) else v) might lead to unexpected behavior if an attribute is a callable that shouldn't be called (e.g., a method that requires arguments, or a property that is a callable object). For the specific use cases in sample_scheduler (e.g., get_scheduler_stats, get_stats), it's likely intended to call them. However, for general attribute access, it's safer to only call if it's known to be a no-argument method or a property getter. Given the context, it's probably fine, but it's a pattern to be aware of.

No specific code change is strictly necessary here, but it's a design choice to keep in mind for future extensions or if _first_attr were to be used more broadly.

gitm/tracer/vllm_stats.py

The _v1_scheduler_stats function attempts to read kv_cache_usage and converts it to float. It's good that it handles int | float.

--- a/gitm/tracer/vllm_stats.py
+++ b/gitm/tracer/vllm_stats.py
@@ -83,6 +83,7 @@
     if stats is None:
         return {}
     out: dict[str, Any] = {}
+    # Use getattr with a default to avoid KeyError if attributes are missing.
     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

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

vllm_sched.py — bad imports (numba.core.types.none, pydantic.fields.PropertyT)

from numba.core.types import none          # line ~19 — unused, wrong module
from pydantic.fields import PropertyT      # removed in Pydantic v2, will crash on import

Neither is used anywhere in the file. Remove both.

vllm_sched.py — total_preemptions is wrong (uses max, not last - first)

"total_preemptions": max(s.num_preemptions for s in self.samples),

num_preemptions is a cumulative counter. The total over the window should be samples[-1].num_preemptions - samples[0].num_preemptions. Using max returns the high-water mark, not the delta.

gpu_generator.py — last chunk silently dropped when n_events % chunk_size != 0

n_chunks = n_events // chunk_size   # integer division — remainder events never generated

Add a trailing partial-chunk pass, or assert n_events % chunk_size == 0 with a clear error.

gpu_generator.py — write errors swallowed in the drain loop

futures.pop(0).result()   # only drains when len > n_writers; last ≤ n_writers futures…
for f in futures:
    f.result()            # …caught here, but exceptions from earlier pops are lost silently

Exceptions from .result() inside the loop body are propagated, which is fine, but the comment is misleading. More importantly, if a write fails mid-run the generator prints "FINISHED" and exits 0. Wrap the drain in a try/except or let the with block handle it.

profile_per_shard.sh — nsys stats CSV output flag is wrong

nsys stats --report cuda_gpu_kern_sum --format csv \
    --output "$OUT/shard_$i" "$OUT/shard_$i.nsys-rep"

--output for nsys stats sets a filename prefix, not a directory. The resulting file will be $OUT/shard_$i_cuda_gpu_kern_sum.csv, but parse_shard_profiles.py globs for shard_*_cuda_gpu_kern_sum.csv — this matches, but only by coincidence. Validate on a real pod; minor nsys version differences change the suffix.

parse_shard_profiles.py — n_shards uses the minimum series length but iterates all ops

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]   # safe only if every op has >= n_shards entries

This is correct if every op appears in every shard profile. If an op is absent from some shards (kernel not triggered), series[shard_i] will silently use a shorter list's valid index from a different shard. Add a per-op length assertion or align by shard index explicitly.

test_vllm_stats_v1.py — batch_occupancy assertion is wrong

assert s.batch_occupancy == 3 / 256   # expects capacity-based occupancy

batch_occupancy in vllm_sched.py is defined as num_running / (num_running + queue_depth), not num_running / max_num_seqs. With running=3, waiting=12, that's 3/15 = 0.2, not 3/256. The test will fail (or the implementation doesn't match the docstring — one of them is wrong).

---

🔒 Security

profile_per_shard.sh — unquoted/interpolated shell variables in nsys invocation

nsys profile ... python3 "$HARNESS" --seed "$SEED" --stage "$OUT/stage_$i" ...

SEED is extracted via grep -oE '[0-9]+$' on an attacker-controlled directory name. If the path contains shell metacharacters the grep result is safe (digits only), but $OUT and $STAGE are user-supplied and unvalidated. Use set -u (already present) but also validate that $OUT doesn't contain spaces or special characters, or quote all expansions consistently.

---

⚡ Performance

gpu_generator.py — six separate GPU kernel launches + six cp.asnumpy calls per chunk

symbol_id = cp.asnumpy(rng.integers(...))
side      = cp.asnumpy(rng.integers(...))
price     = cp.asnumpy(rng.integers(...))
...

Each cp.asnumpy is a synchronous H2D→D2H transfer that blocks the CPU. Batch into a single structured array or use cp.stack + one transfer. This won't change end-to-end time much if the bottleneck is disk, but it adds unnecessary synchronization points that compete with the Numba thread.

measure_compression.py — loads full Parquet metadata for every shard including all row groups

pq.ParquetFile(p).metadata   # opens and reads footer for every file

For 200 shards this is 200 file-open + footer-parse calls. Consider using pq.read_metadata(p) (lighter) or batching with a pool — minor, but worth noting for large shard counts.

---

📊 Reproducibility

gpu_generator.py — CPU and GPU RNGs seeded identically but are independent streams

rng    = cp.random.default_rng(seed)   # GPU
rng_np = np.random.default_rng(seed)   # CPU (Hawkes uniforms)

The same integer seed produces different sequences from the two generators (different algorithms). This is intentional (Hawkes on CPU, fields on GPU), but the documentation and results.md say "same seed always produces identical bytes" — that guarantee depends on both generators being deterministic and the Numba JIT cache being warm (cold JIT can produce different timing, though not different values). Confirm @njit(cache=True) is sufficient and document the two-stream design explicitly.

test_gpu_generator.py / test_hawkes_gpu.py / test_rng_speed.py — not in tests/ and have no assertions
These are ad-hoc scripts committed as test files. They will pass silently in any CI runner that doesn't execute them, and pollute the benchmark directory. Either move to tests/ with proper assert statements, or rename to scripts/ and exclude from test discovery.

---

💡 Suggestions

headroom_kernel_rank.py — covered_kernels and the coverage cov_ns calculation can diverge

cov = covered_kernels(spec, trace)           # returns names
cov_ns = sum(
    (k.end_ns - k.start_ns) for k in trace.kernels() if k.name in cov
)

cov is a list, so k.name in cov is O(n) per kernel. Use a set. More importantly, if two kernels share the same name (impossible here since the dict keys are unique, but fragile), this double-counts. Convert to a set at construction.

hft_library.yaml — applies_to_kernels substrings are ambiguous

applies_to_kernels: [decode_page_data, zstd::decompression]

decode_page_data will match any kernel containing that string. If future kernels are added (e.g. decode_page_data_rle, decode_page_data_dict), they'll all be covered unintentionally. Use full kernel name anchors or add a match_type: prefix|substring|exact field to the spec.

vllm_sched.py — no tests for the batch_size=0 / empty-engine path in summarize()

if not self.samples:
    return {"n_samples": 0}

The early return is correct, but callers that expect keys like mean_batch_occupancy will KeyError. Return a full zero-valued dict or document the sparse contract.

spec.md — em-dash removed from owner line

-> Owner: Ash — baseline + profiling + spec doc.
+> Owner: Ash  baseline + profiling + spec doc.

Looks like an accidental character deletion, not intentional.

[github-actions]

Code Review by Gemini

### `benchmarks/hft/spec.md`

The reproduction steps for building the generator include a hardcoded version for `libparquet.so`. This can be brittle if the `pyarrow` library version changes, leading to `libparquet.so.2400` not being found. A more robust approach would be to dynamically find the latest versioned `libparquet.so` file.

```diff
--- a/benchmarks/hft/spec.md
+++ b/benchmarks/hft/spec.md
@@ -49,7 +49,7 @@
 **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
+ln -sf "$(find "$ARROW_LIB" -name "libparquet.so.*" | sort -V | tail -n 1)" "$ARROW_LIB/libparquet.so"
 cmake .. && make -j$(nproc)

[github-actions]

🤖 Claude Code Review

🐛 Bugs

gpu_generator.py — dual-seed corruption: rng (CuPy) and rng_np (NumPy) are both seeded with the same seed value but advance independently through the generation loop. The Hawkes timestamps and GPU fields will share no common seed progression, meaning the two RNG streams are not reproducibly coordinated. If chunk size ever changes, the field→timestamp pairing shifts arbitrarily.

gpu_generator.py — silent data loss on non-multiple chunk sizes: n_chunks = n_events // chunk_size silently drops the remainder. For --events 100_000_000 with default chunk size this is fine, but with e.g. --events 150_000_000 you silently generate only 100M events with no warning.

gpu_generator.py — future leak on exception: The sliding-window future drain futures.pop(0).result() only fires when len(futures) > n_writers. If write_chunk raises, the remaining futures in the final for f in futures: f.result() loop will re-raise, but intermediate exceptions can be swallowed silently during normal iteration.

parse_shard_profiles.py — n_shards uses minimum series length, silently drops data:

n_shards = min((len(v) for v in obs_series.values()), default=0)

If one op is missing from some shards (e.g. merge_sort doesn't appear in a quiet shard), the series gets truncated to that op's shorter count with no warning, potentially making Granger results unreliable without surfacing why.

vllm_sched.py — total_preemptions calculation is wrong:

"total_preemptions": max(s.num_preemptions for s in self.samples) - min(s.num_preemptions for s in self.samples),

This assumes num_preemptions is a monotonically increasing counter. If the engine resets it, or if it's a per-step count (not cumulative), this gives nonsense. The V0 path in vllm_stats.py doesn't clarify which it is, and the V1 make_stats() path doesn't populate it at all — so it stays 0 silently.

vllm_stats.py — setattr on a Pydantic model with extra="forbid":

setattr(sample, field_name, val)

Pydantic v2 models with extra="forbid" may raise on setattr for fields not in the schema, or silently no-op depending on model config. SchedulerSample should be checked — if it uses model_config = ConfigDict(frozen=True) anywhere upstream this will raise at runtime.

test_vllm_stats_v1.py — truncated test:

assert read_sc

The file is cut off mid-line. This will fail to import entirely, breaking the entire test suite.

hawkes_numba — timestamp overflow risk: ts[i] = int(t * 1e9) converts seconds × 1e9 to int64. At 1B events with MU=100 (baseline ~100 events/sec), t reaches ~10M seconds → ~1e16 ns, well within int64 range. However with high alpha/beta bursting, instantaneous lambda can spike and compress events, but the secular drift is the risk to document.

profile_per_shard.sh — comment/code mismatch: The docstring says "5 per-shard nsys profiles" and "5 single-shard runs", but the default is NSHARDS=20 and the loop variable $3 controls it. The parse_shard_profiles.py docstring also says "5 shards" in its header. This inconsistency will confuse operators about the minimum required.

---

🔒 Security

profile_per_shard.sh — unquoted variable in symlink path:

ln -sf "$STAGE/$SHARD" "$TMP/part-00000.parquet"

STAGE comes directly from $1 with no validation. A path containing spaces or special characters will misbehave. More critically, SEED is extracted via:

SEED="$(basename "$STAGE" | grep -oE '[0-9]+$')"

If STAGE doesn't end in digits (e.g. /workspace/hft_data/), SEED is empty and --seed "" is passed to the harness silently.

spec.md reproduction steps — LD_LIBRARY_PATH injection risk: The doc instructs users to set LD_LIBRARY_PATH=$ARROW_LIB:$LD_LIBRARY_PATH without noting this persists in the shell and can affect subsequent unrelated commands in the same session.

---

⚡ Performance

gpu_generator.py — unnecessary per-field cp.asnumpy calls: Each field (symbol_id, side, price, size, roll) is transferred to CPU separately with individual cp.asnumpy() calls, each triggering a separate H2D→D2H synchronization. Packing all fields into a single structured transfer (or at minimum a single cp.cuda.Stream.null.synchronize() followed by one cp.asnumpy on a stacked array) would reduce PCIe round-trips.

gpu_generator.py — CuPy RNG not on a non-default stream: All GPU generation happens on the null stream, so it serializes with any other GPU work. Using a dedicated cp.cuda.Stream for generation would allow better overlap with the Numba CPU work on the Hawkes path.

measure_compression.py — reads all shard metadata but only prints 5: All 200 shards' metadata is iterated (for i, p in enumerate(paths)) accumulating totals, which is correct, but pq.ParquetFile(p).metadata opens each file. For 200 shards this is 200 file opens — negligible but worth noting if extended to thousands of shards.

---

📊 Reproducibility

gpu_generator.py — CuPy and NumPy RNGs diverge under restarts: If generation is interrupted and restarted mid-run, there's no checkpoint of t_hawkes, lambda_hawkes, or the RNG states. Re-running from scratch with the same seed produces the same output only if run to completion in one pass. Partial restarts silently produce different data for already-written shards if the file isn't overwritten.

gpu_generator.py — same seed for both RNGs creates correlation:

rng = cp.random.default_rng(seed)
rng_np = np.random.default_rng(seed)

Both start at the same seed. For seed=42, both generate from the same initial state but via different algorithms (Philox for CuPy, PCG64 for NumPy). This is not a correctness bug but it means the GPU fields and Hawkes timestamps are not independently seeded — if either RNG is ever replaced, the joint distribution changes in a non-obvious way. Use seed and seed + 1 (or derived seeds) explicitly.

hft_library.yaml — expected_delta_mean is load-bearing for ranking but undocumented sensitivity: The ranking in headroom_kernel_rank.py is predict_delta = coverage × expected_delta_mean. If expected_delta_mean for candidate 1 drops from 0.25 to 0.10 (within the stated lo/hi range), the ranking could change. There's no sensitivity analysis or range-based ranking presented.

hawkes_numba — Numba cache=True with parameter changes: @njit(cache=True) caches the compiled function. If MU, ALPHA, or BETA module-level constants change, Numba's cache may serve the stale compiled version until manually cleared. The constants should either be passed as parameters or the cache invalidation strategy documented.

---

💡 Suggestions

gpu_generator.py — no validation that output files are complete: If the process is killed mid-write, a partial Parquet file exists with the correct name. On restart the generator will skip it (file exists) or overwrite it depending on pq.write_table behavior. Add an atomic write pattern (write to .tmp, rename on success).

parse_shard_profiles.py — Granger requires stationarity, not checked: grangercausalitytests assumes weakly stationary series. Per-shard kernel times may have a trend (e.g. thermal throttling over the 20-shard run). A simple ADF test before calling attribute() would catch this — the current code will return misleading p-values silently.

vllm_sched.py — duplicate of logic in vllm_stats.py: sample_scheduler in vllm_sched.py reimplements _first_attr, _len_or_int, and the scheduler-reading logic that already exists in vllm_stats.py. This creates two code paths for the same concern. Consider whether vllm_sched.py should be a thin wrapper over vllm_stats.py or merged.

headroom_kernel_rank.py — covered_kernels double-iterates the trace:

cov = covered_kernels(spec, trace)   # iterates trace.kernels()
cov_ns = sum(... for k in trace.kernels() if k.name in cov)  # iterates again

With a short trace this is fine, but cov is a list so the in check is O(n). Convert to a set, or combine into one pass.

spec.md — accidental removal of the em-dash:

-> Owner: Ash  baseline + profiling + spec doc.

The — separator was removed, leaving a double-space. Minor but looks unintentional.

[github-actions]

Code Review by Gemini

### `benchmarks/hft/HFT_HEADROOM_TOOLS.md`
```diff
--- a/benchmarks/hft/HFT_HEADROOM_TOOLS.md
+++ b/benchmarks/hft/HFT_HEADROOM_TOOLS.md
@@ -36,7 +36,7 @@
 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:
+The pieces that already exist in the repo (run from repo root):
   - gitm/optimizer/apply.py    — apply_intervention(spec, applicator,
                                  min_keep_delta): snapshot -> apply -> measure
                                  -> keep|rollback. Use this as the gate.

Improvement: Clarified the context for the path to gitm/optimizer/apply.py to avoid confusion about relative paths when running from benchmarks/hft.

benchmarks/hft/parse_shard_profiles.py

--- a/benchmarks/hft/parse_shard_profiles.py
+++ b/benchmarks/hft/parse_shard_profiles.py
@@ -70,7 +70,7 @@
     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:
+    try:  # noqa: BLE001
         from gitm.optimizer.attribution import attribute
         from gitm.optimizer.monitor import KernelResidual, Residuals
         from gitm.planner.hft_graph import HFTDatasetSpec, predict_hft_graph

Issue: The try...except Exception is too broad.
Improvement: Changed to try: # noqa: BLE001 to acknowledge the broad exception catch, as per the project's linting rules, while keeping the existing behavior. A more specific ImportError or ModuleNotFoundError would be ideal if the goal is to catch only missing modules.

gitm/tracer/vllm_sched.py

--- a/gitm/tracer/vllm_sched.py
+++ b/gitm/tracer/vllm_sched.py
@@ -79,6 +79,7 @@
             "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),
+            # Assumes num_preemptions is a cumulative counter.
+            "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,

Improvement: Added a comment to clarify the assumption about num_preemptions being a cumulative counter.

gitm/tracer/vllm_stats.py

--- a/gitm/tracer/vllm_stats.py
+++ b/gitm/tracer/vllm_stats.py
@@ -169,7 +169,7 @@
     if not callable(make):
         return {}
     try:
         stats = make()
-    except Exception:
+    except (AttributeError, TypeError):
         return {}
     if stats is None:

Issue: The except Exception is too broad.
Improvement: Changed to except (AttributeError, TypeError): to catch more specific exceptions that are likely to occur if make_stats or its result is not as expected, without hiding other potential issues.

[github-actions]

🤖 Claude Code Review

Review

🐛 Bugs

gpu_generator.py — dual seeding breaks reproducibility and stream independence

rng = cp.random.default_rng(seed)
rng_np = np.random.default_rng(seed)

Both RNGs get the same seed, but they're advanced independently per chunk. If chunk size or count changes, the Hawkes timestamps (CPU) and GPU fields diverge from each other in unpredictable ways. The same-seed coincidence is also misleading — these should be explicitly documented as independent streams or derived via a seeding scheme (e.g. seed and seed + 1).

gpu_generator.py — trailing chunk silently dropped

n_chunks = n_events // chunk_size

If n_events is not divisible by chunk_size, the remainder is silently lost. For the default 100M / 5M this is fine, but with arbitrary --events values (e.g. --events 123456789) the output row count won't match the requested count and no warning is raised.

gpu_generator.py — write errors silently swallowed during drain

if len(futures) > n_writers:
    futures.pop(0).result()   # raises on write failure

Only the oldest future is drained; after the loop the remaining futures are drained correctly. However, if a write fails mid-loop, the exception propagates and the remaining futures are never .result()-ed, potentially leaving file handles open and the output directory in a partial state. The with ThreadPoolExecutor block will cancel but not drain them. Wrapping with try/finally or using as_completed would be safer.

parse_shard_profiles.py — n_shards truncates silently when series lengths differ

n_shards = min((len(v) for v in obs_series.values()), default=0)

If one op is missing from some CSVs (kernel not triggered on that shard), its series is shorter and later ops use fewer samples without any warning. The residuals for the truncated op won't line up with the others in the Granger input, corrupting the causal ordering. Should log a warning when lengths differ.

vllm_sched.py — total_preemptions calculation assumes monotone counter

"total_preemptions": max(s.num_preemptions for s in self.samples) - min(s.num_preemptions for s in self.samples),

This breaks if the engine resets its preemption counter between samples (e.g. on engine restart or stats reset), producing a negative or underestimated value. A running delta (sum of positive increments between consecutive samples) would be more robust.

vllm_stats.py — setattr on a Pydantic model with extra="forbid" may fail silently or raise

if getattr(sample, field_name) is None:
    setattr(sample, field_name, val)

SchedulerSample (presumably a Pydantic BaseModel) with extra="forbid" may reject setattr for non-declared fields at runtime depending on Pydantic version and model config. If field_name comes from _v1_scheduler_stats and happens to be a key not in the model, this will either silently no-op or raise. The valid field names should be validated against the model schema before calling setattr.

headroom_kernel_rank.py — covered_kernels and coverage calculation use different membership tests

def covered_kernels(spec, trace):
    return [k.name for k in trace.kernels() if _applies(spec, k.name)]

cov_ns = sum(
    (k.end_ns - k.start_ns) for k in trace.kernels() if k.name in cov
)

cov is a list of names; k.name in cov is O(n) per kernel. More importantly, _applies does substring matching, but k.name in cov does exact matching on the already-filtered names — these are consistent here, but if two kernels share a name string this double-iteration is redundant. Minor, but worth a set conversion.

---

🔒 Security

profile_per_shard.sh — unquoted variable interpolation

SEED="$(basename "$STAGE" | grep -oE '[0-9]+$')"

SEED is then used unquoted in:

TMP="$OUT/stage_$i/hft_numba_seed${SEED}"

and passed to Python:

python3 "$HARNESS" --seed "$SEED" ...

If STAGE is provided by the user and contains a path that produces unexpected grep output (empty string, multiple matches), SEED can be empty, causing mkdir -p "$OUT/stage_$i/hft_numba_seed" and --seed "" to be passed to Python. The seed extraction should be validated:

[[ "$SEED" =~ ^[0-9]+$ ]] || { echo "could not parse seed from $STAGE"; exit 1; }

---

⚡ Performance

gpu_generator.py — per-field .asnumpy() calls serialise the GPU

symbol_id = cp.asnumpy(rng.integers(...))
side      = cp.asnumpy(rng.integers(...))
price     = cp.asnumpy(rng.integers(...))
...

Each cp.asnumpy is a synchronising H2D transfer. Generating all fields into a single structured array or transferring after a single cp.cuda.Stream.null.synchronize() call with pinned memory would avoid 5 separate PCIe round-trips per chunk and is consistent with the overlap motivation described in the design docs.

measure_compression.py — loads full Parquet metadata for all 200 shards sequentially

for i, p in enumerate(paths):
    md = pq.ParquetFile(p).metadata

pq.ParquetFile opens the file and reads footer metadata. For 200 shards this is 200 serial file opens. A ThreadPoolExecutor here would be trivially parallel and cut wall time significantly on pod storage.

---

📊 Reproducibility

gpu_generator.py — Numba and CuPy RNG streams are not truly independent

As noted under Bugs: both rng and rng_np start from the same integer seed. Since they are different PRNG algorithms (Philox/XOSHIRO on GPU vs PCG on CPU), they won't collide, but the intent is ambiguous. The Hawkes timestamps are seeded identically to the GPU fields, so any analysis comparing timestamp distributions to field distributions needs to account for this.

gpu_generator.py — Numba JIT cache invalidation

@njit(cache=True)
def hawkes_numba(u, n, t0, lam0):

cache=True means the compiled function is reused from disk on subsequent runs. If the function body changes but the cache isn't cleared (e.g. between pod sessions), the old compiled version runs silently, producing different output from the source. For a reproducibility benchmark this is a silent correctness risk. Worth documenting a numba --clear-cache or cache-directory step in the reproduction instructions.

test_gpu_generator.py and test_hawkes_gpu.py — test scripts write to hardcoded /workspace/ paths

pq.write_table(table, '/workspace/test_nocomp.parquet', ...)

These aren't pytest tests (no assertions, no test_ functions in the pytest sense despite the filename prefix). They'll be collected by pytest, fail on any machine without /workspace/, and produce no signal about correctness. They should either be moved to benchmarks/hft/scripts/ or given proper tmp_path fixtures and assertions.

---

💡 Suggestions

profile_per_shard.sh — nsys stats export uses --output which overwrites the .nsys-rep

nsys stats --report cuda_gpu_kern_sum --format csv \
    --output "$OUT/shard_$i" "$OUT/shard_$i.nsys-rep" || true

The || true suppresses all errors including missing nsys stats binary or a corrupt .nsys-rep. parse_shard_profiles.py will then silently find no CSVs. The || true should at minimum emit a warning: || echo "WARNING: nsys stats failed for shard $i".

headroom_kernel_rank.py — no validation that specs is non-empty

rows.sort(...)
top = rows[0]  # IndexError if library has no interventions

If hft_library.yaml is missing or has an empty interventions list, this crashes with an unhelpful IndexError. A guard with a clear error message would be more useful.

hft_library.yaml — DeviceReduce in applies_to_kernels for candidate 2 is not in OBSERVED_KERNELS_NS

applies_to_kernels: [decode_page_data, DeviceMergeSort, DeviceReduce]

DeviceReduce doesn't appear in the hardcoded kernel table in headroom_kernel_rank.py. The coverage calculation will silently miss it. This should either be added to OBSERVED_KERNELS_NS with a 0 duration (not observed) or removed from the spec.

vllm_sched.py — duplicate of logic already in vllm_stats.py
The new vllm_sched.py reimplements _first_attr, _len_or_int, and sample_scheduler in parallel with the existing vllm_stats.py::read_scheduler_stats. The two files will diverge. The shared helpers should live in one place, or vllm_sched.py should import from vllm_stats.py.