Add CUDA argmax kernel for LLM sampler

[larryliu0820]

Add a CUDA kernel for argmax operation to support GPU-based sampling:

• argmax.cuh: Template kernel using warp-level reductions with __shfl_xor_sync
  for efficient parallel max finding. Supports float, half, and bfloat16.
• argmax.cu: Wrapper function argmax_cuda() that launches the kernel,
  handles device-to-host copy, and synchronization.
• test_argmax.cu: Comprehensive unit tests covering various vocab sizes,
  data types, edge cases, and numerical precision.
• CMakeLists.txt: Build configuration for extension_llm_sampler_cuda library
  and GoogleTest-based unit tests.

[larryliu0820]

Stack from ghstack (oldest at bottom):

• Integrate CUDA sampler into ASR runner and enable skip_copy for decoder #16388
• Add CudaSampler class for GPU-based token sampling #16387
• -> Add CUDA argmax kernel for LLM sampler #16386

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16386

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 2 Unrelated Failures

As of commit f8cd4d2 with merge base c5d66a5 ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner / linux-job (gh)
  >>> Lint for extension/llm/sampler/test/test_argmax.cu:
• pull / test-samsung-models-linux / linux-job (gh)
  test_inception_v3_fp16

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Test Metal Backend / export-model-metal-artifact (openai, whisper-small, non-quantized) / macos-job (gh) (matched macos rule in flaky-rules.json)
  File doesn't exist

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• pull / android / run-emulator (gh) (#16137)
  Timeout waiting for emulator to boot.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Gasoonjia]

consider using nested namespace to follow c++ 17 standard

[Gasoonjia]

perhapes we need to raise error here to avoid silent error?

[mergennachin]

I don't understand why you had to do a completely custom code.

There's already reduce aten kernel for argmax in pytorch (https://docs.pytorch.org/docs/stable/generated/torch.argmax.html)

https://github.com/pytorch/pytorch/blob/e20b76caff04de2c451777ea55e6fa97075ccbbf/aten/src/ATen/native/cuda/ReduceArgMaxKernel.cu

The torch layer has libtorch dependency (e.g., depends on TensorIterator) so you want to implement argmax as a shim layer instead in ExecuTorch, similar to how int4mm is done. That way, any model that uses torch.argmax can be lowered to ExecuTorch cuda backend easily in the future.

You can add files in
backends/cuda/runtime/shims/argmax.cu
backends/cuda/runtime/shims/argmax.cuh
backends/cuda/runtime/shims/argmax.h - header file
update aoti_cuda_shims.lib with the new kernel

Side note: another midsize task we can tackle as a team is to decouple TensorIterator in pytorch/pytorch to be libtorch independent as much as possible -- perhaps part of the header only file.