Fix int32 stride overflow in jagged_to_padded_dense at B*L*D > INT_MAX (#5755)

Summary:
Pull Request resolved: #5755

X-link: https://github.com/facebookresearch/FBGEMM/pull/2686

Fixes a class of int32 stride-overflow bugs in fbgemm::jagged_to_padded_dense.forward (and the dense_to_jagged backward path that calls it).

The kernel jagged_dense_elementwise_dense_output_kernel_ in common.cuh declares its dense `y` and `output` parameters as PackedTensorAccessor32. PTA32 stores strides as int32, so `output[oidx][jidx][iidx]` lowers to `oidx * strides_[0]` in int32. When the dense tensor's numel exceeds INT_MAX (specifically when (B - 1) * max_L * D > INT_MAX), the multiply wraps negative, writes go outside the destination buffer, and the high-oidx rows of the at::empty_symint allocation stay unwritten. Because the outer-loop bound B * max_L need not itself overflow, FBGEMM_LAUNCH_KERNEL's grid sanity check does not fire - the kernel runs to completion silently, and the unwritten memory surfaces as NaN / garbage downstream.

Fix (narrow, contained to a single kernel + its launcher):
- Template walk_down_tensor_storage_tree_ on a new PosT type. The existing int callers continue to work via auto template deduction; the int64 path can now pass int64_t.
- Template jagged_dense_elementwise_dense_output_kernel_ on a new IdxT used both as the PackedTensorAccessor stride width and as the type of the loop / index variables (outer, oidx, jidx, offset, iidx). This keeps the int32 fast path register-width consistent and lets the int64 specialization avoid all int32 wrap points.
- INVOKE_KERNEL_WITH_DIM in jagged_dense_elementwise_dense_output_ now branches on whether x_values, y_reshaped, and output_reshaped all have numel < INT_MAX, dispatching either the existing PTA32+int32_t fast path (unchanged) or a new PTA64+int64_t specialization. Workloads that fit in int32 pay zero extra cost; only over-INT_MAX workloads pay the 64-bit indexing tax.

Scope intentionally limited:
- Does NOT touch the sibling jagged_dense_dense_elementwise_jagged_output_kernel_ or its INVOKE_KERNEL_WITH_DIM (in this file or in jagged_dense_dense_elementwise_add_jagged_output_forward.cu).
- Does NOT touch check_shape_and_partition_, jagged_dense_elementwise_mul_backward.cu, jagged_to_padded_dense_forward.cu, jagged_tensor_ops_autograd.cpp, or cuda_prelude.cuh.
- The kernels and call sites that ARE NOT modified retain their previous PTA32/int32 behavior unchanged. They have the same hazard at sufficiently large shapes; addressing them is out of scope for this diff and can be tracked separately if a need arises.

Context: this bug class was identified during investigation of T264042859 (NaN at pos_encoding.grad in Instagram Reels MTML training at 20k UIH). For the IG Reels prod shapes (B_RO=1280, uih_max_len=20480, pe_emb_dim=64 - numel ~ 1.68B, below INT_MAX), this fix does NOT change kernel behavior - the int32 fast path is still selected. The IG Reels prod failure was resolved separately by D99880049, which avoids the dense_to_jagged path entirely with a model-side gather refactor. This diff lands the kernel-level fix as defense-in-depth against future callers whose shapes do cross INT_MAX (e.g. pe_emb_dim scaled to 128, or any other dense_to_jagged-on-broadcast pattern at scale).

Reviewed By: q10

Differential Revision: D104247303

fbshipit-source-id: f27c91b6c17b529c8b577f3822ef85ac675cfa10

Author: AlbertDachiChen

fbgemm_gpu/src/jagged_tensor_ops/common.cuh

@@ -83,18 +83,22 @@ namespace {
  * @returns true if the flattend jagged idx points to zero'ed (masked out)
  *               portion of the jagged tensor
  */
-template <int NUM_JAGGED_DIM, typename index_t>
+// `PosT` is the position/offset type used for indexing into the jagged
+// storage tree. Templating it lets the int64 specialization of the calling
+// kernel pass int64_t while existing int callers keep their int32 behavior
+// via auto template deduction (no callsite change needed).
+template <int NUM_JAGGED_DIM, typename PosT, typename index_t>
 DEVICE_INLINE bool walk_down_tensor_storage_tree_(
-    int& offset,
-    const int flattened_jagged_idx,
+    PosT& offset,
+    const PosT flattened_jagged_idx,
     const StackArray<int64_t>& jagged_dims,
     const StackArray<index_t*>& x_offsets) {
   // compute coorindates
-  int jagged_coords[NUM_JAGGED_DIM];
-  int j_temp = flattened_jagged_idx;
+  PosT jagged_coords[NUM_JAGGED_DIM];
+  PosT j_temp = flattened_jagged_idx;
 #pragma unroll
   for (int d = NUM_JAGGED_DIM - 1; d >= 0; --d) {
-    const int jagged_size = jagged_dims.vals[d];
+    const PosT jagged_size = jagged_dims.vals[d];
     jagged_coords[d] = j_temp % jagged_size;
     j_temp /= jagged_size;
   }
@@ -103,8 +107,8 @@ DEVICE_INLINE bool walk_down_tensor_storage_tree_(
   bool is_zero = false;
 #pragma unroll
   for (int d = 0; d < NUM_JAGGED_DIM; ++d) {
-    const int begin = x_offsets.vals[d][offset];
-    const int end = x_offsets.vals[d][offset + 1];
+    const PosT begin = x_offsets.vals[d][offset];
+    const PosT end = x_offsets.vals[d][offset + 1];
     if (jagged_coords[d] >= end - begin) {
       is_zero = true;
       break;
@@ -128,34 +132,47 @@ DEVICE_INLINE bool walk_down_tensor_storage_tree_(
 // blockDim.x so the inner dense dimension should be similar to or bigger than
 // warp size.
 // We rely on compiler unrolling the compiler time constant NUM_JAGGED_DIM.
-template <int NUM_JAGGED_DIM, typename index_t, typename scalar_t, typename F>
+// `IdxT` is the index type used for the PackedTensorAccessor stride width
+// and outer-loop counter. The launcher picks int32_t for the common fast
+// path and int64_t when (B - 1) * (max_L * D) > INT_MAX, where the int32
+// PackedTensorAccessor would otherwise wrap `oidx * strides_[0]` and
+// scribble outside the destination buffer (T264042859). Loop / index
+// variables stay in IdxT so the compiler picks consistent register width
+// for the int32 fast path.
+template <
+    int NUM_JAGGED_DIM,
+    typename IdxT,
+    typename index_t,
+    typename scalar_t,
+    typename F>
 __global__
 __launch_bounds__(kMaxThreads) void jagged_dense_elementwise_dense_output_kernel_(
-    const pta::PackedTensorAccessor32<scalar_t, 2, at::RestrictPtrTraits>
+    const pta::PackedTensorAccessor<scalar_t, 2, at::RestrictPtrTraits, IdxT>
         x_values,
     StackArray<index_t*> x_offsets,
-    const pta::PackedTensorAccessor32<scalar_t, 3, at::RestrictPtrTraits> y,
-    pta::PackedTensorAccessor32<scalar_t, 3, at::RestrictPtrTraits> output,
+    const pta::PackedTensorAccessor<scalar_t, 3, at::RestrictPtrTraits, IdxT> y,
+    pta::PackedTensorAccessor<scalar_t, 3, at::RestrictPtrTraits, IdxT> output,
     StackArray<int64_t> jagged_dims,
     F f,
     const scalar_t padding_value) {
-  const int outer_dense_size = y.size(0);
-  const int jagged_folded_size = y.size(1);
-  const int inner_dense_size = y.size(2);
-
-  const auto outer_begin = blockIdx.x * blockDim.y + threadIdx.y;
-  const auto outer_stride = gridDim.x * blockDim.y;
-  for (int outer = outer_begin; outer < outer_dense_size * jagged_folded_size;
-       outer += outer_stride) {
-    const int oidx = outer / jagged_folded_size;
-    const int jidx = outer % jagged_folded_size;
-
-    int offset = oidx;
+  const IdxT outer_dense_size = y.size(0);
+  const IdxT jagged_folded_size = y.size(1);
+  const IdxT inner_dense_size = y.size(2);
+
+  const IdxT outer_begin =
+      static_cast<IdxT>(blockIdx.x) * blockDim.y + threadIdx.y;
+  const IdxT outer_stride = static_cast<IdxT>(gridDim.x) * blockDim.y;
+  const IdxT total_outer = outer_dense_size * jagged_folded_size;
+  for (IdxT outer = outer_begin; outer < total_outer; outer += outer_stride) {
+    const IdxT oidx = outer / jagged_folded_size;
+    const IdxT jidx = outer % jagged_folded_size;
+
+    IdxT offset = oidx;
     const bool is_zero = walk_down_tensor_storage_tree_<NUM_JAGGED_DIM>(
         offset, jidx, jagged_dims, x_offsets);
 
     if (is_zero) {
-      int iidx;
+      IdxT iidx;
       for (iidx = threadIdx.x; iidx * 2 + 1 < inner_dense_size;
            iidx += blockDim.x) {
         output[oidx][jidx][2 * iidx] =
@@ -168,7 +185,7 @@ __launch_bounds__(kMaxThreads) void jagged_dense_elementwise_dense_output_kernel
             f(padding_value, y[oidx][jidx][2 * iidx]);
       }
     } else {
-      int iidx;
+      IdxT iidx;
       for (iidx = threadIdx.x; iidx * 2 + 1 < inner_dense_size;
            iidx += blockDim.x) {
         output[oidx][jidx][2 * iidx] =
@@ -256,35 +273,68 @@ void jagged_dense_elementwise_dense_output_(
   const Tensor y_reshaped = y.view({y.size(0), -1, y.size(-1)});
   Tensor output_reshaped = output.view(y_reshaped.sizes());
 
-#define INVOKE_KERNEL_WITH_DIM(NUM_JAGGED_DIM)              \
-  {                                                         \
-    std::vector<Tensor> x_offsets_contig;                   \
-    x_offsets_contig.resize(num_jagged_dim);                \
-    StackArray<index_t*> x_offset_ptrs;                     \
-    x_offset_ptrs.ndim = num_jagged_dim;                    \
-    for (int d = 0; d < num_jagged_dim; ++d) {              \
-      x_offsets_contig[d] = x_offsets[d].contiguous();      \
-      x_offset_ptrs.vals[d] =                               \
-          x_offsets_contig[d].template data_ptr<index_t>(); \
-    }                                                       \
-                                                            \
-    FBGEMM_LAUNCH_KERNEL(                                   \
-        (jagged_dense_elementwise_dense_output_kernel_<     \
-            NUM_JAGGED_DIM,                                 \
-            index_t,                                        \
-            scalar_t,                                       \
-            F>),                                            \
-        blocks,                                             \
-        threads,                                            \
-        0,                                                  \
-        at::cuda::getCurrentCUDAStream(),                   \
-        PTA_B(x_values, scalar_t, 2, 32),                   \
-        x_offset_ptrs,                                      \
-        PTA_B(y_reshaped, scalar_t, 3, 32),                 \
-        PTA_B(output_reshaped, scalar_t, 3, 32),            \
-        jagged_dims_tensor,                                 \
-        f,                                                  \
-        padding_value);                                     \
+#define INVOKE_KERNEL_WITH_DIM(NUM_JAGGED_DIM)                               \
+  {                                                                          \
+    std::vector<Tensor> x_offsets_contig;                                    \
+    x_offsets_contig.resize(num_jagged_dim);                                 \
+    StackArray<index_t*> x_offset_ptrs;                                      \
+    x_offset_ptrs.ndim = num_jagged_dim;                                     \
+    for (int d = 0; d < num_jagged_dim; ++d) {                               \
+      x_offsets_contig[d] = x_offsets[d].contiguous();                       \
+      x_offset_ptrs.vals[d] =                                                \
+          x_offsets_contig[d].template data_ptr<index_t>();                  \
+    }                                                                        \
+                                                                             \
+    /* Pick int32 fast path when every touched tensor's numel fits in int32; \
+     * otherwise dispatch the int64 specialization. The int64 path defends   \
+     * against `oidx * strides_[0]` overflowing int32 inside                 \
+     * PackedTensorAccessor32, which scribbles outside the buffer for high   \
+     * oidx and surfaces as silent NaN downstream (T264042859). */           \
+    constexpr int64_t kInt32Limit =                                          \
+        static_cast<int64_t>(std::numeric_limits<int32_t>::max());           \
+    const bool use_int32_indexing = x_values.numel() < kInt32Limit &&        \
+        y_reshaped.numel() < kInt32Limit &&                                  \
+        output_reshaped.numel() < kInt32Limit;                               \
+                                                                             \
+    if (use_int32_indexing) {                                                \
+      FBGEMM_LAUNCH_KERNEL(                                                  \
+          (jagged_dense_elementwise_dense_output_kernel_<                    \
+              NUM_JAGGED_DIM,                                                \
+              int32_t,                                                       \
+              index_t,                                                       \
+              scalar_t,                                                      \
+              F>),                                                           \
+          blocks,                                                            \
+          threads,                                                           \
+          0,                                                                 \
+          at::cuda::getCurrentCUDAStream(),                                  \
+          PTA_B(x_values, scalar_t, 2, 32),                                  \
+          x_offset_ptrs,                                                     \
+          PTA_B(y_reshaped, scalar_t, 3, 32),                                \
+          PTA_B(output_reshaped, scalar_t, 3, 32),                           \
+          jagged_dims_tensor,                                                \
+          f,                                                                 \
+          padding_value);                                                    \
+    } else {                                                                 \
+      FBGEMM_LAUNCH_KERNEL(                                                  \
+          (jagged_dense_elementwise_dense_output_kernel_<                    \
+              NUM_JAGGED_DIM,                                                \
+              int64_t,                                                       \
+              index_t,                                                       \
+              scalar_t,                                                      \
+              F>),                                                           \
+          blocks,                                                            \
+          threads,                                                           \
+          0,                                                                 \
+          at::cuda::getCurrentCUDAStream(),                                  \
+          PTA_B(x_values, scalar_t, 2, 64),                                  \
+          x_offset_ptrs,                                                     \
+          PTA_B(y_reshaped, scalar_t, 3, 64),                                \
+          PTA_B(output_reshaped, scalar_t, 3, 64),                           \
+          jagged_dims_tensor,                                                \
+          f,                                                                 \
+          padding_value);                                                    \
+    }                                                                        \
   }
 
   JAGGED_TENSOR_DISPATCH_DIMS();

fbgemm_gpu/test/jagged/dense_to_jagged_test.py

@@ -19,9 +19,14 @@
 
 if open_source:
     # pyre-ignore[21]
-    from test_utils import cpu_and_maybe_gpu, gpu_unavailable, optests
+    from test_utils import cpu_and_maybe_gpu, gpu_memory_lt_gb, gpu_unavailable, optests
 else:
-    from fbgemm_gpu.test.test_utils import cpu_and_maybe_gpu, gpu_unavailable, optests
+    from fbgemm_gpu.test.test_utils import (
+        cpu_and_maybe_gpu,
+        gpu_memory_lt_gb,
+        gpu_unavailable,
+        optests,
+    )
 
 
 @optests.generate_opcheck_tests(additional_decorators=additional_decorators)
@@ -160,6 +165,90 @@ def test_dense_to_jagged_opt_large_batch(
             precompute_total_L,
         )
 
+    @optests.dontGenerateOpCheckTests("regression test, not an op-shape check")
+    @unittest.skipIf(*gpu_unavailable)
+    @unittest.skipIf(*gpu_memory_lt_gb(16))
+    def test_dense_to_jagged_backward_int32_overflow(self) -> None:
+        """Regression: int32 overflow in dense_to_jagged backward at
+        B * max_L * D > INT_MAX, even when B * max_L itself fits in int32.
+
+        dense_to_jagged.backward dispatches to fbgemm::jagged_to_padded_dense
+        forward, which fills a dense (B, max_L, D) gradient using
+        jagged_dense_elementwise_dense_output_kernel_ in common.cuh. Pre-fix,
+        that kernel uses PackedTensorAccessor32 for the dense `y` and
+        `output` parameters: stride[0] is stored as int32 and the access
+        `output[oidx][jidx][iidx]` lowers to `oidx * stride[0]` in int32.
+        For shapes where (B - 1) * (max_L * D) > INT_MAX the multiply wraps
+        negative and writes go to addresses outside the destination buffer.
+        Outer-loop bound and grid sizing are unaffected (B * max_L need not
+        overflow), so there is no FBGEMM_LAUNCH_KERNEL grid abort - the
+        kernel runs to completion silently writing wrong memory, leaving
+        ~25% of high-oidx rows of the dense gradient at the at::empty bit
+        pattern. Downstream NaN. (T264042859, observed in production at
+        Instagram Reels MTML pos_encoding training.)
+
+        Workload: B=1024, max_L=40960, D=64, bf16. B * max_L = 42M < INT_MAX
+        (so no outer-loop / grid hazard), B * max_L * D ~ 2.68B > INT_MAX
+        (so PTA32 stride math wraps for oidx >= 819, i.e. the last ~25% of
+        batches). Peak GPU memory ~ B * max_L * D * 2 bytes ~ 5 GB.
+
+        The single jagged value of 1.0 is placed in the LAST batch (at
+        oidx = B - 1 = 1023, well into the wrap zone). On unfixed code the
+        write to out[1023, 0, :] never happens; it stays at uninitialized
+        memory. After the fix routes this shape to the int64 (PTA64) kernel
+        path the write lands correctly.
+        """
+        device = torch.accelerator.current_accelerator()
+        B = 1024
+        max_L = 40960
+        D = 64
+        dtype = torch.bfloat16
+        assert B * max_L < (1 << 31), (
+            "test must keep B * max_L below INT_MAX (we are testing the PTA32 "
+            "stride-overflow path, not the outer-loop overflow path)"
+        )
+        assert (
+            B * max_L * D > (1 << 31) - 1
+        ), "test must exceed INT_MAX on B * max_L * D"
+        # PTA32 stride wraps for oidx >= ceil(INT_MAX / (max_L * D)).
+        wrap_oidx = (((1 << 31) - 1) + (max_L * D) - 1) // (max_L * D)
+        assert (
+            B - 1 >= wrap_oidx
+        ), "last batch must fall inside the PTA32 stride-overflow zone"
+
+        # Tiny jagged: only the LAST batch has length=1, all others empty.
+        # That keeps total_L = 1 (and thus the (total_L, D) values tensor
+        # microscopic) so the only large allocation is the (B, max_L, D)
+        # padded output.
+        offsets = torch.zeros(B + 1, dtype=torch.long, device=device)
+        offsets[B] = 1
+        grad_jagged = torch.ones((1, D), dtype=dtype, device=device)
+
+        # Exact code path that dense_to_jagged.backward executes.
+        out = torch.ops.fbgemm.jagged_to_padded_dense(
+            grad_jagged, [offsets], [max_L], padding_value=0.0
+        )
+
+        self.assertEqual(tuple(out.shape), (B, max_L, D))
+
+        # The strongest signal: the in-range position MUST equal the jagged
+        # value. On unfixed code oidx=B-1 falls in the PTA32 wrap zone and
+        # the write goes to a wrong address; this position stays as
+        # whatever at::empty left in memory - anything but 1.0.
+        self.assertEqual(out[B - 1, 0, 0].item(), 1.0)
+        self.assertEqual(out[B - 1, 0, D - 1].item(), 1.0)
+
+        # Spot-check padding inside the wrap zone (oidx=B-1) and outside
+        # it (oidx=0). 0.0 here means the kernel correctly wrote
+        # padding_value; an unwritten cell would surface as garbage / NaN
+        # after the autograd reduce.
+        for pos in [1, max_L // 2, max_L - 1]:
+            self.assertEqual(out[B - 1, pos, 0].item(), 0.0)
+            self.assertEqual(out[0, pos, 0].item(), 0.0)
+        # Spot-check a batch right at the wrap boundary.
+        self.assertEqual(out[wrap_oidx, 0, 0].item(), 0.0)
+        self.assertEqual(out[wrap_oidx, max_L - 1, 0].item(), 0.0)
+
     @given(
         num_jagged_dim=st.integers(1, 5),
         outer_dense_size=st.integers(0, 5),