[WebGPU] Support block-quantized FP8 dense MatMul

onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#use .getByOffset .setByOffset
+
+const E4M3_MAX: f32 = 448.0;
+
+$MAIN {
+  if (global_idx >= uniforms.size) {
+    return;
+  }
+  let scale = input_scale.getByOffset(0u);
+  let inverse = select(0.0, 1.0 / scale, scale != 0.0);
+  let input_value = f32(input_a.getByOffset(global_idx));
+  let v = input_value * inverse;
+  let magnitude = min(abs(v), E4M3_MAX);
+  let magnitude_exponent = (bitcast<u32>(magnitude) >> 23u) & 255u;
+  let step = bitcast<f32>((max(magnitude_exponent, 121u) - 3u) << 23u);
+  let units = magnitude / step;
+  let lower = floor(units);
+  let halfway = (lower + 0.5) * step;
+  let round_up = magnitude > halfway ||
+                 (magnitude == halfway && (u32(lower) & 1u) != 0u);
+  let rounded = min((lower + select(0.0, 1.0, round_up)) * step, E4M3_MAX);
+  let q = select(rounded, -rounded, v < 0.0);
+  output.setByOffset(global_idx, f16(q * scale));
+}

onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc

@@ -0,0 +1,216 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h"
+
+#include "contrib_ops/webgpu/webgpu_contrib_kernels.h"
+#include "core/common/narrow.h"
+#include "core/providers/cpu/math/matmul_helper.h"
+#include "core/providers/webgpu/math/subgroup_matrix_config.h"
+#include "core/providers/webgpu/shader_helper.h"
+#include "core/providers/webgpu/webgpu_supported_types.h"
+#include "core/providers/webgpu/webgpu_utils.h"
+#if !defined(DISABLE_FLOAT8_TYPES)
+#include "core/common/float8.h"
+#endif
+
+namespace onnxruntime {
+namespace contrib {
+namespace webgpu {
+
+using namespace onnxruntime::webgpu;
+
+namespace {
+
+class Fp8ActivationProgram final : public Program<Fp8ActivationProgram> {
+ public:
+  Fp8ActivationProgram() : Program{"Fp8MatMulActivationQdq"} {}
+
+  Status GenerateShaderCode(ShaderHelper& shader) const override {
+    const auto& input_a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias);
+    const auto& input_scale = shader.AddInput("input_scale", ShaderUsage::UseValueTypeAlias);
+    const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias);
+    return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_block_quantized_fp8_activation.wgsl.template",
+                               WGSL_TEMPLATE_VARIABLE(input_a, input_a),
+                               WGSL_TEMPLATE_VARIABLE(input_scale, input_scale),
+                               WGSL_TEMPLATE_VARIABLE(output, output));
+  }
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32});
+};
+
+class Fp8EmptyReductionProgram final : public Program<Fp8EmptyReductionProgram> {
+ public:
+  explicit Fp8EmptyReductionProgram(bool has_bias)
+      : Program{"Fp8MatMulEmptyReduction"}, has_bias_(has_bias) {}
+
+  Status GenerateShaderCode(ShaderHelper& shader) const override {
+    if (has_bias_) {
+      const auto& bias = shader.AddInput("bias");
+      const auto& output = shader.AddOutput("output");
+      shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.size")
+                                << output.SetByOffset("global_idx", bias.GetByOffset("global_idx % uniforms.N"));
+    } else {
+      const auto& output = shader.AddOutput("output");
+      shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.size")
+                                << output.SetByOffset("global_idx", "f16(0)");
+    }
+    return Status::OK();
+  }
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32},
+                                          {"N", ProgramUniformVariableDataType::Uint32});
+
+ private:
+  bool has_bias_;
+};
+
+class Fp8MatMulProgram final : public Program<Fp8MatMulProgram> {
+ public:
+  Fp8MatMulProgram(bool has_bias, bool use_matrix)
+      : Program{"MatMulBlockQuantizedFp8Weight"}, has_bias_(has_bias), use_matrix_(use_matrix) {}
+
+  Status GenerateShaderCode(ShaderHelper& shader) const override {
+    const auto& input_a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias);
+    const auto& weights = shader.AddInput("weights", ShaderUsage::UseValueTypeAlias);
+    const auto& scales = shader.AddInput("scales", ShaderUsage::UseValueTypeAlias);
+    const ShaderVariableHelper* bias = nullptr;
+    if (has_bias_) {
+      bias = &shader.AddInput("bias", ShaderUsage::UseValueTypeAlias);
+    }
+    const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias);
+    return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_block_quantized_fp8_weight.wgsl.template",
+                               WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_),
+                               WGSL_TEMPLATE_PARAMETER(use_matrix, use_matrix_),
+                               WGSL_TEMPLATE_OPTIONAL_VARIABLE(bias, bias),
+                               WGSL_TEMPLATE_VARIABLE(input_a, input_a),
+                               WGSL_TEMPLATE_VARIABLE(output, output),
+                               WGSL_TEMPLATE_VARIABLE(scales, scales),
+                               WGSL_TEMPLATE_VARIABLE(weights, weights));
+  }
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32},
+                                          {"N", ProgramUniformVariableDataType::Uint32},
+                                          {"K", ProgramUniformVariableDataType::Uint32},
+                                          {"K_blocks", ProgramUniformVariableDataType::Uint32},
+                                          {"block_size", ProgramUniformVariableDataType::Uint32});
+
+ private:
+  bool has_bias_;
+  bool use_matrix_;
+};
+
+}  // namespace
+
+MatMulBlockQuantizedFp8Weight::MatMulBlockQuantizedFp8Weight(const OpKernelInfo& info)
+    : WebGpuKernel(info),
+      block_size_(narrow<uint32_t>(info.GetAttrOrDefault<int64_t>("block_size", 128))) {
+  ORT_ENFORCE(block_size_ > 0, "block_size must be positive.");
+}
+
+Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const {
+  const Tensor* a = context.Input(0);
+  const Tensor* b = context.Input(1);
+  const Tensor* scales = context.Input(2);
+  const Tensor* a_scale = context.Input(3);
+  const Tensor* bias = context.Input(4);
+
+  const auto& a_shape = a->Shape();
+  const auto& b_shape = b->Shape();
+  ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1 && b_shape.NumDimensions() == 2,
+                    "A must have rank >= 1 and B must have shape [N, K].");
+  const int64_t n_dim = b_shape[0];
+  const int64_t k_dim = b_shape[1];
+  ORT_RETURN_IF_NOT(a_shape[a_shape.NumDimensions() - 1] == k_dim,
+                    "A and B contraction dimensions must match.");
+  const int64_t k_blocks = (k_dim + block_size_ - 1) / block_size_;
+  ORT_RETURN_IF_NOT(scales->Shape() == TensorShape({n_dim, k_blocks}),
+                    "b_scale must have shape [N, ceil(K/block_size)].");
+  ORT_RETURN_IF_NOT(a_scale == nullptr || a_scale->Shape().Size() == 1,
+                    "a_scale must contain one value.");
+  ORT_RETURN_IF_NOT(bias == nullptr || bias->Shape() == TensorShape({n_dim}),
+                    "bias must have shape [N].");
+
+  MatMulComputeHelper helper;
+  ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b_shape, false, true));
+  Tensor* y = context.Output(0, helper.OutputShape());
+  if (y->Shape().Size() == 0) {
+    return Status::OK();
+  }
+
+  const uint32_t K = narrow<uint32_t>(k_dim);
+  const uint32_t N = narrow<uint32_t>(n_dim);
+  const uint32_t M = narrow<uint32_t>(y->Shape().Size() / n_dim);
+  const uint32_t K_blocks = narrow<uint32_t>(k_blocks);
+  if (K == 0) {
+    Fp8EmptyReductionProgram empty_program{bias != nullptr};
+    empty_program.SetWorkgroupSize(64);
+    empty_program.SetDispatchGroupSize((narrow<uint32_t>(y->Shape().Size()) + 63u) / 64u);
+    if (bias) {
+      empty_program.AddInput({bias, ProgramTensorMetadataDependency::Type});
+    }
+    empty_program.AddOutput({y, ProgramTensorMetadataDependency::Type})
+        .AddUniformVariables({{narrow<uint32_t>(y->Shape().Size())}, {N}});
+    return context.RunProgram(empty_program);
+  }
+  const Tensor* activation = a;
+  Tensor qdq;
+  if (a_scale) {
+    qdq = context.CreateGPUTensor(a->DataType(), a_shape);
+    Fp8ActivationProgram qdq_program;
+    qdq_program.SetWorkgroupSize(64);
+    qdq_program.SetDispatchGroupSize((narrow<uint32_t>(a_shape.Size()) + 63u) / 64u);
+    qdq_program.AddInputs({{a, ProgramTensorMetadataDependency::Type},
+                           {a_scale, ProgramTensorMetadataDependency::Type}})
+        .AddOutput({&qdq, ProgramTensorMetadataDependency::Type})
+        .AddUniformVariables({{narrow<uint32_t>(a_shape.Size())}});
+    ORT_RETURN_IF_ERROR(context.RunProgram(qdq_program));
+    activation = &qdq;
+  }
+
+  const bool use_matrix = M >= 8 &&
+                          SelectSubgroupMatrixConfig(context, {{wgpu::SubgroupMatrixComponentType::F16,
+                                                                wgpu::SubgroupMatrixComponentType::F32,
+                                                                16, 16, 16, 32, false}})
+                              .has_value();
+
+  auto memory_info = OrtMemoryInfo{
+      WEBGPU_BUFFER,
+      OrtDeviceAllocator,
+      OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0}};
+  // Tensor views require void*, but the alias only binds B as read-only storage; no copy or mutation occurs.
+  Tensor weight_bytes(DataTypeImpl::GetType<uint8_t>(), b_shape, const_cast<void*>(b->DataRaw()), memory_info);
+  Fp8MatMulProgram program{bias != nullptr, use_matrix};
+  program.SetWorkgroupSize(use_matrix ? 128 : 32);
+  if (use_matrix && context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) {
+    program.SetSubgroupSize(32);
+  }
+  program.SetDispatchGroupSize(use_matrix ? (N + 15u) / 16u : N,
+                               use_matrix ? (M + 63u) / 64u : M);
+  // Flatten rounds the byte count up to u32 words; WebGPU buffers are 16-byte padded by BufferManager.
+  // The shader extracts only bytes whose offset is within N*K, even in the final partial word.
+  program.AddInputs({{activation, ProgramTensorMetadataDependency::Type},
+                     {&weight_bytes, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4},
+                     {scales, ProgramTensorMetadataDependency::Type}})
+      .AddOutput({y, ProgramTensorMetadataDependency::Type})
+      .AddUniformVariables({{M}, {N}, {K}, {K_blocks}, {block_size_}});
+  if (bias) {
+    program.AddInput({bias, ProgramTensorMetadataDependency::Type});
+  }
+  return context.RunProgram(program);
+}
+
+#if !defined(DISABLE_FLOAT8_TYPES)
+ONNX_OPERATOR_KERNEL_EX(
+    MatMulBlockQuantizedFp8Weight,
+    kMSDomain,
+    1,
+    kWebGpuExecutionProvider,
+    (*KernelDefBuilder::Create())
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>())
+        .TypeConstraint("T1", DataTypeImpl::GetTensorType<Float8E4M3FN>())
+        .TypeConstraint("T2", DataTypeImpl::GetTensorType<float>()),
+    MatMulBlockQuantizedFp8Weight);
+#endif
+
+}  // namespace webgpu
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/providers/webgpu/program.h"
+#include "core/providers/webgpu/webgpu_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace webgpu {
+
+class MatMulBlockQuantizedFp8Weight final : public onnxruntime::webgpu::WebGpuKernel {
+ public:
+  explicit MatMulBlockQuantizedFp8Weight(const OpKernelInfo& info);
+  Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override;
+
+ private:
+  uint32_t block_size_;
+};
+
+}  // namespace webgpu
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template

@@ -0,0 +1,107 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#param has_bias
+#param use_matrix
+#use .getByOffset .setByOffset
+
+fn decode_fp8(code: u32) -> f32 {
+  let exponent = (code >> 3u) & 15u;
+  let mantissa = code & 7u;
+  let magnitude = select(f32(mantissa) * 0.001953125,
+                         bitcast<f32>(((exponent + 120u) << 23u) | (mantissa << 20u)),
+                         exponent != 0u);
+  let value = select(magnitude, -magnitude, (code & 128u) != 0u);
+  return select(value, bitcast<f32>(0x7fc00000u | code), exponent == 15u && mantissa == 7u);
+}
+
+fn weight(n: u32, k: u32) -> f16 {
+  let offset = n * uniforms.K + k;
+  let code = (weights.getByOffset(offset / 4u) >> ((offset % 4u) * 8u)) & 255u;
+  return f16(decode_fp8(code) * scales.getByOffset(n * uniforms.K_blocks + k / uniforms.block_size));
+}
+
+#if use_matrix
+// The padded workgroup tiles make every subgroup-matrix load and store in bounds.
+var<workgroup> tile_a: array<f16, 1024>;
+var<workgroup> tile_b: array<f16, 256>;
+var<workgroup> tile_y: array<f32, 1024>;
+#else
+var<workgroup> partial: array<f32, 32>;
+#endif
+$MAIN {
+#if use_matrix
+  let lane = local_idx;
+  let n_tiles = (uniforms.N + 15u) / 16u;
+  if (workgroup_idx >= n_tiles * ((uniforms.M + 63u) / 64u)) {
+    return;
+  }
+  let m_base = (workgroup_idx / n_tiles) * 64u;
+  let n_base = (workgroup_idx % n_tiles) * 16u;
+  var acc: subgroup_matrix_result<f32, 16, 16>;
+  for (var tile = 0u; tile < (uniforms.K + 15u) / 16u; tile++) {
+    let k_base = tile * 16u;
+    for (var i = lane; i < 1024u; i += 128u) {
+      let m = m_base + i / 16u;
+      let k = k_base + i % 16u;
+      tile_a[i] = f16(0);
+      if (m < uniforms.M && k < uniforms.K) {
+        tile_a[i] = input_a.getByOffset(m * uniforms.K + k);
+      }
+    }
+    for (var i = lane; i < 256u; i += 128u) {
+      let n = n_base + i / 16u;
+      let k = k_base + i % 16u;
+      tile_b[i] = f16(0);
+      if (n < uniforms.N && k < uniforms.K) {
+        tile_b[i] = weight(n, k);
+      }
+    }
+    workgroupBarrier();
+    let left = subgroupMatrixLoad<subgroup_matrix_left<f16, 16, 16>, row_major>(
+        &tile_a, subgroup_index * 256u, 16u);
+    let right = subgroupMatrixLoad<subgroup_matrix_right<f16, 16, 16>, col_major>(&tile_b, 0u, 16u);
+    acc = subgroupMatrixMultiplyAccumulate(left, right, acc);
+    workgroupBarrier();
+  }
+  subgroupMatrixStore<row_major>(&tile_y, subgroup_index * 256u, acc, 16u);
+  workgroupBarrier();
+  for (var i = lane; i < 1024u; i += 128u) {
+    let m = m_base + i / 16u;
+    let n = n_base + i % 16u;
+    if (m < uniforms.M && n < uniforms.N) {
+      var value = tile_y[i];
+#if has_bias
+      value += f32(bias.getByOffset(n));
+#endif
+      output.setByOffset(m * uniforms.N + n, f16(value));
+    }
+  }
+#else
+  let lane = local_idx;
+  if (workgroup_idx >= uniforms.M * uniforms.N) {
+    return;
+  }
+  let m = workgroup_idx / uniforms.N;
+  let n = workgroup_idx % uniforms.N;
+  var sum = 0.0;
+  for (var k = lane; k < uniforms.K; k += 32u) {
+    sum += f32(input_a.getByOffset(m * uniforms.K + k)) * f32(weight(n, k));
+  }
+  partial[lane] = sum;
+  workgroupBarrier();
+  for (var stride = 16u; stride > 0u; stride /= 2u) {
+    if (lane < stride) {
+      partial[lane] += partial[lane + stride];
+    }
+    workgroupBarrier();
+  }
+  if (lane == 0u) {
+    var value = partial[0];
+#if has_bias
+    value += f32(bias.getByOffset(n));
+#endif
+    output.setByOffset(m * uniforms.N + n, f16(value));
+  }
+#endif
+}

onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc

@@ -59,6 +59,9 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] =
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, LinearAttention)>,
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, LinearAttentionGate)>,
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MatMulBnb4)>,
+#if !defined(DISABLE_FLOAT8_TYPES)
+    BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MatMulBlockQuantizedFp8Weight)>,
+#endif
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MatMulNBits)>,
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MatMulNBitsQkv)>,
     BuildKernelCreateInfo<class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MatMulNBitsMlp)>,