Skip to content

feat(memory): add atomic unified budget coordinator - #353

Open
PhilipJohnBasile wants to merge 2 commits into
youssofal:mainfrom
PhilipJohnBasile:feat/unified-memory-budgeting
Open

feat(memory): add atomic unified budget coordinator#353
PhilipJohnBasile wants to merge 2 commits into
youssofal:mainfrom
PhilipJohnBasile:feat/unified-memory-budgeting

Conversation

@PhilipJohnBasile

Copy link
Copy Markdown
Contributor

Splits the unified-memory budget coordinator out of #336.

This PR contains only the memory coordination core:

  • one bounded plan for SessionBank, expert warm-set, and KV headroom budgets
  • normal, elevated, warning, and critical pressure policies
  • explicit safe-point gating
  • hysteresis against budget churn
  • ordered application with rollback after partial failure
  • adapters for SessionBank and attribute-backed consumers
  • a standalone no-model control-path measurement

It does not include server polling, memory sampling, expert residency, live KV resizing, dashboard code, or workflows.

Measurement on commit a23b74d, macOS 27.0 arm64, Python 3.14.6, 100,000 plans and 10,000 apply attempts with one injected failure per 10 applies:

Metric Result
Plan p50 4.167 us
Plan p95 4.375 us
Plan p99 5.334 us
Apply p50 3.166 us
Apply p95 3.291 us
Rollback p95 3.000 us
Injected rollbacks 1,000
Rollback failures 0
Maximum budget overage 0 bytes

Reproduce:

python scripts/bench_unified_memory.py \
  --plans 100000 \
  --apply-iterations 10000 \
  --failure-every 10 \
  --output /tmp/unified-memory.json

Validation:

python -m pytest tests/test_unified_memory.py
8 passed in 0.02s

python -m compileall -q mtplx/unified_memory.py tests/test_unified_memory.py scripts/bench_unified_memory.py
PASS

ruff check mtplx/unified_memory.py tests/test_unified_memory.py scripts/bench_unified_memory.py
PASS

git diff --check
PASS

This measures policy and atomic mutation cost against explicit fake consumers. It does not claim reduced RSS, fewer page faults, lower swap, or protection from OOM. The current runtime also has no concrete live KV resize capability, so this PR treats KV headroom as a protected planning partition only.

@PhilipJohnBasile
PhilipJohnBasile force-pushed the feat/unified-memory-budgeting branch from a23b74d to b13288f Compare September 1, 2026 13:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant