Skip to content

[CUDA] Add MatMulBlockQuantizedFp4Weight contrib operator#29818

Open
tianleiwu wants to merge 9 commits into
mainfrom
tlwu/20260721/matmul_block_scaled_fp4
Open

[CUDA] Add MatMulBlockQuantizedFp4Weight contrib operator#29818
tianleiwu wants to merge 9 commits into
mainfrom
tlwu/20260721/matmul_block_scaled_fp4

Conversation

@tianleiwu

@tianleiwu tianleiwu commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

Description

Adds a new com.microsoft CUDA contrib operator MatMulBlockQuantizedFp8Weight : a
weight-only NVFP4 (E2M1) matrix multiplication that computes
Y = A * dequant(B)^T (+ bias).

The weight tensor B is stored as packed NVFP4 (two E2M1 values per byte, low
nibble first). The dequantized weight is
e2m1(B) * weight_scale_2 * e4m3(weight_scale[n, k / block_size]), where
weight_scale holds one E4M3 scale per block_size (default 16) consecutive K
values and weight_scale_2 is a single global fp32 scale. The weight is
dequantized to the activation type (FP16/BF16) and multiplied with the FP16/BF16
activation.

Operator schema

  • Inputs: 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), optional bias ([N]).
  • Attributes: K, N, block_size (default 16).
  • Output: Y ([..., N]) in the activation type.

Implementation

  • Architecture-independent dequant + GEMM/GEMV path that runs on Hopper (SM90)
    and Blackwell.
  • Native SM120 CUTLASS NVFP4 path (guarded by ORT_ENABLE_BLOCKQUANT_SM120).
  • ONNX schema, kernel registration, and CMake build wiring for the CUDA EP and
    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

  • Added C++ unit tests (matmul_block_scaled_fp4_test.cc) covering FP16 and
    BF16 activations, with and without bias.
  • Added a Python profiling/verification script
    (profile_matmul_block_scaled.py).
  • Verified the CUDA EP test binary (onnxruntime_provider_test) builds cleanly
    with ENABLE_FP4 for SM86 + SM120a.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 MatMulBlockScaledFp4 schema + 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.

Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu
Comment thread docs/contrib_ops/cuda/matmul_block_scaled_fp4.md
Comment thread docs/contrib_ops/cuda/matmul_block_scaled_fp4_experiments.md
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
tianleiwu force-pushed the tlwu/20260721/matmul_block_scaled_fp4 branch from fbb7a23 to 8c3a850 Compare July 23, 2026 23:47
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc Fixed
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu Fixed
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h Fixed
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu Fixed

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h Outdated
Comment thread onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h Outdated
@tianleiwu tianleiwu changed the title [CUDA] Add MatMulBlockScaledFp4 contrib operator [CUDA] Add MatMulBlockQuantizedFp4Weight contrib operator Jul 24, 2026
@tianleiwu
tianleiwu marked this pull request as ready for review July 24, 2026 06:54
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.

3 participants