Skip to content

feat(nvidia): add Argmax provider - #918

Open
voltjia wants to merge 2 commits into
masterfrom
feat/nvidia-argmax
Open

feat(nvidia): add Argmax provider#918
voltjia wants to merge 2 commits into
masterfrom
feat/nvidia-argmax

Conversation

@voltjia

@voltjia voltjia commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add a CUB-backed native NVIDIA provider for the canonical Argmax operator.
  • Support flattened float32, float16, and bfloat16 input with an int64 scalar output.
  • Add focused Python coverage against torch.argmax.

Motivation

InfiniCore's basic Llama greedy decoding path needs a canonical InfiniOps replacement for the legacy sampling operator when top_k=1. The existing Argmax schema is aligned with torch.argmax(input, dim=None, keepdim=False), but InfiniOps did not have a native NVIDIA provider for the flattened reduction path.

Type of Change

  • feat - new feature / new operator / new platform
  • fix - bug fix
  • perf - performance improvement (no behavioral change)
  • refactor - code restructuring without behavior change
  • test - adding or fixing tests only
  • docs - documentation only
  • build / ci - build system or CI configuration
  • chore - tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

$ cmake --build /tmp/standard-llama/InfiniOps-build-argmax --target ops -j 16
[79/80] Linking CUDA shared module src/ops.cpython-312-x86_64-linux-gnu.so

$ pytest -q tests/test_argmax.py --devices nvidia
Running 6 items in this shard
......                                                                   [100%]
6 passed in 0.39s

The test build used CMAKE_CUDA_ARCHITECTURES=80 on an NVIDIA A100-SXM4-80GB.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Target and focused test passed f32, f16, and bf16; native and PyTorch implementation slots
Iluvatar No N/A - not affected N/A - not affected
MetaX No N/A - not affected N/A - not affected
Cambricon No N/A - not affected N/A - not affected
Moore No N/A - not affected N/A - not affected
Ascend No N/A - not affected N/A - not affected
InfiniCore / InfiniLM integration

The dependent InfiniCore branch was built against this provider and ran paged-attention TinyLlama inference in accelerator-dev/nvidia:latest:

Query: Hello
Response: Certainly!
total_time: 285.2 ms

The run used top_k=1 and preloaded a trap for the legacy infiniop* C operator API; the process exited successfully without triggering the trap.

Benchmark / Performance Impact

N/A. This adds the previously missing native provider; it is not a performance comparison.

Notes for Reviewers

  • The public schema is unchanged. This PR adds one native implementation and no overloads.
  • The NVIDIA provider currently accepts only dim=None, keepdim=false, contiguous non-empty input, and up to INT_MAX elements. These constraints cover the basic Llama greedy sampling path; other torch.argmax modes can be implemented separately.
  • The dependent InfiniCore integration is tracked in InfiniCore #1483.

@voltjia
voltjia requested a review from a team August 10, 2026 15:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant