CUDA/HIP: partial-top-k argsort — up to 6x faster TOP_K on RDNA3.5/RDNA4

[DeanoC]

Internal review PR. Base is base/hip-draft-topk-488 (upstream Luce-Org#488's commit, @cheese-cakee's "compile DFlash GPU draft top-K for HIP"), so the diff below is only our partial-top-k argsort change. When we're ready, the upstream PR to Luce-Org goes out via the GitHub App, stacked on Luce-Org#488.

Summary

The ggml TOP_K path full-sorts every row via bitonic argsort regardless of k, so latency is independent of k. This makes it partial:

• argsort.cu: shared-memory key cache + partial bitonic bounded by the padded k (work scales with k, not ncols).
• argsort.cu: GGML_ARGSORT_SHFL_XOR — for bitonic inner stages below the wavefront width, both compare-exchange partners live in the same wave, so the swap is register-to-register via __shfl_xor, no shared-mem round-trip / __syncthreads() (wave32 on RDNA3.5/RDNA4).
• top-k.cu: dedicated k=1 argmax kernel that skips the sort.

On HIP the CUB fast path is unavailable (GGML_CUDA_USE_CUB undefined), so the bitonic kernel is the hot path on AMD — where DSpark draft/verify top-K lands.

Impact (test-backend-ops perf, ne=[1000,16])

Scenario	before	after	gain
gfx1201 (RDNA4) k=1	15.59 µs	2.59 µs	6.0×
gfx1201 k=10	15.60 µs	4.58 µs	3.4×
gfx1201 k=400	15.77 µs	6.56 µs	2.4×
gfx1151 (Strix Halo) geomean k=1..400	15.6 µs	~4.5 µs	~3.46×

Validation

• test-backend-ops test -o TOP_K passes on gfx1201 (R9700) and gfx1151 (Strix Halo / Radeon 8060S) — 2/2 backends passed on both.
• Discovered + A/B'd with the isolated test-backend-ops harness on gfx1201, re-validated on gfx1151.

On the N>1024 "not supported" cases (previously flagged — now resolved)

Verified on both gfx1201 and gfx1151: 275 cases report "not supported", identical on both archs, and every one is N>1024 (ne=[1025,…] through ne=[524288,…]); zero N≤1024 cases are unsupported. This is not arch-specific and not affected by this PR — it's a HIP-build-wide gate in ggml_backend_cuda_supports_op:

case GGML_OP_TOP_K:
case GGML_OP_ARGSORT:
#ifndef GGML_CUDA_USE_CUB
    return op->src[0]->ne[0] <= 1024;   // all ROCm/HIP builds (CUB is CUDA-only)
#else
    return true;                         // CUDA+CUB handles arbitrary N
#endif

N>1024 (e.g. full-vocab sampling) falls back to the CPU reference on any HIP build. This PR touches only the kernels (argsort.cu/top-k.cu), not ggml-cuda.cu, so it leaves that envelope unchanged and optimizes within the supported N≤1024 range. Lifting the >1024 cap on HIP (a multi-block sort, or a segmented-radix analog to CUB) is a separate follow-up.

[DeanoC]

Superseded by #15 — collapsed the partial-top-k win (#13) and the hipCUB ncols>1024 support (#14) into a single PR for review; they're complementary changes to the same kernels.