202 lines
6.8 KiB
Markdown
202 lines
6.8 KiB
Markdown
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 —
|