Skip to content

[CUDA] MatMulBlockQuantizedFp8Weight: fold the W8A8 activation QDQ into the decode GEMV #31481

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tianleiwu wants to merge 2 commits into
base: main
Choose a base branch
from
+59 −7
Open

[CUDA] MatMulBlockQuantizedFp8Weight: fold the W8A8 activation QDQ into the decode GEMV #31481

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7e74361
[CUDA] MatMulBlockQuantizedFp8Weight: fold the W8A8 activation QDQ in…
tianleiwu Aug 2, 2026
d828726
fix(cuda): use ParseEnvironmentVariableWithDefault instead of std::ge…
tianleiwu Aug 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,21 @@

#include "core/common/safeint.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/platform/env_var_utils.h"
#include "core/providers/cuda/shared_inc/fpgeneric.h"
#include "core/providers/cpu/math/matmul_helper.h"

namespace onnxruntime::contrib::cuda {
using namespace onnxruntime::cuda;

namespace {
bool FusedFp8ActivationQdqDisabled() {
static const bool disabled =
ParseEnvironmentVariableWithDefault<bool>("ORT_DISABLE_FUSED_FP8_ACT_QDQ", false);
return disabled;
}
} // namespace

#if !defined(DISABLE_FLOAT8_TYPES)
ONNX_OPERATOR_KERNEL_EX(
MatMulBlockQuantizedFp8Weight,
Expand Down Expand Up @@ -87,9 +96,16 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
// Optional W8A8 activation path: statically quantize A to FP8 E4M3 and dequantize back so the
// GEMM sees the same activation rounding as native W8A8 execution. When a_scale is absent the
// activation is kept at full FP16/BF16 precision (weight-only W8A16).
//
// The decode GEMV absorbs this round trip in-register, so the standalone kernel (and its [M, K]
// scratch buffer) is only materialized for the cuBLAS path below.
constexpr int kGemvMaxM = 8;
const bool use_gemv = m_i > 0 && m_i <= kGemvMaxM && (k_i % 16 == 0) && (block_size_ % 16 == 0);
const bool fuse_act_qdq = use_gemv && !FusedFp8ActivationQdqDisabled();

const void* a_ptr = a->DataRaw();
IAllocatorUniquePtr<CudaT> a_dequant;
if (a_scale != nullptr) {
if (a_scale != nullptr && !fuse_act_qdq) {
a_dequant = GetScratchBuffer<CudaT>(SafeInt<size_t>(m_i) * SafeInt<size_t>(k_i),
GetComputeStream(context));
ORT_RETURN_IF_ERROR(LaunchQuantizeDequantizeActivationFp8(
Expand All @@ -106,14 +122,14 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
// Decode fast path: for small M (autoregressive generation) this is a memory-bound GEMV.
// A fused warp-per-column kernel reads the FP8 weight directly, avoiding both the [N, K]
// dequant scratch buffer and the cuBLAS GEMM (which is underutilized at M == 1).
constexpr int kGemvMaxM = 8;
if (m_i > 0 && m_i <= kGemvMaxM && (k_i % 16 == 0) && (block_size_ % 16 == 0)) {
if (use_gemv) {
return LaunchMatMulBlockScaledFp8Gemv(
Y->MutableDataRaw(),
a_ptr,
b->DataRaw(),
b_scale->Data<float>(),
bias != nullptr ? bias->DataRaw() : nullptr,
(a_scale != nullptr && fuse_act_qdq) ? a_scale->Data<float>() : nullptr,
m_i,
n_i,
k_i,
Expand Down
40 changes: 36 additions & 4 deletions onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,19 @@ __global__ void QuantizeDequantizeActivationFp8Kernel(T* __restrict__ out,
out[idx] = from_float<T>(q * scale);
}

// In-register form of QuantizeDequantizeActivationFp8Kernel for the 8 activations packed in one
// uint4. Bit-identical to the standalone kernel, so the GEMV can absorb the W8A8 activation
// rounding instead of round-tripping A through a scratch buffer.
template <typename AType>
__device__ __forceinline__ void Fp8ActQdq16(uint4& raw, float inv_scale, float scale) {
AType* v = reinterpret_cast<AType*>(&raw);
#pragma unroll
for (int i = 0; i < 8; ++i) {
const float q = static_cast<float>(__nv_fp8_e4m3(to_float<AType>(v[i]) * inv_scale));
v[i] = from_float<AType>(q * scale);
}
}

// -----------------------------------------------------------------------------
// Fused FP8 weight-only GEMV fast path for the decode phase (small M).
//
Expand Down Expand Up @@ -202,6 +215,7 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output,
const __nv_fp8_e4m3* __restrict__ input_b,
const float* __restrict__ weight_scale,
const AType* __restrict__ bias,
const float* __restrict__ act_scale,
int m,
int n,
int k,
Expand All @@ -219,6 +233,10 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output,
return;
}

const bool act_qdq = act_scale != nullptr;
const float a_scale = act_qdq ? *act_scale : 0.f;
const float a_inv_scale = a_scale != 0.f ? 1.0f / a_scale : 0.f;

float acc[RowsPerWarp][ColsPerWarp] = {};

for (int base = 0; base < k; base += kStride * Unroll) {
Expand Down Expand Up @@ -247,6 +265,10 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output,
const uint4* ap = reinterpret_cast<const uint4*>(input_a + static_cast<size_t>(row) * k + koff[u]);
a_lo[u][r] = ap[0];
a_hi[u][r] = ap[1];
if (act_qdq) {
Fp8ActQdq16<AType>(a_lo[u][r], a_inv_scale, a_scale);
Fp8ActQdq16<AType>(a_hi[u][r], a_inv_scale, a_scale);
}
} else {
a_lo[u][r] = make_uint4(0, 0, 0, 0);
a_hi[u][r] = make_uint4(0, 0, 0, 0);
Expand Down Expand Up @@ -497,13 +519,18 @@ __global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output,
const __nv_fp8_e4m3* __restrict__ input_b,
const float* __restrict__ weight_scale,
const AType* __restrict__ bias,
const float* __restrict__ act_scale,
int m,
int n,
int k,
int block_size,
int k_blocks) {
using Mma = Fp8GemvMma<AType>;

const bool act_qdq = act_scale != nullptr;
const float a_scale = act_qdq ? *act_scale : 0.f;
const float a_inv_scale = a_scale != 0.f ? 1.0f / a_scale : 0.f;

const int lane = threadIdx.x;
const int warp = threadIdx.y;
// See the fragment-ownership table above: `g` indexes the A rows (output columns col_lo/col_hi)
Expand Down Expand Up @@ -547,6 +574,10 @@ __global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output,
const uint4* ap = reinterpret_cast<const uint4*>(input_a + a_off + k0);
a_raw[0] = ap[0];
a_raw[1] = ap[1];
if (act_qdq) {
Fp8ActQdq16<AType>(a_raw[0], a_inv_scale, a_scale);
Fp8ActQdq16<AType>(a_raw[1], a_inv_scale, a_scale);
}
} else {
a_raw[0] = make_uint4(0, 0, 0, 0);
a_raw[1] = make_uint4(0, 0, 0, 0);
Expand Down Expand Up @@ -766,6 +797,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
const void* b_fp8,
const float* weight_scale,
const void* bias,
const float* act_scale,
int m,
int n,
int k,
Expand Down Expand Up @@ -806,11 +838,11 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
if (is_bf16) {
MatMulBlockScaledFp8MmaGemvKernel<KSplit><<<mma_blocks, mma_threads, 0, stream>>>(
reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b,
weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), m, n, k, block_size, k_blocks);
weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k, block_size, k_blocks);
} else {
MatMulBlockScaledFp8MmaGemvKernel<KSplit><<<mma_blocks, mma_threads, 0, stream>>>(
reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,
weight_scale, reinterpret_cast<const half*>(bias), m, n, k, block_size, k_blocks);
weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k, block_size, k_blocks);
}
};
if (k_split == 16) {
Expand All @@ -832,11 +864,11 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
if (is_bf16) {
MatMulBlockScaledFp8GemvKernel<RowsPerWarp, ColsPerWarp, Unroll><<<blocks, threads, 0, stream>>>(
reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b,
weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), m, n, k, block_size, k_blocks);
weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k, block_size, k_blocks);
} else {
MatMulBlockScaledFp8GemvKernel<RowsPerWarp, ColsPerWarp, Unroll><<<blocks, threads, 0, stream>>>(
reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,
weight_scale, reinterpret_cast<const half*>(bias), m, n, k, block_size, k_blocks);
weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k, block_size, k_blocks);
}
};

Expand Down
4 changes: 4 additions & 0 deletions onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,16 @@ Status LaunchQuantizeDequantizeActivationFp8(void* a_out,
// FP8 E4M3, weight_scale is [N, ceil(K/block_size)] fp32, bias is an optional [N] vector (may be
// null). Output y is [M, N] in the activation type. Requires k % 16 == 0 and block_size % 16 == 0.
// device_prop selects the tensor-core (mma.m16n8k16) variant, which needs SM80+.
// act_scale is an optional device fp32 scalar; when non-null the kernel applies the W8A8
// activation quantize/dequantize inline, bit-identically to LaunchQuantizeDequantizeActivationFp8,
// so no scratch buffer or extra launch is needed.
// Runs on any architecture with FP8 conversion intrinsics (CUDA >= 11.8).
Status LaunchMatMulBlockScaledFp8Gemv(void* y,
const void* a,
const void* b_fp8,
const float* weight_scale,
const void* bias,
const float* act_scale,
int m,
int n,
int k,
Expand Down
Loading