Skip to content

[Feature] Size the hbg graph heap from task lifetimes plus a materialization lookahead cap, not a 256 MB constant #1956

Description

@ChaoWao

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:

  1. 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.
  2. 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.
  3. 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:

  1. 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.
  2. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions