Skip to content

feat(cuda): extend topk_softmax backend coverage - #891

Merged
voltjia merged 1 commit into
masterfrom
feat/extend-topk-softmax-backends
Aug 7, 2026
Merged

feat(cuda): extend topk_softmax backend coverage#891
voltjia merged 1 commit into
masterfrom
feat/extend-topk-softmax-backends

Conversation

@voltjia

@voltjia voltjia commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Extend the canonical TopkSoftmax provider to Iluvatar, MetaX, and Moore.
  • Move the NVIDIA CUDA kernel into the shared CUDA path while retaining NVIDIA device-guard handling.

Motivation

InfiniCore still uses the deprecated TopkSoftmaxInfinilm compatibility operator on CUDA-compatible backends. Shared canonical provider coverage is required before that adapter can migrate to TopkSoftmax.

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

Validated on ssh nvidia in accelerator-dev/nvidia:latest as part of the canonical-provider integration stack:

$ scripts/dev/build.sh nvidia --jobs 16
Build succeeded

$ python -m pytest tests/test_topk_softmax.py --devices nvidia -q
46 passed

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes build passed 46 focused tests passed
Iluvatar Yes not run platform hardware unavailable; platform CI required
MetaX Yes not run platform hardware unavailable; platform CI required
Cambricon No N/A N/A - not affected
Moore Yes not run platform hardware unavailable; platform CI required
Ascend No N/A N/A - not affected

Benchmark / Performance Impact

N/A. This PR relocates the existing kernel for reuse and adds dispatch providers.

Notes for Reviewers

The public schema is unchanged. This PR only shares the existing implementation and fills provider coverage.

InfiniOps API Alignment target Evidence
TopkSoftmax(gating_output, bias, is_padding, renormalize, topk_weights, topk_indices, token_expert_indices) vLLM topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize=False, e_score_correction_bias=None, is_padding=None) vLLM Python API at pinned commit

InfiniOps groups input tensors, attributes, and explicit outputs according to its C++ operator convention. Thus bias maps to vLLM's e_score_correction_bias, topk_indices maps to topk_ids, and the output tensors are trailing instead of leading. No overload is added.

The NVIDIA result validates the relocated shared kernel and NVIDIA wrapper. Iluvatar, MetaX, and Moore platform CI remain required.

@voltjia
voltjia requested a review from a team August 6, 2026 11:37
@voltjia voltjia changed the title feat(cuda): extend topk_softmax backend coverage feat(cuda): extend topk_softmax backend coverage Aug 7, 2026
@voltjia
voltjia merged commit 21b07eb into master Aug 7, 2026
18 of 20 checks passed
@voltjia
voltjia deleted the feat/extend-topk-softmax-backends branch August 7, 2026 11:22
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