Skip to content

feat(nvidia): link vLLM moe_wna16_marlin_gemm - #912

Merged
voltjia merged 1 commit into
masterfrom
feat/linked-moe-wna16-marlin-gemm
Aug 9, 2026
Merged

feat(nvidia): link vLLM moe_wna16_marlin_gemm#912
voltjia merged 1 commit into
masterfrom
feat/linked-moe-wna16-marlin-gemm

Conversation

@voltjia

@voltjia voltjia commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add the NVIDIA linked implementation of vLLM _moe_C::moe_wna16_marlin_gemm in slot 16.
  • Align the complete interface with vLLM v0.26.0 and current main: add b_bias_or_none, a_scales, and thread_k/thread_n/blocks_per_sm; remove is_ep; rename the low-level type id to b_type_id.
  • Preserve provider argument order while adapting its optional output buffer to InfiniOps' required trailing out convention.
  • Cover biased 4-bit/8-bit, activation-order, zero-point, output-alias, non-default-stream, INT8 activation-scale, explicit tuning, and noncanonical top-k-weight shape behavior against the resolved provider.

Motivation

vLLM uses this Marlin grouped GEMM primitive in its Ampere-compatible fused MoE path. Linking the installed provider keeps the kernel in the provider DSO and avoids copying another vendor kernel into InfiniOps.

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

Focused Release/NDEBUG rebuild and smoke revalidation for commit 3173d651 is in progress. The NVIDIA SSH endpoint became unavailable during setup; this section will be replaced with final command output rather than retaining the obsolete vLLM v0.10 result.

The current-schema provider itself has already been probed on A100:

provider image: docker.1ms.run/verlai/verl:vllm020.dev2
vLLM 0.20.2, PyTorch 2.11.0+cu130
exact Dispatcher schema: matched
CUDA kernel: registered
FP16 + non-null bias + tuning (128,128,1): passed, max error 4.88e-4
INT8 + a_scales[1,1] + non-null bias: passed, max error 8.12e-3

This runtime provider is vLLM 0.20.2, not the source alignment baseline; its registered schema is byte-for-byte equivalent to v0.26.0/current for this operator.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Focused Release rebuild pending remote availability Full suite will not be run; validation is 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
Completed local validation
focused resolver and wrapper tests: 49 passed
focused wrapper generation: passed
Ruff 0.15.22 check and format: passed
clang-format 21.1.8 dry run: passed
Python syntax check: passed
git diff --check: passed

Benchmark / Performance Impact

N/A - this PR exposes the installed provider implementation and makes no performance claim.

Notes for Reviewers

  • The public contract is pinned to vLLM v0.26.0 commit 568afb3a and was rechecked against main commit 83ad767e.
  • vLLM's c_or_none is output storage. InfiniOps requires a trailing out, passes it in the provider's second position, and verifies that the returned Tensor aliases it.
  • The upstream Python defaults for thread_k/thread_n/blocks_per_sm remain explicit -1 arguments here; no convenience overload is added.
  • float4_e2m1f, global_scale, and FP8 activation are rejected before dispatch because the current InfiniRT type system cannot represent the required FP8 tensors.
  • This PR adds the stable low-level Dispatcher primitive only. The high-level Python fused_marlin_moe composite is outside this PR and requires separate redesign.

@voltjia
voltjia force-pushed the feat/linked-moe-wna16-marlin-gemm branch from 3f89ad2 to 3173d65 Compare August 9, 2026 11:10
@voltjia
voltjia marked this pull request as ready for review August 9, 2026 13:32
@voltjia
voltjia requested a review from a team August 9, 2026 13:32
@voltjia
voltjia merged commit 47c1c49 into master Aug 9, 2026
18 of 20 checks passed
@voltjia
voltjia deleted the feat/linked-moe-wna16-marlin-gemm branch August 9, 2026 13:32
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