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

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

PlatformAffectedBuild / Smoke ResultFull Result / Notes
NVIDIAYesbuild passed46 focused tests passed
IluvatarYesnot runplatform hardware unavailable; platform CI required
MetaXYesnot runplatform hardware unavailable; platform CI required
CambriconNoN/AN/A - not affected
MooreYesnot runplatform hardware unavailable; platform CI required
AscendNoN/AN/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 APIAlignment targetEvidence
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 teamAugust 6, 2026 11:37
@voltjiavoltjia changed the title feat(cuda): extend topk_softmax backend coveragefeat(cuda): extend topk_softmax backend coverageAug 7, 2026
@voltjia
voltjia merged commit 21b07eb into masterAug 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 freeto 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

@voltjia