Skip to content

Adopt the torch_tpu tensor_buffer API and drop the nobuilder build exclusions. - #877

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_972642255
Open

Adopt the torch_tpu tensor_buffer API and drop the nobuilder build exclusions.#877
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_972642255

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Sep 7, 2026

Copy link
Copy Markdown

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.

  • 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.

…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
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