[Perf][Argreduce] Add adaptive streaming pair reduction

[stelladuyx]

Supersedes #1804.

Summary

• replace the shared-memory + serial-scan argreduce kernel with a single-pass (value, index) pair reduction
• stream inputs directly from global memory with four independent per-thread accumulators
• dispatch between subgroup, CTA, multi-CTA, and stride-aware output-parallel layouts
• preserve first-index and NaN semantics for both argmax and argmin
• remove the staged benchmark skips for LM-head reductions

Why

The implementation in #1804 removed one large fragment, but it still materialized a complete input row in shared memory and scanned the reduction dimension serially. It could not execute the manifest LM-head workload at N=102400, and non-last-axis reductions materialized a transposed contiguous tensor.

This version replaces that design rather than incrementally extending it.

Implementation

• small/medium reductions: power-of-two subgroup pair reduction
• N >= 4096: register -> warp shuffle -> shared memory -> global CTA reduction
• small-M, very large-N: two-stage multi-CTA partial and final pair reduction
• contiguous non-last axes: output-parallel traversal using inner_stride, avoiding the transpose
• unusual non-contiguous views: retain the compatibility compacting fallback

H200 performance

CUPTI kernel time with L2 flushing, physical GPU 1:

workload	dtype	TileOPs	PyTorch	result
argmax (4, 102400)	fp16	0.0120 ms	0.0360 ms	3.00x faster
argmax (4, 102400)	bf16	0.0119 ms	0.0366 ms	3.08x faster
argmax (2048, 4096)	fp16	0.0232 ms	0.0252 ms	1.09x faster
argmax (2048, 4096)	bf16	0.0232 ms	0.0254 ms	1.10x faster
argmax (4, 128, 4096), dim=0	fp16	0.0055 ms	0.0115 ms	2.07x faster
argmin (4, 102400)	fp16	0.0122 ms	0.0359 ms	2.95x faster
argmin (4, 102400)	bf16	0.0122 ms	0.0365 ms	3.00x faster
argmin (2048, 4096)	fp16	0.0232 ms	0.0253 ms	1.09x faster
argmin (2048, 4096)	bf16	0.0232 ms	0.0255 ms	1.10x faster

Compared with the production kernel before #1804, the shared workloads improve by about 10.3x, while the non-last-axis workload improves by about 82x. The four LM-head cases change from unsupported/skipped to passing.

Validation

• pytest tests/ops/test_argreduce.py -q — 90 passed
• pytest benchmarks/ops/bench_argreduce.py -q — 9 passed, 0 skipped
• pytest benchmarks/tests/test_run_benchmarks.py tests/test_ops_manifest.py -q — 12 passed
• ruff check and pre-commit hooks — passed
• torch.compile exact-output check — passed

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.