ciflow/torchtitan/187116: Update on "Add stable tie handling to topk"

The reverted deterministic-topk change made topk pick stable tie indices when

global deterministic algorithms were enabled. That changed default user-visible

tie behavior for existing deterministic-mode callers, including internal model

checks that had recorded the previous unspecified tie order.

This change makes stable topk an explicit opt-in instead. The default

stable=False path keeps the existing topk kernels and their existing tie

semantics. stable=True asks topk to preserve the original index order among

equal values, matching the stable-sort convention, and is exposed as a

keyword-only argument so existing positional callsites are not reinterpreted.

The implementation keeps the default fast paths unchanged. CPU and quantized

CPU switch from the existing partial-sort/nth-element selection to stable-sort

only when requested. CUDA and MPS route the requested stable case through their

sort implementations, then narrow to k. The CuTeDSL override declines the stable

case so the canonical implementation handles the stronger tie-order contract.

ONNX export rejects stable=True because ONNX TopK does not specify stable tie

semantics. AOTI gets versioned topk v2 shims so existing compiled ABI entry

points remain unchanged.

This change was authored with assistance from an AI assistant.

Test Plan:

python torchgen/gen.py --update-aoti-c-shim

lintrunner -a

pip install -e . -v --no-build-isolation --no-deps

python test/test_sort_and_select.py -k topk_stable -v

python test/test_sort_and_select.py -k topk -v

python - <<'PY'

import torch

print(torch.ops.aten.topk.default._schema)

x = torch.tensor([[5., 5., 5., 4.], [0., 1., 1., 1.]])

assert torch.topk(x, 2, stable=True).indices.tolist() == [[0, 1], [1, 2]]

try:

    torch.topk(x, 2, -1, True, True, True)

except TypeError:

    pass

else:

    raise AssertionError('stable unexpectedly accepted positionally')

if torch.cuda.is_available():

    y = x.cuda()

    assert torch.topk(y, 2, stable=True).indices.cpu().tolist() == [[0, 1], [1, 2]]

    def f(t):

        return torch.topk(t, 2, stable=True).indices

    compiled_f = torch.compile(f, fullgraph=True)

    assert compiled_f(y[:1]).cpu().tolist() == [[0, 1]]

PY

[ghstack-poisoned]