OpenXLA PR #49368: Why Raw Float Comparisons Corrupted GPU TopK

OpenXLA PR #49368 fixes a specialized GPU TopK kernel whose comparator did not implement the ordering promised by the source operation.

The old comparator was:

return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key;

For ordinary finite values, this usually behaves like a descending stable sort.

For NaNs, neither equality nor greater-than establishes an order.

For signed zero, equality removes the distinction between +0 and -0.

That was enough to make the bespoke GPU kernel disagree with XLA’s sort-based TopK path. The same HLO operation could therefore produce different values and indices depending on whether its shape crossed the threshold for the optimized kernel.

The patch bitcasts float32 and bfloat16 keys to unsigned integers, transforms those bits into a monotonic total-order key, and compares the integer keys instead.

One shape boundary changed the meaning of TopK

XLA has more than one way to implement TopK.

For smaller or unsupported configurations, it can use a sort followed by a slice.

For a supported GPU shape, TopkSpecializer routes the operation to the bespoke __gpu$TopK kernel.

The relevant optimized range is:

input length n >= 1024
and
k <= 16
and
dtype is float32 or bfloat16

The linked issue demonstrated the boundary directly:

n = 1023
→ sort + slice
→ correct TOTALORDER result

n = 1024
→ bespoke GPU TopK
→ incorrect result

An optimization boundary should change performance and implementation structure. It must not change which elements are considered the largest.

The old comparator used a partial floating-point order

The specialized kernel stores a value and its original index together.

Its comparator attempted to provide stable descending behavior:

if values are equal:
    lower original index wins

otherwise:
    greater value wins

That policy works only when the value comparison defines a consistent order.

NaN breaks the assumption.

For any ordinary number x:

NaN == x    → false
NaN > x     → false
x > NaN     → false

Even two NaNs with identical bits do not compare equal through floating-point equality.

The comparator can therefore say that:

1 is not ahead of NaN
NaN is not ahead of 1

NaN is not ahead of 2
2 is not ahead of NaN

while still saying:

2 is ahead of 1

The apparent equivalence created by “neither compares ahead” is not transitive.

That violates the ordering assumption used by the kernel’s insertion and reduction stages.

The failure propagated through a parallel selection network

The custom kernel does not call a general-purpose sort.

Each thread scans a slice of the input and maintains a small TopK list in registers. Warp lanes then merge their local lists through shuffle-based reductions. A final merge produces the output values and indices.

Every stage assumes that Descending::cmp() preserves one consistent ordering.

When the comparator cannot place a NaN relative to another value, the result becomes dependent on the exact sequence of conditional exchanges across threads and lanes.

The issue’s failing k=16, n=1024 run returned sixteen negative NaNs at periodic indices:

1023, 959, 895, 831, ...

The expected TopK began with:

positive NaNs ordered by payload
positive infinity
large positive finite values
positive subnormals
positive zero
negative zero
negative subnormals

The periodic wrong indices reflect the kernel’s reduction layout, not a defined floating-point order.

Identical NaNs also bypassed the stable-index rule

The comparator intended to use the original index for ties.

That branch began with:

lhs.key == rhs.key

For identical NaN bit patterns, the expression is still false.

The index tie-break therefore never ran.

This mattered because the test input filled most remaining positions with the same negative NaN sentinel. Under a stable total order, equal keys should remain ordered by their original indices.

The old comparator could neither recognize them as equal nor place them consistently relative to other values.

The patch moves equality to the transformed integer keys:

auto l = details::ToOrdered(lhs.key);
auto r = details::ToOrdered(rhs.key);

return l == r ? lhs.idx < rhs.idx : l > r;

Identical payload bits now produce identical integer keys, so the original-index rule becomes reachable.

Signed zero required an order too

Ordinary floating-point comparison treats:

+0 == -0

as true.

XLA’s floating-point TOTALORDER distinguishes them.

For descending TopK:

+0 comes before -0

The old comparator treated them as a value tie and chose solely by source index.

That means the relative order of signed zeros could depend on where they appeared in the input rather than on the TopK comparison contract.

The new bitwise keys place -0 immediately below +0.

The bit transform makes unsigned comparison monotonic

For a 32-bit floating-point key, the patch first reads the raw bits as uint32_t.

It then applies:

negative encoding:
invert every bit

nonnegative encoding:
set the sign bit

The 16-bit bfloat16 path applies the same operation with uint16_t.

Conceptually:

if sign bit is set:
    ordered = ~bits;
else:
    ordered = bits | sign_bit;

Unsigned integer comparison now follows XLA’s total-order direction.

Selected float32 encodings become:

Value class	Original bits	Ordered key
negative NaN with maximal payload	`0xffffffff`	`0x00000000`
negative quiet NaN	`0xffc00000`	`0x003fffff`
negative infinity	`0xff800000`	`0x007fffff`
`-64.5`	`0xc2810000`	`0x3d7effff`
`-0`	`0x80000000`	`0x7fffffff`
`+0`	`0x00000000`	`0x80000000`
`+64.5`	`0x42810000`	`0xc2810000`
positive infinity	`0x7f800000`	`0xff800000`
positive quiet NaN	`0x7fc00000`	`0xffc00000`

In ascending key order, the broad classes are:

negative NaNs
negative infinity
negative finite values
negative zero
positive zero
positive finite values
positive infinity
positive NaNs

The kernel compares the keys in descending order because it is selecting the largest values.

Why the transform reverses only negative encodings

Among positive IEEE-style encodings, larger bit patterns already follow increasing numerical magnitude through finite values and infinity, then continue into positive NaNs.

Setting the sign bit moves the complete nonnegative half above the transformed negative half without changing its internal order.

Negative encodings run in the opposite bit direction.

For example, the raw encoding of a more negative finite number is numerically larger as an unsigned integer than the encoding of a value closer to zero.

Inverting every bit reverses that region.

The result is a one-to-one integer key for each source bit pattern. No arithmetic conversion is performed, so NaN payloads and signed zero survive as ordering information.

CUDA and ROCm receive the same comparator

The pull request title names CUDA, matching the original report.

The merged patch applies the same ToOrdered() helper and comparator change in both:

CUDA TopK common header
ROCm TopK common header

Integral keys pass through unchanged.

The registered bespoke CUDA kernels covered by the source are float and bfloat16, with template variants for k values 1, 2, 4, 8, and 16 and different index widths.

The patch does not add an FP16 custom-kernel registration.

The regression crosses optimized and fallback paths

The new test constructs explicit bit patterns for:

multiple positive NaN payloads
multiple negative NaN payloads
positive and negative infinity
positive and negative zero
positive and negative subnormals
positive and negative finite values
many copies of one negative NaN sentinel

It runs both float32 and bfloat16.

The shape matrix includes:

That matters because the matrix crosses the specialized-kernel boundary instead of testing only one implementation.

For every combination, the test checks both values and original indices across eight identical rows.

The expected indices verify:

NaN payload ordering
positive zero before negative zero
stable lower-index order for duplicate values
the low placement of negative NaNs

For the repeated sentinel NaNs, expected indices continue in ascending order once the non-sentinel values have been exhausted.

One GPU generation is explicitly excluded

The test target supports CPU and GPU backends, but its GPU configuration carries:

requires-gpu-sm80

The BUILD comment says P100 and V100 are skipped because of a separate +0 and -0 interleaving problem in sort.

That exclusion is an important boundary.

The new comparator fixes the bespoke TopK implementation. The regression target does not establish corrected end-to-end behavior on every older NVIDIA GPU.

The production patch also changes the ROCm header, but the public patch does not show a dedicated ROCm-only execution test. Direct ROCm validation cannot be inferred from the source change alone.

A local bit-level check matched the XLA reference order

I separately implemented the merged transform in Python.

For all 65,536 possible bfloat16 bit patterns, the transformed keys were:

unique
covering every integer from 0 through 65535
ordered identically to XLA’s signed total-order reference key

The unsigned key was exactly the reference signed key plus 0x8000.

For one million random float32 bit patterns, the equivalent 0x80000000 offset relationship also held.

Using the issue-style input, the local transform produced the same leading indices committed in the regression:

5, 19, 4, 1, 8, 18, 20, 21,
13, 23, 3, 16, 6, 12, 14, 24

JAX 0.9.0.1 on CPU returned the same indices and value bit patterns for that input.

This local check validates the ordering arithmetic and the CPU reference result. It does not execute the merged GPU kernel.

The public evidence does not include a performance result

The new comparator adds bitcasts and simple integer bit operations before comparison.

The pull request publishes no kernel benchmark, generated-code comparison, register-pressure result, or end-to-end TopK timing.

The bit transform is branch-free after compile-time type selection, but that observation is not a performance measurement.

The supported conclusion is about semantic consistency, not speed.

The compiler contract restored by the patch

A specialized kernel must preserve the comparison mode of the operation it replaces.

The corrected path is:

retain the original floating-point payload bits
        ↓
map each encoding to a monotonic integer key
        ↓
compare those keys under one total order
        ↓
use the source index only when the keys are identical
        ↓
produce the same TopK ordering across implementation paths