Skip to content

feat(nvidia): add moe_sum operator#803

Merged
voltjia merged 1 commit into
masterfrom
feat/nvidia-moe-sum
Jul 23, 2026
Merged

feat(nvidia): add moe_sum operator#803
voltjia merged 1 commit into
masterfrom
feat/nvidia-moe-sum

Conversation

@voltjia

@voltjia voltjia commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add the NVIDIA moe_sum operator and generated Python binding.
  • Support FP32, FP16, and BF16 inputs, optional topk_ids / expert_map routing filters, non-negative strided inputs, non-default streams, and multi-GPU device guards.
  • Add focused correctness, descriptor-reuse, routing, stride, stream, and device tests.

Motivation

moe_sum is a representative vLLM MoE primitive needed while moving third-party serving kernels behind InfiniOps. The public name and parameter semantics follow vLLM's stable custom-op schema; InfiniOps keeps its repository-wide convention of input tensors first and the mutable output last.

Alignment reference (pinned to vLLM commit 2f75e7f):

InfiniOps execution API vLLM schema Evidence
moe_sum(input, output)
moe_sum(input, topk_ids, expert_map, output)
moe_sum(Tensor input, Tensor! output, Tensor? topk_ids=None, Tensor? expert_map=None) -> () Stable schema, CUDA semantics, Python wrapper

N/A - no linked issue.

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

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

Commit 264501126626b7f3dbf0a05c9cc3ec33dc8a4852 was exported as a clean Git archive and tested on ssh nvidia using accelerator-dev/nvidia:latest (sha256:dd94fce2f83a...), NVIDIA A100-SXM4-80GB, and PyTorch 2.10.0a0+b4e4ee8.

INFINI_RT_ROOT=/opt/infinirt scripts/dev/build.sh nvidia --smoke --jobs 16 --torch-jobs 2 --binding-jobs 2
[build] done

PYTHONPATH=build-nvidia/install python3 -m pytest tests/test_moe_sum.py -q
31 passed, 6 skipped in 21.81s

PYTHONPATH=build-nvidia/install python3 -m pytest tests -m smoke -q
77 passed, 14 skipped, 12331 deselected in 9.63s

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes operator-pruned build and smoke passed focused moe_sum: 31 passed, 6 skipped; full suite intentionally not run
Iluvatar No N/A - not affected N/A
MetaX No N/A - not affected N/A
Cambricon No N/A - not affected N/A
Moore No N/A - not affected N/A
Ascend No N/A - not affected N/A
Focused coverage

The focused suite covers both schema overloads, FP32/FP16/BF16 accumulation, int32/int64 routing IDs, negative and remote expert filtering, positive non-contiguous strides, non-default streams, multi-GPU device selection, and safe descriptor reuse. Descriptor reuse permits fresh data pointers with identical metadata and rejects changed shape, stride, dtype, device, or optional-argument presence before launch.

Benchmark / Performance Impact

N/A - this PR establishes the aligned interface and correctness baseline; it makes no performance claim.

Notes for Reviewers

  • Review the API mapping and routing semantics against the pinned vLLM schema and CUDA implementation above.
  • InfiniOps places mutable output tensors last; this is the only ordering adaptation from the vLLM dispatcher schema.
  • Kernel indexing is signed 64-bit. The descriptor rejects negative strides and overflow-prone metadata while preserving non-negative non-contiguous layouts.
  • Scope is intentionally limited to one representative NVIDIA MoE operator. moe_align_block_size is being proposed independently.

@voltjia voltjia changed the title feat(nvidia): add moe_sum operator feat(nvidia): add moe_sum operator Jul 23, 2026
@voltjia
voltjia marked this pull request as ready for review July 23, 2026 02:06
@voltjia
voltjia requested a review from a team July 23, 2026 02:06
@voltjia
voltjia merged commit 4cffc76 into master Jul 23, 2026
20 checks passed
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