kv_cache: reject non-row-addressable layouts and mixed-slice managers

tpu_raiden/core/raw_transfer_core.h

@@ -92,6 +92,41 @@ inline const PJRT_RawBuffer_Extension* GetRawBufferExtension(
       PJRT_Extension_Type::PJRT_Extension_Type_RawBuffer);
 }
 
+// Returns OK when `buffer` decomposes into contiguous per-major-row byte
+// slices (row i addressable at offset i * GetMajorSliceByteSize(buffer)).
+// Layouts that interleave the major dimension — a tiled buffer whose major
+// dim is not the outermost physical dim, or any tiled buffer of rank < 3
+// (both dims are then inside the tile pair) — have no such slice;
+// GetMajorSliceByteSize on them yields a value that is NOT a per-row
+// stride (up to the whole buffer), so callers that index rows with it
+// must reject these layouts first.
+inline absl::Status ValidateMajorSliceAddressable(
+    const xla::PjRtBuffer* buffer) {
+  const xla::Shape& shape = buffer->on_device_shape();
+  if (shape.dimensions_size() == 0) {
+    return absl::InvalidArgumentError(
+        "scalar buffer has no major dimension to slice");
+  }
+  auto pjrt_layout = buffer->layout();
+  const xla::Layout* xla_layout =
+      pjrt_layout ? &pjrt_layout->xla_layout() : nullptr;
+  if (!xla_layout) return absl::OkStatus();
+  if (!xla_layout->minor_to_major().empty() &&
+      xla_layout->minor_to_major().back() != 0) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "buffer ", shape.ToString(), " (layout ", xla_layout->ToString(),
+        ") does not keep its major dimension outermost in device memory; "
+        "major rows interleave and no per-row byte slice exists"));
+  }
+  if (!xla_layout->tiles().empty() && shape.dimensions_size() < 3) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "buffer ", shape.ToString(), " (layout ", xla_layout->ToString(),
+        ") is tiled with rank < 3; rows sit inside the tile pair and no "
+        "per-row byte slice exists"));
+  }
+  return absl::OkStatus();
+}
+
 inline int64_t GetMajorSliceByteSize(const xla::PjRtBuffer* buffer) {
   const xla::Shape& shape = buffer->on_device_shape();
   if (shape.dimensions_size() == 0) return 0;

tpu_raiden/frameworks/torch/BUILD

@@ -40,6 +40,13 @@ header_only_cc_info(
     ],
 )
 
+header_only_cc_info(
+    name = "torch_tpu_materialize_headers",
+    deps = [
+        "@torch_tpu//torch_tpu/eager:materialize",
+    ],
+)
+
 cc_library(
     name = "weight_synchronizer_torch",
     srcs = ["weight_synchronizer.cc"],
@@ -82,6 +89,7 @@ cc_library(
     ],
     deps = [
         ":torch_tpu_device_buffer_headers",
+        ":torch_tpu_materialize_headers",
         ":torch_tpu_structured_log_buffer_headers",
         ":torch_tpu_tensor_to_buffer_headers",
         "//tpu_raiden/core:xla_raw_transfer_headers",
@@ -465,6 +473,8 @@ py_test(
         "@pypi//numpy",
         "@torch_tpu//shims/torch:pytorch",
         "@torch_tpu//torch_tpu",
+        "@torch_tpu//torch_tpu/_internal/execution_mode",
+        "@torch_tpu//torch_tpu/_internal/sync",
     ],
 )
 

tpu_raiden/frameworks/torch/kv_cache_manager_test.py

@@ -14,13 +14,26 @@
 
 """E2E physical unit tests for PyTorch KVCacheManager on XLA TPUs."""
 
+import faulthandler
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
 import torch
+from torch_tpu._internal import execution_mode
+from torch_tpu._internal import sync
 
 from tpu_raiden.frameworks.torch import _tpu_raiden_torch as _kv_cache_manager
 
+EagerMode = execution_mode.EagerMode
+
+# Upper bound on how long constructing a manager over lazy KV caches may take.
+# The constructor holds the GIL, so a regressed deadlock freezes the whole
+# process and no Python-thread watchdog can fire; faulthandler dumps every
+# thread stack and hard-exits the process past this bound so a regression fails
+# loudly instead of hanging the test runner.
+_CONSTRUCTION_DEADLOCK_TIMEOUT_S = 60.0
+
 
 class KVCacheManagerTorchTest(parameterized.TestCase):
 
@@ -33,14 +46,49 @@ def setUp(self):
     self.block_size = 2
     self.slice_byte_size = 16384 // 4  # float32 capacity
 
+  def _new_manager(self, tensors, num_blocks):
+    """Constructs a manager, hard-failing if construction hangs.
+
+    Construction unpacks each cache and awaits its device buffer; a cache still
+    carrying a deferred op whose graph is never materialized would block here
+    forever. The constructor holds the GIL, so a hang freezes the process and no
+    Python-thread watchdog could fire -- guard it with faulthandler instead.
+    """
+    faulthandler.dump_traceback_later(
+        _CONSTRUCTION_DEADLOCK_TIMEOUT_S, exit=True
+    )
+    try:
+      return _kv_cache_manager.KVCacheManager(
+          tensors,
+          local_port=0,
+          host_blocks_to_allocate=num_blocks * self.block_size,
+          parallelism=1,
+      )
+    finally:
+      faulthandler.cancel_dump_traceback_later()
+
   @parameterized.named_parameters(
-      ("fp32", torch.float32),
-      ("int32", torch.int32),
+      ("fp32", torch.float32, False),
+      ("int32", torch.int32, False),
+      ("fp32_deferred", torch.float32, True),
+      ("int32_deferred", torch.int32, True),
   )
-  def test_e2e_distributed_staging_offloads_and_socket_transceives(self, dtype):
+  def test_e2e_distributed_staging_offloads_and_socket_transceives(
+      self, dtype, defer
+  ):
     num_blocks = 2
     shape = (num_blocks * self.block_size, 128, 8)  # 2 blocks capacity
 
+    if defer:
+      # Create and register the caches while they still carry a pending deferred
+      # op (vLLM's DEFER_AND_FUSE mode), exercising the in-place materialization
+      # that construction performs. If that materialization forked a separate
+      # buffer instead of filling the live one, the H2d in step 3 would write
+      # into the copy and the parity check in step 4 would fail.
+      cm = execution_mode.set_eager_mode(EagerMode.DEFER_AND_FUSE)
+      cm.__enter__()
+      self.addCleanup(cm.__exit__, None, None, None)
+
     # 1. Allocate sharded eager tensors directly on the physical XLA TPU devices!
     # Layer 0 and Layer 1 for Source (Trainer)
     src_tensors = []
@@ -62,19 +110,15 @@ def test_e2e_distributed_staging_offloads_and_socket_transceives(self, dtype):
         shards.append(t)
       dst_tensors.append(shards)
 
+    if defer:
+      # Guard: the caches must still be deferred here, or this variant would not
+      # exercise construction-time materialization at all.
+      self.assertFalse(sync.is_materialized(src_tensors[0][0]))
+      self.assertFalse(sync.is_materialized(dst_tensors[0][0]))
+
     # 2. Instantiate two real managers locally on ephemeral loopback ports!
-    ws_source = _kv_cache_manager.KVCacheManager(
-        src_tensors,
-        local_port=0,
-        host_blocks_to_allocate=num_blocks * self.block_size,
-        parallelism=1,
-    )
-    ws_dest = _kv_cache_manager.KVCacheManager(
-        dst_tensors,
-        local_port=0,
-        host_blocks_to_allocate=num_blocks * self.block_size,
-        parallelism=1,
-    )
+    ws_source = self._new_manager(src_tensors, num_blocks)
+    ws_dest = self._new_manager(dst_tensors, num_blocks)
 
     self.assertIsNotNone(ws_source.local_port)
     self.assertIsNotNone(ws_dest.local_port)
@@ -146,6 +190,62 @@ def test_e2e_distributed_staging_offloads_and_socket_transceives(self, dtype):
         np.testing.assert_allclose(actual_data[0:2], expected_val, atol=1e-5)
         np.testing.assert_allclose(actual_data[2:4], expected_val, atol=1e-5)
 
+  def test_construction_over_lazy_kv_caches_does_not_deadlock(self):
+    """Constructing a manager over still-deferred KV caches must not hang.
+
+    Under DEFER_AND_FUSE (vLLM's execution mode) a KV cache registered right
+    after allocation still carries a pending deferred op. UnpackTorchTensor
+    resolves each tensor's base buffer by awaiting it, and the await only
+    completes once the deferred graph has been executed. The unpack forces that
+    execution before awaiting; without it the constructor waits on a graph that
+    only this (now-blocked) thread could flush.
+    """
+    num_blocks = 2
+    shape = (num_blocks * self.block_size, 128, 8)
+
+    with execution_mode.set_eager_mode(EagerMode.DEFER_AND_FUSE):
+      lazy_tensors = [
+          [torch.ones(shape, dtype=torch.float32, device=self.device)]
+          for _ in range(self.num_layers)
+      ]
+      # The regression only exists for unmaterialized tensors; assert the
+      # tensors are actually deferred so a future eager-materialization change
+      # can't silently turn this into a no-op test.
+      for layer in lazy_tensors:
+        for shard in layer:
+          self.assertFalse(sync.is_materialized(shard))
+
+      manager = self._new_manager(lazy_tensors, num_blocks)
+
+    # Unpack awaited each base buffer, so the tensors are materialized now.
+    for layer in lazy_tensors:
+      for shard in layer:
+        self.assertTrue(sync.is_materialized(shard))
+    self.assertIsNotNone(manager.local_port)
+
+  def test_construction_over_host_transferred_kv_caches(self):
+    """KV caches created as `cpu_tensor.to(device)` must unpack.
+
+    vLLM allocates KV caches as `torch.zeros(shape, dtype).to(device)`. Under
+    DEFER_AND_FUSE the resulting tensor's TensorImpl may not expose a device
+    until the transfer materializes, so unpack must not gate on
+    `tensor.device()` before resolving and materializing the base buffer.
+    """
+    num_blocks = 2
+    shape = (num_blocks * self.block_size, 128, 8)
+
+    with execution_mode.set_eager_mode(EagerMode.DEFER_AND_FUSE):
+      transferred = [
+          [torch.zeros(shape, dtype=torch.float32).to(self.device)]
+          for _ in range(self.num_layers)
+      ]
+      manager = self._new_manager(transferred, num_blocks)
+
+    for layer in transferred:
+      for shard in layer:
+        self.assertTrue(sync.is_materialized(shard))
+    self.assertIsNotNone(manager.local_port)
+
 
 if __name__ == "__main__":
   absltest.main()

tpu_raiden/frameworks/torch/torch_tpu_utils.cc

@@ -21,6 +21,7 @@
 #include "ATen/core/TensorBody.h"
 #include "torch/headeronly/core/DeviceType.h"
 
+#include "torch_tpu/eager/materialize.h"
 #include "torch_tpu/eager/structured_log_buffer.h"
 #include "torch_tpu/eager/tensor_to_buffer.h"
 
@@ -79,6 +80,15 @@ UnpackedTensor UnpackTorchTensor(const at::Tensor& tensor) {
         "requires a full reinterpret of the storage, not a partial view.");
   }
 
+  // Materialize deferred tensor so AwaitBuffer() won't hang. No-op if already
+  // materialized.
+  if (auto status = torch_tpu::Materialize(
+          base_ref, torch_tpu::MaterializationReason::kExplicitSync);
+      !status.ok()) {
+    throw std::runtime_error("Failed to materialize base device buffer: " +
+                             std::string(status.message()));
+  }
+
   auto status_or_buf = base_ref.AwaitBuffer();
   if (!status_or_buf.ok()) {
     throw std::runtime_error("Failed to fetch PjRtBuffer from TPU reference: " +

tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc

@@ -144,9 +144,10 @@ NB_MODULE(_tpu_raiden_torch, m) {
       .def(
           "register_active_plan",
           [](KVCacheManager& self, uint64_t uuid,
-             const std::string& request_bytes, bool is_sender) {
+             const nb::bytes& request_bytes, bool is_sender) {
             tpu_raiden::rpc::StartTransferRequest request;
-            if (!request.ParseFromString(request_bytes)) {
+            if (!request.ParseFromArray(request_bytes.c_str(),
+                                        request_bytes.size())) {
               throw std::runtime_error(
                   "KVCacheManager register_active_plan failed: invalid "
                   "StartTransferRequest bytes");

tpu_raiden/kv_cache/kv_cache_manager_base.cc

@@ -155,6 +155,29 @@ KVCacheManagerBase::KVCacheManagerBase(
     // may differ (e.g. mamba conv_state bf16 vs ssm_state f32).
     device_info.physical_size =
         layer_buffers[layer_idx][0]->GetOnDeviceSizeInBytes().value();
+
+    // Host allocation and host-side block offsets are both derived from
+    // layer 0's slice (bytes_per_block()), and per-block device offsets
+    // require a contiguous per-row slice. A layer whose layout breaks
+    // either assumption corrupts host data or over-allocates by the full
+    // buffer per block, so reject at registration.
+    absl::Status addressable = raiden::ValidateMajorSliceAddressable(
+        layer_buffers[layer_idx][0]);
+    if (!addressable.ok()) {
+      throw std::runtime_error(absl::StrCat(
+          "KVCacheManager: layer ", layer_idx,
+          " is not per-block DMA-addressable: ",
+          addressable.ToString()));
+    }
+    int64_t layer_slice =
+        raiden::GetMajorSliceByteSize(layer_buffers[layer_idx][0]);
+    if (layer_slice != static_cast<int64_t>(bytes_per_block())) {
+      throw std::runtime_error(absl::StrCat(
+          "KVCacheManager: layer ", layer_idx, " slice byte size ",
+          layer_slice, " differs from layer 0's ", bytes_per_block(),
+          "; register layers with distinct slice sizes in separate "
+          "managers"));
+    }
     max_physical_size_ =
         std::max(max_physical_size_, device_info.physical_size);
     VLOG(1) << "KVCacheManagerBase: layer " << layer_idx << " on_device_shape: "