Skip to content

Commit f2c6c1a

Browse files
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

File tree

0 commit comments

Comments
 (0)