jgrusewski 87ca1f5f55 perf(ml-alpha): CUDA Graph capture of training step (#162 finale)
Captures the full Mamba2 fwd → K-loop fwd → BCE → K-loop bwd →
Mamba2 bwd → AdamW × 7 → loss DtoD chain into a single CUDA Graph.
First step_batched call runs uncaptured (cuBLAS warmup); second
call captures; third+ replays. Replaces ~155 individual kernel
launches per step with one graph launch.

Four root-cause issues had to be fixed in concert to make the
capture region capture-compatible:

1. cuBLAS lazy workspace allocation
   crates/ml-alpha/src/mamba2_block.rs — pre-allocate an 8 MiB
   workspace via cublasSetWorkspace_v2 at Mamba2Block::new.
   cuBLAS would otherwise allocate on first gemm call with each
   new shape, breaking capture.

2. cuBLAS heuristic plan-cache lookup
   crates/ml-core/src/cuda_autograd/linear.rs — switch
   gemm_ex_f32 from CUBLAS_GEMM_DEFAULT_TENSOR_OP to
   CUBLAS_GEMM_DFALT. The heuristic algo path triggers
   plan-cache allocs; DFALT is deterministic with negligible
   perf delta for our shapes (Mamba2: 128×32, 128×state_dim).

3. Per-call cuModuleLoadData in bias kernels
   crates/ml-core/src/cuda_autograd/linear.rs — add a
   `BiasKernels` struct (add_bias_2d_kernel + reduce_sum_axis0
   handles) cached at construction. `BiasKernels::shared(stream)`
   uses a per-context OnceLock cache so the cubin loads exactly
   once per CUDA context for the process lifetime. Every
   `GpuLinear` and `OwnedGpuLinear` constructor now stores its
   `BiasKernels`. Removed the per-call `get_bias_kernels`
   helper entirely (no legacy aliases — greenfield).

4. Per-call CudaSlice::clone() in Mamba2 fwd + bwd
   crates/ml-alpha/src/mamba2_block.rs — three sites cloned
   the input slice to build a fresh `GpuTensor` view. Each
   `CudaSlice::clone()` does cuMemAlloc + dtod copy
   (vendor/cudarc/src/driver/safe/core.rs:1437); cuMemAlloc
   is forbidden during capture.

   Refactored `forward_with_slices_into` and
   `backward_with_slices_into` to take raw `x_data: &CudaSlice<f32>,
   batch: usize` instead of `&GpuTensor` / `&LinearActivations`.
   All three Mamba2 fwd call sites + three bwd call sites now
   pass the underlying CudaSlice directly.

Trainer changes (trainer/perception.rs):
  - Three-state machine in step_batched: warmup → capture → replay.
  - Labels staging fill moved BEFORE the captured region (host writes
    only; replays read whatever the host last wrote).
  - Removed the post-snap_batched sync that was inside dispatch (it
    would trip STREAM_CAPTURE_ISOLATION; kernels are stream-ordered
    so the sync was unnecessary).
  - Dispatch extracted into `dispatch_train_step` method; the
    captured region brackets exactly this method.
  - Final sync + mapped-pinned loss read happens once per step in
    step_batched, outside the captured region.

Validation:
  - Synthetic overfit smoke: 0.397 → 0.0007 (matches pre-refactor
    trajectory; capture+replay produces equivalent loss).
  - 26 ml-alpha lib tests + 23 integration tests pass.
  - 306 ml-core lib tests pass.

Also fixed (orthogonal but required for green workspace):
  crates/ml-core/src/action_space.rs — tests assumed 7-exposure
  layout (63 actions). Production `ExposureLevel` enum has 8
  variants (Hold inserted at idx 3, Flat moved to idx 7 → 72
  total actions). Updated tests + `get_valid_action_mask` to
  match the 8-level layout.

Honors:
  - feedback_no_legacy_aliases.md (no `add_bias_2d_with_fn` /
    legacy `get_bias_kernels` fallback; one canonical API).
  - feedback_no_partial_refactor.md (signature change propagated
    to every caller atomically; no half-migrated state).
  - feedback_wire_everything_up.md (BiasKernels wired into all
    GpuLinear/OwnedGpuLinear constructors in same commit).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-17 15:21:00 +02:00

Foxhunt

Production HFT trading system in Rust.

Architecture

The workspace contains 32 crates organized as follows:

Core Libraries (16)

Crate Purpose
trading_engine Order processing, FIX 4.4, IB TWS, SIMD, RDTSC timing
risk VaR, Kelly, circuit breakers, kill switches, compliance
risk-data Risk data types and shared structures
trading-data Trading data types
ml DQN Rainbow, PPO, TFT, Mamba2, ensemble inference
ml-data ML data types and feature definitions
data Market data ingestion and storage
backtesting Replay engine, strategy tester
adaptive-strategy Ensemble execution, microstructure analysis
common Shared types, resilience, error handling
storage S3 and local model storage
model_loader Model serialization and loading
market-data Market data feed handlers
database PostgreSQL access layer (SQLx)
config Configuration management
tli CLI commands and tooling

Services (8)

Service Purpose
backtesting_service gRPC backtesting service
broker_gateway_service FIX routing, broker connectivity
trading_service Core trading operations
ml_training_service Model training orchestration
data_acquisition_service Market data acquisition
trading_agent_service Autonomous trading agents
api_gateway gRPC API gateway with auth
web-gateway Axum REST + WebSocket gateway

Frontend

web-dashboard/ -- React 19 + TypeScript + Vite + TradingView charts.

Building

# Check compilation (no PostgreSQL required)
SQLX_OFFLINE=true cargo check --workspace

# Run tests for a specific crate
SQLX_OFFLINE=true cargo test -p <crate> --lib

# Clippy
SQLX_OFFLINE=true cargo clippy --workspace

ML Models

Four production model architectures on Candle v0.9.1 with CUDA:

  • DQN Rainbow -- Deep Q-Network with prioritized replay, dueling heads, noisy nets
  • PPO -- Proximal Policy Optimization with GAE and LSTM policies
  • TFT -- Temporal Fusion Transformer for multi-horizon forecasting
  • Mamba2 -- State space model for sequence prediction

Each model has a standalone trainer and a UnifiedTrainable adapter for the hyperopt pipeline.

Infrastructure

  • Git: Gitea at git.fxhnt.ai (Tailscale-only), Scaleway DEV1-S
  • Observability: OpenTelemetry OTLP (env OTEL_EXPORTER_OTLP_ENDPOINT)
  • Database: PostgreSQL with SQLx offline mode for CI

License

Proprietary. All rights reserved.

Description
No description provided
Readme 849 MiB
Languages
Rust 88.2%
Cuda 7.7%
Python 1.3%
Shell 1.1%
PLpgSQL 0.8%
Other 0.8%