From b60d098dacbe4fed9d350881d1b8ab0f7774bec8 Mon Sep 17 00:00:00 2001 From: fhzhang Date: Tue, 4 Aug 2026 21:22:55 -0700 Subject: [PATCH] Migration of JAX integration to use public API for extracting raw buffers. PiperOrigin-RevId: 959399603 --- tpu_raiden/api/jax/BUILD | 3 +- tpu_raiden/frameworks/jax/BUILD | 28 +++++++-------- .../frameworks/jax/raw_transfer_test.cc | 36 ++++++++++++++++++- 3 files changed, 48 insertions(+), 19 deletions(-) diff --git a/tpu_raiden/api/jax/BUILD b/tpu_raiden/api/jax/BUILD index 1a7c7e18..ad3d2ac1 100644 --- a/tpu_raiden/api/jax/BUILD +++ b/tpu_raiden/api/jax/BUILD @@ -42,7 +42,6 @@ py_test( ], tags = [ "manual", - "notap", ], deps = [ ":kv_cache_store", @@ -85,7 +84,7 @@ py_library( py_test( name = "weight_synchronizer_test", srcs = ["weight_synchronizer_test.py"], - tags = ["notap"], + tags = [], deps = [ ":weight_synchronizer_jax_py", "//tpu_raiden/rpc:raiden_service_py_pb2", diff --git a/tpu_raiden/frameworks/jax/BUILD b/tpu_raiden/frameworks/jax/BUILD index 41bbc321..7e811e70 100644 --- a/tpu_raiden/frameworks/jax/BUILD +++ b/tpu_raiden/frameworks/jax/BUILD @@ -81,10 +81,9 @@ cc_library( "@xla//xla:shape_util", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt:pjrt_layout", + "@xla//xla/pjrt:raw_buffer", "@xla//xla/python:nb_numpy", "@xla//xla/python:types", - "@xla//xla/python/ifrt", - "@xla//xla/python/pjrt_ifrt", ], ) @@ -430,7 +429,6 @@ cc_library( name = "raw_transfer_jax", srcs = ["raw_transfer.cc"], hdrs = [ - "jax_utils.h", "raw_transfer_internal.h", ], copts = [ @@ -440,6 +438,7 @@ cc_library( ], features = ["-use_header_modules"], deps = [ + ":jax_utils", ":mock_nanobind", "//tpu_raiden/core:raw_transfer_core", "//tpu_raiden/core:raw_transfer_impl", @@ -449,14 +448,15 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", - "@jax//jaxlib:py_client", "@nanobind", "@rules_python//python/cc:current_py_cc_headers", + "@xla//xla:shape_util", "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_layout", "@xla//xla/pjrt:status_casters", "@xla//xla/pjrt/c:pjrt_c_api_raw_buffer_external", # buildcleaner: keep - "@xla//xla/python/ifrt", - "@xla//xla/python/pjrt_ifrt", + "@xla//xla/python:nb_numpy", + "@xla//xla/python:types", ], ) @@ -483,12 +483,11 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", - "@xla//xla/pjrt:pjrt_client", + "@xla//xla:shape_util", + "@xla//xla/pjrt:pjrt_layout", + "@xla//xla/pjrt:raw_buffer", "@xla//xla/pjrt:status_casters", "@xla//xla/pjrt/c:pjrt_c_api_raw_buffer_external", # buildcleaner: keep - "@xla//xla/python/ifrt", - "@xla//xla/python/pjrt_ifrt", - "@xla//xla/tsl/concurrency:ref_count", ], ) @@ -500,8 +499,10 @@ cc_test( deps = [ ":raw_transfer_jax_mock", "//tpu_raiden/core:raw_transfer_core", + "@com_google_absl//absl/log:check", "@xla//xla/pjrt:common_pjrt_client", "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:raw_buffer", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/stream_executor/host:host_platform", "@xla//xla/tsl/platform:statusor", @@ -518,7 +519,6 @@ py_test( main = "kv_cache_manager_test.py", tags = [ "cpu:4", - "notap", ], deps = [ ":_tpu_raiden_jax", @@ -581,7 +581,6 @@ py_test( main = "raw_transfer_perf_test.py", tags = [ "cpu:8", - "notap", "requires-accel-hbm", "requires-ghostlite:8", "requires-net:external", @@ -604,7 +603,6 @@ py_test( main = "raw_transfer_perf_test.py", tags = [ "cpu:8", - "notap", "requires-accel-hbm", "requires-ghostfish:4", "requires-net:external", @@ -627,7 +625,6 @@ py_test( main = "raw_transfer_test.py", tags = [ "cpu:8", - "notap", "requires-accel-hbm", "requires-ghostlite:8", "requires-net:external", @@ -648,7 +645,6 @@ py_test( main = "raw_transfer_test.py", tags = [ "cpu:8", - "notap", "requires-accel-hbm", "requires-ghostfish:4", "requires-net:external", @@ -707,7 +703,7 @@ cc_test( "-fexceptions", ], features = ["-use_header_modules"], - tags = ["notap"], + tags = [], deps = [ ":kv_cache_manager_core_lib", "//tpu_raiden/core:kv_cache_manager_with_transfer", diff --git a/tpu_raiden/frameworks/jax/raw_transfer_test.cc b/tpu_raiden/frameworks/jax/raw_transfer_test.cc index 317ce5c8..250072e6 100644 --- a/tpu_raiden/frameworks/jax/raw_transfer_test.cc +++ b/tpu_raiden/frameworks/jax/raw_transfer_test.cc @@ -22,8 +22,11 @@ #include #include +#include "absl/log/check.h" +#include "xla/pjrt/abstract_tracked_device_buffer.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" +#include "xla/pjrt/raw_buffer.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" #include "tpu_raiden/core/raw_transfer_core.h" @@ -63,7 +66,38 @@ nb::object CreateMockDeviceArray(xla::PjRtBuffer* pjrt_buffer) { nb::object arr; nb::list shards; nb::object shard; - nb::object shard_data(reinterpret_cast(pjrt_buffer)); + nb::object shard_data; + + xla::PjRtRawBufferInterface* raw_buf = nullptr; + auto* common_buf = dynamic_cast(pjrt_buffer); + if (common_buf) { + auto hold = common_buf->GetBufferWithHold( + xla::CommonPjRtBuffer::ScopedHold::kUsage); + if (hold.ok()) { + raw_buf = hold.buffer()->raw_buffer().get(); + } + } + + // We must have a valid raw buffer for this test path. + CHECK_NE(raw_buf, nullptr) << "Failed to get raw buffer from pjrt_buffer"; + + nb::object unsafe_raw_buffer_fn; + unsafe_raw_buffer_fn.set_callable([raw_buf]() { + nb::object raw_buf_obj(0); // Dummy value to not be None + raw_buf_obj.set_attr("ptr", nb::object(reinterpret_cast(raw_buf))); + return raw_buf_obj; + }); + shard_data.set_attr("unsafe_raw_buffer", unsafe_raw_buffer_fn); + + nb::list shape_list; + for (int64_t dim : pjrt_buffer->on_device_shape().dimensions()) { + shape_list.add_element(nb::object(dim)); + } + shard_data.set_attr("shape", shape_list); + + shard_data.set_attr("dtype", nb::object()); + shard_data.set_attr("_pjrt_layout", nb::object()); + shard_data.set_attr("ptr", nb::object(reinterpret_cast(pjrt_buffer))); shard.set_attr("data", shard_data); shards.add_element(shard);