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

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc

@@ -3,6 +3,7 @@
 
 #include "contrib_ops/cuda/math/matmul_block_scaled_fp8.h"
 
+#include <cstdlib>
 #include <type_traits>
 
 #include "core/common/safeint.h"
@@ -13,6 +14,16 @@
 namespace onnxruntime::contrib::cuda {
 using namespace onnxruntime::cuda;
 
+namespace {
+bool FusedFp8ActivationQdqDisabled() {
+  static const bool disabled = [] {
+    const char* v = std::getenv("ORT_DISABLE_FUSED_FP8_ACT_QDQ");
+    return v != nullptr && v[0] != '\0' && v[0] != '0';
+  }();
+  return disabled;
+}
+}  // namespace
+
 #if !defined(DISABLE_FLOAT8_TYPES)
 ONNX_OPERATOR_KERNEL_EX(
     MatMulBlockQuantizedFp8Weight,
@@ -87,9 +98,16 @@
   // 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(
@@ -106,14 +124,14 @@
   // 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,

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu

@@ -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).
 //
@@ -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,
@@ -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) {
@@ -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);
@@ -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)
@@ -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);
@@ -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,
@@ -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) {
@@ -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);
     }
   };
 

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h

@@ -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,