Summary
host_build_graph reserves its graph heap at a compile-time constant —
#define PTO2_HEAP_SIZE (256 * 1024 * 1024) in both
src/{a2a3,a5}/runtime/host_build_graph/runtime/pto_runtime2_types.h:87 — while the
host knows, by the end of orchestration, exactly which tasks are live at the same
time. Derive the reservation from the tasks' reuse relationships instead.
The value is three terms, not one, and each has a different derivation. That is
the part worth stating up front, because a lifetime-only computation sizes the heap
far too small:
- Outputs (
required_heap) — written by a node at execution and read by its
consumers, so its lifetime is DAG-bounded. Interval overlap over the
dependency graph is the correct model, and this is the term lifetime analysis
computes.
- Graph execution storage — written when the scheduler materializes a
replay, which is deliberately decoupled from dependency readiness. The DAG does
not bound its liveness, so this term is lookahead × per-window bytes and is
only finite once the scheduler caps how many replays it materializes
concurrently.
- Long-lived allocations (
alloc_tensors and friends) — live for the whole
run, so neither lever touches them. On qwen3-14b decode this is the majority
of the heap.
So: heap = peak_live_outputs(lifetime analysis) + lookahead × execution_storage + long_lived.
Motivation / Use Case
Measured on qwen3-14b decode, a2a3, one die:
|
bytes |
|
share |
heap high-water (host_orch heap_used) |
127,673,344 |
121.8 MB |
nothing is reclaimed |
| replay windows × 40 |
55,316,480 |
52.8 MB |
43.3% |
long-lived (alloc_tensors + others) |
72,356,864 |
69.0 MB |
56.7% |
With a materialization lookahead of 3 the floor is
72,356,864 + 3 × 1,382,912 = 76,505,600 B (73.0 MB).
Against the reserved 256 MB that is a ~3.5x over-reservation on a case that has
already been squeezed hard: #1952 brought actual use down to heap_used=80,405,504
(76.7 MB) on current main (1b637ef0), so the gap between what is reserved and
what is used is now the largest remaining item in hbg's device-memory footprint.
The constant also has to cover the worst case any workload might need, so it cannot
be lowered safely without computing the requirement — which is the point of this
request.
This is the same defect shape as #1920 (ready queues sized by a fixed constant),
one level up: there the fix is to size each queue by the tasks that can reach it,
here it is to size the heap by the bytes that can be simultaneously live.
Proposed API / Behavior
Two changes, and the second is a precondition for the first being a bound:
- Bind-time computation. The host already walks every task and every Graph
node during orchestration. Accumulate the interval-overlap peak of the outputs
term there, add lookahead × execution_storage, add the long-lived total, and
pass the result to setup_static_arena instead of PTO2_HEAP_SIZE. 6864edc7
on a local branch is a first step (reserve at the size the graph needs) but sizes
from the total, not the live peak.
- Scheduler lookahead cap. Bound how many replays are materialized
concurrently — a small pool of expansion buffers rather than one window per
occurrence. Without this, term 2 is unbounded by anything the host can see and a
computed size is not an upper bound.
Order matters: shipping (1) without (2) computes a number the scheduler is free to
exceed, and the failure mode is heap exhaustion mid-run.
Alternatives Considered
- Reuse one replay window's bytes for another replay at runtime. Dropped, and
the reason is structural rather than a matter of tuning: execution storage is
written at materialization, and the dispatch loop's own contract is that
"External dependency readiness and bounded definition materialization progress
independently, then meet at the outer task payload's graph_activation_gate."
A later replay can therefore write into a window an earlier replay is still
using, so interval overlap on submission order does not model it and reuse is
unsafe regardless of what the DAG says. Capping the lookahead is the lever for
this term, which is why it appears above as a precondition rather than an
optimization.
- Just lower the constant. It has to cover the worst case across every
workload, so lowering it without a computed requirement trades a memory win for
an exhaustion risk on some other case.
- Survey what reuse would be worth first. A tool that measures the potential
(hbg_heap_live_set.py) was written and is being dropped: if the value is
computed at bind time anyway, the survey is redundant.
Additional Context
Two things constrain how this can be validated:
- qwen3-14b decode cannot validate the lifetime analysis. Its
required_heap
is 0 bytes — the whole replay window is execution storage — so term 1
evaluates to zero and a correct and an incorrect implementation are
indistinguishable on it. Coverage needs a case whose Graph bodies create
intermediates.
- Term 3 caps the achievable win. Even a perfect lifetime analysis leaves
56.7% of qwen's heap untouched, because those tensors are held for the whole run.
Pushing past ~73 MB on this case means attacking that, not the replay windows.
An earlier pass on this claimed 13.3x headroom (a 9.13 MB lifetime peak against
121.8 MB). That was wrong twice over: it attributed the entire heap to the replay
windows, and it mismeasured a window at 3,191,808 B where logging its two halves
shows 1,382,912 B. The corrected write-up (three-term decomposition, the
decoupling contract, the required_heap == 0 coverage gap) is not yet landed under
docs/investigations/; it should be, since the numbers above are what make the
design reviewable.
Related: #1920, #1706
Summary
host_build_graphreserves its graph heap at a compile-time constant —#define PTO2_HEAP_SIZE (256 * 1024 * 1024)in bothsrc/{a2a3,a5}/runtime/host_build_graph/runtime/pto_runtime2_types.h:87— while thehost knows, by the end of orchestration, exactly which tasks are live at the same
time. Derive the reservation from the tasks' reuse relationships instead.
The value is three terms, not one, and each has a different derivation. That is
the part worth stating up front, because a lifetime-only computation sizes the heap
far too small:
required_heap) — written by a node at execution and read by itsconsumers, so its lifetime is DAG-bounded. Interval overlap over the
dependency graph is the correct model, and this is the term lifetime analysis
computes.
replay, which is deliberately decoupled from dependency readiness. The DAG does
not bound its liveness, so this term is
lookahead × per-window bytesand isonly finite once the scheduler caps how many replays it materializes
concurrently.
alloc_tensorsand friends) — live for the wholerun, so neither lever touches them. On qwen3-14b decode this is the majority
of the heap.
So:
heap = peak_live_outputs(lifetime analysis) + lookahead × execution_storage + long_lived.Motivation / Use Case
Measured on qwen3-14b decode, a2a3, one die:
host_orch heap_used)alloc_tensors+ others)With a materialization lookahead of 3 the floor is
72,356,864 + 3 × 1,382,912 = 76,505,600B (73.0 MB).Against the reserved 256 MB that is a ~3.5x over-reservation on a case that has
already been squeezed hard: #1952 brought actual use down to
heap_used=80,405,504(76.7 MB) on current
main(1b637ef0), so the gap between what is reserved andwhat is used is now the largest remaining item in hbg's device-memory footprint.
The constant also has to cover the worst case any workload might need, so it cannot
be lowered safely without computing the requirement — which is the point of this
request.
This is the same defect shape as #1920 (ready queues sized by a fixed constant),
one level up: there the fix is to size each queue by the tasks that can reach it,
here it is to size the heap by the bytes that can be simultaneously live.
Proposed API / Behavior
Two changes, and the second is a precondition for the first being a bound:
node during orchestration. Accumulate the interval-overlap peak of the outputs
term there, add
lookahead × execution_storage, add the long-lived total, andpass the result to
setup_static_arenainstead ofPTO2_HEAP_SIZE.6864edc7on a local branch is a first step (reserve at the size the graph needs) but sizes
from the total, not the live peak.
concurrently — a small pool of expansion buffers rather than one window per
occurrence. Without this, term 2 is unbounded by anything the host can see and a
computed size is not an upper bound.
Order matters: shipping (1) without (2) computes a number the scheduler is free to
exceed, and the failure mode is heap exhaustion mid-run.
Alternatives Considered
the reason is structural rather than a matter of tuning: execution storage is
written at materialization, and the dispatch loop's own contract is that
"External dependency readiness and bounded definition materialization progress
independently, then meet at the outer task payload's
graph_activation_gate."A later replay can therefore write into a window an earlier replay is still
using, so interval overlap on submission order does not model it and reuse is
unsafe regardless of what the DAG says. Capping the lookahead is the lever for
this term, which is why it appears above as a precondition rather than an
optimization.
workload, so lowering it without a computed requirement trades a memory win for
an exhaustion risk on some other case.
(
hbg_heap_live_set.py) was written and is being dropped: if the value iscomputed at bind time anyway, the survey is redundant.
Additional Context
Two things constrain how this can be validated:
required_heapis 0 bytes — the whole replay window is execution storage — so term 1
evaluates to zero and a correct and an incorrect implementation are
indistinguishable on it. Coverage needs a case whose Graph bodies create
intermediates.
56.7% of qwen's heap untouched, because those tensors are held for the whole run.
Pushing past ~73 MB on this case means attacking that, not the replay windows.
An earlier pass on this claimed 13.3x headroom (a 9.13 MB lifetime peak against
121.8 MB). That was wrong twice over: it attributed the entire heap to the replay
windows, and it mismeasured a window at 3,191,808 B where logging its two halves
shows 1,382,912 B. The corrected write-up (three-term decomposition, the
decoupling contract, the
required_heap == 0coverage gap) is not yet landed underdocs/investigations/; it should be, since the numbers above are what make thedesign reviewable.
Related: #1920, #1706