Skip to content

feat(nvidia): link vLLM fused_marlin_moe - #914

Draft
voltjia wants to merge 1 commit into
fix/linked-multiple-operator-schemasfrom
feat/linked-fused-marlin-moe
Draft

feat(nvidia): link vLLM fused_marlin_moe#914
voltjia wants to merge 1 commit into
fix/linked-multiple-operator-schemasfrom
feat/linked-fused-marlin-moe

Conversation

@voltjia

@voltjia voltjia commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add the NVIDIA linked slot-16 implementation of vLLM fused_marlin_moe with one complete upstream-aligned overload and InfiniOps' trailing out convention.
  • Compose the provider's registered _moe_C::moe_align_block_size and _moe_C::moe_wna16_marlin_gemm operators with ATen SiLU/multiply and reduction, without copying provider kernels.
  • Cover out-of-place, in-place aliasing, expert mapping, empty activation-order metadata, non-default streams, and linked capability gates in the existing single operator test module.

Motivation

On Ampere, vLLM uses fused_marlin_moe as a complete quantized MoE path. Its public operator is registered by Python as a composite rather than exported by the provider DSO, so InfiniOps cannot honestly bind a same-named shared-library symbol. This PR reconstructs the pinned composite from the two exact Dispatcher contracts supplied by the installed vLLM _moe_C library and standard ATen operations.

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

Remote environment: A100-SXM4-80GB, vllm/vllm-openai:v0.10.0 (vLLM 0.10.1.dev1+gbcc0a3cbe, PyTorch 2.7.1+cu128), CUDA 12.8, CUDA-12-compatible InfiniRT.

# Clean Release/NDEBUG, SM80, WITH_LINKED=ON, smoke union plus fused_marlin_moe
linked resolver tests: 32 passed
tests/test_fused_marlin_moe.py: 7 passed
pytest tests -m smoke -q: 78 passed, 22 skipped, 2 failed, 7809 deselected

The two smoke failures are the same float32 GEMM tolerance cases already reproduced on a fresh origin/master build in the identical environment while validating the preceding Marlin primitive PR. No additional smoke failure remains after running from the correct generated-header working directory.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Clean Release linked build and focused operator tests passed; smoke matched the established master baseline Full suite not run; validation was intentionally operator-pruned
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
Additional focused validation
clang-format 21.1.8 dry run: passed
focused strict wrapper generation: passed
real provider schemas and CUDA dispatch kernels: passed
FP4, global_scale, and undersized workspace rejection: passed
out-of-place, in-place alias, identity expert_map, empty [0] act-order, and non-default stream: passed
git diff --check: passed

Benchmark / Performance Impact

N/A - this PR establishes the composite interface and correctness path and makes no performance claim.

Notes for Reviewers

  • The public argument order follows vLLM commit bcc0a3cb, with the InfiniOps-required output moved to the end.
  • vLLM's high-level vllm::fused_marlin_moe appears only after Python registration; it is not a DSO contract. The binding therefore declares the exact low-level align and Marlin schemas that its C++ source actually calls.
  • inplace=true requires out to alias hidden_states; inplace=false requires a distinct output. The linked boundary checks this in Release builds.
  • The pinned high-level qtype set is enforced independently from the broader low-level primitive. FP4 and global scales are rejected before tensor conversion because InfiniRT cannot currently represent the required float8 scales.
  • This PR is stacked on fix(linked): support multiple operator schemas #913, which extends the existing operator_schema YAML field to validate both low-level Dispatcher contracts from one DSO.

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