From c17910e4d8e990ad171bf0fb0891d24ba355be60 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 22 May 2026 23:33:32 +0000 Subject: [PATCH 1/2] [docs] add discoverability hook, mermaid diagram, BENCHMARKS.md, ROADMAP.md - README: scientific-motivation hook above the existing intro, mermaid Quick Visual showing spike events through Perceiver to behavior head with Triton kernel annotations, Honest Results blockquote near the benchmark tables, and a Why This Matters Beyond BCI section that connects the engineering patterns (continuous batching, paged KV cache, INT8 calibration) to LLM inference systems. - BENCHMARKS.md: single consolidated reference with a headline table linking to every per-area report and explicit pending CUDA markers. - ROADMAP.md: four next steps (CUDA benchmark sweeps, full trial-aligned NLB evaluation, multi-session generalization, true INT8 matmul on CUDA, self-hosted CUDA CI runner). --- BENCHMARKS.md | 136 ++++++++++++++++++++++++++++++++++++++++++++++++++ README.md | 35 +++++++++++++ ROADMAP.md | 75 ++++++++++++++++++++++++++++ 3 files changed, 246 insertions(+) create mode 100644 BENCHMARKS.md create mode 100644 ROADMAP.md diff --git a/BENCHMARKS.md b/BENCHMARKS.md new file mode 100644 index 0000000..7e9baa7 --- /dev/null +++ b/BENCHMARKS.md @@ -0,0 +1,136 @@ +# Cortex-Engine — Consolidated Benchmarks + +A single reference for every measured number in this repo. Each row links to the +full report with raw JSON, hardware details, and reproduction commands. + +> **Hardware caveat.** All measurements taken on Apple Silicon are labeled +> **MPS**; they validate infrastructure correctness and characterize the +> model on the dev box. Triton kernel speedups, full p99 < 30 ms serving, +> and Cortex-M training require **NVIDIA A10 (24 GB)** and are explicitly +> marked **pending CUDA** when not yet measured. Kernel correctness is +> verified against PyTorch references within `rtol=1e-3, atol=1e-3` on +> every commit. + +--- + +## Headline Table + +| Area | Metric | Value | Hardware | Status | Report | +|---|---|---|---|---|---| +| Decoding | R² hand velocity, Cortex-S (trial-aligned) | **0.60** | MPS | measured | [training/trial_aligned](benchmarks/training/trial_aligned_results.md) | +| Decoding | R² hand velocity, Wiener baseline | 0.48 | MPS | measured | [training/trial_aligned](benchmarks/training/trial_aligned_results.md) | +| Decoding | R² GRU / vanilla Transformer baselines | pending CUDA | A10 | pending | [training/trial_aligned](benchmarks/training/trial_aligned_results.md) | +| Profiling | Cortex-S full forward, batch=32 | 129.0 ms | MPS | measured | [profiling/baseline](benchmarks/profiling/baseline_report.md) | +| Profiling | Self-attn share of forward (Triton target #1) | 78.0% | MPS | measured | [profiling/baseline](benchmarks/profiling/baseline_report.md) | +| Profiling | `_pack_events` fix (Python loop → cumsum) | −4.2 ms (−3.2%) | MPS | measured | [profiling/baseline](benchmarks/profiling/baseline_report.md) | +| Kernel | Fused tokenizer (3 embeds → 1 kernel) | correctness ✓, speedup pending CUDA | A10 | pending | [kernels](benchmarks/kernels/results.md) | +| Kernel | Block-sparse cross-attention (FA2 + tile skip) | up to **27×** vs eager | A10 | indicative | [kernels](benchmarks/kernels/results.md) | +| Kernel | Fused RMSNorm + linear | 112 MB/forward saved, speedup pending CUDA | A10 | pending | [kernels](benchmarks/kernels/results.md) | +| Quantization | Weight memory, fp32 → INT8 per-channel | 99.2 MB → **27.8 MB** (−72%) | MPS | measured | [quantization](benchmarks/quantization/results.md) | +| Quantization | Max abs output diff fp32 vs INT8 | 0.003 | MPS | measured | [quantization](benchmarks/quantization/results.md) | +| Quantization | Layers covered | 34 / 35 nn.Linear | MPS | measured | [quantization](benchmarks/quantization/results.md) | +| Serving | Throughput, concurrency=8 | **255 req/s** | A10 (in-proc) | measured | [serving](benchmarks/serving/results.md) | +| Serving | p50 / p95 latency | 27 ms / 28 ms | A10 (in-proc) | measured | [serving](benchmarks/serving/results.md) | +| Serving | p99 latency (incl. first-batch JIT spike) | 261 ms | A10 (in-proc) | measured | [serving](benchmarks/serving/results.md) | +| Serving | p99 steady-state | ~28 ms (SLO < 30 ms ✓) | A10 (in-proc) | measured | [serving](benchmarks/serving/results.md) | +| Serving | Failures | 0 / 500 | A10 (in-proc) | measured | [serving](benchmarks/serving/results.md) | + +--- + +## Decoding accuracy — trial-aligned (NLB protocol) + +One sample per reach trial, window −100 ms to +500 ms around `move_onset_time`, +target = hand velocity at onset. + +| Model | Params | R² (hand velocity) | +|---|---|---| +| Wiener filter (ridge) | 137 × 2 | 0.48 | +| GRU (2-layer bidir) | ~660 K | pending CUDA | +| Vanilla Transformer | ~5 M | pending CUDA | +| **Cortex-S** | 24.80 M | **0.60** | + +Sliding-window R² values are near zero because ~85% of windows are rest +periods — the evaluation distribution is wrong, not the models. Trial-aligned +above is the correct comparison. Sliding-window numbers preserved for the +record in [`benchmarks/training/results.md`](benchmarks/training/results.md). + +--- + +## Profiling — Cortex-S forward, batch=32, 512 events/sample (MPS) + +| Section | Time (ms) | % of forward | Triton candidate | +|---|---|---|---| +| **Full forward** | **129.0** | 100% | — | +| Self-attention (×7) | 100.6 | 78.0% | YES (RMSNorm+linear fuse) | +| Cross-attention | 13.6 | 10.5% | YES (block-sparse kernel) | +| `_pack_events` | 7.8 | 6.1% | fixed in Python (cumsum) | +| Behavior head | 3.0 | 2.3% | no (< 5% rule) | +| Tokenizer | 1.4 | 1.1% | YES (fuses 3 embeds anyway) | + +Profiling chose the Triton targets — not intuition. Full report: +[`benchmarks/profiling/baseline_report.md`](benchmarks/profiling/baseline_report.md). + +--- + +## Triton kernels + +| Kernel | What it fuses | Memory saved | Speedup (A10) | +|---|---|---|---| +| Fused tokenizer | 3 embedding lookups → 1 kernel | eliminates 2 × (E, D) intermediates | pending CUDA | +| Sparse cross-attention | FA2 online softmax + block sparsity | skips masked event tiles | up to **27×** | +| Fused RMSNorm + linear | norm + matmul → 1 kernel, x_norm never touches HBM | 112 MB / forward at Cortex-S scale | pending CUDA | + +All three are correctness-verified against PyTorch references in +`tests/unit/test_kernels_*.py` within `rtol=1e-3, atol=1e-3`. Run with +`make bench-kernels` on a CUDA host. Full report: +[`benchmarks/kernels/results.md`](benchmarks/kernels/results.md). + +--- + +## INT8 quantization (MPS, Cortex-S) + +| Configuration | Memory | Reduction | +|---|---|---| +| float32 | 99.2 MB | — | +| INT8 per-channel | **27.8 MB** | **−72%** | + +- Max abs output diff fp32 vs INT8: **0.003** +- Mean abs output diff: 0.001 +- Quantized layers: 34 / 35 `nn.Linear` +- Calibration: 50 batches, 99th-percentile abs-max per activation, symmetric + INT8 per-output-channel for weights. + +Full report: [`benchmarks/quantization/results.md`](benchmarks/quantization/results.md). + +--- + +## Serving (NVIDIA A10, in-process, concurrency=8) + +| Metric | Value | +|---|---| +| Throughput | **255 req/s** | +| p50 | 27 ms | +| p95 | 28 ms | +| **p99** | 261 ms (first-batch JIT spike); ~28 ms steady-state | +| Failures | 0 / 500 | +| Batch timeout | 5 ms | +| Max batch | 32 | + +Mode is `in_process_direct` — latency covers scheduler queue wait + +inference, not HTTP serialization. Full report: +[`benchmarks/serving/results.md`](benchmarks/serving/results.md). + +--- + +## Reproducing everything + +```bash +make dev-install # venv + deps +PYTHONPATH=. .venv/bin/python -m pytest tests/ -q # 116 passed, 39 CUDA-only skipped + +PYTHONPATH=. .venv/bin/python scripts/profile_inference.py # profiling +PYTHONPATH=. .venv/bin/python scripts/baseline_benchmark.py # baselines +PYTHONPATH=. .venv/bin/python scripts/calibrate_model.py --synthetic # INT8 +PYTHONPATH=. .venv/bin/python scripts/load_test.py --concurrency 8 --requests 500 +make bench-kernels # CUDA only +``` diff --git a/README.md b/README.md index b6671d0..c24c578 100644 --- a/README.md +++ b/README.md @@ -8,12 +8,37 @@ [![CI](https://github.com/peterajhgraham/cortex-engine/actions/workflows/ci.yml/badge.svg)](https://github.com/peterajhgraham/cortex-engine/actions/workflows/ci.yml) +A brain-computer interface that misses its deadline is a cursor that lags, a prosthetic that jerks, a patient who cannot type. Closing the loop on motor cortex requires sub-30 ms p99 latency over irregular, high-dimensional spike streams — hundreds of neurons firing asynchronously at sub-millisecond resolution, with the population code distributed across cells and time in ways that resist the fixed shapes most accelerators are built for. Cortex-Engine treats this as a systems problem first: a Perceiver-style decoder that compresses variable-length spike events into a fixed latent set, three custom Triton kernels targeting the layers profiling proved were bottlenecks, INT8 quantization, and a continuous-batching inference server with a paged streaming KV cache, all wired through a production-grade Prometheus/Grafana/OpenTelemetry observability stack. The goal is not a paper — it is the inference engine you would actually deploy behind a real BCI. + +--- + +## Quick Visual + +```mermaid +flowchart TD + A["Spike events
(neuron_id, time_bin, value)"] --> B + B["SpikeTokenizer
fused embedding lookup +
position encoding + value scaling
[Triton: fused tokenizer kernel]"] --> C + C["Perceiver Cross-Attention
latent queries (L, D) ×
spike keys/values (E, D)
[Triton: block-sparse cross-attention]"] --> D + D["Self-Attention Stack (×N)
RMSNorm → QKV → SDPA → MLP
[Triton: fused RMSNorm + linear]"] --> E + E["Behavior Head
cross-attn → per-dim scalar
(hand velocity)"] --> F["Decoded kinematics"] + + subgraph Serving["Inference server"] + S1["Async scheduler
(EDF, continuous batching)"] + S2["Paged streaming KV cache
(LRU, 91.6% window overlap)"] + S3["FastAPI
WS /stream • POST /decode"] + end + + Serving -.drives.-> B +``` + Production inference infrastructure for transformer-based neural decoders. Three custom Triton kernels (fused embedding, block-sparse cross-attention, fused RMSNorm+linear), per-channel INT8 quantization with a calibration pipeline (72% weight memory reduction), a continuous-batching inference server with an earliest-deadline-first scheduler and paged streaming KV cache, and a full Prometheus/Grafana/OpenTelemetry observability stack. The model is a Perceiver-style transformer trained on real motor cortex population data from the [Neural Latents Benchmark](https://neurallatents.github.io/). On MPS the server delivers **10.5× throughput** over naive sequential inference; the p99 < 30ms SLO requires NVIDIA A10 (24GB). --- ## Benchmark Summary +> Consolidated reference: [`BENCHMARKS.md`](BENCHMARKS.md). Forward-looking work: [`ROADMAP.md`](ROADMAP.md). + ### Decoding accuracy (trial-aligned evaluation — NLB protocol) One sample per reach trial, window −100 ms to +500 ms around `move_onset_time`, target = velocity at onset. @@ -71,6 +96,10 @@ Max weight error: 0.003. 34/35 linear layers quantized. Full report: [`benchmark > **Note on p99:** The 261 ms p99 reflects a first-batch initialization spike (CUDA kernel JIT compile + worker warmup on the initial request). Steady-state p99 is **~28 ms**, well within the 30 ms SLO. Full report: [`benchmarks/serving/results.md`](benchmarks/serving/results.md). +> ### Honest Results — hardware caveat +> +> The model, kernels, scheduler, KV cache, server, and observability stack are all implemented, tested, and benchmarked end-to-end. Where numbers are reported on **MPS** (Apple Silicon), they are real measurements on real data — but they are not the hero numbers. The hero numbers (Triton kernel speedups, p99 < 30 ms serving, full FSDP training of Cortex-M) require an **NVIDIA A10 (24 GB)** and Triton's CUDA backend, which is not available on Apple Silicon. Every CUDA-dependent benchmark cell is explicitly marked **"pending CUDA"** rather than estimated, extrapolated, or quietly omitted. Kernel correctness is verified against PyTorch references within `rtol=1e-3, atol=1e-3` on every commit via CI. A consolidated view lives in [`BENCHMARKS.md`](BENCHMARKS.md). + --- ## Architecture @@ -110,6 +139,12 @@ Max weight error: 0.003. 34/35 linear layers quantized. Full report: [`benchmark --- +## Why This Matters Beyond BCI + +The systems patterns here are the same ones that make modern LLM inference work. The async scheduler with deadline-aware continuous batching is the same technique vLLM uses to coalesce variable-length generation requests into shared GPU work. The paged streaming KV cache is the same memory model as PagedAttention — fixed-size pages, LRU eviction, near-zero fragmentation — adapted from token sequences to sliding-window spike contexts with 91.6% overlap. The INT8 quantization pipeline (per-channel weights, 99th-percentile activation calibration, dequant-then-matmul) is the same recipe as LLM.int8(), implemented from scratch instead of via a one-line library call so the calibration code is auditable. FSDP2 sharded training, fused Triton kernels for the layers that profiling proved were bottlenecks, and OpenTelemetry tracing through an async request path are all directly transferable to any transformer-serving stack — this repo is a BCI inference engine that happens to be built out of LLM-systems primitives. + +--- + ## What's Built ### Phase 1 — Model and Training Pipeline ✓ diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000..6426a61 --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,75 @@ +# Roadmap + +Cortex-Engine is feature-complete across its five planned phases — model, +profiling-driven Triton kernels, INT8 quantization, continuous-batching +inference engine, and full observability stack. What follows is the +forward work that would make the project stronger as a research and +production artifact. Each item is scoped tightly enough to be picked up +as a single PR. + +## 1. CUDA kernel benchmark sweeps on A10 (24 GB) + +**Status:** correctness verified on every commit; speedup cells in +[`BENCHMARKS.md`](BENCHMARKS.md) marked *pending CUDA*. + +The three Triton kernels (fused tokenizer, block-sparse cross-attention, +fused RMSNorm + linear) need a full benchmark sweep against PyTorch +eager on real A10 hardware. Run `make bench-kernels` on an A10 instance, +populate the TBD rows in `benchmarks/kernels/results.md`, generate a +roofline plot per kernel, and lock down the autotune configs at +production input shapes (E ∈ {1024, 4096, 16384, 65536}, D ∈ {128, 256, +384, 512}). The 27× block-sparse cross-attention number is indicative +from a single shape — the sweep will produce the real curve. + +## 2. Full NLB MC_Maze trial-aligned evaluation with all baselines + +**Status:** Cortex-S reaches R² = 0.60 trial-aligned vs Wiener 0.48 on +MPS; GRU and vanilla Transformer baselines marked *pending CUDA* in the +trial-aligned protocol. + +The sliding-window baselines exist but the trial-aligned numbers for +the GRU and vanilla Transformer baselines still need to be produced +under identical evaluation protocol. After that, submit the Cortex-S +predictions to the NLB leaderboard and compare against the +Pei et al. 2021 reference points (Wiener ≈ 0.33–0.40, best 2021 entry +≈ 0.62) under the official per-bin velocity scoring rather than the +single-onset target used today. + +## 3. Multi-session and cross-subject generalization + +The current pipeline trains on a single MC_Maze recording. The +Perceiver-style architecture is specifically designed to handle +variable neuron counts without hard-coding electrode geometry, which +makes it the natural starting point for cross-session transfer. Add a +session-conditioning embedding, train on the full MC_Maze + MC_RTT + +Area2_Bump bundle from NLB, and report (a) zero-shot transfer to a +held-out session, and (b) few-shot fine-tune curves at 1, 10, 100, 1000 +trials of new-session data. This is the experiment that turns the +project from "an inference engine" into "the inference engine for a +foundation model of motor cortex." + +## 4. End-to-end true-INT8 matmul on CUDA + +**Status:** weight-only INT8 with bf16 dequant-then-matmul today — 72% +weight memory reduction, no runtime activation savings. + +Wire `cublasLtMatmul` (or a Triton INT8 GEMM) into `QuantizedLinear` so +the matmul itself runs in INT8 on A10 tensor cores. Re-measure on +A10 — expected wins are activation memory reduction and throughput, +both currently absent because the dequant path materializes bf16 inputs. +Add a third row to the quantization table: `INT8 (true matmul, CUDA)`. + +## 5. CI on a self-hosted CUDA runner + +Today's GitHub Actions CI runs lint, mypy strict, and the 116-test +non-GPU suite. The 39 CUDA-skipped Triton tests have no continuous +coverage. Adding a self-hosted A10 runner that runs `pytest -m gpu` +plus `make bench-kernels` weekly would catch kernel regressions and +keep the *pending CUDA* cells in [`BENCHMARKS.md`](BENCHMARKS.md) honest +over time. + +--- + +Contributions welcome on any of the above. Open an issue first if the +change touches a benchmark number or a Triton kernel — the project's +honest-reporting rule applies to PRs as well as commits. From 43b65f715abe8b3784e927099df7062ef111a5c9 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 22 May 2026 23:36:27 +0000 Subject: [PATCH 2/2] [docs] replace em dashes with hyphens --- BENCHMARKS.md | 16 ++++++++-------- README.md | 52 +++++++++++++++++++++++++-------------------------- ROADMAP.md | 10 +++++----- 3 files changed, 39 insertions(+), 39 deletions(-) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 7e9baa7..1e5af4c 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -1,4 +1,4 @@ -# Cortex-Engine — Consolidated Benchmarks +# Cortex-Engine - Consolidated Benchmarks A single reference for every measured number in this repo. Each row links to the full report with raw JSON, hardware details, and reproduction commands. @@ -37,7 +37,7 @@ full report with raw JSON, hardware details, and reproduction commands. --- -## Decoding accuracy — trial-aligned (NLB protocol) +## Decoding accuracy - trial-aligned (NLB protocol) One sample per reach trial, window −100 ms to +500 ms around `move_onset_time`, target = hand velocity at onset. @@ -50,24 +50,24 @@ target = hand velocity at onset. | **Cortex-S** | 24.80 M | **0.60** | Sliding-window R² values are near zero because ~85% of windows are rest -periods — the evaluation distribution is wrong, not the models. Trial-aligned +periods - the evaluation distribution is wrong, not the models. Trial-aligned above is the correct comparison. Sliding-window numbers preserved for the record in [`benchmarks/training/results.md`](benchmarks/training/results.md). --- -## Profiling — Cortex-S forward, batch=32, 512 events/sample (MPS) +## Profiling - Cortex-S forward, batch=32, 512 events/sample (MPS) | Section | Time (ms) | % of forward | Triton candidate | |---|---|---|---| -| **Full forward** | **129.0** | 100% | — | +| **Full forward** | **129.0** | 100% | - | | Self-attention (×7) | 100.6 | 78.0% | YES (RMSNorm+linear fuse) | | Cross-attention | 13.6 | 10.5% | YES (block-sparse kernel) | | `_pack_events` | 7.8 | 6.1% | fixed in Python (cumsum) | | Behavior head | 3.0 | 2.3% | no (< 5% rule) | | Tokenizer | 1.4 | 1.1% | YES (fuses 3 embeds anyway) | -Profiling chose the Triton targets — not intuition. Full report: +Profiling chose the Triton targets - not intuition. Full report: [`benchmarks/profiling/baseline_report.md`](benchmarks/profiling/baseline_report.md). --- @@ -91,7 +91,7 @@ All three are correctness-verified against PyTorch references in | Configuration | Memory | Reduction | |---|---|---| -| float32 | 99.2 MB | — | +| float32 | 99.2 MB | - | | INT8 per-channel | **27.8 MB** | **−72%** | - Max abs output diff fp32 vs INT8: **0.003** @@ -116,7 +116,7 @@ Full report: [`benchmarks/quantization/results.md`](benchmarks/quantization/resu | Batch timeout | 5 ms | | Max batch | 32 | -Mode is `in_process_direct` — latency covers scheduler queue wait + +Mode is `in_process_direct` - latency covers scheduler queue wait + inference, not HTTP serialization. Full report: [`benchmarks/serving/results.md`](benchmarks/serving/results.md). diff --git a/README.md b/README.md index c24c578..194b93b 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ [![CI](https://github.com/peterajhgraham/cortex-engine/actions/workflows/ci.yml/badge.svg)](https://github.com/peterajhgraham/cortex-engine/actions/workflows/ci.yml) -A brain-computer interface that misses its deadline is a cursor that lags, a prosthetic that jerks, a patient who cannot type. Closing the loop on motor cortex requires sub-30 ms p99 latency over irregular, high-dimensional spike streams — hundreds of neurons firing asynchronously at sub-millisecond resolution, with the population code distributed across cells and time in ways that resist the fixed shapes most accelerators are built for. Cortex-Engine treats this as a systems problem first: a Perceiver-style decoder that compresses variable-length spike events into a fixed latent set, three custom Triton kernels targeting the layers profiling proved were bottlenecks, INT8 quantization, and a continuous-batching inference server with a paged streaming KV cache, all wired through a production-grade Prometheus/Grafana/OpenTelemetry observability stack. The goal is not a paper — it is the inference engine you would actually deploy behind a real BCI. +A brain-computer interface that misses its deadline is a cursor that lags, a prosthetic that jerks, a patient who cannot type. Closing the loop on motor cortex requires sub-30 ms p99 latency over irregular, high-dimensional spike streams - hundreds of neurons firing asynchronously at sub-millisecond resolution, with the population code distributed across cells and time in ways that resist the fixed shapes most accelerators are built for. Cortex-Engine treats this as a systems problem first: a Perceiver-style decoder that compresses variable-length spike events into a fixed latent set, three custom Triton kernels targeting the layers profiling proved were bottlenecks, INT8 quantization, and a continuous-batching inference server with a paged streaming KV cache, all wired through a production-grade Prometheus/Grafana/OpenTelemetry observability stack. The goal is not a paper - it is the inference engine you would actually deploy behind a real BCI. --- @@ -39,7 +39,7 @@ Production inference infrastructure for transformer-based neural decoders. Three > Consolidated reference: [`BENCHMARKS.md`](BENCHMARKS.md). Forward-looking work: [`ROADMAP.md`](ROADMAP.md). -### Decoding accuracy (trial-aligned evaluation — NLB protocol) +### Decoding accuracy (trial-aligned evaluation - NLB protocol) One sample per reach trial, window −100 ms to +500 ms around `move_onset_time`, target = velocity at onset. Full report: [`benchmarks/training/trial_aligned_results.md`](benchmarks/training/trial_aligned_results.md). @@ -51,7 +51,7 @@ Full report: [`benchmarks/training/trial_aligned_results.md`](benchmarks/trainin | Vanilla Transformer | ~5 M | pending CUDA | | Cortex-S | 24.80 M | **0.60** | -*Sliding-window R² values (published in [`benchmarks/training/results.md`](benchmarks/training/results.md)) are near zero because ~85% of windows are rest periods — the evaluation distribution is wrong, not the models. Trial-aligned evaluation above is the correct comparison.* +*Sliding-window R² values (published in [`benchmarks/training/results.md`](benchmarks/training/results.md)) are near zero because ~85% of windows are rest periods - the evaluation distribution is wrong, not the models. Trial-aligned evaluation above is the correct comparison.* ### Inference (MPS, Cortex-S, batch=32, 512 events/sample) @@ -69,9 +69,9 @@ Full report: [`benchmarks/profiling/baseline_report.md`](benchmarks/profiling/ba | Kernel | What it fuses | Memory saved | Speedup (A10) | |---|---|---|---| -| Fused tokenizer | 3 embedding lookups → 1 kernel | Eliminates 2 × (E, D) intermediates | — | +| Fused tokenizer | 3 embedding lookups → 1 kernel | Eliminates 2 × (E, D) intermediates | - | | Sparse cross-attention | FA2 online softmax + block sparsity | Skips masked event tiles | **up to 27×** | -| Fused RMSNorm + linear | Norm + matmul → 1 kernel | 112 MB / forward at Cortex-S scale | — | +| Fused RMSNorm + linear | Norm + matmul → 1 kernel | 112 MB / forward at Cortex-S scale | - | Run with `make bench-kernels` on an NVIDIA A10 (24GB) host. @@ -79,7 +79,7 @@ Run with `make bench-kernels` on an NVIDIA A10 (24GB) host. | Configuration | Memory | Reduction | |---|---|---| -| float32 | 99.2 MB | — | +| float32 | 99.2 MB | - | | INT8 per-channel | **27.8 MB** | **−72%** | Max weight error: 0.003. 34/35 linear layers quantized. Full report: [`benchmarks/quantization/results.md`](benchmarks/quantization/results.md). @@ -96,9 +96,9 @@ Max weight error: 0.003. 34/35 linear layers quantized. Full report: [`benchmark > **Note on p99:** The 261 ms p99 reflects a first-batch initialization spike (CUDA kernel JIT compile + worker warmup on the initial request). Steady-state p99 is **~28 ms**, well within the 30 ms SLO. Full report: [`benchmarks/serving/results.md`](benchmarks/serving/results.md). -> ### Honest Results — hardware caveat +> ### Honest Results - hardware caveat > -> The model, kernels, scheduler, KV cache, server, and observability stack are all implemented, tested, and benchmarked end-to-end. Where numbers are reported on **MPS** (Apple Silicon), they are real measurements on real data — but they are not the hero numbers. The hero numbers (Triton kernel speedups, p99 < 30 ms serving, full FSDP training of Cortex-M) require an **NVIDIA A10 (24 GB)** and Triton's CUDA backend, which is not available on Apple Silicon. Every CUDA-dependent benchmark cell is explicitly marked **"pending CUDA"** rather than estimated, extrapolated, or quietly omitted. Kernel correctness is verified against PyTorch references within `rtol=1e-3, atol=1e-3` on every commit via CI. A consolidated view lives in [`BENCHMARKS.md`](BENCHMARKS.md). +> The model, kernels, scheduler, KV cache, server, and observability stack are all implemented, tested, and benchmarked end-to-end. Where numbers are reported on **MPS** (Apple Silicon), they are real measurements on real data - but they are not the hero numbers. The hero numbers (Triton kernel speedups, p99 < 30 ms serving, full FSDP training of Cortex-M) require an **NVIDIA A10 (24 GB)** and Triton's CUDA backend, which is not available on Apple Silicon. Every CUDA-dependent benchmark cell is explicitly marked **"pending CUDA"** rather than estimated, extrapolated, or quietly omitted. Kernel correctness is verified against PyTorch references within `rtol=1e-3, atol=1e-3` on every commit via CI. A consolidated view lives in [`BENCHMARKS.md`](BENCHMARKS.md). --- @@ -141,47 +141,47 @@ Max weight error: 0.003. 34/35 linear layers quantized. Full report: [`benchmark ## Why This Matters Beyond BCI -The systems patterns here are the same ones that make modern LLM inference work. The async scheduler with deadline-aware continuous batching is the same technique vLLM uses to coalesce variable-length generation requests into shared GPU work. The paged streaming KV cache is the same memory model as PagedAttention — fixed-size pages, LRU eviction, near-zero fragmentation — adapted from token sequences to sliding-window spike contexts with 91.6% overlap. The INT8 quantization pipeline (per-channel weights, 99th-percentile activation calibration, dequant-then-matmul) is the same recipe as LLM.int8(), implemented from scratch instead of via a one-line library call so the calibration code is auditable. FSDP2 sharded training, fused Triton kernels for the layers that profiling proved were bottlenecks, and OpenTelemetry tracing through an async request path are all directly transferable to any transformer-serving stack — this repo is a BCI inference engine that happens to be built out of LLM-systems primitives. +The systems patterns here are the same ones that make modern LLM inference work. The async scheduler with deadline-aware continuous batching is the same technique vLLM uses to coalesce variable-length generation requests into shared GPU work. The paged streaming KV cache is the same memory model as PagedAttention - fixed-size pages, LRU eviction, near-zero fragmentation - adapted from token sequences to sliding-window spike contexts with 91.6% overlap. The INT8 quantization pipeline (per-channel weights, 99th-percentile activation calibration, dequant-then-matmul) is the same recipe as LLM.int8(), implemented from scratch instead of via a one-line library call so the calibration code is auditable. FSDP2 sharded training, fused Triton kernels for the layers that profiling proved were bottlenecks, and OpenTelemetry tracing through an async request path are all directly transferable to any transformer-serving stack - this repo is a BCI inference engine that happens to be built out of LLM-systems primitives. --- ## What's Built -### Phase 1 — Model and Training Pipeline ✓ +### Phase 1 - Model and Training Pipeline ✓ - **Three model sizes:** Cortex-XS (4.83M), Cortex-S (24.80M), Cortex-M (83.51M). Perceiver cross-attention encoder with a fixed latent array (L=256 for Cortex-S); handles variable neuron counts without hard-coding electrode geometry. - **FSDP training loop:** Mixed-precision bfloat16, sharded checkpointing via `torch.distributed.checkpoint`, cosine LR schedule with warmup. Single-GPU mode on MPS. - **Three baselines + MC_Maze loader:** Wiener filter, GRU, vanilla Transformer evaluated under identical protocol. pynwb/DANDI pipeline, 137 heldin units, 5 ms bins, trial-aligned and sliding-window dataset modes. -### Phase 2.1 — Profiling ✓ +### Phase 2.1 - Profiling ✓ - **Bottleneck found:** `_pack_events` Python loop forced 32 CPU↔MPS sync stalls per forward pass. Replaced with vectorized `cumsum`; saves 4.2 ms / forward (−3.2%). - **Bottleneck hierarchy:** Self-attention 78%, cross-attention 10.5%, pack_events 6.1%. Triton targets chosen from this, not intuition. -### Phase 2.2 — Three Triton Kernels ✓ +### Phase 2.2 - Three Triton Kernels ✓ - **Fused tokenizer:** 2D kernel fuses three embedding lookups into one pass; eliminates two intermediate (E, D) tensors. 9-config autotuner keyed on (E, D). - **Block-sparse cross-attention:** FlashAttention-2 online softmax (running m/l/o). Skips entire event tiles where the block mask is false. External mask API keeps sparsity policy decoupled from the kernel. -- **Fused RMSNorm+linear:** Two-pass kernel — accumulate x² for RMS in pass 1, apply norm × gamma inline during `tl.dot()` in pass 2. x_norm never touches HBM. +- **Fused RMSNorm+linear:** Two-pass kernel - accumulate x² for RMS in pass 1, apply norm × gamma inline during `tl.dot()` in pass 2. x_norm never touches HBM. -### Phase 2.6 — INT8 Quantization ✓ +### Phase 2.6 - INT8 Quantization ✓ - **Per-channel calibration:** Activation scales use 99th-percentile abs-max across calibration batches. Weight scales: absmax per output neuron. -- **`QuantizedLinear`:** Stores INT8 weights + float32 scales; dequantizes to bf16 before matmul — device-agnostic, no quantized CUDA kernel dependency. +- **`QuantizedLinear`:** Stores INT8 weights + float32 scales; dequantizes to bf16 before matmul - device-agnostic, no quantized CUDA kernel dependency. -### Phase 3 — Inference Engine ✓ +### Phase 3 - Inference Engine ✓ - **Scheduler + worker:** `asyncio.PriorityQueue` ordered by deadline (EDF). Worker runs in a `ThreadPoolExecutor`; on CUDA, `compute_stream` and `copy_stream` overlap H2D transfer with the previous batch's compute. - **`StreamingKVCache`:** Paged embedding cache `(num_pages, page_size, hidden_dim)`. LRU eviction via `OrderedDict`. Exploits the 91.6% overlap between consecutive 600 ms / 50 ms-stride BCI windows. - **FastAPI server:** `POST /decode`, `WS /stream`, `GET /metrics` (Prometheus sub-app). Admission control raises HTTP 429 when the queue is full. -### Phase 4 — Operations and Observability ✓ +### Phase 4 - Operations and Observability ✓ - **Three Grafana dashboards** auto-provisioned at startup: traffic (req/s, error rate, queue depth), latency (p50–p99.9, SLO burn gauge, batch-size heatmap), resources (GPU memory, utilization, KV cache hit rate). - **docker-compose stack:** cortex-engine + Prometheus + Grafana + OTel Collector + k6 loadgen. CPU override in `docker-compose.cpu.yml`. `docker compose up` brings up the full stack. - **Helm chart + Alertmanager:** GPU node selector/toleration, ServiceMonitor, autoscaling on `cortex_queue_depth`. Alerts for p99 > 50ms, error rate > 0.1%, queue saturation. -### Phase 5 — Writeup ✓ +### Phase 5 - Writeup ✓ - Engineering postmortem (`docs/writeup.md`, ~3 500 words): architecture decisions, profiling methodology, kernel design, honest results with hardware caveats, what I'd do differently. @@ -193,7 +193,7 @@ The systems patterns here are the same ones that make modern LLM inference work. |---|---| | **PyTorch 2.2+** | FSDP2, `scaled_dot_product_attention` → FlashAttention dispatch, MPS backend for dev on Apple Silicon. | | **Triton** | Custom GPU kernels in Python; generates PTX directly. CUDA-only. | -| **einops** | Readable tensor reshapes — `rearrange` instead of chains of `.view()`. | +| **einops** | Readable tensor reshapes - `rearrange` instead of chains of `.view()`. | | **Hydra** | Hierarchical config composition with CLI overrides. | | **Pydantic v2** | Typed config and I/O schemas; `model_validator` for cross-field constraints. | | **pynwb / DANDI** | NLB data is NWB-native; no conversion step needed. | @@ -318,10 +318,10 @@ make docker-build # builds cortex-engine:latest from Dockerfile make docker-up # docker compose up -d # Services: -# http://localhost:8080 — inference API -# http://localhost:9090 — Prometheus -# http://localhost:3000 — Grafana (admin / admin) -# localhost:4317 — OTel Collector (gRPC) +# http://localhost:8080 - inference API +# http://localhost:9090 - Prometheus +# http://localhost:3000 - Grafana (admin / admin) +# localhost:4317 - OTel Collector (gRPC) # Run k6 load test against the live stack: docker compose -f ops/docker/docker-compose.yml --profile loadtest run loadgen @@ -340,11 +340,11 @@ docker compose \ Three dashboards are auto-provisioned at startup (Grafana → Dashboards → Cortex): -**Traffic** (`cortex-traffic`) — request rate by endpoint, error rate by type, queue depth, Little's Law in-flight estimate. +**Traffic** (`cortex-traffic`) - request rate by endpoint, error rate by type, queue depth, Little's Law in-flight estimate. -**Latency** (`cortex-latency`) — p50/p95/p99/p99.9 time series with 30 ms / 50 ms thresholds. SLO burn gauge spikes red when the k6 ramping scenario saturates the server. Batch-size heatmap shows continuous batching coalescing requests into batches of 16–32. +**Latency** (`cortex-latency`) - p50/p95/p99/p99.9 time series with 30 ms / 50 ms thresholds. SLO burn gauge spikes red when the k6 ramping scenario saturates the server. Batch-size heatmap shows continuous batching coalescing requests into batches of 16–32. -**Resources** (`cortex-resources`) — GPU memory and utilization, KV cache pages and hit rate. Hit rate should stay above 80% during streaming sessions given the 91.6% sliding-window overlap. +**Resources** (`cortex-resources`) - GPU memory and utilization, KV cache pages and hit rate. Hit rate should stay above 80% during streaming sessions given the 91.6% sliding-window overlap. --- diff --git a/ROADMAP.md b/ROADMAP.md index 6426a61..505f5a9 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,6 +1,6 @@ # Roadmap -Cortex-Engine is feature-complete across its five planned phases — model, +Cortex-Engine is feature-complete across its five planned phases - model, profiling-driven Triton kernels, INT8 quantization, continuous-batching inference engine, and full observability stack. What follows is the forward work that would make the project stronger as a research and @@ -19,7 +19,7 @@ populate the TBD rows in `benchmarks/kernels/results.md`, generate a roofline plot per kernel, and lock down the autotune configs at production input shapes (E ∈ {1024, 4096, 16384, 65536}, D ∈ {128, 256, 384, 512}). The 27× block-sparse cross-attention number is indicative -from a single shape — the sweep will produce the real curve. +from a single shape - the sweep will produce the real curve. ## 2. Full NLB MC_Maze trial-aligned evaluation with all baselines @@ -50,12 +50,12 @@ foundation model of motor cortex." ## 4. End-to-end true-INT8 matmul on CUDA -**Status:** weight-only INT8 with bf16 dequant-then-matmul today — 72% +**Status:** weight-only INT8 with bf16 dequant-then-matmul today - 72% weight memory reduction, no runtime activation savings. Wire `cublasLtMatmul` (or a Triton INT8 GEMM) into `QuantizedLinear` so the matmul itself runs in INT8 on A10 tensor cores. Re-measure on -A10 — expected wins are activation memory reduction and throughput, +A10 - expected wins are activation memory reduction and throughput, both currently absent because the dequant path materializes bf16 inputs. Add a third row to the quantization table: `INT8 (true matmul, CUDA)`. @@ -71,5 +71,5 @@ over time. --- Contributions welcome on any of the above. Open an issue first if the -change touches a benchmark number or a Triton kernel — the project's +change touches a benchmark number or a Triton kernel - the project's honest-reporting rule applies to PRs as well as commits.