Skip to content

feat(nvidia): add moe_align_block_size operator#804

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

feat(nvidia): add moe_align_block_size operator#804
voltjia merged 1 commit into
masterfrom
feat/nvidia-moe-align-block-size

Conversation

@voltjia

@voltjia voltjia commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add the NVIDIA moe_align_block_size operator and generated Python binding.
  • Support the vLLM low-level ABI with optional global-to-local expert_map, deterministic expert grouping, block padding, non-default streams, and multi-GPU device guards.
  • Add focused correctness, invalid-ID, sparse/boundary, repeatability, descriptor-reuse, stream, and device tests.

Motivation

moe_align_block_size prepares routed token indices and expert block IDs for blockwise MoE GEMM. It is an independent vLLM MoE primitive needed while moving serving kernels behind InfiniOps.

This PR aligns with the low-level stable vLLM operator, not the higher-level Python allocation helper. InfiniOps keeps its repository-wide execution convention of input tensors first, scalar attributes next, and mutable outputs last.

Alignment reference (pinned to vLLM commit 2f75e7f):

InfiniOps execution API vLLM low-level interface Evidence
moe_align_block_size(topk_ids, num_experts, block_size, sorted_token_ids, experts_ids, num_tokens_post_pad)
moe_align_block_size(topk_ids, expert_map, num_experts, block_size, sorted_token_ids, experts_ids, num_tokens_post_pad)
moe_align_block_size(topk_ids, num_experts, block_size, sorted_token_ids, experts_ids, num_tokens_post_pad, maybe_expert_map) Stable C++ declaration, CUDA semantics, High-level allocation helper

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

The operator implementation and focused test at rebased commit 63296a2b are byte-for-byte identical to validated commit 53a686af. That tree was tested from a clean archive on ssh nvidia using accelerator-dev/nvidia:latest (sha256:dd94fce2f83a...), NVIDIA A100-SXM4-80GB, and PyTorch 2.10.0a0+b4e4ee8.

After #803 merged, this PR no longer adds moe_align_block_size to the permanent smoke build or smoke matcher. The existing moe_sum case represents the vLLM MoE integration in the routine smoke suite; this operator remains covered by its focused test file.

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_align_block_size.py -q
18 passed, 1 skipped in 36.58s

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

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes operator-pruned focused build passed; vLLM MoE smoke is represented by moe_sum from #803 focused moe_align_block_size: 18 passed, 1 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 checks exact deterministic token ordering, five repeated launches, invalid raw and mapped expert IDs, sparse and zero-token experts, more than 256 routes, the num_experts=257 / block_size=128 boundary, non-default streams, multi-GPU device selection, and safe descriptor reuse. Descriptor reuse permits fresh data pointers with identical metadata and rejects changed attributes, tensor metadata, or expert_map presence before launch.

Benchmark / Performance Impact

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

Notes for Reviewers

  • Review the low-level API mapping and output capacity rules against the pinned vLLM implementation above.
  • InfiniOps places optional tensor inputs before scalar attributes and mutable outputs last; the semantic arguments and defaults are otherwise the low-level vLLM contract.
  • Invalid global or mapped expert IDs are ignored before indexing. Output tails are initialized to the vLLM-compatible token sentinel / -1 expert sentinel.
  • The stable scatter is a single-writer O(N) phase. It avoids nondeterministic atomic placement, but very large route counts may warrant a separately benchmarked parallel optimization later.
  • Scope is intentionally limited to this one NVIDIA operator. moe_sum was merged independently in feat(nvidia): add moe_sum operator #803 and supplies the vLLM MoE smoke coverage.

@voltjia
voltjia force-pushed the feat/nvidia-moe-align-block-size branch 3 times, most recently from b04ef7c to 53a686a Compare July 23, 2026 00:15
@voltjia voltjia changed the title feat(nvidia): add moe_align_block_size operator feat(nvidia): add moe_align_block_size operator Jul 23, 2026
@voltjia
voltjia marked this pull request as ready for review July 23, 2026 02:22
@voltjia
voltjia requested a review from a team July 23, 2026 02:22
@voltjia
voltjia force-pushed the feat/nvidia-moe-align-block-size branch from 53a686a to 63296a2 Compare July 23, 2026 02:33
@voltjia
voltjia merged commit fd15321 into master Jul 23, 2026
20 checks passed
@voltjia
voltjia deleted the feat/nvidia-moe-align-block-size branch July 23, 2026 02:47
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