pytorch · iuliur-meta · Jul 21, 2026
diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu
@@ -51,6 +51,15 @@ inline __device__ void embedding_inplace_update_kernel_impl(
         lxu_cache_locations,
     const int64_t i) {
   const int32_t table_idx = update_table_idx[i];
+  // Defense-in-depth backstop (callers are expected to validate up front). An
+  // out-of-range table_idx would index D_offsets / weights_tys /
+  // weights_offsets out of bounds and then scatter to a wild address.
+  // weights_offsets has one entry per table; D_offsets has num_tables + 1
+  // entries, so table_idx + 1 is in range once table_idx is.
+  const int64_t num_tables = weights_offsets.size(0);
+  CUDA_KERNEL_ASSERT(
+      table_idx >= 0 && table_idx < num_tables &&
+      "embedding_inplace_update: table_idx out of range");
   const auto row_idx = update_row_idx[i];
 
   // TODO: We don't need to compute these here.
@@ -62,17 +71,27 @@ inline __device__ void embedding_inplace_update_kernel_impl(
       nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment);
 
   const int64_t weight_offset = weights_offsets[table_idx];
+  const int64_t byte_offset = weight_offset +
+      static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx);
+  const int64_t row_end = weight_offset +
+      static_cast<int64_t>(D_bytes) * (static_cast<int64_t>(row_idx) + 1);
   uint8_t* __restrict__ weight_row;
   if (weights_placement == PlacementType::DEVICE) {
-    weight_row =
-        &dev_weights
-            [weight_offset +
-             static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx)];
+    // Bound the row's byte range against the dev_weights allocation. A tighter
+    // per-table upper bound (weights_offsets[table_idx + 1]) is NOT safe here:
+    // weights_offsets is per-placement-buffer, so under mixed DEVICE/UVM
+    // placement (kernel_1) the next entry can belong to uvm_weights and is not
+    // a valid end for this table. The exact per-table [0, num_embeddings) range
+    // is enforced host-side before dispatch; this is the allocation backstop.
+    CUDA_KERNEL_ASSERT(
+        row_idx >= 0 && byte_offset >= 0 && row_end <= dev_weights.size(0) &&
+        "embedding_inplace_update: row scatters past dev_weights");
+    weight_row = &dev_weights[byte_offset];
   } else {
-    weight_row =
-        &uvm_weights
-            [weight_offset +
-             static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx)];
+    CUDA_KERNEL_ASSERT(
+        row_idx >= 0 && byte_offset >= 0 && row_end <= uvm_weights.size(0) &&
+        "embedding_inplace_update: row scatters past uvm_weights");
+    weight_row = &uvm_weights[byte_offset];
   }
 
   // padded_row_size_in_bytes pad each row with row_alignment (16 bytes on GPUs)

diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp
@@ -6,10 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <algorithm>
-#include <cmath>
-#include <functional>
-
 #include <ATen/ATen.h>
 #include <c10/util/irange.h>
 #include <torch/library.h>
@@ -35,8 +31,20 @@ void embedding_inplace_update_cpu_kernel(
     const at::TensorAccessor<index_t, 1>& update_row_idx,
     const at::TensorAccessor<int64_t, 1>& update_offsets,
     int64_t row_alignment) {
+  const int64_t num_tables = weights_offsets.size(0);
   for (const auto n : c10::irange(update_row_idx.size(0))) {
     int32_t t = update_table_idx[n];
+    // Defense-in-depth backstop. An out-of-range table_idx would index the
+    // per-table metadata (D_offsets / weights_tys / weights_offsets) out of
+    // bounds and then scatter to a wild address. weights_offsets has one entry
+    // per table; D_offsets has num_tables + 1 entries.
+    TORCH_CHECK(
+        t >= 0 && t < num_tables,
+        "embedding_inplace_update: table index ",
+        t,
+        " out of range [0, ",
+        num_tables,
+        ")");
     auto row_idx = update_row_idx[n];
 
     SparseType weight_ty = static_cast<SparseType>(weights_tys[t]);
@@ -46,19 +54,41 @@ void embedding_inplace_update_cpu_kernel(
     const int32_t D_bytes =
         nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment);
     int64_t weight_offset = weights_offsets[t];
+    const int64_t byte_offset = weight_offset +
+        static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx);
+    const int64_t row_end = weight_offset +
+        static_cast<int64_t>(D_bytes) * (static_cast<int64_t>(row_idx) + 1);
 
     uint8_t* __restrict__ weight_row = nullptr;
     const auto placement = static_cast<PlacementType>(weights_placements[t]);
+    // Bound the row's byte range against the destination allocation. A tighter
+    // per-table bound (weights_offsets[t + 1]) is not safe: weights_offsets is
+    // per-placement-buffer, so under mixed HOST/UVM placement the next entry
+    // can belong to the other buffer. The exact per-table [0, num_embeddings)
+    // range is enforced by callers before dispatch; this is the allocation
+    // backstop.
     if (placement == PlacementType::HOST) {
-      weight_row =
-          &dev_weights
-              [weight_offset +
-               static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx)];
+      TORCH_CHECK(
+          row_idx >= 0 && byte_offset >= 0 && row_end <= dev_weights.size(0),
+          "embedding_inplace_update: row_idx ",
+          row_idx,
+          " (table ",
+          t,
+          ") scatters past dev_weights of ",
+          dev_weights.size(0),
+          " bytes");
+      weight_row = &dev_weights[byte_offset];
     } else {
-      weight_row =
-          &uvm_weights
-              [weight_offset +
-               static_cast<int64_t>(D_bytes) * static_cast<int64_t>(row_idx)];
+      TORCH_CHECK(
+          row_idx >= 0 && byte_offset >= 0 && row_end <= uvm_weights.size(0),
+          "embedding_inplace_update: row_idx ",
+          row_idx,
+          " (table ",
+          t,
+          ") scatters past uvm_weights of ",
+          uvm_weights.size(0),
+          " bytes");
+      weight_row = &uvm_weights[byte_offset];
     }
 
     int64_t update_weight_offset = update_offsets[n];

diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp
@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <ATen/ATen.h>
 #include <c10/util/irange.h>
 #include <folly/Random.h>
 #include <gtest/gtest.h>
@@ -28,7 +29,7 @@ int32_t get_D_bytes(
 }
 
 template <typename index_t>
-void test_embedding_inplace_update() {
+void test_embedding_inplace_update(bool use_cpu) {
   int T = folly::Random::rand32() % 3 + 2; // total number of tables
   std::vector<SparseType> weight_ty_list = {
       SparseType::FP32,
@@ -45,18 +46,34 @@ void test_embedding_inplace_update() {
   std::vector<index_t> update_row_idx;
   int64_t dev_weights_offset = 0;
   int64_t uvm_weights_offset = 0;
+  // dev_weights vs uvm_weights routing differs by kernel: the CUDA kernel
+  // treats PlacementType::DEVICE as dev_weights, while the CPU kernel treats
+  // PlacementType::HOST as dev_weights. Pick the matching enum per device so
+  // the harness drives both buffers correctly on either path (the old
+  // hard-coded 0/1 convention only matched the CUDA kernel).
+  const int32_t devPlacement = use_cpu
+      ? static_cast<int32_t>(PlacementType::HOST)
+      : static_cast<int32_t>(PlacementType::DEVICE);
+  const int32_t uvmPlacement = static_cast<int32_t>(PlacementType::MANAGED);
+  // Row alignment differs by kernel (16B on device, 1B on CPU). D_bytes below
+  // MUST use the same alignment the kernel will use, otherwise the harness lays
+  // out buffers/offsets differently than the kernel scatters (the old
+  // hard-coded 16 only matched the device path).
+  const int64_t row_alignment = use_cpu ? 1L : 16L;
   for (const auto i : c10::irange(T)) {
     SparseType weight_ty =
         weight_ty_list[folly::Random::rand32() % weight_ty_list.size()];
     weights_tys.push_back(uint8_t(weight_ty));
     int D = 1 << (1 + folly::Random::rand32() % 5); // table dimension
-    int32_t D_bytes = nbit::padded_row_size_in_bytes(D, weight_ty, 16);
+    int32_t D_bytes =
+        nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment);
     int total_rows = 10 + folly::Random::rand32() % 50;
     D_offsets.push_back(D_offsets.back() + D);
 
-    int32_t weights_placement = folly::Random::rand32() % 2;
+    int32_t weights_placement =
+        (folly::Random::rand32() % 2 == 0) ? devPlacement : uvmPlacement;
     weights_placements.push_back(weights_placement);
-    if (weights_placement == 0) {
+    if (weights_placement == devPlacement) {
       weights_offsets.push_back(dev_weights_offset);
       dev_weights_offset += D_bytes * total_rows;
     } else {
@@ -83,9 +100,7 @@ void test_embedding_inplace_update() {
               << ", update rows: " << update_row_idx_str;
   }
 
-  bool use_cpu = folly::Random::rand32() % 2;
   auto device = use_cpu ? at::kCPU : at::kCUDA;
-  int64_t row_alignment = use_cpu ? 1L : 16L;
 
   auto dev_weight = at::randint(
       0, 255, {dev_weights_offset}, at::device(device).dtype(at::kByte));
@@ -157,11 +172,11 @@ void test_embedding_inplace_update() {
     auto weight_placement = weights_placements[table_idx];
     auto D = D_offsets[table_idx + 1] - D_offsets[table_idx];
     SparseType ty = static_cast<SparseType>(weights_tys[table_idx]);
-    int32_t D_bytes = nbit::padded_row_size_in_bytes(D, ty, 16);
+    int32_t D_bytes = nbit::padded_row_size_in_bytes(D, ty, row_alignment);
     auto dev_weight_acc = dev_weight_cpu.const_data_ptr<uint8_t>();
     auto uvm_weight_acc = uvm_weight_cpu.const_data_ptr<uint8_t>();
     auto update_weight_acc = update_weight_cpu.const_data_ptr<uint8_t>();
-    if (weight_placement == 0) {
+    if (weight_placement == devPlacement) {
       for (const auto j : c10::irange(D_bytes)) {
         ASSERT_EQ(
             dev_weight_acc[weight_offset + D_bytes * row_idx + j],
@@ -178,10 +193,157 @@ void test_embedding_inplace_update() {
   }
 }
 
-TEST(EmbeddingInplaceUpdateTest, random_update) {
-  // TODO: Skipping tests for now because it is
-  // unreliable and crashes occasionally.  This should be fixed and re-enabled.
-  //
-  // test_embedding_inplace_update<int32_t>();
-  // test_embedding_inplace_update<int64_t>();
+// Positive coverage via the shared harness: a valid (in-range) random update
+// must scatter correctly and must NOT trip the new bounds asserts. The harness
+// randomly assigns per-table DEVICE/UVM placement, so this also exercises the
+// mixed-placement path that the allocation-size row bound must not false-abort.
+TEST(EmbeddingInplaceUpdateTest, CpuRandomValidUpdate) {
+  test_embedding_inplace_update<int32_t>(/*use_cpu=*/true);
+  test_embedding_inplace_update<int64_t>(/*use_cpu=*/true);
+}
+
+// Real device-kernel (.cu) coverage: proves the CUDA_KERNEL_ASSERT guards do
+// not false-abort valid callers. Skipped when no GPU is present.
+//
+// There is intentionally NO device NEGATIVE (assert-fires) test: on ROCm
+// CUDA_KERNEL_ASSERT compiles to a bare abort(), which SIGABRTs the whole
+// process and poisons the GPU/HIP context; gtest death tests fork and CUDA/HIP
+// is not fork-safe, so an out-of-range device index cannot be reliably caught
+// in-process (it would need the DSA framework / TORCH_USE_CUDA_DSA). The
+// out-of-range negative path is instead covered at the CPU-kernel layer
+// (Cpu*OutOfRange* below, TORCH_CHECK -> catchable) and the host layer
+// (parent diff D112649885).
+TEST(EmbeddingInplaceUpdateTest, GpuRandomValidUpdate) {
+  if (!at::hasCUDA()) {
+    GTEST_SKIP() << "No CUDA/HIP device available";
+  }
+  test_embedding_inplace_update<int32_t>(/*use_cpu=*/false);
+  test_embedding_inplace_update<int64_t>(/*use_cpu=*/false);
+}
+
+// Tests for the defense-in-depth bounds asserts added to the inplace-update
+// kernels. Only the CPU kernel is exercised here: it uses TORCH_CHECK, which
+// throws a catchable c10::Error. The CUDA kernel uses CUDA_KERNEL_ASSERT, which
+// on ROCm compiles to a bare abort() that crashes the whole process, so the
+// device out-of-range case cannot be a catchable/death gtest in-process; it is
+// compile-verified only.
+
+namespace {
+
+// Apply one delta update to (tableIdx, rowIdx) of a single HOST-placement FP32
+// table with `totalRows` rows of dimension `D`, via the CPU kernel. Throws (via
+// TORCH_CHECK) if a bounds assert fires.
+void runSingleTableCpuUpdate(
+    int64_t totalRows,
+    int32_t D,
+    int32_t tableIdx,
+    int64_t rowIdx) {
+  const SparseType weight_ty = SparseType::FP32;
+  const int64_t row_alignment = 1;
+  const int32_t D_bytes =
+      nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment);
+
+  auto dev_weights =
+      at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte));
+  auto uvm_weights = at::zeros({0}, at::device(at::kCPU).dtype(at::kByte));
+  // CPU kernel routes PlacementType::HOST to dev_weights.
+  auto weights_placements = at::tensor(
+      {static_cast<int32_t>(PlacementType::HOST)}, at::dtype(at::kInt));
+  auto weights_offsets = at::tensor({int64_t(0)}, at::dtype(at::kLong));
+  auto weights_tys =
+      at::tensor({static_cast<uint8_t>(weight_ty)}, at::dtype(at::kByte));
+  auto D_offsets = at::tensor({0, D}, at::dtype(at::kInt));
+  auto update_weights =
+      at::zeros({D_bytes}, at::device(at::kCPU).dtype(at::kByte));
+  auto update_table_idx = at::tensor({tableIdx}, at::dtype(at::kInt));
+  auto update_row_idx = at::tensor({rowIdx}, at::dtype(at::kLong));
+  auto update_offsets =
+      at::tensor({int64_t(0), int64_t(D_bytes)}, at::dtype(at::kLong));
+
+  embedding_inplace_update_cpu(
+      dev_weights,
+      uvm_weights,
+      weights_placements,
+      weights_offsets,
+      weights_tys,
+      D_offsets,
+      update_weights,
+      update_table_idx,
+      update_row_idx,
+      update_offsets,
+      row_alignment);
+}
+
+} // namespace
+
+TEST(EmbeddingInplaceUpdateTest, CpuRejectsOutOfRangeTableIdx) {
+  // Single table (num_tables == 1); table index 5 is out of range and would
+  // index per-table metadata out of bounds. Must throw, not read wild memory.
+  EXPECT_ANY_THROW(runSingleTableCpuUpdate(
+      /*totalRows=*/8, /*D=*/4, /*tableIdx=*/5, /*rowIdx=*/0));
+}
+
+TEST(EmbeddingInplaceUpdateTest, CpuRejectsOutOfRangeRowIdx) {
+  // Row 100 in an 8-row table scatters past dev_weights. Must throw.
+  EXPECT_ANY_THROW(runSingleTableCpuUpdate(
+      /*totalRows=*/8, /*D=*/4, /*tableIdx=*/0, /*rowIdx=*/100));
+}
+
+TEST(EmbeddingInplaceUpdateTest, CpuValidUpdateSucceeds) {
+  // In-range update must not false-abort.
+  EXPECT_NO_THROW(runSingleTableCpuUpdate(
+      /*totalRows=*/8, /*D=*/4, /*tableIdx=*/0, /*rowIdx=*/5));
+}
+
+TEST(EmbeddingInplaceUpdateTest, CpuMixedPlacementValidUpdateSucceeds) {
+  // Two tables under MIXED placement: table 0 HOST (dev_weights), table 1
+  // non-HOST (uvm_weights). weights_offsets is per-placement-buffer, so both
+  // start at offset 0 in their own buffer. A valid in-range update to each
+  // table must NOT abort. This locks in the allocation-size row bound: a
+  // per-table weights_offsets[table_idx + 1] bound would treat table 1's
+  // offset (0, in the uvm buffer) as table 0's end and false-abort this
+  // valid caller.
+  const SparseType weight_ty = SparseType::FP32;
+  const int64_t row_alignment = 1;
+  const int32_t D = 4;
+  const int32_t D_bytes =
+      nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment);
+  const int64_t totalRows = 8;
+
+  auto dev_weights =
+      at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte));
+  auto uvm_weights =
+      at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte));
+  // table 0 -> dev_weights (HOST), table 1 -> uvm_weights (non-HOST).
+  auto weights_placements = at::tensor(
+      {static_cast<int32_t>(PlacementType::HOST),
+       static_cast<int32_t>(PlacementType::DEVICE)},
+      at::dtype(at::kInt));
+  auto weights_offsets =
+      at::tensor({int64_t(0), int64_t(0)}, at::dtype(at::kLong));
+  auto weights_tys = at::tensor(
+      {static_cast<uint8_t>(weight_ty), static_cast<uint8_t>(weight_ty)},
+      at::dtype(at::kByte));
+  auto D_offsets = at::tensor({0, D, 2 * D}, at::dtype(at::kInt));
+  auto update_weights =
+      at::zeros({2 * D_bytes}, at::device(at::kCPU).dtype(at::kByte));
+  auto update_table_idx = at::tensor({0, 1}, at::dtype(at::kInt));
+  auto update_row_idx =
+      at::tensor({int64_t(5), int64_t(6)}, at::dtype(at::kLong));
+  auto update_offsets = at::tensor(
+      {int64_t(0), int64_t(D_bytes), int64_t(2 * D_bytes)},
+      at::dtype(at::kLong));
+
+  EXPECT_NO_THROW(embedding_inplace_update_cpu(
+      dev_weights,
+      uvm_weights,
+      weights_placements,
+      weights_offsets,
+      weights_tys,
+      D_offsets,
+      update_weights,
+      update_table_idx,
+      update_row_idx,
+      update_offsets,
+      row_alignment));
 }