PyTorch #188031 — Why Did the Existing 64-Bit Test Miss Spatial Softmax Backward?

PyTorch Spatial Softmax, Part 3 of 3 — The test boundary between inner_size == 1 and inner_size != 1, and between INT_MAX and UINT32_MAX

PyTorch already had a test named:

test_softmax_backward_64bit_indexing

At first glance, that sounds as though 64-bit indexing in softmax backward had already been covered.

The tensor used by the test was certainly large:

roughly 20GB of GPU memory
more than 2,147,483,650 FP16 elements
an input beyond INT_MAX

Yet the spatial softmax backward path still contained address calculations fixed to uint32_t.

How could a large-tensor 64-bit test already exist while this path remained untested?

The answer is not that the existing tensor was too small.

The existing test never entered the kernel that contained the defect.

The same API call does not imply the same CUDA kernel

Both the existing test and the new regression test call the same internal API:

torch._softmax_backward_data(...)

But PyTorch examines the tensor shape and the softmax dimension before selecting a kernel.

It first reduces the tensor logically into three regions:

outer_size
dim_size
inner_size

Suppose the tensor has the shape:

[A, B, C, D]

and softmax is applied along dimension B.

Then:

outer_size = A
dim_size   = B
inner_size = C × D

The key branch in the backward implementation is:

if (inner_size == 1) {
    // last-dimension-oriented softmax backward
} else {
    // spatial softmax backward
}

The selected implementation therefore depends on more than the total number of elements.

It depends on how many elements remain after the softmax dimension.

The shape used by the existing test

The existing 64-bit backward test creates a tensor like this:

x = torch.ones(
    [1, 1, numel],
    device=device,
    dtype=torch.float16,
)

It applies softmax backward along dimension 2:

out = torch._softmax_backward_data(
    x,
    x,
    2,
    x.dtype,
)

For:

shape = [1, 1, numel]
softmax dimension = 2

the logical decomposition is:

outer_size
=
1 × 1
=
1

dim_size
=
numel

There are no dimensions after dimension 2, so:

inner_size
=
1

The final values are:

outer_size = 1
dim_size   = numel
inner_size = 1

That sends execution into the first branch:

inner_size == 1
→ last-dimension softmax backward path

For a large dim_size, the call proceeds through the general host dispatch and eventually reaches the ordinary softmax backward kernel family.

Conceptually:

test_softmax_backward_64bit_indexing
        ↓
inner_size == 1
        ↓
dispatch_host_softmax_backward()
        ↓
cunn_SoftMaxBackward

The defective kernel was not on this path.

The spatial path required a different shape and axis

Spatial softmax backward is selected when:

inner_size != 1

That happens when softmax is applied to a non-final dimension and one or more dimensions remain after it.

The new regression test uses:

inner_size = 2147483649

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

It applies softmax backward along dimension 1:

gI = torch._softmax_backward_data(
    grad,
    out,
    1,
    out.dtype,
)

For:

shape = [1, 2, inner_size]
softmax dimension = 1

the decomposition becomes:

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

Now:

inner_size != 1

so PyTorch enters the other branch:

spatial softmax backward
→ cunn_SpatialSoftMaxBackward

Both tests invoke the same Python-facing operation.

They do not execute the same CUDA kernel.

The two tests side by side

Existing test

Shape:
[1, 1, N]

Softmax dimension:
2

Decomposition:
outer = 1
dim   = N
inner = 1

Kernel path:
cunn_SoftMaxBackward

New regression test

Shape:
[1, 2, N]

Softmax dimension:
1

Decomposition:
outer = 1
dim   = 2
inner = N

Kernel path:
cunn_SpatialSoftMaxBackward

At the API level, both are enormous softmax backward operations.

At the implementation level, they are almost opposite workloads.

Existing test:
one extremely long softmax group

New test:
an enormous number of two-element softmax groups

The spatial kernel exists specifically for the second memory pattern.

Why PyTorch keeps separate kernel paths

When softmax operates along the last dimension, the values in each softmax group are laid out contiguously:

[..., x0, x1, x2, x3]

The kernel can read adjacent values while reducing one long group.

When softmax operates along a non-final dimension, the values belonging to one group are separated by a stride.

For the shape:

[1, 2, N]

with softmax along dimension 1, each column contains two values at:

Column i:

first value  → offset i
second value → offset i + N

Each reduction contains only two elements.

But those elements can be separated by billions of positions.

The spatial implementation treats the tensor as:

outer × dim × inner

and uses a two-dimensional block and grid to parallelize:

the outer dimension

the inner dimension

the reduction across dim

The separate path is therefore not incidental duplication.

The memory layout and parallelization strategy are fundamentally different.

The existing test did cross `INT_MAX`

The first large value used by the existing test is:

2,147,483,650

The largest signed 32-bit integer is:

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

So the test does cross the signed 32-bit boundary:

numel
=
INT_MAX + 3

That makes the existing test meaningful.

It verifies that the last-dimension backward path can operate beyond the signed 32-bit indexing range.

The test was not pointless or incorrectly constructed.

It covered the boundary of the implementation it actually reached.

But the spatial kernel used `uint32_t`

The defective spatial backward kernel did not use signed int32_t.

It used uint32_t.

Those types have different limits:

Maximum signed int32 value:
2**31 - 1
=
2,147,483,647

Maximum uint32_t value:
2**32 - 1
=
4,294,967,295

The element count in the existing test exceeded INT_MAX, but it remained well below UINT32_MAX:

2,147,483,650
<
4,294,967,295

Even if the same number had somehow been passed through the spatial kernel, it would not yet have crossed the direct unsigned wraparound boundary.

To expose the actual spatial defect, two conditions were required:

1. Execution must reach cunn_SpatialSoftMaxBackward.

2. A calculated element offset must exceed 2**32 - 1.

The existing test satisfied neither condition.

The new test crosses both boundaries

The new test chooses:

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

and:

dim_size = 2

The total number of elements is:

2 × (2**31 + 1)

=
2**32 + 2

=
4,294,967,298

That is three values beyond UINT32_MAX:

UINT32_MAX
=
4,294,967,295

The test therefore exercises both relevant boundaries:

inner_size > INT_MAX
→ repaired dispatch must choose 64-bit index_t

total spatial offset > UINT32_MAX
→ old uint32_t code would actually wrap

The existing test verified a signed 32-bit transition in one kernel family.

The new test reaches the unsigned 32-bit wraparound boundary in another.

One large tensor does not cover every large-tensor path

This case demonstrates an important distinction:

Large input
≠
coverage of every large-input implementation

No matter how large the tensor becomes, applying softmax along the last dimension still produces:

inner_size = 1

The spatial kernel remains unreachable.

Using the same total number of elements but moving the softmax dimension into the middle can produce:

inner_size > 1

and select a completely different kernel.

Actual test reachability depended on the combination of:

total element count
+
tensor shape
+
softmax axis
+
inner_size
+
selected branch
+
launched kernel

Element count alone could not describe the coverage.

A test name does not establish implementation coverage

The name:

test_softmax_backward_64bit_indexing

sounds broad.

But a test name is written by a person.

Actual coverage is determined by control flow.

Test name:
softmax backward 64-bit indexing

Actual path:
one backward implementation under inner_size == 1

The following two statements are therefore not equivalent:

A 64-bit softmax backward test exists.

Every softmax backward kernel has been exercised at a 64-bit boundary.

The first can be confirmed by looking at the test file.

The second requires following:

shape
→ dimension selection
→ branch condition
→ dispatch
→ actual kernel launch

The existing PASS was still valid

The existing test passed.

That PASS correctly established:

The large last-dimension softmax backward path produced the expected result under that input.

The new regression test asks a different question:

Can the non-last-dimension spatial backward path address every element after its offsets exceed the old unsigned 32-bit range?

These are separate implementation contracts:

Existing test
→ 64-bit boundary of the last-dimension kernel

New test
→ 64-bit boundary of the spatial kernel

The weakness was not that the earlier test was wrong.

It was treating one passing implementation path as evidence for the whole operation name.

Why the new test uses `dim_size = 2`

The regression test fixes the softmax dimension at two elements:

dim_size = 2

That makes each reduction as simple as possible.

The softmax output in every column is:

[0.25, 0.75]

The incoming gradient is:

[1.0, 0.0]

The expected result is exactly:

[+0.1875, -0.1875]

for every column.

This removes several unrelated variables:

long reductions

complex value distributions

position-dependent expected values

floating-point tolerance ambiguity

The logical computation stays small.

Only the address space becomes enormous.

Why the test also fixes `outer_size = 1`

The first dimension is kept at 1:

outer_size = 1

That removes another component of the address calculation.

The kernel no longer needs a meaningful movement across outer groups.

The central offset reduces to:

inner_index
+
d × inner_size

For the second softmax value:

d = 1

so the test needs only to determine whether:

inner_index + inner_size

can cross 2**32.

The shape is not merely large.

It is deliberately arranged to isolate the spatial address calculation.

No outer movement
+
two-element reduction
+
enormous inner dimension

This is close to a minimal counterexample for the indexing contract, even though the allocation itself is huge.

The test compares the beginning and end of the tensor

The regression test checks gradients near the start:

gI[0, 0, 0]
gI[0, 1, 0]

It also checks the final column:

gI[0, 0, -1]
gI[0, 1, -1]

The beginning uses low offsets:

well within uint32_t

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

2**32 + 1

which exceeds the old range.

Every column contains the same input values, so both ends must produce the same gradients:

Beginning:
+0.1875 / -0.1875

End:
+0.1875 / -0.1875

This gives one test both a normal-address reference and an overflow-address probe.

Why the new test needs approximately 30GB

The existing test is marked as requiring roughly 20GB.

The new spatial regression is marked for roughly 30GB.

That difference follows from the shape and number of tensors involved.

The spatial test contains:

2**32 + 2 FP16 elements

in each large tensor.

At two bytes per FP16 element:

one tensor
≈ 8 GiB

The test needs at least:

softmax output

incoming gradient

result gradient

The raw tensor storage is therefore approximately:

8 GiB × 3
≈ 24 GiB

The CUDA context, allocator, kernel execution, and safety margin raise the practical requirement to around 30GB.

The memory cost comes directly from testing the real pointer path beyond the unsigned 32-bit element boundary.

The 64-bit kernel introduced another boundary: launch viability

Changing the spatial backward kernel from fixed uint32_t offsets to an index_t template creates a 64-bit kernel variant for large tensors.

A 64-bit variant can consume more registers than the 32-bit version.

Higher register pressure can reduce:

the usable block size
the number of blocks that can reside simultaneously
occupancy on a particular GPU

A launch configuration that works for the 32-bit variant may not be valid for the 64-bit instantiation.

The patch therefore had to reinforce the launch-selection logic as well.

Correct address type
does not automatically imply
a launchable kernel configuration

The old launch helper asked the occupancy question only once

The spatial helper chooses a block shape and queries the driver for occupancy.

Conceptually, the previous flow was:

Select block dimensions
        ↓
Query occupancy once
        ↓
Construct the grid from that result

But the 64-bit kernel may have a lower practical thread limit because of its register usage.

The patch notes a further backend-specific behavior observed on gfx950.

For an oversized block, the occupancy query can return:

error code:
cudaSuccess

resident block count:
0

The API call itself reports success.

But the configuration cannot keep even one block resident.

Checking only:

err == cudaSuccess

would therefore be insufficient.

A valid result also requires:

max_active_blocks > 0

The repaired launch helper shrinks the block until it is viable

The updated logic repeatedly probes the selected configuration:

Query occupancy for the current block
        ↓
No error and at least one resident block?
        ↓
Yes → accept
        ↓
No → reduce the block
        ↓
query again

It first reduces block.y.

When block.y can no longer be reduced, it reduces block.x.

if (block.y > 1) {
    block.y /= 2;
} else {
    block.x /= 2;
}

After every reduction, it recalculates:

total threads
shared-memory requirements
occupancy

The helper accepts a configuration only when:

the driver accepts it
+
at least one block can actually reside

This was not an unrelated cleanup.

It was part of making the newly introduced 64-bit kernel variant executable across devices with different register and occupancy limits.

Address correctness and launch viability had to be closed together

The full repair crosses several layers:

Layer 1:
old uint32_t offsets wrap

Layer 2:
runtime selects a 64-bit index variant

Layer 3:
the 64-bit variant uses more registers

Layer 4:
the original block shape may no longer be resident

Layer 5:
the launch helper finds a viable configuration

Correct offsets are not useful if the kernel cannot launch.

A successful launch is not useful if the offsets are still wrong.

The complete correctness chain is:

correct branch
→ sufficient index width
→ correct kernel variant
→ valid launch configuration
→ correct gradient values

Test presence and test execution are different facts

The new test is guarded by:

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

An environment without sufficient GPU memory skips it.

The change author also stated that their local GPU was not large enough to execute the approximately 30GB path.

They verified the expected numerical values with a smaller CPU example, and the index dispatch mirrors the already-deployed forward implementation.

The public evidence therefore establishes:

Regression-test code exists
✓

The exact boundary shape is encoded
✓

The source change landed in main
✓

It does not automatically establish:

Every CI environment executed the 30GB case
✗

A test being committed and a test being exercised on the required hardware are separate claims.

A stronger test matrix starts with implementation branches

A complete large-index softmax backward matrix would need to distinguish more than tensor size.

Relevant dimensions include:

Softmax axis:
last / non-last

inner_size:
1 / greater than 1

Signed index range:
at or below INT_MAX / above INT_MAX

Unsigned spatial offset:
at or below 2**32 - 1 / above it

Selected implementation:
ordinary softmax / spatial softmax

Observed location:
beginning / end of tensor

Index variant:
32-bit / 64-bit

Launch configuration:
immediately resident / requires block reduction

This does not mean testing every possible combination exhaustively.

It means selecting a minimal set of inputs that actually crosses each implementation and type boundary.

API coverage and implementation coverage are not the same

The structure of this issue can be reduced to one statement:

A test for the softmax backward API exists
≠
Every softmax backward kernel is covered

At the API level, the operation is the same.

Inside PyTorch, the route splits:

Python API
        ↓
Analyze shape and softmax dimension
        ↓
Calculate outer / dim / inner
        ↓
Is inner_size equal to 1?
       ↙                     ↘
ordinary path             spatial path

Without following the actual arrow taken by the test, the implementation coverage remains unknown.

A PASS speaks only for the path that was reached

The existing test passed.

That result proved:

The inner_size == 1 large softmax backward implementation produced the expected result for that input.

It did not prove:

The inner_size != 1 spatial backward implementation can represent the same address range safely.

The scope of test evidence does not automatically expand beyond the executed control flow.

Test PASS
→ evidence for the reached path

Unreached path
→ no evidence yet

The patch added more than another test

At the surface, #188031 added one approximately 30GB regression test.

The deeper contract is broader:

Regardless of the softmax axis, forward and backward must interpret the same tensor size consistently, use 64-bit addressing when required, and launch the resulting kernel under a valid device configuration.

The previous test covered an operation name.

The new test pins down an implementation branch.

That is the important difference.

Did the test actually reach the defect?

The final lesson from PyTorch #188031 is not about the number of tests.

It is not even about whether a test already existed.

The important question is:

Did the test pass through the producer, branch, kernel, and address boundary where the defect actually lived?

The existing test was large enough.

Its name explicitly mentioned 64-bit indexing.

It passed correctly.

But its shape and axis selected another CUDA kernel.

The defective path was never reached.

A test does not gain coverage from its name.

Coverage exists only along the control flow it actually executes.

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?
PyTorch Softmax Backward — How an Offset Beyond 2**32 Wrapped Back to the Start of the Tensor

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

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

Related material

Patch status: Landed in PyTorch main
Evidence: Actual branch dispatch, source diff, and existing and new large-tensor tests
Boundary note: The approximately 30GB regression test is skipped on devices without sufficient memory, and the change author stated that they could not execute that path on their local GPU.

This is Part 3 and the final article in the PyTorch spatial softmax 64-bit indexing series.

Part 1 examined the asymmetry between forward and backward index widths.

Part 2 followed the exact regression shape and calculated how offsets beyond 2**32 wrapped to positions 0 and 1.

This final article showed why an existing large 64-bit test still missed the defect: inner_size == 1 and inner_size != 1 selected different CUDA kernel paths.

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

Search This Blog

Resonetic Lab | AI · Code · Market