Commit f2c6c1a
committed
feat(relax): PoC + roadmap for graph-level (TVM Relax) memory planning vs eager
Proves on real hardware (CPU LLVM Relax VM, tvm 0.25.dev0) that lowering a
multi-op step to ONE Relax graph and running StaticPlanBlockMemory gives a
strictly lower peak than eager mx.eval.
Measured (relax_memory_plan_poc.py, fail-loud, RULE #1):
fwd-only 8L: all-live 256MB -> planned 48MB = 5.33x lower
fwd+bwd 6L: all-live 170MB -> planned 77MB = 2.21x; strict peak 52->45MB
fwd+bwd 8L: all-live 57MB -> planned 25.25MB = 2.26x; strict peak 17->13.25MB
Strict-peak win grows with depth (1.16x@6L -> 1.28x@8L): the cross-layer
fwd/bwd liveness eager cannot exploit. relax.transform.Gradient produces the
joint fwd+bwd adjoint in one dataflow block and it runs under planning.
docs/RELAX-GRAPH-MEMORY-PATH.md: result, autodiff coverage, MLX->Relax
boundary (path_c regions as call_tir leaves in one R.function), first PR,
118GB->~26-30GB pass pipeline, honest multi-quarter risks (DPS/physical-ABI
adapter, no auto-remat, optimizer co-planning single-block ICHECK).1 parent 1b021e8 commit f2c6c1a
2 files changed
Lines changed: 713 additions & 0 deletions
0 commit comments