feat(runtime): always-on TRB temporary variable buffer (internal, no user switch)

[puddingfjz]

Summary

Cut the per-run device_malloc / device_free of ordinary non-child tensors on the tensormap_and_ringbuffer (TRB) bind path. TRB now stages them through a runner-scoped retained temporary buffer that grows to the high-water plan size and is reused across runs — steady-state (stable decode shapes) converges to zero temporary allocations.

This is an internal allocation optimization, always on for TRB — there is no user-facing switch. It is always result-equivalent to the old malloc/free path (only faster), and TRB runs are already serial per runner (enforced by the static arena pool), so there is no correctness reason to gate it. The buffer is a single retained allocation per DeviceRunnerBase; begin_temporary_buffer_run(plan) sizes it from the run's non-child tensors, acquire slices from it, validate ends the run. Grow = free-old + alloc-new (no data preserved — per-run scratch); it never shrinks; on-exhaustion the run fails (no mid-run malloc fallback).

Design notes

• No budget, no user API. The buffer self-sizes from each run's plan; the caller configures nothing. There is no max_temporary_buffer_bytes, no temporary_buffer_mode, no configure_* entry across Python / pybind / C-ABI / ChipWorker.
• Ownership. TensorLease + TensorReleaseKind replace the implicit "every recorded dev_ptr was malloc'd this run" assumption: buffer slices release as no-ops (whole buffer reused at run end); the Free kind covers any backend that leaves the temp callbacks unwired and falls back to device_malloc.
• child-memory preserved. Child-memory and zero-byte tensors are skipped by the plan and pass through unchanged.
• hbg unaffected — host-orchestration bind does not use the temp buffer callbacks.

See the runtime's RUNTIME_LOGIC.md §2.4 and docs/task-flow.md for mechanics.

Testing

• a2a3 + a5 build clean (all runtimes, sim + onboard).
• Python UT: 92 passed (test_chip_worker, test_host_worker).
• Sim ST: tensormap_and_ringbuffer/prepared_callable a2a3sim — 7 passed (end-to-end register→run→validate with the buffer always on).
• C++ UT (test_temporary_variable_buffer, test_trb_runtime_temp_buffer) cover grow / reuse / shrink-noop / grow-failure / plan-sizing / child-memory pass-through / copy-back / lease-release — to be exercised by CI.

Follow-up on the optimization direction discussed in #1186.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 49186d93-d58d-472e-a041-edb1461ad604

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds docs/trb-serial-tensor-buffer-pool-plan.md, a 590-line design document specifying a lease-based bump allocator for reusing pre-retained device memory across decode steps in the tensormap_and_ringbuffer path, covering ownership model, DeviceRunnerBase placement, HostApi wiring, bind/validate path changes, Qwen3 tensor classification, metrics, tests, acceptance criteria, and deferred work.

Changes

TRB Serial Tensor Buffer Pool Plan

Layer / File(s)	Summary
Scope, non-goals, and lease ownership model docs/trb-serial-tensor-buffer-pool-plan.md	Establishes the optimization's decision scope, explicit non-goals, the current bind/validate allocation/free flow for non-child tensors, and replaces implicit TensorPair ownership with an explicit lease model using release kinds.
DeviceRunnerBase placement, budget contract, and bump allocator lifecycle docs/trb-serial-tensor-buffer-pool-plan.md	Specifies relocating buffer management to DeviceRunnerBase, defines the aggregate max_temporary_buffer_bytes budget contract with alignment/padding and error rules, configuration ingress rules, segmented-chunk retained memory approach, HostApi wiring, and the per-run configure/begin/acquire/end/clear lifecycle.
Concurrency assumptions and bind/validate path changes docs/trb-serial-tensor-buffer-pool-plan.md	States the single-active-run concurrency assumption and enumerates non-additions (no locking, no fallback malloc, no double buffering), then details bind path changes (begin run, acquire slices, record leases), validate path changes (release-kind dispatch replacing device_free, clear leases, end run), cleanup contract, and preserved copy semantics including Qwen3 tensor classification.
Metrics, tests, acceptance criteria, and deferred work docs/trb-serial-tensor-buffer-pool-plan.md	Defines lightweight counter/metrics requirements, minimum unit and integration tests, acceptance criteria for correctness/compatibility/performance, and explicitly deferred items.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~15 minutes

Poem

🐇 A plan arrives, fresh from the burrow below,
Bump allocators reuse memory's flow,
Leases replace frees on the hot decode path,
One runner, one run — no concurrency wrath,
The doc hops in, and the rabbit says: go! 🥕

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly captures the main change: an always-on TRB temporary variable buffer with no user switch.
Description check	✅ Passed	The description is directly about the TRB temporary buffer optimization and matches the documented implementation plan.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces an implementation plan for a TRB Temporary Variable Buffer to optimize memory allocation overhead. The review feedback highlights three key areas for improvement: simplifying the cleanup contract by leveraging the existing validate_runtime_impl gateway, resolving a design contradiction between the aggregate budget contract and segmented chunk allocation, and adding a lightweight active-run guard to prevent silent data corruption from concurrent runs.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

docs/trb-serial-tensor-buffer-pool-plan.md (3)

11-11: 📐 Maintainability & Code Quality | 🔵 Trivial

Consider using English consistently throughout the document.

Line 11 uses Chinese characters "临时变量" while the rest of the document is entirely in English. For consistency and to avoid encoding issues in some toolchains, consider using only English: "This plan uses 'temporary variable buffer' for the same concept as 'linshi biancun buffer'."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/trb-serial-tensor-buffer-pool-plan.md` at line 11, The document mixes
Chinese and English, so update the referenced text in the plan to use English
consistently. Replace the Chinese phrase in the buffer description with an
English equivalent in the same section of
docs/trb-serial-tensor-buffer-pool-plan.md, keeping the wording aligned with the
existing terminology used elsewhere in the document.

---

154-157: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Clarify finalize cleanup ordering when a run is active.

The text says to "log it, release retained chunks" if finalize sees an active run. For defensive correctness, explicitly state whether end_temporary_buffer_run() should be called before clear() to reset the active_ flag, or if clear() alone is sufficient because the runner is being destroyed. This prevents implementers from leaving the flag in an inconsistent state if any post-finalize diagnostics check it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/trb-serial-tensor-buffer-pool-plan.md` around lines 154 - 157, Clarify
the finalize cleanup order in finalize_common()/sim finalize when an active
temporary-buffer run is present: specify whether end_temporary_buffer_run() must
be called before clear() to reset active_, or whether clear() alone is
sufficient because the runner is being torn down. Update the contract wording so
implementers know how to handle the active run state consistently and avoid
leaving diagnostics-visible flags stale.

---

526-552: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add unit tests for configuration edge cases.

Consider adding focused unit tests for:

• configure_temporary_buffer with same budget twice is a no-op (no reallocation);
• configure_temporary_buffer while a run is active fails with clear error;
• begin_temporary_buffer_run failure path (e.g., OOM during chunk allocation) is handled correctly.

These complement the existing test list and guard against reconfiguration bugs.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/trb-serial-tensor-buffer-pool-plan.md` around lines 526 - 552, Add
focused unit tests for the temporary-buffer configuration edge cases: verify
configure_temporary_buffer is a no-op when called twice with the same budget,
verify configure_temporary_buffer rejects changes while a run is active with a
clear error, and verify begin_temporary_buffer_run handles allocation
failure/OOM correctly. Place these alongside the existing buffer unit tests and
use the same temporary-buffer pool APIs and run lifecycle symbols so the new
coverage matches the current single-active-run model.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/trb-serial-tensor-buffer-pool-plan.md`:
- Around line 448-460: The bind pseudocode currently assumes
begin_temporary_buffer_run() always succeeds and sets temp_run_active=true
unconditionally, but the HostApi callback returns a bool. Update the bind flow
to check the return value from begin_temporary_buffer_run() and fail the bind
path immediately with a clear error if it returns false, before any TensorLease
recording or device memory work; keep temp_run_active=true only after a
successful begin_temporary_buffer_run() call.

---

Nitpick comments:
In `@docs/trb-serial-tensor-buffer-pool-plan.md`:
- Line 11: The document mixes Chinese and English, so update the referenced text
in the plan to use English consistently. Replace the Chinese phrase in the
buffer description with an English equivalent in the same section of
docs/trb-serial-tensor-buffer-pool-plan.md, keeping the wording aligned with the
existing terminology used elsewhere in the document.
- Around line 154-157: Clarify the finalize cleanup order in
finalize_common()/sim finalize when an active temporary-buffer run is present:
specify whether end_temporary_buffer_run() must be called before clear() to
reset active_, or whether clear() alone is sufficient because the runner is
being torn down. Update the contract wording so implementers know how to handle
the active run state consistently and avoid leaving diagnostics-visible flags
stale.
- Around line 526-552: Add focused unit tests for the temporary-buffer
configuration edge cases: verify configure_temporary_buffer is a no-op when
called twice with the same budget, verify configure_temporary_buffer rejects
changes while a run is active with a clear error, and verify
begin_temporary_buffer_run handles allocation failure/OOM correctly. Place these
alongside the existing buffer unit tests and use the same temporary-buffer pool
APIs and run lifecycle symbols so the new coverage matches the current
single-active-run model.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: e992fafc-1d81-49ff-a3d7-3cea3aeb9bfd

📥 Commits

Reviewing files that changed from the base of the PR and between 86bd60a and 185bc6e.

📒 Files selected for processing (1)

• docs/trb-serial-tensor-buffer-pool-plan.md

[ChaoWao]

Two things: (1) replace explicit budget with a self-sizing AUTO buffer, (2) title/body vs. actual diff

1. Reconsider the core Decision: drop the explicit budget, make the buffer self-sizing.

This is a direction-level suggestion, not a tweak — it changes the plan's Decision, so flagging it up front.

The current design pushes model-shape knowledge (hidden size, vocab pad, max batch/token) onto the serving/runner owner via max_temporary_buffer_bytes, and fail-fasts when a run exceeds it. I'd argue the buffer should instead be user-invisible and self-sizing: remember the high-water size, and only device_malloc a bigger chunk when a run actually needs more — PyTorch-caching-allocator style. Steady-state converges to zero allocations, and the caller configures nothing.

Concretely:

• Drop max_temporary_buffer_bytes from the public API and the fail-fast "budget exceeded" path.
• configure no longer pre-allocates; the buffer starts empty.
• On an acquire miss, grow: device_malloc a new retained chunk (reuse the existing allocate_chunks machinery) and continue. Track high-water at the whole-run/arena granularity, not per-tensor, to keep aggregate semantics and avoid fragmentation.
• Modes collapse to two: off (budget == 0, current malloc/free path) and AUTO (enabled → self-sizing). No numeric budget.

I want to be honest about what this gives up, because these were presumably why you chose explicit budget:

• It violates the plan's acquire() must not device_malloc() on the hot path invariant until the size converges. For qwen3 decode (stable shapes) that's 1–2 runs then never again — near-ideal. For prefill (growing seq len) it can re-trigger growth repeatedly, adding hot-path malloc jitter in exactly the window where 507018 (slow prefill killed by short timeouts) lives. Worth a caveat / measurement on the prefill path.
• It loses fail-fast diagnosis — AUTO grows until true OOM, whose error is less actionable than "required X / configured Y". If bounded/predictable memory is a hard serving requirement, AUTO can't provide it. If you believe serving needs that guarantee, this whole suggestion may not fit — that's the trade to decide.

The upside is that "zero-config, just make the churn go away" is a real product win, and it keeps model-shape knowledge out of the runtime entirely — which the plan already lists as a goal.

2. Title/body vs. diff. The PR says "documentation-only... Not run", but the diff carries ~2400 lines of implementation + C++/Python UTs beyond the 592-line plan doc. Those tests should be actually run / CI-green before merge, and the title probably wants feat: rather than docs:.

[ChaoWao]

Re-review after the AUTO rework — looks good, two loose ends

The AUTO switch landed well, and the begin_run(plan) design is actually better than what I suggested: pre-scanning the run's non-child tensors and reallocating once up front means acquire() is pure slicing and never mallocs mid-run. That directly addresses my main worry — the "hot-path malloc jitter on prefill / 507018 window" concern is gone, since growth happens at most once per run in begin_run, and not at all once shapes converge (RepeatedSamePlanReusesRetainedBuffer / SmallerPlanDoesNotShrinkRetainedBuffer cover that). The align_up_checked / align_ptr_checked overflow guards and the reworked UTs are solid too. Nice.

Two loose ends before merge:

1. Two plan docs now coexist and one is stale. docs/trb-serial-tensor-buffer-pool-plan.md still describes the explicit-budget design in full ("requires a caller-provided max_temporary_buffer_bytes", the budget-exceeded fail-fast), while the new docs/trb-auto-temporary-buffer-modification-plan.md describes AUTO. The old doc now documents a design that no longer exists — per .claude/rules/doc-consistency.md it should be deleted or rewritten, not left contradicting the AUTO doc.

2. Title/CI (still open from my last comment). Diff is ~2878 lines of implementation + C++/Python UTs; title is still docs:. Should be feat:, and the UTs need to actually run / be CI-green rather than treated as a docs-only change.

Neither is a design issue — the mechanism itself is in good shape now.

[ChaoWao]

Two more cleanups, both about "make it the cleanest version":

1. Drop the max_temporary_buffer_bytes migration error entirely.

worker.py currently keeps:

if "max_temporary_buffer_bytes" in config:
    raise ValueError("max_temporary_buffer_bytes has been removed; use temporary_buffer_mode='auto'")

plus its UT in test_host_worker.py. But max_temporary_buffer_bytes never landed on main — git log --all -S / git grep HEAD both come up empty; it only ever existed in this PR's own earlier commits. So this is a back-compat shim for a public API that was never released: nobody has code to migrate. It just plants the false impression that the param once existed. Cleanest is to delete the branch and its UT — the identifier shouldn't appear in the PR at all.

2. realloc naming is misleading — CANN has no rtRealloc.

I checked the CANN 9.0.0 runtime headers: the only device-memory primitives are rtMalloc / rtFree (+ variants) — there is no rtRealloc. So begin_run's "realloc" is necessarily release_buffer() (rtFree) + allocate_buffer() (rtMalloc), i.e. free-then-alloc that does not preserve data. That's correct for a per-run scratch buffer (last run's contents are dead anyway), but the names realloc_count_ / realloc_failed_count_ / "AUTO realloc failed" invite the C-realloc() reading of "in-place grow that keeps data". Suggest renaming to something like regrow_count_ / "buffer grow failed" so nobody assumes a data-preserving resize that the hardware layer can't even provide.

Neither blocks the design — both are polish toward the minimal/cleanest form.

[ChaoWao]

Pushed a commit directly (maintainer edit, 9e8d1a8b) taking the "cleanest form" a step further: made the temporary buffer an unconditional TRB internal default and removed the user-facing switch entirely.

Rationale: the AUTO buffer is always result-equivalent to the old per-run malloc/free path — only faster — and TRB runs are already serial per runner (the static arena pool enforces this), so there was no correctness reason to gate it behind opt-in. Exposing it through Python / pybind / C-ABI / ChipWorker pushed an internal allocation decision up the whole stack for no benefit.

Removed: Worker.configure_temporary_buffer_auto / temporary_buffer_mode / _normalize_temporary_buffer_mode, the pybind binding, ChipWorker::configure_temporary_buffer_auto + dlsym, configure_temporary_buffer_auto_ctx (C ABI), DeviceRunnerBase::configure_temporary_buffer_auto / temporary_buffer_enabled, TemporaryVariableBuffer::configure_auto / enabled() / Stats.enabled, and the HostApi.temporary_buffer_enabled callback. runtime_maker now gates purely on whether the begin/acquire/end callbacks are wired (trb always wires them; null → transparent per-tensor device_malloc fallback). Net −228 lines.

Docs (RUNTIME_LOGIC.md §2.4, task-flow.md) and tests retargeted to the always-on behavior; OFF-path test coverage now nulls the callbacks instead of toggling an enable flag.

Verified locally: a2a3+a5 sim+onboard runtimes build clean; Python UT 92 passed; trb prepared_callable a2a3sim ST 7 passed. C++ cpput not run locally (pre-existing gtest link-env issue on this box, unrelated — the .o's compile); leaving that to CI.

[ChaoWao]

Pushed a follow-up maintainer commit (e09e591b) that collapses the abstraction one more level, per the observation that this layer never needed to be an allocator.

The prior commit kept a ~330-line TemporaryVariableBuffer class in the platform, driven through three HostApi callbacks (begin/acquire/end) + a TemporaryBufferPlanItem wire type. But the runtime already has device_malloc/device_free and the 1024-byte packing logic — so all the platform actually needs to provide is somewhere to remember one {addr, size} across runs.

• Deleted TemporaryVariableBuffer (header + its unit test) entirely.
• HostApi: the three temp callbacks + TemporaryBufferPlanItem → two slot accessors get_retained_temp_buffer / set_retained_temp_buffer.
• DeviceRunnerBase (onboard+sim): the buffer object → a plain void* retained_temp_addr_ + size_t retained_temp_size_; freed once at finalize.
• runtime_maker (a2a3+a5): a file-local RetainedTempBump does the whole thing — pack non-child tensors to a 1024-aligned required size, read the slot, grow (device_free old + device_malloc new + write back) only when too small, then bump-slice. Grow failure fails the run; a backend leaving the slot callbacks null falls back to per-tensor device_malloc transparently.
• Dropped the now-unused Runtime::temporary_buffer_run_active_ (bump offset is bind-local; only the slot spans runs).

Net −442 lines on top of the previous commit. No behavior change — same always-on retained-buffer reuse, just without the allocator-class indirection. Verified: a2a3+a5 sim+onboard build clean; Python UT 92 passed; prepared_callable + orch_so_cache a2a3sim ST pass. C++ cpput left to CI (pre-existing local gtest link-env issue, unrelated).