rustytorch/docs/archive/legacy/phase4.md

RustyTorch++ — Phase 4 Plan (Advanced Compilation & Auto‑Kernel Synthesis) Objective (5–6 months): Outperform hand‑tuned kernels by introducing auto‑kernel synthesis, hardware‑profile‑driven autotuning, and AOT graph compilation across CUDA/ROCm/Metal. Preserve determinism, stability, and safety, while expanding operator coverage and reducing TCO.

Reference paths

RUSTYTORCH_ROOT: /home/osobh/projects/rustytorch

RUSTG_COMPILER: /home/osobh/projects/rust/rustg (RTX 5090 sm_120 primary target)

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

Phase 4 builds on Phase 1–3 cores (compiler/runtime, tensor+autograd+graph, and ops/model zoo). Distributed training and serving remain separate tracks.

1. Scope & Non‑Goals In scope

Hardware Profile DB (per‑SKU): latency/bandwidth, occupancy sweet spots, cache behavior, launch constraints

Auto‑Kernel Synthesis: generate specialized kernels from IR patterns (tile/warp/vector widths, epilogue fusions)

Autotuning Engine: bandit/grid hybrid; persistent kernel cache keyed by {op, layout, shape, arch, driver}

AOT Graph Compilation: serialize compiled graphs & kernel bundles; fast load with compatibility checks

Cross‑backend parity: CUDA↔ROCm↔Metal feature alignment for supported ops

Operator coverage expansion: Flash‑class attention forward (and Phase 4.5: backward), FFT/convolution primitives

Safety: numerical parity checks, OOB fuzz, deterministic mode support

Out of scope (Phase 5+)

Multi‑tenant inference scheduler and KV‑paging (serving)

Full distributed optimizer sharding (handled in distributed track)

2. Success Criteria (Exit / Phase‑Gate) Throughput wins: ≥20–40% step‑time reduction on Phase‑3 macrobench via synthesized/fused kernels.

Inference speedup: ≥1.5× tokens/sec single‑GPU for LLM‑block with synthesized attention/MLP.

Kernel cache hit‑rate: ≥80% across repeated runs with same shapes/layouts.

AOT load time: <100ms load of compiled graph + kernels on identical SKU/driver.

Parity & determinism: fp32 ≤1e‑6; bf16/fp16 ≤1e‑3; strict mode bit‑stable where declared.

Cross‑backend: CUDA and ROCm produce parity on smoke suites; Metal reaches feature‑parity subset.

3. Architecture Work A. Hardware Profile DB

Capture: microbenches for GEMM bandwidth, L2 behavior, register pressure thresholds, warp occupancy vs. tile

Storage: versioned profiles by {arch, driver, rustg_sha}; Prometheus export + JSON artifacts

B. Auto‑Kernel Synthesis

Pattern library: IR → kernel templates (attention, MLP, norm epilogues, conv/FFT)

Parameter space: tile sizes, block dim, warp count, vector width, epilogue fusion choices

Synthesis flow: generate → compile (nvcc/hipcc/metal) → validate numerics → bench → select → cache

C. Autotuning Engine

Strategies: warm grid search → bandit refinement → exploitation cache

Constraints: thermal/power from NVML; wall‑time budget per op; early‑stop

Cache keys: {op, shape, dtype, layout, arch, driver, pass_ver} → chosen params + perf summary

D. AOT Graph Compilation

Serialization: *.rtxgc (graph + schedule + kernel bundle manifest)

Compatibility: verify against {arch, driver, rustg_sha}; fallback to JIT if mismatch

Pre‑capture CUDA/HIP graphs for hot regions; lazy‑init remaining

E. Cross‑Backend Strategy

Feature flags per op; portability layer for epilogues

ROCm parity plan; Metal subset (mac dev runners)

4. Operator & Kernel Targets (Phase 4) Attention family: fused QKV projection → attention core (Flash‑class forward) → output proj (epilogue fusion)

MLP family: GEMM + bias + activation (GELU/SiLU) fused

Norms: LayerNorm/RMSNorm epilogues

Conv/FFT (pilot): 1D/2D FFT; small‑kernel convs with autotuned tiles

Bench shapes: sequence 1k–8k; hidden 4k–16k; heads 32–64; bf16/fp16 primary, fp32 reference.

5. Benchmarks & Baselines Micro: synthesized vs. baseline kernels (latency, GB/s, occupancy)

Macro: LLM block forward (and training fwd+bwd where implemented); step time, tokens/sec, mem peak

Determinism: fixed seeds; parity to reference kernels; graph capture hash stability; kernel cache version pinning

6. CI Gates (Phase 4) Perf Gate: ≥20% macro step‑time win vs. Phase‑3 baseline; no weekly regression >5%

Cache Gate: kernel cache hit‑rate ≥80% on standard benches

Determinism Gate: parity thresholds hold; strict tests pass where enabled

Security Gate: supply chain scans, SBOM/sign; no unsigned kernel bundles

7. Observability & Telemetry Nsight/rocprof traces with synthesized‑kernel annotations and selected params

Metrics: synth_win_pct, cache_hit_pct, aot_load_ms, occ_pct, dram_gbps, l2_hit_pct

Dashboards: per‑SKU perf deltas; cache effectiveness; AOT load times

8. Deliverables rtx-synth module: pattern library, template emitters, validators

Autotuning engine with NVML‑aware constraints; persistent cache

AOT graph compiler/loader + compatibility checks

Synthesized kernels: attention fwd, fused MLP, norm epilogues; conv/FFT pilot

Bench suite & CI gates; metrics dashboards

Documentation: docs/phase4.md (synthesis, tuning, AOT); ADR updates (kernel cache, graph format)

9. RACI — Phase 4 Agent Ownership Conductor (Orchestrator) — Accountable: phase delivery, dependency orchestration

IR Rewrite & Fusion Agent — Responsible: pattern detection & rewrite rules

Kernel Synthesizer Agent — Responsible: template emitters; compile; numeric parity; occupancy reports

Auto‑Tuner Agent — Responsible: search strategies; profile DB; cache management

Rust Engineer — Responsible: runtime hooks for AOT load/verify; kernel cache I/O; error handling

Performance Engineer — Accountable: micro/macro benches; thresholds; dashboards

Security Engineer & Auditor — Accountable: provenance for kernel bundles; scan results; signing

DevOps Engineer — Responsible: CI matrices; artifact storage; runner provisioning

Governance & Provenance Agent — Responsible: SBOM/signature pipeline; manifest policies

Agent Organizer — Accountable: tool allowlists; routing; freeze on gate failure

10. Risks & Mitigations Search explosion → bandit pruning; parameter templates; profile seeding

Numeric drift → robust reference kernels; tolerance checks; soak tests

Cache invalidation → version keys include driver + rustg_sha; fallback JIT path

Cross‑backend gaps → feature flags; ROCm/Metal parity backlog; transparent fallsbacks

AOT mismatch → strict compatibility checks; informative downgrade path

11. Timeline (suggested) Weeks 1–2: Hardware profile DB; microbench capture; perf dashboards

Weeks 3–4: Pattern library + template emitters; compile/validate baseline templates

Weeks 5–6: Autotuning engine; persistent cache; Nsight/rocprof wiring

Weeks 7–8: AOT compiler/loader; synthesized attention/MLP; CI gates; docs & ADRs; phase‑gate review

12. Phase‑5 Handover Seeds Inference scheduler design (continuous batching, KV paging, SLA lanes)

Cross‑node AOT caching & distribution for Stratoswarm deployments

Online autotuning: workload‑aware live refinement under safety guardrails

— End of Phase 4 Plan —