[CUDA] Add MatMulBlockQuantizedFp4Weight contrib operator#29818
Open
tianleiwu wants to merge 9 commits into
Open
[CUDA] Add MatMulBlockQuantizedFp4Weight contrib operator#29818tianleiwu wants to merge 9 commits into
tianleiwu wants to merge 9 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new CUDA contrib operator com.microsoft::MatMulBlockScaledFp4 to ONNX Runtime, implementing weight-only NVFP4 (E2M1) matmul with per-block E4M3 scales and an optional opt-in native SM120 (Blackwell) CUTLASS path. It also adds C++ tests, profiling tooling, documentation, and CMake wiring for both the CUDA EP and CUDA plugin EP.
Changes:
- Introduce
MatMulBlockScaledFp4schema + CUDA kernel implementation (dequant+cuBLAS, small-M GEMV, optional native SM120 path). - Register the operator in the MS opset and CUDA contrib kernel registry.
- Add unit tests, a Python profiling harness, and operator documentation/experiments notes.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/python/contrib_ops/profile_matmul_block_scaled.py | Opt-in Python accuracy/latency harness for the new CUDA contrib op. |
| onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc | New CUDA-focused unit tests for FP16/BF16, bias, and GEMV decode path. |
| onnxruntime/core/graph/contrib_ops/ms_opset.h | Adds the operator to the Microsoft opset v1 schema list. |
| onnxruntime/core/graph/contrib_ops/contrib_defs.cc | Defines the MatMulBlockScaledFp4 schema and shape inference. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h | Declares the CUDA kernel class and launcher APIs. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu | Implements NVFP4 dequant, bias add, and small-M fused GEMV kernels. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc | Implements kernel registration, validation, dispatch, and native-path prepacking. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu | Implements the native SM120 CUTLASS NVFP4 x NVFP4 GEMM path. |
| onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc | Registers the new CUDA contrib kernel. |
| docs/contrib_ops/cuda/matmul_block_scaled_fp4.md | Adds operator documentation (format, dispatch chain, env vars, workflows). |
| docs/contrib_ops/cuda/matmul_block_scaled_fp4_experiments.md | Records performance experiments and rationale for rejected variants. |
| cmake/onnxruntime_providers_cuda.cmake | Enables SM120 build define needed for native path in the CUDA EP target(s). |
| cmake/onnxruntime_providers_cuda_plugin.cmake | Enables the same SM120 build define for the CUDA plugin EP target. |
| cmake/onnxruntime_cuda_source_filters.cmake | Ensures the SM120-specific .cu file is compiled only when arch 120 is enabled. |
tianleiwu
marked this pull request as draft
July 22, 2026 21:10
Add the com.microsoft MatMulBlockScaledFp4 CUDA contrib operator, a weight-only NVFP4 (E2M1) matmul that computes Y = A * dequant(B)^T (+bias) with packed 4-bit weights and per-block E4M3 scales plus a global fp32 scale. Includes an architecture-independent dequant + GEMM/GEMV path and a native SM120 CUTLASS NVFP4 path, the ONNX schema, kernel registration, build wiring, unit tests, docs and a profiling script.
… GEMV guards, doc env vars - Replace __builtin_clz with a portable next-power-of-two loop so MSVC host builds compile. - Validate exact ranks/dims of B, weight_scale, and bias in ComputeImpl instead of only total element counts. - Enforce rank-2 [N, K/16] weight_scale shape in PrePack. - Add block_size==16 and K%32==0 runtime guards in LaunchMatMulBlockScaledFp4Gemv. - Replace hard-coded absolute paths in docs with ORT_REPO/ORT_BUILD variables.
tianleiwu
force-pushed
the
tlwu/20260721/matmul_block_scaled_fp4
branch
from
July 23, 2026 23:47
fbb7a23 to
8c3a850
Compare
tianleiwu
marked this pull request as ready for review
July 24, 2026 06:54
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.
Description
Adds a new
com.microsoftCUDA contrib operatorMatMulBlockQuantizedFp8Weight: aweight-only NVFP4 (E2M1) matrix multiplication that computes
Y = A * dequant(B)^T (+ bias).The weight tensor
Bis stored as packed NVFP4 (two E2M1 values per byte, lownibble first). The dequantized weight is
e2m1(B) * weight_scale_2 * e4m3(weight_scale[n, k / block_size]), whereweight_scaleholds one E4M3 scale perblock_size(default 16) consecutive Kvalues and
weight_scale_2is a single global fp32 scale. The weight isdequantized to the activation type (FP16/BF16) and multiplied with the FP16/BF16
activation.
Operator schema
A(FP16/BF16[..., K]),B(packed uint8 NVFP4[N, K/2]),weight_scale(uint8 E4M3[N, ceil(K/block_size)]),weight_scale_2(scalar fp32), optional
input_scale(scalar fp32), optionalbias([N]).K,N,block_size(default 16).Y([..., N]) in the activation type.Implementation
and Blackwell.
ORT_ENABLE_BLOCKQUANT_SM120).the CUDA plugin EP.
Motivation and Context
Enables weight-only NVFP4 quantized matmul for LLM inference, reducing weight
memory footprint while keeping FP16/BF16 activation precision.
Testing
matmul_block_scaled_fp4_test.cc) covering FP16 andBF16 activations, with and without bias.
(
profile_matmul_block_scaled.py).onnxruntime_provider_test) builds cleanlywith
ENABLE_FP4for SM86 + SM120a.