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

PyTorch Spatial Softmax, Part 1 of 3 — An asymmetric indexing contract and the patch that aligned it

A recent PyTorch change fixed a 64-bit indexing problem in the spatial softmax backward path.

The bug was not in the softmax formula itself.

It was in the width of the integer used to calculate memory locations.

For the same large tensor, the forward path already selected either 32-bit or 64-bit indexing based on the input size.

One backward path, however, continued to calculate sizes, strides, and offsets using uint32_t.

That created a possible split in behavior:

Forward
→ 64-bit address calculation
→ accesses the intended elements

Backward
→ 32-bit address calculation
→ offset wraparound
→ gradients read from or written to different locations

A tensor could therefore complete the forward pass correctly while the corresponding backward kernel could no longer represent every address in that tensor.

The changes described in PyTorch #188031 were later landed in PyTorch’s main branch.

Softmax does not have only one CUDA kernel

Softmax is usually explained as a mathematical sequence:

Input values
→ exponentiation
→ reduction along one dimension
→ normalization by the sum

But a GPU implementation must answer another question:

Along which dimension is softmax being computed?

Consider a tensor with the shape:

[B, C, H, W]

If softmax is applied along the final dimension, W, the values in each softmax group are usually contiguous in memory.

There are no dimensions after the softmax axis, so:

inner_size = 1

If softmax is applied along the channel dimension, C, the memory layout is different.

Values belonging to the same softmax group are separated by H × W elements.

The tensor can then be interpreted as:

outer_size = B
dim_size   = C
inner_size = H × W

PyTorch uses a separate spatial softmax path for this non-last-dimension case.

Conceptually:

inner_size == 1
→ last-dimension-oriented kernel path

inner_size != 1
→ cunn_SpatialSoftMaxForward
→ cunn_SpatialSoftMaxBackward

The indexing problem appeared in the second path: spatial softmax backward.

The address has to be calculated before the gradient can be written

Spatial softmax backward interprets the tensor through three sizes:

outer_size
dim_size
inner_size

It then derives strides and offsets to locate each element.

A simplified version of the address calculation looks like this:

outer_stride = inner_size × dim_size
dim_stride   = inner_size

data_offset =
    outer_index × outer_stride
    + inner_index

final_offset =
    data_offset
    + d × dim_stride

If those calculations are correct, the backward kernel reads the relevant softmax output and incoming gradient, computes the result, and writes it to the matching position in the gradient tensor.

The problem was the type used to store those values.

For a sufficiently large tensor, outer_stride, data_offset, and final_offset can exceed several billion.

Before the fix, the spatial backward kernel calculated them with uint32_t.

const uint32_t outer_stride =
    inner_size * dim_size;

const uint32_t dim_stride =
    inner_size;

Its loop variables and intermediate offsets also used the same 32-bit unsigned type.

uint32_t outer_index;
uint32_t inner_index;
uint32_t data_offset;
uint32_t d;

`uint32_t` returns to zero after roughly 4.29 billion

The largest value representable by uint32_t is:

4,294,967,295
=
2**32 - 1

Adding one does not produce a wider integer.

Within 32-bit unsigned arithmetic, the value wraps back to zero:

4,294,967,295 + 1
→ 0

The same rule applies to element offsets.

Suppose the correct logical position is:

4,294,967,297

Stored in uint32_t, it becomes:

4,294,967,297 mod 2**32
= 1

The kernel may believe that it is processing an element near the end of the tensor while the calculated offset points near the beginning.

The failure does not have to produce an immediate crash

It is tempting to assume that an incorrect address must immediately trigger a GPU fault.

That is not always what happens.

After wraparound, the calculated offset may still point inside the allocated tensor.

Intended location
→ near the end of the tensor

Wrapped location
→ a valid address near the beginning

In that case, the kernel can continue running.

The output tensor is still produced.

Python may receive no exception.

But a gradient can be written into another element’s location.

This is more difficult to detect than a clear launch failure:

Kernel launch succeeds
+
output exists
+
some gradients are stored at the wrong offsets

Depending on the shape and memory access, the same defect could also lead to a fault.

But failure was not guaranteed to be explicit.

Forward already changed index width according to tensor size

The surprising part is that spatial softmax as a whole was not unaware of 64-bit indexing.

The forward kernel already accepted index_t as a template parameter.

Conceptually, its definition looked like this:

template <
    typename scalar_t,
    typename accscalar_t,
    typename outscalar_t,
    typename index_t,
    ...
>
__global__ void cunn_SpatialSoftMaxForward(...);

At launch time, PyTorch checked whether the tensor could be handled safely with 32-bit index arithmetic.

It then selected the appropriate kernel variant:

Tensor fits in the 32-bit indexing contract
→ 32-bit index_t
→ common fast path

Tensor exceeds that contract
→ 64-bit index_t
→ wide address calculations

The forward path therefore already had an explicit boundary for large spatial softmax tensors.

The backward path did not share it.

Backward remained fixed to `uint32_t`

Before the patch, the backward kernel template did not include index_t.

template <
    typename scalar_t,
    typename accscalar_t,
    typename outscalar_t,
    ...
>
__global__ void cunn_SpatialSoftMaxBackward(...);

Its sizes and offsets were hard-coded to uint32_t.

The same tensor could therefore be handled under two different indexing policies:

Forward:
Does this tensor exceed the 32-bit indexing range?
→ If so, use 64-bit indexing.

Backward:
Use uint32_t regardless of tensor size.

Forward and backward compute different mathematical operations.

But they still interpret the same tensor’s memory.

If forward declares that every element in a tensor can be addressed safely, backward must also be able to represent every corresponding gradient position.

That contract had been implemented on only one side.

Why could this asymmetry remain hidden?

Most tensors are far smaller than 2**32 elements.

Even a tensor containing 100 million elements is large in many practical workloads, but it remains far below the unsigned 32-bit boundary:

100,000,000
<
4,294,967,296

Under ordinary tensor sizes, uint32_t does not wrap.

The implementation can therefore appear correct under all of the following:

small or medium tensor
+
non-last-dimension softmax
+
spatial backward execution

The asymmetry still exists.

The input simply has not crossed the range where the difference becomes observable.

There was also an existing large softmax backward test, but it exercised a different kernel path.

That test used:

inner_size == 1

The defective spatial kernel required:

inner_size != 1

The public operation was still called softmax backward, but the implementation underneath was not the same.

The fix passed `index_t` through the backward kernel

The central change was to give spatial backward the same index-width abstraction already used by forward.

Before:

template <
    typename scalar_t,
    typename accscalar_t,
    typename outscalar_t,
    ...
>

After:

template <
    typename scalar_t,
    typename accscalar_t,
    typename outscalar_t,
    typename index_t,
    ...
>

Sizes and offsets were then changed from uint32_t to index_t.

const index_t outer_stride =
    inner_size * dim_size;

const index_t dim_stride =
    inner_size;

The loop indices and intermediate address calculations were updated as well.

index_t outer_index;
index_t inner_index;
index_t data_offset;
index_t d;

Changing only the final variable would not have been enough.

If an intermediate value first overflows in 32-bit arithmetic and is then copied into a 64-bit variable, the original address has already been lost.

32-bit intermediate overflows
        ↓
incorrect value is produced
        ↓
value is widened to 64 bits
        ↓
the incorrect value is preserved

The size, stride, loop index, intermediate offset, and final pointer offset all need to share the same indexing contract.

PyTorch did not force every backward call onto 64-bit indexing

The simplest possible fix would have been to use 64-bit indexing for every spatial backward launch.

But 64-bit integer arithmetic can consume additional registers and instructions.

Most tensors do not require it.

PyTorch instead mirrored the forward dispatch and selected the index width at runtime.

The patch added a dispatch equivalent to:

AT_DISPATCH_INDEX_TYPES(
    canUse32BitIndexMath(grad, INT_MAX)
        ? ScalarType::Int
        : ScalarType::Long,
    ...
);

Conceptually:

Safe under the signed 32-bit indexing contract
→ compile and launch the int32 variant

Not safe
→ compile and launch the int64 variant

The common case keeps the lighter 32-bit path.

Only large tensors use the wider index type.

Make every input pay for 64-bit indexing
✗

Switch to 64-bit indexing only at the required boundary
✓

Why switch at `INT_MAX` rather than waiting for `2**32`?

The direct wraparound boundary of the old uint32_t implementation was:

2**32

But the new dispatch uses INT_MAX:

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

This means PyTorch moves to the 64-bit path before unsigned wraparound is reached.

The reason is that the shared 32-bit indexing fast path is based on signed int.

Safe signed 32-bit range
→ 0 through 2**31 - 1

Larger indexing space required
→ use a 64-bit index

The old backward implementation happened to use uint32_t.

The repaired dispatch follows the same signed 32-bit safety boundary already used elsewhere.

Forward and backward once again agree on when the tensor requires a wider address type.

A 64-bit kernel can also change the launch constraints

The patch did more than replace the index type.

A 64-bit kernel variant may use more registers than its 32-bit equivalent.

Higher register usage can reduce:

the maximum practical block size
the number of resident blocks
occupancy on a given device

A block configuration that works for the 32-bit kernel may not be launchable for the 64-bit instantiation.

The patch therefore strengthened the occupancy-selection logic as well.

Conceptually:

Test the current block configuration
        ↓
Can at least one block reside on the device?
        ↓
Yes → use it
        ↓
No → reduce the block dimensions
        ↓
test again

A comment in the patch notes that gfx950 can report an oversized configuration as:

cudaSuccess
+
0 resident blocks

So checking the error code alone is not enough.

The launch helper also verifies that the number of active blocks is greater than zero.

This launch-boundary change becomes especially important once the kernel has a 64-bit variant with different register pressure.

The new regression test crosses the real boundary

The patch added a dedicated large-tensor test for spatial softmax backward.

Its key shape is:

inner_size = 2,147,483,649
dim_size   = 2

The total number of elements is:

2 × 2,147,483,649
=
4,294,967,298

That is two elements larger than:

2**32
=
4,294,967,296

The final part of the tensor therefore lies beyond the old uint32_t element-offset range.

The test fixes the two softmax output values at:

0.25
0.75

The incoming gradients are:

1
0

Under those conditions, the expected backward results are exactly:

+0.1875
-0.1875

The test checks those values both near the beginning of the tensor and at its final column.

Low-address column
→ +0.1875 / -0.1875

Column beyond the old 32-bit offset boundary
→ +0.1875 / -0.1875

The old implementation could wrap the final locations back toward the beginning of the tensor.

The new test was constructed specifically to make that boundary observable.

The public validation scope still needs to be stated carefully

The source change landed in PyTorch main, and the repository now contains a large-tensor regression test for the spatial path.

However, the author of #188031 stated that the test required roughly 30GB of GPU memory and could not be exercised locally on the hardware available to them.

They verified the expected +0.1875 and -0.1875 values with a smaller CPU example.

The indexing change itself also mirrors the already-shipping forward dispatch.

The public evidence therefore supports the following:

Forward/backward indexing asymmetry
→ confirmed in source

Potential uint32_t wraparound
→ confirmed from the address arithmetic

Existing 64-bit forward dispatch
→ confirmed

Backward implementation updated
→ landed in main

Large boundary-focused regression test
→ added

It does not support the broader claim that the 30GB path was executed locally by the PR author on every relevant GPU environment.

A test being present in the source and a test being executed under every required environment are different facts.

The patch restored the tensor-size contract between forward and backward

At the diff level, the change can be summarized as:

Backward uint32_t offsets
→ replaced by index_t

Runtime size check
→ selects int32 or int64 kernel

But the actual contract restored by the patch is broader:

Any tensor size that forward accepts must also be representable by the corresponding backward address calculations.

The failure was not discovered inside an unusually complex softmax equation.

It became visible by comparing two corresponding execution paths.

Forward switches to 64-bit indexing when required.

Why does backward remain fixed to uint32_t?

Each path can look reasonable when read independently.

Most ordinary inputs also pass.

The asymmetry appears only when both sides are compared under the same tensor-size contract.

A successful forward pass is not evidence that backward is safe

A neural-network operation is typically executed in two phases:

Forward
→ compute outputs

Backward
→ compute gradients

When forward succeeds, it is natural to assume that the shape, allocation, and operation are valid.

But forward and backward may use:

different kernels
different indexing implementations
different launch configurations
different memory access patterns

The following statements are not equivalent:

The forward pass succeeded.

Every backward offset can be represented safely.

Nor are these:

Same public operation name

Same underlying kernel

PyTorch #188031 illustrates the distinction.

The large tensor did not create a new softmax formula bug.

It exposed an index-width asymmetry that had already been present.

The same tensor requires the same addressable range

The contradiction at the center of the issue was simple:

Forward:
This tensor can be processed.

Backward:
Some positions in the same tensor cannot be represented.

Both statements should not be true at once.

Backward performs a different mathematical computation, but it must still be able to address every gradient corresponding to the forward tensor.

The 32-bit fast path remains valuable.

But once the input exceeds its addressable range, the implementation must move to a wider contract.

Small input
→ 32-bit performance

Large input
→ 64-bit correctness

The patch did not choose performance over correctness or correctness over performance.

It defined the boundary between them explicitly.

That is the boundary repaired by PyTorch #188031.

Related material

Patch status: Landed in PyTorch main
Evidence: Public change description, source diff, and a newly added large-tensor regression test
Boundary note: The change author stated that the approximately 30GB test path could not be executed on their local GPU hardware.

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

Part 2 follows the exact regression-test shape and calculates how an offset beyond 2**32 could wrap from the end of the tensor back to positions 0 and 1.

https://resoneticlab.blogspot.com/2026/08/pytorch-softmax-backward-how-offset.html

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

Search This Blog

Resonetic Lab | AI Systems Code Review