From cbcd5e90d857e22dcfa118224bc63bbf03b41202 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:06:13 +0000 Subject: [PATCH 1/8] Add WebGPU block-quantized FP8 dense matmul Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ock_quantized_fp8_activation.wgsl.template | 18 ++ .../matmul_block_quantized_fp8_weight.cc | 213 ++++++++++++++++++ .../matmul_block_quantized_fp8_weight.h | 24 ++ ...l_block_quantized_fp8_weight.wgsl.template | 100 ++++++++ .../webgpu/webgpu_contrib_kernels.cc | 3 + .../matmul_block_quantized_fp8_weight_test.cc | 194 ++++++++++++++++ 6 files changed, 552 insertions(+) create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template create mode 100644 onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template new file mode 100644 index 0000000000000..5efd451136276 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#use .getByOffset .setByOffset + +$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 v = f32(input_a.getByOffset(global_idx)) * inverse; + let magnitude = min(abs(v), 448.0); + let step = exp2(max(floor(log2(magnitude)), -6.0) - 3.0); + let rounded = min(round(magnitude / step) * step, 448.0); + let q = select(rounded, -rounded, v < 0.0); + output.setByOffset(global_idx, f16(q * scale)); +} diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc new file mode 100644 index 0000000000000..f02831c6f0bca --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc @@ -0,0 +1,213 @@ +// 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 { + 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 { + 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 { + 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(info.GetAttrOrDefault("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(k_dim); + const uint32_t N = narrow(n_dim); + const uint32_t M = narrow(y->Shape().Size() / n_dim); + const uint32_t K_blocks = narrow(k_blocks); + if (K == 0) { + Fp8EmptyReductionProgram empty_program{bias != nullptr}; + empty_program.SetWorkgroupSize(64); + empty_program.SetDispatchGroupSize((narrow(y->Shape().Size()) + 63u) / 64u); + if (bias) { + empty_program.AddInput({bias, ProgramTensorMetadataDependency::Type}); + } + empty_program.AddOutput({y, ProgramTensorMetadataDependency::Type}) + .AddUniformVariables({{narrow(y->Shape().Size())}, {N}}); + return context.RunProgram(empty_program); + } + const Tensor* activation = a; + Tensor qdq; + if (a_scale && K != 0) { + qdq = context.CreateGPUTensor(a->DataType(), a_shape); + Fp8ActivationProgram qdq_program; + qdq_program.SetWorkgroupSize(64); + qdq_program.SetDispatchGroupSize((narrow(a_shape.Size()) + 63u) / 64u); + qdq_program.AddInputs({{a, ProgramTensorMetadataDependency::Type}, + {a_scale, ProgramTensorMetadataDependency::None}}) + .AddOutput({&qdq, ProgramTensorMetadataDependency::None}) + .AddUniformVariables({{narrow(a_shape.Size())}}); + ORT_RETURN_IF_ERROR(context.RunProgram(qdq_program)); + activation = &qdq; + } + + const bool use_matrix = M >= 8 && K != 0 && + 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 weight_bytes(DataTypeImpl::GetType(), b_shape, const_cast(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); + 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()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + MatMulBlockQuantizedFp8Weight); +#endif + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h new file mode 100644 index 0000000000000..db8954e5d5670 --- /dev/null +++ b/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 diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template new file mode 100644 index 0000000000000..97aadc622f00d --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template @@ -0,0 +1,100 @@ +// 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(((exponent + 120u) << 23u) | (mantissa << 20u)), + exponent != 0u); + let value = select(magnitude, -magnitude, (code & 128u) != 0u); + return select(value, bitcast(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 tile_a: array; +var tile_b: array; +var tile_y: array; +#else +var partial: array; +#endif +$MAIN { +#if use_matrix + let lane = local_idx; + let m_base = workgroup_id.y * 64u; + let n_base = workgroup_id.x * 16u; + var acc: subgroup_matrix_result; + 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, row_major>( + &tile_a, (lane / 32u) * 256u, 16u); + let right = subgroupMatrixLoad, col_major>(&tile_b, 0u, 16u); + acc = subgroupMatrixMultiplyAccumulate(left, right, acc); + workgroupBarrier(); + } + subgroupMatrixStore(&tile_y, (lane / 32u) * 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; + let m = workgroup_id.y; + let n = workgroup_id.x; + 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 +} diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 9fd1262c33c3a..d96b057663e97 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -55,6 +55,9 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, +#if !defined(DISABLE_FLOAT8_TYPES) + BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc new file mode 100644 index 0000000000000..4955544b55e89 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -0,0 +1,194 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/common/float8.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +#if defined(USE_WEBGPU) && !defined(DISABLE_FLOAT8_TYPES) +namespace { +void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_size, + bool with_a_scale, bool with_bias) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + + const auto k = a_dims.back(); + int64_t m = 1; + for (size_t i = 0; i + 1 < a_dims.size(); ++i) { + m *= a_dims[i]; + } + const auto blocks = (k + block_size - 1) / block_size; + std::vector a; + std::vector b; + std::vector scales; + std::vector biases; + for (int64_t i = 0; i < m * k; ++i) { + a.emplace_back(0.125f * static_cast((i % 11) - 5)); + } + for (int64_t i = 0; i < n; ++i) { + biases.emplace_back(static_cast((i % 5) - 2) * 0.25f); + for (int64_t j = 0; j < blocks; ++j) { + scales.push_back(0.125f * static_cast(1 + (i + j) % 4)); + } + for (int64_t j = 0; j < k; ++j) { + b.emplace_back(0.25f * static_cast(((i * 3 + j) % 9) - 4), true); + } + } + + std::vector expected; + expected.reserve(m * n); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < n; ++col) { + float sum = 0.0f; + for (int64_t j = 0; j < k; ++j) { + float activation = a[row * k + j].ToFloat(); + if (with_a_scale) { + activation = MLFloat16(Float8E4M3FN(activation / 0.375f, true).ToFloat() * 0.375f).ToFloat(); + } + const auto weight = MLFloat16(b[col * k + j].ToFloat() * + scales[col * blocks + j / block_size]) + .ToFloat(); + sum += activation * weight; + } + expected.emplace_back(MLFloat16(sum + (with_bias ? biases[col].ToFloat() : 0.0f))); + } + } + + std::vector output_dims(a_dims.begin(), a_dims.end()); + output_dims.back() = n; + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + test.AddInput("A", a_dims, a); + test.AddInput("B", {n, k}, b, /*is_initializer=*/true); + test.AddInput("b_scale", {n, blocks}, scales, /*is_initializer=*/true); + if (with_a_scale) { + test.AddInput("a_scale", {}, {0.375f}, /*is_initializer=*/true); + } else { + test.AddOptionalInputEdge(); + } + if (with_bias) { + test.AddInput("bias", {n}, biases); + } + test.AddOutput("Y", output_dims, expected); + test.SetOutputTolerance(0.03f); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} +} // namespace + +TEST(MatMulBlockQuantizedFp8WeightTest, DecodeRankOneOddBytes) { + RunFp8MatMul({3}, 1, 2, false, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, DecodeWithBlockScalesAndBias) { + RunFp8MatMul({1, 33}, 19, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqBatched) { + RunFp8MatMul({2, 3, 33}, 19, 16, true, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixTailWithBlockScales) { + RunFp8MatMul({9, 33}, 19, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixWithActivationQdq) { + RunFp8MatMul({8, 32}, 16, 32, true, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, SubnormalSignedWeightsAndBias) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const std::vector b{ + Float8E4M3FN(uint8_t{0x01}, Float8E4M3FN::FromBits()), + Float8E4M3FN(uint8_t{0x81}, Float8E4M3FN::FromBits()), + Float8E4M3FN(0.5625f, true), + Float8E4M3FN(-1.0f, true)}; + float expected = 0.25f; + for (const auto weight : b) { + expected += weight.ToFloat(); + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{4}); + test.AddInput("A", {1, 4}, std::vector(4, MLFloat16(1.0f))); + test.AddInput("B", {1, 4}, b, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddOptionalInputEdge(); + test.AddInput("bias", {1}, {MLFloat16(0.25f)}); + test.AddOutput("Y", {1, 1}, {MLFloat16(expected)}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqTieAndSaturation) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {5, 1}, + ToFloat16({0.3984375f, -0.3984375f, 200.0f, -200.0f, 0.0f})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {0.375f}); + test.AddOutput("Y", {5, 1}, ToFloat16({0.375f, -0.375f, 168.0f, -168.0f, 0.0f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, EmptyReductionWithBias) { + RunFp8MatMul({2, 0}, 3, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, EmptyOutput) { + RunFp8MatMul({0, 16}, 3, 16, false, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, InvalidScaleShape) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{16}); + test.AddInput("A", {1, 17}, std::vector(17, MLFloat16(1.0f))); + test.AddInput("B", {1, 17}, std::vector(17, Float8E4M3FN(1.0f, true))); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddOutput("Y", {1, 1}, {MLFloat16(17.0f)}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectFailure, "b_scale must have shape", {}, nullptr, &providers); +} + +#endif + +} // namespace test +} // namespace onnxruntime From 255d4b375a48680da3e3fd1576f0fb7ea42f4ed6 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:34:44 +0000 Subject: [PATCH 2/8] Preserve FP16-only ops and select registered WebGPU plugin in Python Avoid casting FP16-only operators to unsupported FP32 before EP assignment and prefer the registered WebGPU plugin when initializing a Python session. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/optimizer/insert_cast_transformer.cc | 28 +++++++++++++++-- .../python/onnxruntime_pybind_state.cc | 5 ++++ .../framework/insert_cast_transformer_test.cc | 30 +++++++++++++++++++ 3 files changed, 60 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/optimizer/insert_cast_transformer.cc b/onnxruntime/core/optimizer/insert_cast_transformer.cc index 05f53273447b9..f2983b18d9def 100644 --- a/onnxruntime/core/optimizer/insert_cast_transformer.cc +++ b/onnxruntime/core/optimizer/insert_cast_transformer.cc @@ -25,9 +25,31 @@ bool InsertCastTransformer::NeedInsertCast(const onnxruntime::Node* node, const // We don't cast a node with a subgraph as we'd need to do a lot more checking of the subgraph inputs // (both explicit and implicit) and contents to determine if it was safe to do so. // TODO: a better check is to check does the CPU kernel with float exist or not. - return node->GetExecutionProviderType().empty() && - !node->ContainsSubgraph() && - IsMLFloat16Tensor(*input); + if (!node->GetExecutionProviderType().empty() || node->ContainsSubgraph() || !IsMLFloat16Tensor(*input)) { + return false; + } + + if (const auto* schema = node->Op()) { + const auto& formal_inputs = schema->inputs(); + const auto& input_counts = node->InputArgCount(); + const auto& inputs = node->InputDefs(); + const auto& constraints = schema->typeConstraintMap(); + const auto float_type = ONNX_NAMESPACE::Utils::DataTypeUtils::ToType("tensor(float)"); + size_t actual_index = 0; + for (size_t formal_index = 0; formal_index < formal_inputs.size() && formal_index < input_counts.size(); + ++formal_index) { + const auto constraint = constraints.find(formal_inputs[formal_index].GetTypeStr()); + if (constraint != constraints.end() && constraint->second.first.count(float_type) == 0) { + for (int count = 0; count < input_counts[formal_index]; ++count) { + if (inputs[actual_index + count] && IsMLFloat16Tensor(*inputs[actual_index + count])) { + return false; + } + } + } + actual_index += input_counts[formal_index]; + } + } + return true; } onnxruntime::NodeArg* AddCastNode(onnxruntime::Graph& graph, diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 0c210269704c3..a33f822de1517 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1473,6 +1473,11 @@ static std::shared_ptr CreateExecutionProviderFactory cit == provider_options_map.end() ? ProviderOptions{} : cit->second, &session_options); #endif } else if (type == kWebGpuExecutionProvider) { +#if !defined(ORT_MINIMAL_BUILD) + if (auto ep_factory = try_create_registered_plugin_factory(); ep_factory) { + return ep_factory; + } +#endif #if defined(USE_WEBGPU) return onnxruntime::WebGpuProviderFactoryCreator::Create(session_options.config_options); #endif diff --git a/onnxruntime/test/framework/insert_cast_transformer_test.cc b/onnxruntime/test/framework/insert_cast_transformer_test.cc index ae9fccf8fedd3..21bd7ba8ff622 100644 --- a/onnxruntime/test/framework/insert_cast_transformer_test.cc +++ b/onnxruntime/test/framework/insert_cast_transformer_test.cc @@ -20,6 +20,36 @@ namespace test { #define MODEL_FOLDER ORT_TSTR("testdata/transform/") typedef std::vector ArgMap; + +TEST(TransformerTest, Fp16OnlyContribOpIsNotCastToUnsupportedFloat32) { +#if defined(DISABLE_CONTRIB_OPS) || defined(DISABLE_FLOAT8_TYPES) + GTEST_SKIP() << "Block-quantized FP8 MatMul is unavailable in this build."; +#else + auto model = std::make_shared( + "test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + std::unordered_map{{onnxruntime::kOnnxDomain, 22}, {onnxruntime::kMSDomain, 1}}, + std::vector(), DefaultLoggingManager().DefaultLogger()); + onnxruntime::Graph& graph = model->MainGraph(); + + TypeProto fp16, fp8, float32; + fp16.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT16); + fp8.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT8E4M3FN); + float32.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + onnxruntime::NodeArg a("A", &fp16), b("B", &fp8), scale("b_scale", &float32), y("Y", &fp16); + auto& node = graph.AddNode("fp8_matmul", "MatMulBlockQuantizedFp8Weight", "", + ArgMap{&a, &b, &scale}, ArgMap{&y}, nullptr, kMSDomain); + graph.SetOutputs({node.OutputDefs()[0]}); + + ASSERT_STATUS_OK(graph.Resolve()); + InsertCastTransformer transformer("Test", DefaultCpuExecutionProvider()->GetKernelRegistry().get()); + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, DefaultLoggingManager().DefaultLogger())); + ASSERT_STATUS_OK(graph.Resolve()); + EXPECT_FALSE(modified); + EXPECT_EQ(graph.NumberOfNodes(), 1); +#endif +} + TEST(TransformerTest, InsertCastGPUTest) { auto model = std::make_shared("test", false, DefaultLoggingManager().DefaultLogger()); onnxruntime::Graph& graph = model->MainGraph(); From 3e9903689783ef15c243929ea21e3554e4442075 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:50:21 +0000 Subject: [PATCH 3/8] Make FP8 activation rounding portable across WebGPU backends Apply explicit E4M3 ties-to-even instead of backend-dependent WGSL round. Document the read-only FP8 byte alias and padded binding and extend tie coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ul_block_quantized_fp8_activation.wgsl.template | 14 +++++++++++--- .../matmul_block_quantized_fp8_weight.cc | 7 +++++-- .../matmul_block_quantized_fp8_weight_test.cc | 9 ++++++--- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template index 5efd451136276..f4d978093ed60 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -3,6 +3,10 @@ #use .getByOffset .setByOffset +const E4M3_MAX: f32 = 448.0; +const E4M3_MIN_NORMAL_EXP: f32 = -6.0; +const E4M3_MANTISSA_BITS: f32 = 3.0; + $MAIN { if (global_idx >= uniforms.size) { return; @@ -10,9 +14,13 @@ $MAIN { let scale = input_scale.getByOffset(0u); let inverse = select(0.0, 1.0 / scale, scale != 0.0); let v = f32(input_a.getByOffset(global_idx)) * inverse; - let magnitude = min(abs(v), 448.0); - let step = exp2(max(floor(log2(magnitude)), -6.0) - 3.0); - let rounded = min(round(magnitude / step) * step, 448.0); + let magnitude = min(abs(v), E4M3_MAX); + let step = exp2(max(floor(log2(magnitude)), E4M3_MIN_NORMAL_EXP) - E4M3_MANTISSA_BITS); + let units = magnitude / step; + let lower = floor(units); + let fraction = units - lower; + let round_up = fraction > 0.5 || (fraction == 0.5 && (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)); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc index f02831c6f0bca..656a70ebf572d 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc @@ -153,7 +153,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu } const Tensor* activation = a; Tensor qdq; - if (a_scale && K != 0) { + if (a_scale) { qdq = context.CreateGPUTensor(a->DataType(), a_shape); Fp8ActivationProgram qdq_program; qdq_program.SetWorkgroupSize(64); @@ -166,7 +166,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu activation = &qdq; } - const bool use_matrix = M >= 8 && K != 0 && + const bool use_matrix = M >= 8 && SelectSubgroupMatrixConfig(context, {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, 16, 16, 16, 32, false}}) @@ -176,6 +176,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu 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(), b_shape, const_cast(b->DataRaw()), memory_info); Fp8MatMulProgram program{bias != nullptr, use_matrix}; program.SetWorkgroupSize(use_matrix ? 128 : 32); @@ -184,6 +185,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu } 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}}) diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc index 4955544b55e89..9be5c8757adf0 100644 --- a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -149,12 +149,15 @@ TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqTieAndSaturation) { } OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); test.AddAttribute("block_size", int64_t{1}); - test.AddInput("A", {5, 1}, - ToFloat16({0.3984375f, -0.3984375f, 200.0f, -200.0f, 0.0f})); + test.AddInput("A", {7, 1}, + ToFloat16({0.3984375f, -0.3984375f, 0.4453125f, -0.4453125f, + 200.0f, -200.0f, 0.0f})); test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); test.AddInput("b_scale", {1, 1}, {1.0f}); test.AddInput("a_scale", {}, {0.375f}); - test.AddOutput("Y", {5, 1}, ToFloat16({0.375f, -0.375f, 168.0f, -168.0f, 0.0f})); + test.AddOutput("Y", {7, 1}, + ToFloat16({0.375f, -0.375f, 0.46875f, -0.46875f, + 168.0f, -168.0f, 0.0f})); SessionOptions options; ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); std::vector> providers; From c6ca2e2470531c5e58d1ec51d8e112e2838a0390 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:08:04 +0000 Subject: [PATCH 4/8] Declare FP8 activation shader type dependencies Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../webgpu/quantization/matmul_block_quantized_fp8_weight.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc index 656a70ebf572d..754331942f880 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc @@ -159,8 +159,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu qdq_program.SetWorkgroupSize(64); qdq_program.SetDispatchGroupSize((narrow(a_shape.Size()) + 63u) / 64u); qdq_program.AddInputs({{a, ProgramTensorMetadataDependency::Type}, - {a_scale, ProgramTensorMetadataDependency::None}}) - .AddOutput({&qdq, ProgramTensorMetadataDependency::None}) + {a_scale, ProgramTensorMetadataDependency::Type}}) + .AddOutput({&qdq, ProgramTensorMetadataDependency::Type}) .AddUniformVariables({{narrow(a_shape.Size())}}); ORT_RETURN_IF_ERROR(context.RunProgram(qdq_program)); activation = &qdq; From 2f12bcdb08d2aefc278022e473d90a3fdb325dd6 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:47:36 +0000 Subject: [PATCH 5/8] Compare FP8 activation ties in input space Avoid backend-dependent reciprocal rounding at halfway cases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../matmul_block_quantized_fp8_activation.wgsl.template | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template index f4d978093ed60..cedcedba1c163 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -13,13 +13,15 @@ $MAIN { } let scale = input_scale.getByOffset(0u); let inverse = select(0.0, 1.0 / scale, scale != 0.0); - let v = f32(input_a.getByOffset(global_idx)) * inverse; + let input_value = f32(input_a.getByOffset(global_idx)); + let v = input_value * inverse; let magnitude = min(abs(v), E4M3_MAX); let step = exp2(max(floor(log2(magnitude)), E4M3_MIN_NORMAL_EXP) - E4M3_MANTISSA_BITS); let units = magnitude / step; let lower = floor(units); - let fraction = units - lower; - let round_up = fraction > 0.5 || (fraction == 0.5 && (u32(lower) & 1u) != 0u); + let halfway_input = (lower + 0.5) * step * abs(scale); + let round_up = abs(input_value) > halfway_input || + (abs(input_value) == halfway_input && (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)); From 38a12629ef4a6487f16e929dccb244c311f7d8c1 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Tue, 6 Oct 2026 17:04:44 +0000 Subject: [PATCH 6/8] Fix FP8 WebGPU dispatch, subgroup ownership and quantization ties Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ock_quantized_fp8_activation.wgsl.template | 5 +- ...l_block_quantized_fp8_weight.wgsl.template | 19 +++-- .../core/providers/webgpu/shader_helper.cc | 1 + .../matmul_block_quantized_fp8_weight_test.cc | 80 ++++++++++++++++++- 4 files changed, 94 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template index cedcedba1c163..00d069b21c658 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -4,8 +4,6 @@ #use .getByOffset .setByOffset const E4M3_MAX: f32 = 448.0; -const E4M3_MIN_NORMAL_EXP: f32 = -6.0; -const E4M3_MANTISSA_BITS: f32 = 3.0; $MAIN { if (global_idx >= uniforms.size) { @@ -16,7 +14,8 @@ $MAIN { let input_value = f32(input_a.getByOffset(global_idx)); let v = input_value * inverse; let magnitude = min(abs(v), E4M3_MAX); - let step = exp2(max(floor(log2(magnitude)), E4M3_MIN_NORMAL_EXP) - E4M3_MANTISSA_BITS); + let magnitude_exponent = (bitcast(magnitude) >> 23u) & 255u; + let step = bitcast((max(magnitude_exponent, 121u) - 3u) << 23u); let units = magnitude / step; let lower = floor(units); let halfway_input = (lower + 0.5) * step * abs(scale); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template index 97aadc622f00d..ba0321cf88832 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template @@ -32,8 +32,12 @@ var partial: array; $MAIN { #if use_matrix let lane = local_idx; - let m_base = workgroup_id.y * 64u; - let n_base = workgroup_id.x * 16u; + 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; for (var tile = 0u; tile < (uniforms.K + 15u) / 16u; tile++) { let k_base = tile * 16u; @@ -55,12 +59,12 @@ $MAIN { } workgroupBarrier(); let left = subgroupMatrixLoad, row_major>( - &tile_a, (lane / 32u) * 256u, 16u); + &tile_a, subgroup_index * 256u, 16u); let right = subgroupMatrixLoad, col_major>(&tile_b, 0u, 16u); acc = subgroupMatrixMultiplyAccumulate(left, right, acc); workgroupBarrier(); } - subgroupMatrixStore(&tile_y, (lane / 32u) * 256u, acc, 16u); + subgroupMatrixStore(&tile_y, subgroup_index * 256u, acc, 16u); workgroupBarrier(); for (var i = lane; i < 1024u; i += 128u) { let m = m_base + i / 16u; @@ -75,8 +79,11 @@ $MAIN { } #else let lane = local_idx; - let m = workgroup_id.y; - let n = workgroup_id.x; + 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)); diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index 2b68a75a450be..edfc8227b198f 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -88,6 +88,7 @@ Status ShaderHelper::Init() { if (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups)) { body_ss_ << ",\n" " @builtin(subgroup_invocation_id) sg_id : u32,\n" + " @builtin(subgroup_id) subgroup_index : u32,\n" " @builtin(subgroup_size) sg_size : u32"; } // When using indirect dispatch, avoid @builtin(num_workgroups) to skip Dawn's validation diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc index 9be5c8757adf0..336605228db87 100644 --- a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -9,6 +9,8 @@ #include "gtest/gtest.h" #include "core/common/float8.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/webgpu_context.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" @@ -20,7 +22,7 @@ namespace test { #if defined(USE_WEBGPU) && !defined(DISABLE_FLOAT8_TYPES) namespace { void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_size, - bool with_a_scale, bool with_bias) { + bool with_a_scale, bool with_bias, bool distinct_rows = false) { auto ep = DefaultWebGpuExecutionProvider(); if (!ep) { GTEST_SKIP() << "WebGPU execution provider is not available"; @@ -37,7 +39,8 @@ void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_s std::vector scales; std::vector biases; for (int64_t i = 0; i < m * k; ++i) { - a.emplace_back(0.125f * static_cast((i % 11) - 5)); + a.emplace_back(0.125f * static_cast((i % 11) - 5) + + (distinct_rows ? 0.015625f * static_cast(i / k) : 0.0f)); } for (int64_t i = 0; i < n; ++i) { biases.emplace_back(static_cast((i % 5) - 2) * 0.25f); @@ -91,12 +94,33 @@ void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_s providers.push_back(std::move(ep)); test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); } + +bool SupportsFp8MatrixPath() { + const auto& context = webgpu::WebGpuContextFactory::DefaultContext(); + bool supported = false; + if (context.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + const auto& configs = context.SubgroupMatrixConfigs(); + const auto& adapter = context.AdapterInfo(); + supported = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {configs.configs, configs.configCount}, adapter.subgroupMinSize, adapter.subgroupMaxSize, + context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, + 16, 16, 16, 32, false}}) + .has_value(); + } + webgpu::WebGpuContextFactory::ReleaseContext(0); + return supported; +} } // namespace TEST(MatMulBlockQuantizedFp8WeightTest, DecodeRankOneOddBytes) { RunFp8MatMul({3}, 1, 2, false, false); } +TEST(MatMulBlockQuantizedFp8WeightTest, ScalarDispatchBeyondWebGpuDimensionLimit) { + RunFp8MatMul({1, 1}, 65537, 1, false, false); +} + TEST(MatMulBlockQuantizedFp8WeightTest, DecodeWithBlockScalesAndBias) { RunFp8MatMul({1, 33}, 19, 16, false, true); } @@ -106,13 +130,36 @@ TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqBatched) { } TEST(MatMulBlockQuantizedFp8WeightTest, MatrixTailWithBlockScales) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } RunFp8MatMul({9, 33}, 19, 16, false, true); } TEST(MatMulBlockQuantizedFp8WeightTest, MatrixWithActivationQdq) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } RunFp8MatMul({8, 32}, 16, 32, true, false); } +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixSubgroupTilesAndSecondWorkgroup) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } + RunFp8MatMul({64, 16}, 16, 16, false, false, true); + RunFp8MatMul({65, 16}, 16, 16, false, false, true); +} + TEST(MatMulBlockQuantizedFp8WeightTest, SubnormalSignedWeightsAndBias) { auto ep = DefaultWebGpuExecutionProvider(); if (!ep) { @@ -165,6 +212,35 @@ TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqTieAndSaturation) { test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); } +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqAdjacentEvenAndOddTies) { + for (float scale : {0.375f, 0.75f, 1.0f}) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const float factor = scale / 0.375f; + const float delta = factor * 0.000244140625f; + const float even_tie = factor * 0.3984375f; + const float odd_tie = factor * 0.4453125f; + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {6, 1}, + ToFloat16({even_tie - delta, even_tie, even_tie + delta, + odd_tie - delta, odd_tie, odd_tie + delta})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {scale}); + test.AddOutput("Y", {6, 1}, + ToFloat16({factor * 0.375f, factor * 0.375f, factor * 0.421875f, + factor * 0.421875f, factor * 0.46875f, factor * 0.46875f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); + } +} + TEST(MatMulBlockQuantizedFp8WeightTest, EmptyReductionWithBias) { RunFp8MatMul({2, 0}, 3, 16, false, true); } From dc9a8e15dae4e86470568fc7b0ceca7634c3ec11 Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Tue, 6 Oct 2026 21:15:24 +0000 Subject: [PATCH 7/8] Compare FP8 activation midpoints in reciprocal domain Keep FP8 rounding decisions consistent with CUDA reciprocal quantization and test adjacent FP16 values at a non-power-of-two scale. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ock_quantized_fp8_activation.wgsl.template | 6 ++--- .../matmul_block_quantized_fp8_weight_test.cc | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template index 00d069b21c658..6f9f1e556c975 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -18,9 +18,9 @@ $MAIN { let step = bitcast((max(magnitude_exponent, 121u) - 3u) << 23u); let units = magnitude / step; let lower = floor(units); - let halfway_input = (lower + 0.5) * step * abs(scale); - let round_up = abs(input_value) > halfway_input || - (abs(input_value) == halfway_input && (u32(lower) & 1u) != 0u); + 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)); diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc index 336605228db87..28230be3c111c 100644 --- a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -241,6 +241,30 @@ TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqAdjacentEvenAndOddTies) { } } +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqReciprocalMidpoint) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + constexpr float scale = 1.21603262424469f; + const auto input = ToFloat16({0.0272979736328125f, 0.027313232421875f, 0.0273284912109375f}); + ASSERT_LT(input[1].ToFloat() * (1.0f / scale), 0.0224609375f); + + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {3, 1}, input); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {scale}); + test.AddOutput("Y", {3, 1}, + ToFloat16({0.021484375f * scale, 0.021484375f * scale, 0.0234375f * scale})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + TEST(MatMulBlockQuantizedFp8WeightTest, EmptyReductionWithBias) { RunFp8MatMul({2, 0}, 3, 16, false, true); } From 4905cd01f250d5d9ab8aedc58d523d4ee9cd6b9c Mon Sep 17 00:00:00 2001 From: Sayan Shaw <52221015+sayanshaw24@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:41:53 +0000 Subject: [PATCH 8/8] Preserve FP8 activation NaNs and gate subgroup ID shader inputs Strengthen non-canceling subnormal coverage and explicitly request subgroup_id only for supported dense and routed FP8 matrix programs. Cache the shader opt-in separately and document the capability fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../matmul_block_quantized_fp8_weight.md | 27 +++++++++ onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 + ...ock_quantized_fp8_activation.wgsl.template | 5 ++ .../matmul_block_quantized_fp8_weight.cc | 2 + .../core/providers/webgpu/compute_context.h | 3 + onnxruntime/core/providers/webgpu/program.h | 8 +++ .../providers/webgpu/program_cache_key.cc | 3 + .../core/providers/webgpu/shader_helper.cc | 12 +++- .../matmul_block_quantized_fp8_weight_test.cc | 59 ++++++++++++++++++- 9 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md diff --git a/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md b/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md new file mode 100644 index 0000000000000..9a728f17edfed --- /dev/null +++ b/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md @@ -0,0 +1,27 @@ +# WebGPU block-quantized FP8 dense MatMul + +`com.microsoft::MatMulBlockQuantizedFp8Weight` supports FP16 activations and +row-major E4M3FN weights with FP32 block scales shaped +`[N, ceil(K / block_size)]`. It supports optional activation quantize/dequantize, +bias, rank-one and batched inputs, partial blocks, and non-word-aligned weight rows. +BF16 activation storage is not supported by this WebGPU kernel. + +When activation scaling is supplied, quantization uses saturating E4M3FN +round-to-nearest, ties-to-even semantics. NaN activations remain NaN rather than +being replaced by the maximum finite value. Weight decoding preserves finite +E4M3FN subnormals before multiplication. + +Prefill can use 16x16x16 FP16 subgroup matrices with FP32 accumulation when the +adapter exposes the required configuration, subgroup-size requirements, and the +instance exposes the WGSL `subgroup_id` language feature. Other configurations +use the scalar FP32-reduction path. This is FP8 storage with FP16 widening, not +native FP8 matrix multiplication. + +Shared shader generation includes `@builtin(subgroup_id)` only for programs +that explicitly request it. Both dense FP8 and routed block-FP8 QMoE matrix +programs check the language feature before selecting that path. Ordinary +operators and scalar fallbacks do not acquire this language-feature requirement. + +The focused `MatMulBlockQuantizedFp8WeightTest.*` tests disable CPU EP fallback. +Matrix-specific cases check capabilities and skip explicitly on unsupported +adapters; scalar cases remain separate. diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index f6e691cb1d7f4..93182e9e0fdbf 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -136,6 +136,7 @@ Status ApplyBlockFp8ExpertMatMul(ComputeContext& context, const uint32_t scale_k_blocks = (inner - 1) / 128 + 1; uint32_t byte_offset = 0; const bool use_matrix = !indirect_experts && rows >= 128 && cols >= 256 && inner >= 256 && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && BlockFp8MatrixDispatchFits(rows, cols, context.DeviceLimits().maxComputeWorkgroupsPerDimension) && input->DataType() == DataTypeImpl::GetType() && @@ -144,6 +145,7 @@ Status ApplyBlockFp8ExpertMatMul(ComputeContext& context, 16, 16, 16, 32, false}}) .has_value(); BlockFp8ExpertMatMulProgram program{bias != nullptr, indirect_experts != nullptr, broadcast_input, use_matrix}; + program.SetUseSubgroupId(use_matrix); program.AddInputs({{input, ProgramTensorMetadataDependency::Type}}); if (indirect_experts) { program.AddInputs({{&raw_weights, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template index 6f9f1e556c975..64387a727c0e4 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -13,6 +13,11 @@ $MAIN { 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; + // Classify bits so finite-math optimizations cannot fold a NaN self-comparison away. + if ((bitcast(v) & 0x7fffffffu) > 0x7f800000u) { + output.setByOffset(global_idx, f16(v)); + return; + } let magnitude = min(abs(v), E4M3_MAX); let magnitude_exponent = (bitcast(magnitude) >> 23u) & 255u; let step = bitcast((max(magnitude_exponent, 121u) - 3u) << 23u); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc index 754331942f880..75c5c8cde31d2 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc @@ -167,6 +167,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu } const bool use_matrix = M >= 8 && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && SelectSubgroupMatrixConfig(context, {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, 16, 16, 16, 32, false}}) @@ -179,6 +180,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::Compu // Tensor views require void*, but the alias only binds B as read-only storage; no copy or mutation occurs. Tensor weight_bytes(DataTypeImpl::GetType(), b_shape, const_cast(b->DataRaw()), memory_info); Fp8MatMulProgram program{bias != nullptr, use_matrix}; + program.SetUseSubgroupId(use_matrix); program.SetWorkgroupSize(use_matrix ? 128 : 32); if (use_matrix && context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { program.SetSubgroupSize(32); diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index bdcbafccace57..4030e66d53064 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -77,6 +77,9 @@ class ComputeContextBase { inline bool HasFeature(wgpu::FeatureName feature) const { return webgpu_context_.DeviceHasFeature(feature); } + inline bool HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName feature) const { + return webgpu_context_.Instance().HasWGSLLanguageFeature(feature); + } inline const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return webgpu_context_.SubgroupMatrixConfigs(); } diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index 67e94d3035616..8b62c69da3115 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -353,6 +353,12 @@ class ProgramBase { // requires the device to support wgpu::FeatureName::SubgroupSizeControl; see ShaderHelper::GenerateSourceCode. ProgramBase& SetSubgroupSize(uint32_t size); + // Opt in only when the shader uses subgroup_index and subgroup_id is supported by the instance. + ProgramBase& SetUseSubgroupId(bool value) { + use_subgroup_id_ = value; + return *this; + } + // add a uniform variable. // // the specified uniform variable should match the uniform definition in the class, @@ -398,6 +404,7 @@ class ProgramBase { inline uint32_t WorkgroupSizeY() const { return workgroup_size_y_; } inline uint32_t WorkgroupSizeZ() const { return workgroup_size_z_; } inline uint32_t SubgroupSize() const { return subgroup_size_; } + inline bool UseSubgroupId() const { return use_subgroup_id_; } inline const std::vector& UniformVariables() const { return variables_; } inline const std::vector& OverridableConstants() const { return overridable_constants_; } @@ -427,6 +434,7 @@ class ProgramBase { uint32_t workgroup_size_y_; uint32_t workgroup_size_z_; uint32_t subgroup_size_; + bool use_subgroup_id_{false}; std::vector variables_; std::vector overridable_constants_; diff --git a/onnxruntime/core/providers/webgpu/program_cache_key.cc b/onnxruntime/core/providers/webgpu/program_cache_key.cc index ae83918a77376..029da45cdcf03 100644 --- a/onnxruntime/core/providers/webgpu/program_cache_key.cc +++ b/onnxruntime/core/providers/webgpu/program_cache_key.cc @@ -71,6 +71,9 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, // = ||... // = ;; ss << program.Name(); + if (program.UseSubgroupId()) { + ss << "[SubgroupId]"; + } // append custom cache hint if any if (auto& hint = program.CacheHint(); !hint.empty()) { diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index edfc8227b198f..f1e7b78db8f5f 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -39,6 +39,10 @@ ShaderHelper::ShaderHelper(const ProgramBase& program, Status ShaderHelper::Init() { // dispatch group size is normalized so no need to validate it here + ORT_RETURN_IF_NOT(!program_.UseSubgroupId() || + (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups) && + webgpu_context_.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId)), + "Program ", program_.Name(), " requires subgroups and the subgroup_id WGSL language feature."); // validate workgroup size // normalize 0 (meaning "use default") the same way GenerateSourceCode() does, so validation @@ -87,9 +91,11 @@ Status ShaderHelper::Init() { " @builtin(local_invocation_id) local_id : vec3"; if (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups)) { body_ss_ << ",\n" - " @builtin(subgroup_invocation_id) sg_id : u32,\n" - " @builtin(subgroup_id) subgroup_index : u32,\n" - " @builtin(subgroup_size) sg_size : u32"; + " @builtin(subgroup_invocation_id) sg_id : u32,\n"; + if (program_.UseSubgroupId()) { + body_ss_ << " @builtin(subgroup_id) subgroup_index : u32,\n"; + } + body_ss_ << " @builtin(subgroup_size) sg_size : u32"; } // When using indirect dispatch, avoid @builtin(num_workgroups) to skip Dawn's validation // and duplication overhead in TransformIndirectDispatchBuffer. diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc index 28230be3c111c..7cb3986c71d04 100644 --- a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -3,13 +3,17 @@ #include #include +#include #include +#include #include #include "gtest/gtest.h" #include "core/common/float8.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/program_cache_key.h" +#include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" @@ -21,6 +25,12 @@ namespace test { #if defined(USE_WEBGPU) && !defined(DISABLE_FLOAT8_TYPES) namespace { +class SubgroupIdTestProgram final : public webgpu::Program { + public: + SubgroupIdTestProgram() : Program{"SubgroupIdTest"} {} + Status GenerateShaderCode(webgpu::ShaderHelper&) const override { return Status::OK(); } +}; + void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_size, bool with_a_scale, bool with_bias, bool distinct_rows = false) { auto ep = DefaultWebGpuExecutionProvider(); @@ -98,7 +108,8 @@ void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_s bool SupportsFp8MatrixPath() { const auto& context = webgpu::WebGpuContextFactory::DefaultContext(); bool supported = false; - if (context.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + if (context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && + context.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { const auto& configs = context.SubgroupMatrixConfigs(); const auto& adapter = context.AdapterInfo(); supported = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( @@ -113,6 +124,30 @@ bool SupportsFp8MatrixPath() { } } // namespace +TEST(MatMulBlockQuantizedFp8WeightTest, SubgroupIdIsOptInAndCacheKeyed) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const auto& context = webgpu::WebGpuContextFactory::DefaultContext(); + SubgroupIdTestProgram program; + const auto original_key = webgpu::CalculateProgramCacheKey(program, {}, {}); + webgpu::ShaderHelper ordinary_shader(program, program.Metadata(), context, {}, {}, 1, 1, 1); + ASSERT_STATUS_OK(ordinary_shader.Init()); + EXPECT_EQ(std::move(ordinary_shader.MainFunctionBody()).str().find("@builtin(subgroup_id)"), std::string::npos); + + program.SetUseSubgroupId(true); + EXPECT_NE(original_key, webgpu::CalculateProgramCacheKey(program, {}, {})); + webgpu::ShaderHelper matrix_shader(program, program.Metadata(), context, {}, {}, 1, 1, 1); + const bool supported = context.DeviceHasFeature(wgpu::FeatureName::Subgroups) && + context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId); + const auto status = matrix_shader.Init(); + EXPECT_EQ(status.IsOK(), supported); + if (supported) { + EXPECT_NE(std::move(matrix_shader.MainFunctionBody()).str().find("@builtin(subgroup_id)"), std::string::npos); + } + webgpu::WebGpuContextFactory::ReleaseContext(0); +} + TEST(MatMulBlockQuantizedFp8WeightTest, DecodeRankOneOddBytes) { RunFp8MatMul({3}, 1, 2, false, false); } @@ -167,7 +202,7 @@ TEST(MatMulBlockQuantizedFp8WeightTest, SubnormalSignedWeightsAndBias) { } const std::vector b{ Float8E4M3FN(uint8_t{0x01}, Float8E4M3FN::FromBits()), - Float8E4M3FN(uint8_t{0x81}, Float8E4M3FN::FromBits()), + Float8E4M3FN(uint8_t{0x87}, Float8E4M3FN::FromBits()), Float8E4M3FN(0.5625f, true), Float8E4M3FN(-1.0f, true)}; float expected = 0.25f; @@ -189,6 +224,26 @@ TEST(MatMulBlockQuantizedFp8WeightTest, SubnormalSignedWeightsAndBias) { test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); } +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqPreservesNaN) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const float nan = std::numeric_limits::quiet_NaN(); + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {2, 1}, ToFloat16({nan, 1.0f})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {0.375f}); + test.AddOutput("Y", {2, 1}, ToFloat16({nan, 1.03125f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqTieAndSaturation) { auto ep = DefaultWebGpuExecutionProvider(); if (!ep) {