Adopt the torch_tpu tensor_buffer API and drop the nobuilder build exclusions. - #877
Open
copybara-service[bot] wants to merge 1 commit into
Open
Adopt the torch_tpu tensor_buffer API and drop the nobuilder build exclusions.#877copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_972642255
branch
from
September 7, 2026 22:30
9dd7acc to
3fa7854
Compare
…clusions. The torch integration compiled against three torch_tpu internal headers that only a source checkout provides, so every torch-dependent target carried a "nobuilder" tag. tensor_buffer now covers everything the integration uses. - Tensors are resolved, inspected, and materialized through GetBaseTensorBuffer / TensorBufferHandle; every keep-alive holder stores the handle instead of DeviceBufferRef. - BUILD: targets depend on tensor_buffer through a header_only_cc_info wrapper; the three internal-header forwarding targets and the "nobuilder" tags are removed. Exception: kv_cache_manager_test_py still imports torch_tpu._internal Python modules with no public equivalent, so it keeps "nobuilder" and is tagged "manual". - The external wheel build (build.sh) passes --define raiden_wheel_build=true, which makes that wrapper header-only: the torch extension is compiled against the API but does not link its objects, and binds the symbols at import to the installed torch_tpu wheel, which exports them from its per-torch-version common library. Every other build links the API as a normal dependency. Previously the extension embedded its own copy of the API together with its allocator state, and a running torch_tpu rejected the extension's tensors on the allocator-identity check. - The Copybara transform that stripped the device_buffer.h include from torch_tpu_utils.cc on export is removed; the include no longer exists. Validation: - Both torch packages build with no torch_tpu source checkout (26 + 25 targets), and the package unit tests pass on CPU (4/4). - Wheel: the torch extension leaves the seven torch_tpu API symbols undefined in all three torch ABI variants (nm); built without the wrapper, they are embedded. - On a TPU host, against a torch_tpu wheel that exports the API, the extension resolves GetBaseTensorBuffer from torch_tpu's common library and constructs a KVCacheManager over device tensors. The same probe with the previous link model aborts on the allocator-identity check. PiperOrigin-RevId: 972642255
copybara-service
Bot
force-pushed
the
test_972642255
branch
from
September 7, 2026 22:40
3fa7854 to
93e1a77
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adopt the torch_tpu tensor_buffer API and drop the nobuilder build exclusions.
The torch integration compiled against three torch_tpu internal headers
that only a source checkout provides, so every torch-dependent target
carried a "nobuilder" tag. tensor_buffer now covers everything the
integration uses.
GetBaseTensorBuffer / TensorBufferHandle; every keep-alive holder
stores the handle instead of DeviceBufferRef.
wrapper; the three internal-header forwarding targets and the
"nobuilder" tags are removed. Exception: kv_cache_manager_test_py still
imports torch_tpu._internal Python modules with no public equivalent,
so it keeps "nobuilder" and is tagged "manual".
raiden_wheel_build=true, which makes that wrapper header-only: the
torch extension is compiled against the API but does not link its
objects, and binds the symbols at import to the installed torch_tpu
wheel, which exports them from its per-torch-version common library.
Every other build links the API as a normal dependency. Previously the
extension embedded its own copy of the API together with its allocator
state, and a running torch_tpu rejected the extension's tensors on the
allocator-identity check.
torch_tpu_utils.cc on export is removed; the include no longer exists.
Validation:
(26 + 25 targets), and the package unit tests pass on CPU (4/4).
undefined in all three torch ABI variants (nm); built without the
wrapper, they are embedded.
extension resolves GetBaseTensorBuffer from torch_tpu's common library
and constructs a KVCacheManager over device tensors. The same probe
with the previous link model aborts on the allocator-identity check.