Linked issue or supporting maintainer

Part of #N

BC-breaking?

No

Summary (human written only)

The original FlyDSL top-k implementation supports fp32 when largest=True. This PR extends the kernels to fp16, bfloat16, int32, int16, and int8, and adds support for largest=False.

This is done by introducing TopKSpec, a compile-time specialization for each dtype and ordering. It selects the dtype-specific bit transformation used to encode values as order-preserving ordinals and decode them, and carries the element type and width used by the load path. A comprehensive proof of why this construction is sound is included in Why it works below.

• derive vector width, radix passes, and value storage from the dtype while retaining the signed-int32 ordinal plus inverted-index key
• include dtype and largest in the compilation cache keys
• leave operator dispatch restricted to fp32, largest=True, and sorted=True until the policy windows are retuned

Agent details (if applicable)

The proof below was outlined, written, and reviewed by the author. An AI assistant helped verify the bit-level reasoning, edit the presentation, and implement the corresponding refactor. The author reviewed and takes responsibility for both the code and the proof.

Testing

• lintrunner --take RUFF,PYFMT -- torch/_native/ops/topk/flydsl_kernels.py
• Python AST syntax check
• exhaustive host formula validation for all fp16 and bfloat16 bit patterns, plus fp32 boundary cases
• GPU kernel compilation and execution: not run yet (draft)

Why it works

Each selected pair (v, i) is packed into one signed int64 key, and the kernel keeps the greater key. ord(v) is the signed int32 in the high half, and the low 32 bits are ~i:

[ ord(v) | ~i ]

A narrow value occupies the top of ord(v) and the bits underneath are zero. Those zeros are part of ord(v), not a third field. fp32 and int32 fill the int32. fp16, bf16, and int16 use the top 16 bits. int8 uses the top 8.

The current kernel already uses this key for fp32 with largest=True: the high half is the fp32 ordinal and the low half is ~i, with no pad. This note records the order that implementation assumes, and the same key is what extends it to the other dtypes and to largest=False.

Define < on pairs as "loses to":

• largest=True: (v1, i1) < (v2, i2) when v1 < v2, or the values tie and i2 < i1. This is lexicographic order on (v, -i).
• largest=False: (v1, i1) < (v2, i2) when v2 < v1, or the values tie and i2 < i1. This is lexicographic order on (-v, -i).

Here -v denotes the value order reversed, not a value computed by arithmetic negation. Likewise, -i denotes the index order reversed. The kernel does not negate either field.

The value tie rule is the same in both: the smaller index wins. The < on values agrees with ordinary comparison wherever two numbers compare unequal. It also places -0 before +0, and every NaN above +inf, tied with every other NaN.

Overview

We want to show that packaging (v, i) pairs as [ ord(v) | ~i ] for largest=True, and as [ ~ord(v) | ~i ] for largest=False, then choosing the pairs with the largest keys, gives the correct semantics. We do so by showing three separate facts:

1. ord preserves the value order. This is dtype-specific. The signed int32 ord(v) compares exactly as v does: v1 < v2 if and only if ord(v1) < ord(v2), and the ordinals are equal exactly when the values tie. The fp32 case is the ordinal the current kernel already builds. The same claim is what embeds fp16, bf16, int32, int16, and int8 into that int32. For this scheme to work, v1 and v2 have the same dtype.
2. The full key preserves lexicographic order on (v, -i). Representing a pair as [ ord(v) | ~i ] keeps that order. Sorting these keys sorts (v, -i).
3. Complementing ord(v) reverses the value order. The current kernel never does this: it only builds the largest=True key. largest=False is fact 2 applied to the key [ ~ord(v) | ~i ], with the same ~i.

From the above facts we can infer:

• largest=True semantics are correct: (v, -i) -> (ord(v), -i) and (ord(v), -i) -> [ ord(v) | ~i ] are order preserving because of points 1 and 2 respectively.
• largest=False semantics are correct: (-v, -i) -> (~ord(v), -i) and (~ord(v), -i) -> [ ~ord(v) | ~i ] are order preserving because of points 1 and 3, and point 2, respectively. Point 3 reverses the ordinal; point 1 is what ties that ordinal back to v.

Side note. Before the reversible value transform, step 0 collapses every NaN of either sign to one sentinel. The current fp32 kernel already uses 0x7FFFFFFF; its ordinal sits above +inf at 0x7F800000. This PR keeps that collapse and uses 0x7FFF for fp16 and bf16. After the left shift, their stored sentinel is 0x7FFF0000, above fp16 +inf at 0x7C000000 and bf16 +inf at 0x7F800000. largest=True therefore returns NaNs first, tied with each other, and the index breaks the tie. largest=False complements after the collapse, so the maximum sentinel becomes the minimum code and a smallest-k returns NaNs last.

The collapse is not injective, so deconversion cannot recover the original NaN payload. After undoing the complement, alignment, and value transform, the caller compares the recovered narrow bits with the narrow sentinel and reloads row[i] on a match.

The three facts

1. ord preserves the value order

[ fp32 ]

This is the ordinal the current kernel already uses, unchanged. Bitcast the float to int32 and apply

sign = (bits >> 31) & 1
mask = (0 - sign) & 0x7FFFFFFF
bits ^ mask

For this width the mask equals (bits >> 31) & 0x7FFFFFFF.

As a background, positive fp32 is a clear sign bit, then an 8-bit exponent, then a 23-bit fraction. +1 is 0x3F800000. A negative fp32 is that same pattern with bit 31 set, so -1 is 0xBF800000 and the exponent and fraction are unchanged.

On a non-negative pattern the formula is a no-op. The sign bit is clear, so sign is 0, 0 - sign is 0, and the XOR changes nothing. The ordinal is the raw bits. The exponent sits above the fraction, so a larger exponent is a larger integer, and the same exponent with a larger fraction differs only in the low bits. +0 is the all-zero pattern, so its ordinal is 0 and every positive ordinal is greater than 0.

On a negative pattern the sign bit is set, so sign is 1 and 0 - sign is all ones. The mask keeps the low 31 of those ones, and the formula collapses to bits ^ 0x7FFFFFFF. The exponent and fraction are stored exactly as they are for a positive value, so a larger magnitude still has a larger field in the low 31 bits. A signed compare walks the bits from bit 31 downward, and the first bit that differs decides, with a 1 there making the larger integer. Both patterns have bit 31 set, so the walk falls through to the exponent. A larger exponent is therefore a larger int32, the same rule that ordered the positive side. Here it points the wrong way: -inf has a larger exponent than -1, so the raw pattern ranks -inf above -1. The XOR flips that field and leaves bit 31 set. The walk now meets a 0 first on the larger magnitude, so that value is the smaller int32, which is the numeric order, and the sign bit keeps the result negative.

The sign bit keeps every negative ordinal below ord(+0) = 0, and within that half a larger magnitude is a smaller ordinal. Every positive ordinal is above ord(+0). That is the total order.

Deconversion applies the same three lines:

sign = (bits >> 31) & 1
mask = (0 - sign) & 0x7FFFFFFF
bits ^ mask

The mask depends only on bit 31. 0x7FFFFFFF has bit 31 clear, so the sign bit is the same before and after the original transform and the second application builds the same mask. For a non-NaN, xoring the same mask twice returns the original bits.

NaN handling: before this transform, step 0 collapses every NaN to 0x7FFFFFFF. Its ordinal sits above +inf at 0x7F800000.

[ fp16 and bf16 ]

This is the same three lines at 16 bits. The sign is bit 15, and the operations run on that 16-bit signed integer:

sign = (bits >> 15) & 1
mask = (0 - sign) & 0x7FFF
bits ^ mask

The 16-bit result is shifted left by 16 inside an int32. The shift keeps those 16 bits, moves bit 15 to bit 31, and writes zeros under the field, so the signed walk from the fp32 case applies. Bits above the field are shifted off the top, so it does not matter whether sign extension had filled them with ones.

Deconversion shifts the int32 right by 16 and truncates to 16 bits. An arithmetic shift fills ones above a negative ordinal, and the truncate keeps the low 16:

field = low 16 bits of (ord >> 16)
sign = (field >> 15) & 1
mask = (0 - sign) & 0x7FFF
field ^ mask

Bit 15 of that field is the bit the forward map left unchanged, so the mask built here is the mask applied to the original input. For a non-NaN, xoring it again returns the original bits.

NaN handling: before this transform, step 0 collapses every NaN to 0x7FFF. The stored ordinal after the shift is 0x7FFF0000, which sits above fp16 +inf at 0x7C000000 and bf16 +inf at 0x7F800000.

int32, int16, and int8

This PR adds int32, int16, and int8. A signed integer needs no sign flip: its bits are already in signed order. The only step is the left shift of the next fact. int32 already fills the high half, so its shift is zero. int16 and int8 are shifted left so their sign lands on bit 31.

Deconversion is the arithmetic right shift by the same amount, then a truncate to width bits. int32 stays as it is. int16 shifts by 16 and int8 by 24, which sign-fills. The low width bits are the original value, and the truncate reads them as that signed integer:

low width bits of (ord >> (32 - width))

NaN handling: these dtypes have no NaN.

Zero-padding preserves order

Zero-padding on the right is the shift that places a signed width-bit value into the high bits of an int32:

z << (32 - width)

The shift keeps the low width bits and moves them to the top. Bits above the field are shifted off, whether sign extension had written ones there or not, and the bits under the field are zero.

It preserves signed order when the shift does not overflow. Let z be a signed width-bit integer, width ≤ 32, so -2^(width-1) ≤ z ≤ 2^(width-1) - 1. The shift is multiplication by 2^(32-width), which is positive, so the order is preserved as long as the product fits. The largest width-bit value is 2^(width-1) - 1, so the largest product is 2^31 - 2^(32-width). That is strictly smaller than 2^31, hence at most 2^31 - 1. The smallest value is -2^(width-1), and its product with 2^(32-width) is -2^31, which is INT32_MIN. Both products fit in a signed int32. The bits below the field are zero. For fp32, width = 32, so the shift is zero and the current kernel has no pad.

Deconversion is the arithmetic right shift by the same amount, then a truncate to width bits:

low width bits of (ord >> (32 - width))

A negative value sign-fills above the field. The truncate keeps the low width bits, and their top bit is the original sign. The filled ones sit above those bits, so the sign fill does not change the recovered value.

Reference implementation

The value maps in point 1 are not proposed here for the first time. CUDA radix top-k already defines that order for float, half, bfloat16, int8, int16, and int32 in TopKTypeConfig. A float keeps its order by an xor of the sign, with NaN sent to all-ones. A signed integer is biased by 2^(width-1). Both are then compared as unsigned. This kernel uses the signed form of that same order, because the existing fp32 path already turns the value into a signed int32 and packs it into a signed int64. The new dtypes extend that path rather than switching the compare to the unsigned CUDA codes. For a float the two codes differ by one sign-bit xor: ^ 0x80000000 for fp32, and ^ 0x8000 on the 16-bit field for fp16 and bf16, which is ^ 0x80000000 after the left shift. That xor preserves order between the two spaces: a signed compare of this code and an unsigned compare of the reference code rank values the same way. It flips the sign bit, so the order is not the same inside either compare alone. For an integer the reference adds 2^(width-1) and obtains a non-negative code. This kernel multiplies by 2^(32-width), a left shift that puts the sign on bit 31, and a negative value stays negative. The signed order of that result is the unsigned order of the biased code.

2. Packing preserves the order of (v, -i)

The current fp32 kernel already assumes this and implements it by sorting the int64 keys descending. Nothing below depends on the high half being an fp32 ordinal, so every new dtype inherits it once lemma 1 has produced an int32.

The packed key is [ ord(v) | ~i ]. The high 32 bits are ord(v) and the low 32 bits are the inverted index.

A signed int64 compare walks from bit 63 downward. The key with bit 63 set is smaller. If that bit matches, the walk continues, and at the first later difference the key with the 1 is larger. The first 32 steps are the same walk as a signed compare of the high half. When the ordinals differ, the keys therefore compare as the ordinals do, and the index bits are not reached.

The walk reaches bit 31 only when the ordinals are equal, including two NaNs. From there it compares indexes. Take i1 > i2 ≥ 0. Walking the raw indexes from bit 31 down, the first bit that differs is 1 in i1 and 0 in i2, so the larger index would be the larger key. The low half stores ~i instead. Inversion flips that bit, so i2 has the 1 there and the smaller index wins.

The bitonic sort pads with INT64_MIN, whose bits are [ INT32_MIN | 0 ]. Every real key is greater: either its ordinal is greater than INT32_MIN, or its ordinal is INT32_MIN and its low half is the nonzero ~i of a valid index.

3. Complementing ord(v) reverses the value order

largest=False is the new half of the scheme. The current kernel only builds the largest=True key, so this reversal is not something it already does. The low half stays ~i.

ord(v) is the signed int32 the key already walks from bit 63. At bit 31, a 1 is the smaller ordinal. Below that, a 1 is the larger ordinal. Two different ordinals first differ in one of those bits. Complementing both swaps who has the 1 there, so the comparison reverses, and equal ordinals stay equal. Fact 2 applied to [ ~ord(v) | ~i ] is therefore the order of (-v, -i), which is largest=False.

Complementing also turns the zero bits under a narrow field into ones. Those ones are the same for every value of the dtype, so the first difference is still in the field. The low half is still ~i, so a tie still goes to the smaller index.

The NaN sentinel is the largest code of its field. For fp32 that sentinel is the 0x7FFFFFFF the current kernel already writes; complementing it is the new step. Its complement is strictly smaller than every numeric ordinal. fp32 0x7FFFFFFF becomes 0x80000000, below -inf at 0x807FFFFF. A smallest-k therefore returns NaNs last, still tied with each other.

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @jataylo @hongxiayang @naromero77amd @pragupta @jerrymannil @xinyazhang