feat(nvidia): link vLLM fused_marlin_moe - #914
Draft
voltjia wants to merge 1 commit into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
fused_marlin_moewith one complete upstream-aligned overload and InfiniOps' trailingoutconvention._moe_C::moe_align_block_sizeand_moe_C::moe_wna16_marlin_gemmoperators with ATen SiLU/multiply and reduction, without copying provider kernels.Motivation
On Ampere, vLLM uses
fused_marlin_moeas 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_Clibrary and standard ATen operations.N/A - no linked issue.
Type of Change
feat- new feature / new operator / new platformfix- bug fixperf- performance improvement (no behavioral change)refactor- code restructuring without behavior changetest- adding or fixing tests onlydocs- documentation onlybuild/ci- build system or CI configurationchore- tooling, formatting, or other non-code changesPlatforms Affected
WITH_CPU)WITH_NVIDIA)WITH_ILUVATAR)WITH_METAX)WITH_CAMBRICON)WITH_MOORE)WITH_ASCEND)WITH_TORCH)Smoke Test Result
Remote environment: A100-SXM4-80GB,
vllm/vllm-openai:v0.10.0(vLLM0.10.1.dev1+gbcc0a3cbe, PyTorch2.7.1+cu128), CUDA 12.8, CUDA-12-compatible InfiniRT.The two smoke failures are the same float32 GEMM tolerance cases already reproduced on a fresh
origin/masterbuild 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
Additional focused validation
Benchmark / Performance Impact
N/A - this PR establishes the composite interface and correctness path and makes no performance claim.
Notes for Reviewers
bcc0a3cb, with the InfiniOps-required output moved to the end.vllm::fused_marlin_moeappears 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=truerequiresoutto aliashidden_states;inplace=falserequires a distinct output. The linked boundary checks this in Release builds.operator_schemaYAML field to validate both low-level Dispatcher contracts from one DSO.