[ROCm] Fix GPU memory fault in batched_unary_embeddings backward

.github/scripts/fbgemm_gpu_test.bash

@@ -168,7 +168,6 @@ __configure_fbgemm_gpu_test_rocm () {
 
   # https://github.com/pytorch/FBGEMM/issues/1559
   export ignored_tests=(
-    ./batched_unary_embeddings_test.py
     ./sll/triton_sll_test.py
     ./gather_scatter/gather_scatter_test.py
     ./moe/layers_test.py  # Not a python unittest file

fbgemm_gpu/src/sparse_ops/sparse_batched_unary_embeddings.cu

@@ -61,10 +61,21 @@ Tensor batched_unary_embeddings_forward_cuda(
   // N: number of tasks, T: number of tables, B: batch size
   const int32_t N = weight.size(0);
   const int32_t T = table_offsets.numel() - 1;
-  const int32_t B = (offsets.numel() - 1) / T;
   TORCH_CHECK(N > 0);
-  TORCH_CHECK(B > 0);
   TORCH_CHECK(T > 0);
+  // offsets delimits a flattened [T, B] grid of segments, so it must hold
+  // exactly T*B+1 entries. Reject inputs whose length is not T*B+1 for any
+  // integer B: otherwise the trailing partial row is silently dropped here
+  // while the backward pass walks it with an out-of-range table id.
+  TORCH_CHECK(
+      (offsets.numel() - 1) % T == 0,
+      "offsets.numel() - 1 (",
+      offsets.numel() - 1,
+      ") must be divisible by the number of tables T (",
+      T,
+      ")");
+  const int32_t B = (offsets.numel() - 1) / T;
+  TORCH_CHECK(B > 0);
   TORCH_CHECK(T <= 65535);
   TORCH_CHECK(N <= 65535);
   int32_t threads = std::min<int32_t>(B, 512);
@@ -203,6 +214,20 @@ DLL_PUBLIC Tensor batched_unary_embeddings_backward_cuda(
   TORCH_CHECK(N > 0);
   TORCH_CHECK(B > 0);
   TORCH_CHECK(T > 0);
+  // offsets delimits a flattened [T, B] grid of segments, so it must hold
+  // exactly T*B+1 entries. A shorter/longer offsets tensor makes the index
+  // linearization walk segments with an out-of-range table id, writing
+  // grad_weight out of bounds.
+  TORCH_CHECK(
+      offsets.numel() == static_cast<int64_t>(T) * B + 1,
+      "offsets.numel() (",
+      offsets.numel(),
+      ") must equal T*B+1 (",
+      static_cast<int64_t>(T) * B + 1,
+      ") for T=",
+      T,
+      ", B=",
+      B);
 
   int32_t info_B_num_bits;
   uint32_t info_B_mask;

fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp

@@ -2373,9 +2373,20 @@ Tensor batched_unary_embeddings_forward_cpu(
   // N: number of tasks, T: number of tables, B: batch size
   const int64_t N = weight.sizes()[0];
   const int64_t T = table_offsets.numel() - 1;
-  const int64_t B = (offsets.numel() - 1) / T;
   TORCH_CHECK(N > 0);
   TORCH_CHECK(T > 0);
+  // offsets delimits a flattened [T, B] grid of segments, so it must hold
+  // exactly T*B+1 entries. Reject inputs whose length is not T*B+1 for any
+  // integer B: otherwise the trailing partial row is silently dropped here
+  // while the backward pass walks it with an out-of-range table id.
+  TORCH_CHECK(
+      (offsets.numel() - 1) % T == 0,
+      "offsets.numel() - 1 (",
+      offsets.numel() - 1,
+      ") must be divisible by the number of tables T (",
+      T,
+      ")");
+  const int64_t B = (offsets.numel() - 1) / T;
   TORCH_CHECK(B > 0);
 
   // Make sure the index_t are consistent among table_offsets, offsets and

fbgemm_gpu/src/sparse_ops/sparse_ops_meta.cpp

@@ -106,7 +106,12 @@ Tensor batched_unary_embeddings_forward_meta(
     const Tensor& /* indices */) {
   at::SymInt N = weight.sym_sizes()[0];
   at::SymInt T = table_offsets.sym_numel() - 1;
+  TORCH_SYM_CHECK(T.sym_gt(0), "number of tables T must be positive");
+  TORCH_SYM_CHECK(
+      ((offsets.sym_numel() - 1) % T).sym_eq(0),
+      "offsets.sym_numel() - 1 must be divisible by the number of tables T");
   at::SymInt B = (offsets.sym_numel() - 1) / T;
+  TORCH_SYM_CHECK(B.sym_gt(0), "batch size B must be positive");
   return at::empty_symint({N, B, T}, weight.options());
 }
 

fbgemm_gpu/test/batched_unary_embeddings_test.py

@@ -216,7 +216,9 @@ def _test_main(
             hash_sizes=[71, 107],
             long_index=True,
         ).to(device)
-        offsets = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.long).to(device)
+        # T=2 tables, so offsets must hold T*B+1 entries. With B=4 that is 9
+        # offsets delimiting 8 unit-length segments over the 8 values.
+        offsets = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.long).to(device)
         values = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.long).to(device)
         for _ in range(10):
             output = unary_embedding_module(offsets, values).transpose(1, 0)
@@ -427,6 +429,34 @@ def test_batched_unary_embeddings_forward_backward_small_repro(self) -> None:
                 for r in range(1, small_hash_sizes[t]):
                     self.assertEqual(gpu_grad[n, row0_idx + r].item(), 0.0)
 
+    def test_batched_unary_embeddings_meta(self) -> None:
+        # The meta kernel must mirror the eager forward validation so that
+        # FakeTensor/export reject the same malformed offsets eager rejects
+        # (rather than floor-dividing and reporting a misleading B). Runs on
+        # the meta device, so no GPU is required.
+        T = 2  # tables -> offsets must hold T*B+1 entries
+        B = 4
+        N = 3
+        weight = torch.empty(N, sum([71, 107]), 1, device="meta")
+        table_offsets = torch.tensor([0, 71, 178], dtype=torch.long, device="meta")
+
+        # Well-formed: T*B+1 = 9 offsets -> output shape [N, B, T].
+        good_offsets = torch.empty(T * B + 1, dtype=torch.long, device="meta")
+        indices = torch.empty(T * B, dtype=torch.long, device="meta")
+        out = torch.ops.fbgemm.batched_unary_embeddings(
+            weight, table_offsets, good_offsets, indices
+        )
+        self.assertEqual(list(out.shape), [N, B, T])
+        self.assertEqual(out.device.type, "meta")
+
+        # Malformed: length T*B gives numel-1 = 7, which is not a multiple of
+        # T=2. Eager now rejects this; the meta kernel must too.
+        bad_offsets = torch.empty(T * B, dtype=torch.long, device="meta")
+        with self.assertRaises(RuntimeError):
+            torch.ops.fbgemm.batched_unary_embeddings(
+                weight, table_offsets, bad_offsets, indices
+            )
+
 
 if __name__ == "__main__":
     unittest.main()