183 lines
7.1 KiB
Markdown
183 lines
7.1 KiB
Markdown
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 —
|