8000 Add model.Exec.DonateStore() and support for backend SharedConstants by janpfeifer · Pull Request #476 · gomlx/gomlx · GitHub
[go: up one dir, main page]

Skip to content

Add model.Exec.DonateStore() and support for backend SharedConstants - #476

Merged
janpfeifer merged 3 commits into
mainfrom
gobackend-memusage
Nov 1, 2026
Merged

janpfeifer merged 3 commits into
mainfrom
gobackend-memusage

Conversation

@janpfeifer
Copy link
Copy Markdown
Contributor

Summary

Optimize memory usage for inference by allowing model.Store variables to be donated and embedded as shared backend constants across dynamic graph compilations:

  • Added model.Exec.DonateStore(): Marks the model.Store as donated to the model.Exec for inference.
    • When the backend supports Capabilities.SharedConstants, variables used by the computation graph are converted into backend constants and their host/device Tensor values in the Store are immediately finalized to minimize memory usage.
    • Any variables in the Store not used by the graph during compilation are also immediately finalized to reclaim memory.
    • Prevents external access (Variable.Value()) or mutation (Variable.SetNodeValue()) on donated variables.
  • Added support for Capabilities.SharedConstants:
    • Added graph.SharedConstant, Node.BackendValue, and graph.FinalizeConstantValue in core/graph.
    • Allows cached backend constants (and any derived packed weight caches) in model.Variable to be shared and reused across dynamic graph recompilations (e.g., for different input shapes) within the same backend without keeping duplicate copies or preserving the original Tensor.
  • Dependencies & Docs:
    • Updated github.com/gomlx/compute dependency.
    • Updated docs/CHANGELOG.md and added unit tests in ml/model/exec_test.go.

Added support for backend Capabilities.SharedConstants: allow constants to be shared across graphs.
@janpfeifer
janpfeifer merged commit 5835769 into main Nov 1, 2026
4 checks passed
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

0