[Model] Learn the retention budget directly from observed peaks

Drop the staged-bytes seed per review direction: the budget is now
min(observed_peak * 1.25, 25% of device capacity), nothing else. A cold
start pays one release per observation step while the peak climbs (two on
the validated MiniMax workload), then boundaries retain steadily; the 5%
free floor and capacity cap are unchanged. The staged-byte ledger remains
only to detect full detachment, which now resets the observed peak so one
workload cannot permanently influence another. Trade-off accepted per
review: lower reserved memory and a simpler policy over a flush-free
first request.

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
hsliuustc0106 committed 2026-08-26 13:43:00 +08:00
1 parent b78db6036a
commit c9bb65902e
3 files changed
+89 -86

No files matched your search

@@ -177,18 +177,19 @@ component can be reused by the next. The DiT prefetch path and its two shared
device buffers are unchanged.
After a component is offloaded, cached-but-unallocated memory is retained only
while it stays within a budget derived from the workload itself and at least
5% of device capacity remains physically free. The budget starts from the
model topology — the staged component byte total reported by the attached
`PinnedModuleStager` instances times a seed multiplier — and then tracks the
largest cached value observed at any component boundary: if a workload's
transient footprint outgrows the seed, the first over-budget boundary
releases once, the observation raises the budget to the observed peak times
a margin, and later boundaries retain. The budget is capped at a quarter of
device capacity so neither a large staged total nor a spurious observation
can monopolize a small device, and a detached stager returns its bytes to the
seed. No user-chosen fraction and no warmup flag are needed. Each decision
logs the staged byte total, observed peak, effective budget and cap,
while it stays within a budget learned from the running workload and at least
5% of device capacity remains physically free. The budget is the largest
cached value observed at any component boundary times a margin, so it tracks
the weight-reuse payload plus the transient footprint whatever the
resolution. It starts at zero: a cold start pays one release per observation
step while the peak climbs, after which boundaries retain steadily, and a
workload change whose cached footprint grows by more than the margin pays one
further release while the new peak lands. The budget is capped at a quarter
of device capacity so a spurious observation cannot monopolize a small
device, and detaching the last attached stager resets the learned peak so one
workload never permanently influences another. No user-chosen fraction, no
multiplier constant, and no warmup flag are needed. Each decision logs the
staged byte total, observed peak, effective budget and cap,
cached-but-unallocated bytes, and free memory.
Missing allocator telemetry also releases it conservatively. Component or
staging failure forces a release, and an out-of-memory allocation gets one