RustyTorch++ — Phase 1 Plan (Core Compiler & Runtime)

Objective (3–4 months): Deliver a functioning core of RustyTorch++ capable of compiling a minimal model graph to GPU code (via rustg) and executing it with a safe runtime: device/stream/event abstractions, pooled allocator, graph capture/replay, and initial fused kernels. Establish deterministic benchmarks and CI gates beyond Phase 0.

Reference paths

RUSTYTORCH_ROOT: /home/osobh/projects/rustytorch

RUSTG_COMPILER: /home/osobh/projects/rust/rustg (GPU codegen backend; RTX 5090 sm_120 primary)

STRATOSWARM_ROOT: /home/osobh/projects/stratoswarm (deployment later)

1. Scope & Non‑Goals

In scope

Front‑end → IR → rustg codegen integration for a minimal operator set (matmul, layernorm, GELU, softmax, elementwise)

Runtime MVP: Device/Stream/Event/Graph/Module/TensorStorage traits + pooled allocator

CUDA Graphs (capture/replay) around forward path; fallback to eager launches

Initial fused kernels (fused MLP, RMSNorm/LayerNorm, rotary embeddings)

AMP (fp16/bf16) with loss‑scaler and overflow recovery

Determinism suite + seed discipline; fragment/leak checks

Perf smoke suites and baselines (micro + macro) with Nsight/rocprof traces

Out of scope (Phase 2+)

Full autograd; full operator coverage

Distributed training; NCCL/RCCL collectives (design only in Phase 1)

Advanced autotuning and auto‑kernel synthesis (basic param sweeps only)

2. Success Criteria (Exit / Phase‑Gate)

Compiler path works: Selected ops lower to IR and generate GPU binaries with rustg; runs on RTX 5090 (sm_120).

Runtime stability: Multi‑stream scheduling + pooled allocator pass soak tests (≥ 6h) without leaks; frag% < 15%.

Determinism: Fixed‑seed runs produce stable outputs within tolerances (fp32 ≤ 1e‑6; bf16/fp16 ≤ 1e‑3).

Performance: ≥ 20% step‑time improvement vs. naive eager launches on the Phase 1 macrobench.

Observability: Nsight/rocprof traces captured; metrics exported to dashboards; graph capture hit‑rate ≥ 70%.

Governance: SBOM + signed artifacts; CI gates enforced (perf/determinism/security).

3. Architecture Work

Compiler integration

Wire rtx-compiler passes: Canonicalize → Shape/Type check → Layout Normalize → Early Fusion

Lower supported ops to rustg backends; embed specialization constants; record kernel cache key {arch, layout, tile}

Golden IR snapshots + round‑trip tests

Runtime MVP

Traits: Device, Stream, Event, Graph, Module, TensorStorage

Pooled arena allocator + page‑locked host pool; deterministic allocation order

Stream scheduler with dependency DAG; memcpy/compute overlap; priority streams

Graph capture/replay API (CUDA/HIP), toggleable; eager fallback path

Module loader: fatbin/hsaco; versioned on arch and specialization

Precision & Numerics

AMP (fp16/bf16) loss‑scaler with overflow counters; automatic recovery policies

Strict determinism mode: stable reductions; reproducible checkpoints (WAL)

4. Operator & Kernel Targets (Phase 1)

Core: matmul, add, mul, gelu, softmax, layer_norm, rms_norm

Embeddings: rotary position embedding (RoPE)

Fusions: MLP (GEMM + bias + activation), LayerNorm epilogue fusions

Utilities: dtype casts, view/reshape, transpose/permute

Bench reference shapes (LLM‑ish):

GEMM: (batch=16..64, seq=1024..4096, hidden=4096..8192)

LN/RMSNorm: (N=hidden); RoPE: (batch, seq, heads, head_dim)

5. Benchmarks & Baselines

Micro

GEMM throughput (TFLOPs), achieved occupancy, DRAM GB/s, L2 hit rate

LN/RMSNorm latency, kernel fusion coverage

Macro

Encoder/decoder block forward step time and peak memory

Graph capture hit‑rate, memcpy/compute overlap %, allocator fragmentation %

Determinism

Fixed seed, fixed inputs; record outputs, graph hash, kernel cache key; check tolerances

6. CI Gates (Phase 1)

Determinism Gate: bit‑stable (declared) or tolerance‑stable outputs; seed & graph hash must match

Perf Gate: macrobench ≥ 20% improvement vs. eager baseline, no >5% regressions week‑over‑week

Memory Gate: frag% ≤ 15%, no leaks in 6h soak

Security Gate: pinned deps, cargo-deny/cargo-audit clean, SBOM + signed artifacts

7. Observability & Telemetry

nsys/ncu or rocprof traces attached to CI artifacts; normalized summaries

Prometheus metrics: step_time_ms, tokens_per_s (if applicable), alloc_frag_pct, occ_pct, graph_hit_pct, dram_gbps

Dashboards: per‑SKU view (RTX 5090 sm_120), trend lines, regression alerts

8. Deliverables

Compiler: IR passes + rustg lowering for core ops; golden tests

Runtime: allocator + stream scheduler + graph capture (CUDA), eager fallback, module loader

Kernels: fused MLP, LN/RMSNorm, RoPE; AMP scaler

Benches: micro + macro; deterministic harness; baseline JSONs

CI: gates for determinism/perf/memory/security; artifact uploads (traces, SBOM)

Docs: docs/phase1.md (how to run benches, interpret traces), ADRs updated

9. RACI — Phase 1 Agent Ownership

Conductor (Orchestrator) — Accountable: critical path, phase‑gate readiness, dependency unblocking

Rust Engineer — Responsible: runtime traits, allocator, graph capture, module loader

Frontend Compiler Agent — Responsible: IR passes, lowering to rustg, golden IR

IR Rewrite & Fusion Agent — Responsible: early fusions; cost model annotations; epilogue fusion specs

Kernel Synthesizer Agent — Responsible: compile fused kernels; occupancy reports; fallbacks

Performance Engineer — Accountable: bench design, profiler harness, baseline capture, perf gates

Memory Planner Agent — Consulted: allocator policies; fragmentation tests

Security Engineer & Security Auditor — Accountable: dependency policies; scans; SBOM/sign

DevOps Engineer — Responsible: CI matrix; GPU runners; artifact storage; reproducibility

Governance & Provenance Agent — Responsible: signatures, provenance manifests

Agent Organizer — Accountable: tool allowlists; routing; freeze switch if gates fail

10. Risk Register & Mitigations

rustg ABI churn → pin compiler SHA per build; compatibility tests

Graph capture instability → robust fallback path; capture boundary tests

Profiler variance → warm‑up protocol; driver/toolchain pinning; result smoothing

Memory fragmentation → arena growth policy; periodic defrag windows (doc only in P1)

Numeric drift with AMP → strict reference runs; overflow logging; scaler tuning

11. Timeline (suggested)

Weeks 1–2: IR integration & golden tests; runtime trait scaffolding; allocator MVP

Weeks 3–4: Graph capture API + eager fallback; fused kernels v0; AMP scaler

Weeks 5–6: Micro/macro benches; profiler wiring; baseline capture; CI gates online

Weeks 7–8: Soak tests; tuning; docs & ADRs; phase‑gate review

12. Phase‑2 Handover Seeds

Autograd engine design notes and gradient numerics tests scaffold

NCCL/RCCL design doc: topology discovery, overlap strategy, retry/timeouts

Autotuning roadmap: tile/warp/vector sweeps → bandit search → kernel cache semantics

— End of Phase 1 Plan —