PyTorch Softmax Backward — How an Offset Beyond 2**32 Wrapped Back to the Start of the Tensor

PyTorch Spatial Softmax, Part 2 of 3 — Following the uint32_t wraparound with the exact regression-test shape

Part 1 examined an indexing asymmetry in PyTorch’s spatial softmax implementation.

The forward path already switched to 64-bit indexing when a tensor became too large for the 32-bit fast path.

One backward path, however, continued to calculate sizes, strides, and final element offsets with uint32_t.

The defect became observable when the tensor crossed 2**32 elements.

An unsigned 32-bit integer cannot represent a larger element offset.

Instead of raising an error, the arithmetic wraps back to zero.

This article uses the exact shape added by the regression test to answer one specific question:

How could a gradient near the end of the tensor resolve to an address near the beginning?

The shape chosen by the regression test

The new test creates tensors with the following shape:

inner_size = 2147483649

out = torch.empty(
    [1, 2, inner_size],
    device="cuda",
    dtype=torch.float16,
)

The value of inner_size is:

2,147,483,649
=
2**31 + 1

The full tensor shape is therefore:

[1, 2, 2**31 + 1]

Softmax backward is applied along dimension 1, the middle dimension.

From the spatial softmax kernel’s perspective:

outer_size = 1
dim_size   = 2
inner_size = 2**31 + 1

The total number of elements is:

1 × 2 × (2**31 + 1)
=
2**32 + 2
=
4,294,967,298

That is exactly two elements beyond 2**32.

The tensor can be viewed as two extremely long rows

Flattening the tensor conceptually gives us two rows, each containing inner_size elements.

Row 0:
offset 0
offset 1
offset 2
...
offset inner_size - 1

Row 1:
offset inner_size
offset inner_size + 1
...
offset 2 × inner_size - 1

Each column contains the two values belonging to one softmax group.

For a column index i:

Row 0 element offset = i

Row 1 element offset = i + inner_size

At the first column:

i = 0

Row 0 offset = 0

Row 1 offset =
    inner_size
    =
    2,147,483,649

Both offsets still fit inside uint32_t.

The failure appears only near the end of the tensor.

The address calculation used by spatial backward

Before the fix, the spatial softmax backward kernel used calculations equivalent to:

const uint32_t outer_stride =
    inner_size * dim_size;

const uint32_t dim_stride =
    inner_size;

const uint32_t data_offset =
    outer_index * outer_stride
    + inner_index;

const uint32_t final_offset =
    data_offset
    + d * dim_stride;

Because the regression test uses:

outer_size = 1

the only valid outer_index is 0.

The relevant calculation therefore reduces to:

data_offset = inner_index

and:

final_offset =
    inner_index
    + d × inner_size

When d = 0, the kernel accesses the first row:

final_offset = inner_index

When d = 1, it accesses the second row:

final_offset =
    inner_index + inner_size

Mathematically, this formula is correct.

The problem was that its result was stored in uint32_t.

The last value representable by `uint32_t`

An unsigned 32-bit integer can represent:

0
through
4,294,967,295

or:

0
through
2**32 - 1

The next value does not become a wider integer:

2**32

Within uint32_t, it becomes:

2**32 mod 2**32
=
0

The value after that becomes 1:

(2**32 + 1) mod 2**32
=
1

Unsigned wraparound in C++ is defined modulo 2**32.

It does not automatically raise an exception.

Everything remains correct until just before the boundary

Let:

N = inner_size
  = 2**31 + 1

The second-row element for column i is located at:

i + N

Consider the third-to-last column:

i = N - 3
  = 2**31 - 2

Its second-row offset is:

i + N

=
(2**31 - 2)
+
(2**31 + 1)

=
2**32 - 1

=
4,294,967,295

This is exactly the largest value that uint32_t can represent.

The address is still correct.

Third-to-last column
→ last safely representable uint32_t offset

The second-to-last column wraps to offset 0

Now move one column forward:

i = N - 2
  = 2**31 - 1

The correct second-row offset is:

i + N

=
(2**31 - 1)
+
(2**31 + 1)

=
2**32

But in uint32_t:

2**32
→ 0

The intended element is near the end of the tensor.

The calculated offset points to its very first element.

Correct logical offset:
4,294,967,296

Offset represented by uint32_t:
0

The second-row element in the second-to-last column now aliases the first element of the first row.

The final column wraps to offset 1

The final column index is:

i = N - 1
  = 2**31

Its correct second-row offset is:

i + N

=
2**31
+
(2**31 + 1)

=
2**32 + 1

=
4,294,967,297

Stored in uint32_t, it becomes:

2**32 + 1
→ 1

So the final second-row element resolves to the second element of the entire tensor.

Correct logical offset:
4,294,967,297

Offset represented by uint32_t:
1

The final two elements of the second row therefore wrap as follows:

Second-to-last element of Row 1
→ tensor offset 0

Last element of Row 1
→ tensor offset 1

The aliasing relationship becomes explicit

The tensor contains logically distinct coordinates.

But under the wrapped 32-bit offsets, some of those coordinates resolve to the same physical element location.

Coordinate:
[0, 0, 0]

Logical offset:
0

Coordinate:
[0, 1, N - 2]

Correct logical offset:
2**32

Wrapped uint32_t offset:
0

Those two different coordinates now alias the same element.

The next pair behaves similarly:

Coordinate:
[0, 0, 1]

Logical offset:
1

Coordinate:
[0, 1, N - 1]

Correct logical offset:
2**32 + 1

Wrapped uint32_t offset:
1

The original one-to-one mapping between tensor coordinates and element offsets has been broken.

The problem affects reads as well as writes

The backward kernel does not use the calculated offset only when writing gradInput.

It uses the same expression when reading the incoming gradient and the softmax output.

The relevant accesses have the form:

gradOutput[
    data_offset + d * dim_stride
]

output[
    data_offset + d * dim_stride
]

gradInput[
    data_offset + d * dim_stride
]

Once the offset wraps, three things can go wrong:

The kernel reads gradOutput from the wrong element.

The kernel reads the softmax output from the wrong element.

The kernel writes the result into the wrong gradInput element.

The final column may therefore be calculated using values taken from the beginning of the tensor and then written back into another low-address location.

This is broader than a single incorrect destination write.

The inputs to the gradient calculation can also come from the wrong coordinates.

Different threads can converge on the same address

A thread processing the first column legitimately accesses offset 0.

A different thread processing a column near the end can wrap around and also access offset 0.

Thread handling an early column
→ writes to offset 0

Thread handling a late column
→ wrapped offset also becomes 0

Logically independent columns can therefore collide on the same memory location.

The final value can become dependent on execution ordering.

GPU threads do not execute under a simple globally sequential order, so this collision may create behavior that is more complicated than a single deterministic replacement.

The central failure is already established before considering the exact schedule:

Two distinct logical coordinates have been mapped to the same element address.

`outer_stride` had already overflowed as well

For the regression-test shape:

outer_stride
=
inner_size × dim_size

Substituting the actual values:

(2**31 + 1) × 2

=
2**32 + 2

Stored in uint32_t, this becomes:

2**32 + 2
→ 2

So:

Correct outer_stride:
4,294,967,298

uint32_t outer_stride:
2

In this specific test:

outer_size = 1

and therefore:

outer_index = 0

The incorrect outer_stride does not directly affect the final offset because:

0 × outer_stride = 0

But if the same kernel processed multiple outer groups, the distance between those groups could also be miscalculated.

The regression test deliberately removes that additional variable and exposes the defect through the d × dim_stride path alone.

The `2**32` boundary refers to elements, not bytes

The kernel uses the calculated value as an element offset:

gradInput[final_offset]

Because gradInput is a typed pointer, C++ later converts the element index into a byte displacement.

For FP16:

byte displacement
=
final_offset × 2 bytes

But the wraparound has already occurred before that multiplication.

The intended element offset:

2**32 + 1

first becomes:

The pointer access is then:

gradInput[1]

It is not:

gradInput[2**32 + 1]

The 64-bit device pointer cannot recover the original address because the element index supplied to it has already been corrupted.

The values used by the regression test

The test fills every column with the same softmax output:

y0 = 0.25
y1 = 0.75

The incoming gradient is:

g0 = 1.0
g1 = 0.0

Softmax backward can be written as:

dx_i
=
y_i × (
    g_i - Σ(y_j × g_j)
)

The weighted sum is:

Σ(y_j × g_j)

=
0.25 × 1
+
0.75 × 0

=
0.25

For the first element:

dx0

=
0.25 × (1 - 0.25)

=
0.25 × 0.75

=
0.1875

For the second:

dx1

=
0.75 × (0 - 0.25)

=
-0.1875

Every column should therefore produce:

[+0.1875, -0.1875]

Why choose `0.25`, `0.75`, and `0.1875`?

All three values are exactly representable in binary floating-point.

0.25
=
1 / 4

0.75
=
3 / 4

0.1875
=
3 / 16

Their denominators are powers of two.

That makes the expected values exactly representable in FP16 as well.

If the test fails, the failure is difficult to dismiss as ordinary floating-point tolerance noise.

Expected-value rounding ambiguity
→ minimized

Primary variable under test
→ element-address calculation

The numerical setup isolates the indexing boundary.

Why check both the first and final columns?

The test verifies values near the beginning of the tensor:

self.assertEqual(
    gI[0, 0, 0],
    0.1875,
)

self.assertEqual(
    gI[0, 1, 0],
    -0.1875,
)

It also verifies the final column:

self.assertEqual(
    gI[0, 0, -1],
    0.1875,
)

self.assertEqual(
    gI[0, 1, -1],
    -0.1875,
)

The first column uses low element offsets that fit easily within 32 bits.

The final element of the second row has the logical offset:

2**32 + 1

The same mathematical values are repeated in every column.

The gradient should therefore be identical at both ends.

Beginning of tensor:
+0.1875 / -0.1875

Beyond the old 32-bit offset boundary:
+0.1875 / -0.1875

A difference points directly toward address-range handling rather than a difference in the softmax data.

Why does the test require roughly 30GB of GPU memory?

Each tensor contains:

2**32 + 2 FP16 elements

Each FP16 element occupies 2 bytes.

The raw size of one tensor is therefore approximately:

(2**32 + 2) × 2 bytes
≈ 8 GiB

The test needs at least three large tensors:

softmax output

incoming gradient

result gradient

The raw tensor storage alone is therefore approximately:

8 GiB × 3
≈ 24 GiB

Additional memory is needed for:

the CUDA context
the caching allocator
kernel execution
temporary runtime state
test-environment overhead

The test is consequently marked as:

@largeTensorTest("30GB", "cuda")

A small synthetic calculation can demonstrate the modular arithmetic.

But exercising the actual GPU pointer path beyond 2**32 element positions requires a genuinely enormous allocation.

The new implementation does not wait for unsigned wraparound

The old uint32_t implementation visibly wraps at:

2**32

The new dispatch switches to 64-bit indexing at the lower signed boundary:

INT_MAX
=
2**31 - 1
=
2,147,483,647

Its decision is based on logic equivalent to:

canUse32BitIndexMath(
    grad,
    INT_MAX
)

Conceptually:

Safe under signed 32-bit indexing
→ use the 32-bit variant

Not safe
→ use the 64-bit variant

The regression test uses:

inner_size
=
2**31 + 1

which already exceeds INT_MAX.

Under the repaired code, the 64-bit index_t kernel is selected before the old unsigned arithmetic can approach its wraparound boundary.

Widening only the final variable would not be enough

The final address is produced through several intermediate calculations:

inner_size
→ dim_stride

inner_size × dim_size
→ outer_stride

outer_index × outer_stride
→ outer_offset

outer_offset + inner_index
→ data_offset

data_offset + d × dim_stride
→ final_offset

If any intermediate step still uses 32-bit arithmetic, the value can be corrupted before it reaches a 64-bit destination.

For example:

uint32_t wrapped =
    data_offset + d * dim_stride;

int64_t final_offset =
    wrapped;

The 64-bit variable does not restore the original value.

2**32 + 1
→ wraps to 1 in uint32_t
→ widened to int64_t as 1

The wider type preserves the wrong result perfectly.

That is why the patch changes the full chain:

kernel size arguments
strides
loop counters
intermediate offsets
final element offsets

They now share the same index_t contract.

A valid address is not necessarily the correct address

The wrapped offsets 0 and 1 are still inside the allocated tensor.

Address is accessible
✓

But they do not correspond to the intended logical coordinates.

Address is logically correct
✗

A bounds check may therefore fail to detect the problem.

Did the access leave the allocation?
→ No

Did the access reach the intended tensor element?
→ No

Memory safety and numerical correctness are related, but they are not identical properties.

This defect could remain within allocated memory while still corrupting gradient semantics.

The large input did not create the bug

The address formula was the same for smaller tensors:

data_offset
+
d × dim_stride

What changed was whether the result still fit in the chosen integer type.

Small tensor
→ the old type happens to represent every offset

Large tensor
→ the pre-existing width mismatch becomes observable

The large input was not the source of the defect.

It crossed the boundary of an implementation that had always been limited to 32-bit unsigned offsets.

Just as kElementsPerAccess = 1 hid the CUTLASS asymmetry, ordinary tensor sizes hid the PyTorch indexing asymmetry.

A boundary test is more than a very large number

A useful boundary test removes unrelated variables.

This test fixes:

outer_size = 1
→ removes outer-group movement

dim_size = 2
→ minimizes the softmax reduction

same values in every column
→ removes position-dependent mathematics

exactly representable FP16 values
→ removes tolerance ambiguity

final element beyond 2**32
→ directly exercises address width

The shape is enormous, but the logical computation is deliberately simple.

If the final result differs from the first result, the possible causes have already been narrowed substantially.

Complex softmax distribution?
→ No

Different values at different positions?
→ No

FP16 approximation noise?
→ No

Crossing the indexing boundary?
→ Directly testable

The allocation is large.

The counterexample itself is structurally minimal.

The core calculation

The final address in the old backward path can be reduced to:

last inner index
+
stride to the second softmax value

Substituting the test values:

2**31
+
(2**31 + 1)

gives:

2**32 + 1

That is the correct logical element offset.

But in uint32_t:

2**32 + 1
→ 1

The result is:

Gradient for the final tensor element
→ address of the tensor’s second element

The defect was not that the formula produced the wrong mathematical number.

The defect was that the selected integer type could not represent the number the formula produced.

An integer type is therefore not merely an implementation detail.

It declares the largest address space the kernel can express.

Address width must be shared across the entire operation path

A 64-bit tensor size at the host level does not make the kernel safe automatically.

A 64-bit device pointer does not help if the element offset has already wrapped.

A correct forward kernel does not make a separate backward kernel correct.

The full address path must agree:

shape
→ stride
→ loop index
→ intermediate offset
→ final element index
→ pointer access

PyTorch #188031 repaired the spatial backward path by reconnecting those stages through one index_t contract.

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

https://resoneticlab.blogspot.com/2026/08/pytorch-188031-why-did-forward-use-64.html

Related material

This is Part 2 of a three-part series on PyTorch spatial softmax 64-bit indexing.

Part 3 examines why an existing test named test_softmax_backward_64bit_indexing did not detect this defect: inner_size == 1 and inner_size != 1 select different CUDA kernel paths.

https://resoneticlab.blogspot.com/2026/08/pytorch-188031-why-did-existing-64-bit.html

#PyTorch #CUDA #Softmax #GPUProgramming #IntegerOverflow #64BitIndexing #GPUCorrectness #CodeAnalysis

Search This Blog

Resonetic Lab | AI · Code · Market