From 5578aabb814678e4188da3ebde4d5a1ea33a177e Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 12 Oct 2026 11:53:44 +0800 Subject: [PATCH 01/35] Fix 1DS telemetry deadlock during Windows process termination Skip uploader teardown before acquiring telemetry locks when Windows is detaching DLLs for process termination. WinHTTP callbacks and worker threads are already gone at that point, so draining them can wait forever. Keep normal object destruction and explicit DLL unload cleanup. Validated against captured native hang stacks and identical machine code at all resolved frames. A pending-uploader fault injection hangs the unchanged runtime and exits in 20 of 20 fixed trials. Both WebGPU modes completed 1024-prompt, 32-output-token native profiles with LoRA on and off; explicit unload also passed. --- onnxruntime/core/platform/telemetry_1ds.cc | 8 ++++++++ onnxruntime/core/platform/telemetry_1ds_platform.h | 4 ++++ onnxruntime/core/platform/windows/telemetry_1ds.cc | 9 +++++++++ 3 files changed, 21 insertions(+) diff --git a/onnxruntime/core/platform/telemetry_1ds.cc b/onnxruntime/core/platform/telemetry_1ds.cc index 3d9a9d04cbbd8..fa11a2f1bde68 100644 --- a/onnxruntime/core/platform/telemetry_1ds.cc +++ b/onnxruntime/core/platform/telemetry_1ds.cc @@ -342,6 +342,14 @@ OneDsTelemetry::OneDsTelemetry(const Telemetry& local_telemetry) : OneDsTelemetr } OneDsTelemetry::~OneDsTelemetry() { +#if defined(_WIN32) + // ExitProcess terminates the other threads before running DLL destructors. + // The uploader cannot drain WinHTTP callbacks here, and a terminated thread + // may still own one of our locks. Normal destruction/FreeLibrary still cleans up. + if (telemetry_internal::IsDllShutdownInProgress()) { + return; + } +#endif std::lock_guard lock(global_mutex_); global_register_count_--; diff --git a/onnxruntime/core/platform/telemetry_1ds_platform.h b/onnxruntime/core/platform/telemetry_1ds_platform.h index 2cfdbd99dd0a8..311b80b19c553 100644 --- a/onnxruntime/core/platform/telemetry_1ds_platform.h +++ b/onnxruntime/core/platform/telemetry_1ds_platform.h @@ -13,4 +13,8 @@ const char* GetDefaultEncodedToken(); int32_t GetProcessorCount(); std::string GetCachePath(const std::string& directory); +#if defined(_WIN32) +bool IsDllShutdownInProgress(); +#endif + } // namespace onnxruntime::telemetry_internal diff --git a/onnxruntime/core/platform/windows/telemetry_1ds.cc b/onnxruntime/core/platform/windows/telemetry_1ds.cc index 88ad276a796a0..5fb8bd45bacb5 100644 --- a/onnxruntime/core/platform/windows/telemetry_1ds.cc +++ b/onnxruntime/core/platform/windows/telemetry_1ds.cc @@ -54,6 +54,15 @@ bool PrepareSampledProcessEvent(EventProperties& event) { namespace telemetry_internal { +bool IsDllShutdownInProgress() { + using RtlDllShutdownInProgressFn = BOOLEAN(NTAPI*)(); + const HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll"); + const auto is_shutting_down = + ntdll == nullptr ? nullptr + : reinterpret_cast(::GetProcAddress(ntdll, "RtlDllShutdownInProgress")); + return is_shutting_down != nullptr && is_shutting_down(); +} + const char* GetDefaultEncodedToken() { return "fllXSmJHWBYMX1cqWldIYU0KTA0IVnpWD0lkRloSWwhILF5WGmBDDUdEC1csC0NMZhZYWVEPBH9DWEhlR15MCl1UfgwPVWVDX0U="; } From 625b47c46453b586d21cc291a7aa8dfa022c04b5 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 12 Oct 2026 15:28:49 +0800 Subject: [PATCH 02/35] Add Windows child-process coverage for 1DS teardown --- cmake/onnxruntime_unittests.cmake | 20 +++ onnxruntime/core/platform/telemetry_1ds.h | 2 + .../lifecycle/telemetry_shutdown_test.cc | 110 ++++++++++++++++ .../telemetry_shutdown_test_library.cc | 123 ++++++++++++++++++ .../telemetry_shutdown_test_library.def | 3 + 5 files changed, 258 insertions(+) create mode 100644 onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc create mode 100644 onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc create mode 100644 onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index d7f82262d3a5e..c9777525f6c28 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -474,6 +474,26 @@ if(WIN32) if(onnxruntime_USE_1DS_TELEMETRY) list(APPEND onnxruntime_test_framework_src_patterns "${TEST_SRC_DIR}/platform/windows_telemetry/*.cc" ) + + if(NOT GDK_PLATFORM AND NOT CMAKE_SYSTEM_NAME STREQUAL "WindowsStore") + set(telemetry_lifecycle_test_dir "${TEST_SRC_DIR}/platform/windows_telemetry/lifecycle") + onnxruntime_add_shared_library(onnxruntime_telemetry_shutdown_test_library + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.cc" + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.def") + target_link_libraries(onnxruntime_telemetry_shutdown_test_library PRIVATE onnxruntime_common mat) + onnxruntime_add_executable(onnxruntime_telemetry_shutdown_test + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test.cc") + add_dependencies(onnxruntime_telemetry_shutdown_test onnxruntime_telemetry_shutdown_test_library) + foreach(mode process_exit normal_destruction explicit_unload) + add_test(NAME onnxruntime_telemetry_${mode}_test + COMMAND onnxruntime_telemetry_shutdown_test + "$" ${mode} + WORKING_DIRECTORY "$") + set_tests_properties(onnxruntime_telemetry_${mode}_test PROPERTIES TIMEOUT 15) + endforeach() + set_target_properties(onnxruntime_telemetry_shutdown_test_library onnxruntime_telemetry_shutdown_test + PROPERTIES FOLDER "ONNXRuntimeTest") + endif() endif() endif() diff --git a/onnxruntime/core/platform/telemetry_1ds.h b/onnxruntime/core/platform/telemetry_1ds.h index 2d0eabfd48dc2..665d38acf2ffe 100644 --- a/onnxruntime/core/platform/telemetry_1ds.h +++ b/onnxruntime/core/platform/telemetry_1ds.h @@ -140,6 +140,8 @@ class OneDsTelemetry : public Telemetry { const std::string& lib_path) const override; private: + friend class OneDsTelemetryTestPeer; + static void ReportFailure(const char* operation_name, const char* message); template diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc new file mode 100644 index 0000000000000..e04fd6357fc22 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct CloseWin32Handle { + void operator()(HANDLE handle) const { CloseHandle(handle); } +}; +using UniqueHandle = std::unique_ptr; + +int Fail(const char* message) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + return EXIT_FAILURE; +} + +int RunChild(const wchar_t* library_path, std::wstring_view mode, HANDLE verified) { + if (!SetEnvironmentVariableW(L"ORT_DISABLE_TELEMETRY", L"1")) { + return Fail("Could not disable network telemetry"); + } + const HMODULE library = LoadLibraryW(library_path); + if (!library) { + return Fail("LoadLibraryW failed"); + } + using RunTest = void (*)(bool, bool, HANDLE); + const auto run_test = reinterpret_cast(GetProcAddress(library, "OrtTestTelemetryShutdown")); + if (!run_test) { + return Fail("GetProcAddress failed"); + } + const bool process_exit = mode == L"process_exit"; + run_test(process_exit, mode == L"normal_destruction", verified); + const DWORD expected_signal = mode == L"normal_destruction" ? WAIT_OBJECT_0 : WAIT_TIMEOUT; + if (WaitForSingleObject(verified, 0) != expected_signal) { + return Fail("Telemetry destruction happened at the wrong lifecycle stage"); + } + if (process_exit) { + ExitProcess(EXIT_SUCCESS); + } + if (!FreeLibrary(library)) { + return Fail("FreeLibrary failed"); + } + return EXIT_SUCCESS; +} + +} // namespace + +int wmain(int argc, wchar_t** argv) { + if (argc == 5 && std::wstring_view(argv[1]) == L"--child") { + const auto verified = reinterpret_cast(static_cast(std::wcstoull(argv[4], nullptr, 10))); + return RunChild(argv[2], argv[3], verified); + } + if (argc != 3) { + return Fail("Expected test DLL path and lifecycle mode"); + } + const std::wstring_view mode(argv[2]); + if (mode != L"process_exit" && mode != L"normal_destruction" && mode != L"explicit_unload") { + return Fail("Unknown lifecycle mode"); + } + + // The parent checks a signal from the DLL destructor, so an early child exit + // cannot pass without actually exercising and verifying telemetry teardown. + SECURITY_ATTRIBUTES attributes{sizeof(SECURITY_ATTRIBUTES), nullptr, TRUE}; + UniqueHandle verified(CreateEventW(&attributes, TRUE, FALSE, nullptr)); + if (!verified) { + return Fail("CreateEventW failed"); + } + std::wstring executable(32768, L'\0'); + const DWORD length = GetModuleFileNameW(nullptr, executable.data(), static_cast(executable.size())); + if (length == 0 || length >= executable.size()) { + return Fail("GetModuleFileNameW failed"); + } + executable.resize(length); + std::wstring command = L"\"" + executable + L"\" --child \"" + argv[1] + L"\" " + argv[2] + L" " + + std::to_wstring(reinterpret_cast(verified.get())); + STARTUPINFOW startup{}; + startup.cb = sizeof(startup); + PROCESS_INFORMATION info{}; + if (!CreateProcessW(executable.c_str(), command.data(), nullptr, nullptr, TRUE, CREATE_NO_WINDOW, + nullptr, nullptr, &startup, &info)) { + return Fail("CreateProcessW failed"); + } + UniqueHandle process(info.hProcess); + UniqueHandle thread(info.hThread); + if (WaitForSingleObject(process.get(), 10000) != WAIT_OBJECT_0) { + TerminateProcess(process.get(), EXIT_FAILURE); + WaitForSingleObject(process.get(), 2000); + return Fail("Child process did not exit within 10 seconds"); + } + DWORD exit_code = EXIT_FAILURE; + if (!GetExitCodeProcess(process.get(), &exit_code)) { + return Fail("GetExitCodeProcess failed"); + } + if (exit_code != EXIT_SUCCESS) { + std::fprintf(stderr, "Child process failed with exit code %lu\n", exit_code); + return EXIT_FAILURE; + } + if (WaitForSingleObject(verified.get(), 0) != WAIT_OBJECT_0) { + return Fail("Child did not verify telemetry destruction"); + } + return EXIT_SUCCESS; +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc new file mode 100644 index 0000000000000..0b7b15235ff48 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include + +#include +#include + +#include "core/platform/telemetry_1ds.h" +#include "core/platform/telemetry_1ds_platform.h" + +namespace onnxruntime { +namespace { + +void Check(bool condition, const char* message) { + if (!condition) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + std::fflush(stderr); + // Also fail reliably from a DLL process-detach destructor. + TerminateProcess(GetCurrentProcess(), EXIT_FAILURE); + std::abort(); + } +} + +} // namespace + +class OneDsTelemetryTestPeer { + public: + static void Install(Microsoft::Applications::Events::ILogManager& manager) { + Check(OneDsTelemetry::global_register_count_ == 1, "Expected exactly one telemetry instance"); + Check(!OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_, "Real uploader must be disabled"); + OneDsTelemetry::config_ = std::make_unique(); + OneDsTelemetry::log_manager_ = &manager; + } + + static bool IsClean() { + return OneDsTelemetry::global_register_count_ == 0 && + !OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_; + } +}; + +namespace { + +class MockLogManager : public Microsoft::Applications::Events::NullLogManager { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MockLogManager); + + explicit MockLogManager(bool pending) : callback_(CreateEventW(nullptr, TRUE, !pending, nullptr)) { + Check(callback_ != nullptr, "CreateEventW failed"); + } + + ~MockLogManager() override { CloseHandle(callback_); } + + Microsoft::Applications::Events::status_t Flush() override { + ++flush_count; + return Microsoft::Applications::Events::STATUS_SUCCESS; + } + + void FlushAndTeardown() override { + Check(flush_count == 1, "Expected Flush before FlushAndTeardown"); + ++teardown_count; + // A pending callback cannot arrive after ExitProcess has killed the other + // threads. Bound the mock's wait so a regression does not leave a stuck DLL. + Check(WaitForSingleObject(callback_, 2000) == WAIT_OBJECT_0, "Timed out draining the pending uploader"); + } + + int flush_count = 0; + int teardown_count = 0; + + private: + HANDLE callback_; +}; + +class TelemetryLifetime { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TelemetryLifetime); + + TelemetryLifetime(bool process_exit, HANDLE verified) + : manager_(process_exit), process_exit_(process_exit), verified_(verified) { + telemetry_.emplace(); + OneDsTelemetryTestPeer::Install(manager_); + } + + ~TelemetryLifetime() { + if (telemetry_) { + Destroy(); + } + } + + void Destroy() { + Check(telemetry_internal::IsDllShutdownInProgress() == process_exit_, "Unexpected DLL shutdown state"); + telemetry_.reset(); + const int expected_calls = process_exit_ ? 0 : 1; + Check(manager_.flush_count == expected_calls, "Unexpected Flush count"); + Check(manager_.teardown_count == expected_calls, "Unexpected FlushAndTeardown count"); + if (!process_exit_) { + Check(OneDsTelemetryTestPeer::IsClean(), "Telemetry state was not cleaned up"); + } + Check(SetEvent(verified_) != FALSE, "Could not signal verified destruction"); + } + + private: + MockLogManager manager_; + std::optional telemetry_; + bool process_exit_; + HANDLE verified_; +}; + +} // namespace +} // namespace onnxruntime + +extern "C" void OrtTestTelemetryShutdown(bool process_exit, bool destroy_now, HANDLE verified) { + // A DLL-owned static exercises CRT destruction during real Windows shutdown. + static onnxruntime::TelemetryLifetime lifetime(process_exit, verified); + if (destroy_now) { + lifetime.Destroy(); + } +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def new file mode 100644 index 0000000000000..517b8cf037088 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def @@ -0,0 +1,3 @@ +LIBRARY onnxruntime_telemetry_shutdown_test_library +EXPORTS + OrtTestTelemetryShutdown From 897462aa9ad4c041ae4feca477bbcb9457bfc70c Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:10:31 +0800 Subject: [PATCH 03/35] [WebGPU] Fuse rank-32 FP16 LoRA with MatMulNBits Add an opt-in MatMulNBitsLoRA operator and graph fusion for INT4/block-32 base weights with replaceable FP16 rank-32 adapters. Empty adapters execute only the base projection. Match original and offline-transposed layouts, absorb private adapter Transpose nodes, and preserve overridable inputs. Use two fusion implementations on the same transposed model: - fuse_down: base + T, followed by base + T * Bt^T. - fuse_up: T, followed by base with an integrated T * Bt^T epilogue. Default session.webgpu.lora_fusion_mode to auto: select fuse_up for M > 1 prefill and fuse_down for M == 1 decode on each invocation. Use X's per-matrix row dimension, independent of batch size. Retain explicit fuse_down/fuse_up overrides for comparisons. A one-token prompt uses the decode path. Mode selection applies to transposed adapters only. Both implementations read transposed B and T as vec4 FP16 components. Share the decode GEMV reduction and prefill down projection, preserving FP16 rounding, 128-thread decode producers and 64-thread standalone output. Keep DP4A activation quantization for the base, with original FP16 X for T. Name internal flags has_lora_down/has_lora_up by the fused projection. Retain compatible original-layout and base-kernel fallback paths. Validation on Windows/D3D12, RTX 5080, built-in and dynamic-plugin WebGPU: - 26 focused native tests pass with no skips; cover omitted mode, auto, both forced paths, decode, prefill, batches and empty/partial adapters. - 192 auto/default full-model cases and 4024 full-vocabulary logit vectors match the saved original unfused FP16 reference bitwise; rewind matches. - Cover adapter on/off and graph capture on/off for both providers. - Traces/shader dumps confirm 196 fused projections switching from prefill fuse_up to decode fuse_down, vec4 bindings, no adapter transpose, and base-only work with LoRA off. - All 72 native 1024-input/300-output benchmark sequences match exactly. --- docs/WebGPU_MatMulNBitsLoRA_Prototype.md | 116 +++++++ .../quantization/dp4a_matmul.wgsl.template | 63 +++- .../webgpu/quantization/dp4a_matmul_nbits.cc | 46 ++- .../webgpu/quantization/dp4a_matmul_nbits.h | 55 ++-- .../dp4a_matmul_small_m.wgsl.template | 32 +- .../webgpu/quantization/matmul_nbits.cc | 98 ++++-- .../webgpu/quantization/matmul_nbits.h | 39 ++- .../quantization/matmul_nbits.wgsl.template | 21 ++ .../webgpu/quantization/matmul_nbits_lora.cc | 303 ++++++++++++++++++ .../matmul_nbits_wide_tile.wgsl.template | 21 ++ .../webgpu/webgpu_contrib_kernels.cc | 1 + .../core/graph/contrib_ops/contrib_defs.cc | 50 +++ .../core/optimizer/graph_transformer_utils.cc | 7 + .../optimizer/matmul_nbits_lora_fusion.cc | 188 +++++++++++ .../core/optimizer/matmul_nbits_lora_fusion.h | 20 ++ .../contrib_ops/matmul_nbits_lora_test.cc | 66 ++++ .../matmul_nbits_lora_fusion_test.cc | 294 +++++++++++++++++ 17 files changed, 1343 insertions(+), 77 deletions(-) create mode 100644 docs/WebGPU_MatMulNBitsLoRA_Prototype.md create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc create mode 100644 onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc create mode 100644 onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc create mode 100644 onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc diff --git a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md new file mode 100644 index 0000000000000..e802852b6079d --- /dev/null +++ b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md @@ -0,0 +1,116 @@ +# WebGPU MatMulNBitsLoRA prototype + +This opt-in prototype replaces `MatMulNBits(X, W) + MatMul(MatMul(X, A), B)` with +one `com.microsoft.MatMulNBitsLoRA` node during session optimization. It retains +the packed base weights and the replaceable FP16 adapter inputs. It does not +merge or requantize the adapter into the base weights, or change the source +model file. + +The current matcher supports FP16 rank-32 adapters, INT4 base weights with block +size 32, an implicit zero point, no base bias, and no separate adapter scaling +node. Both MatMuls must use the same activation as the base projection, and +their intermediate outputs must have no other consumers or graph outputs. + +## Enable in native GenAI + +Add this entry to the decoder's `session_options`, or supply it in the native +benchmark's configuration overlay: + +```json +{ + "model": { + "decoder": { + "session_options": { + "session.enable_matmul_nbits_lora_fusion": "1" + } + } + } +} +``` + +Use the original model directory and original `.onnx_adapter` file. The +transformer is registered only when the selected WebGPU provider advertises +the new kernel. Both the ORT host and the WebGPU plugin must come from this +branch for dynamic plugin use. Setting the option to `0` restores the original +graph for a matched comparison. + +## Execution + +With an adapter, the kernel computes `T = X * A`, then incorporates `T * B` +and the addition into the base kernel's output computation. The small +intermediate `T` remains FP16. The base result and adapter result are rounded +separately before addition. Decode uses two GPU dispatches per projection. +DP4A prefill retains the existing activation quantization dispatch, so it uses +three. Subgroup matrix base kernels use a separate output epilogue fallback. + +For offline-transposed adapters, both fusion modes consume the same `At[32,K]` +and `Bt[N,32]` tensors directly. The optimizer absorbs the adapter transpose +wrappers; neither mode dispatches an adapter transpose. Select the mode with +`session.webgpu.lora_fusion_mode` in the same decoder session options: + +| Mode | First dispatch | Second dispatch | +|---|---|---| +| `fuse_down` | Base + `T = X * At^T` | `Y = base + T * Bt^T` | +| `fuse_up` | `T = X * At^T` | Base + `Y = base + T * Bt^T` | + +The default setting, `auto`, selects `fuse_up` for prefill (`M > 1`) and +`fuse_down` for decode (`M == 1`) on each invocation. Here `M` is the +second-to-last dimension of X, so batched one-token decode also selects +`fuse_down`. A one-token prompt uses the same path as decode. This setting +chooses between the two existing implementations; it adds no shader or GPU +dispatch. Explicit `fuse_down` and `fuse_up` settings force that path for all +input sizes and remain available for comparisons. + +On NVIDIA decode (one input row, K from 2048 through 8192), `fuse_down` uses +`MatMulNBitsLoRADown` with eight additional workgroups that calculate T once, +followed by `LoRAOutputVec4`. `fuse_up` uses `LoRADownGemv`, followed by +`MatMulNBitsLoRAUp`. Both use the same FP32 GEMV reduction and FP16 T, with +128 threads per producer workgroup. The standalone output shader retains +64 threads and one output per thread. Both output implementations load Bt and +T as vec4 along rank, preserving scalar rank accumulation order and separate +FP16 rounding before the base/adapter addition. + +DP4A prefill retains its activation quantization dispatch. `fuse_down` adds +the tiled down projection to the base dispatch; `fuse_up` runs that same down +projection separately and integrates the vectorized up projection into DP4A. +Other shapes use the existing base kernels, with a separate output fallback +when the selected kernel has no integrated up projection. The mode selection +applies to transposed adapters; original-layout models retain their compatible +down-then-base/up implementation. Empty adapters execute only the base kernel +in either mode. Unsupported mode strings produce an error at initialization. + +For an overridable pair of private, all-zero default initializers, the +transformer replaces the defaults in the session graph with empty rank +dimensions and relaxes those graph-input dimensions. Input names, element +types, and non-rank dimensions remain intact. A zero rank in either adapter +input selects the ordinary base kernel directly. The zero check happens once +during optimization; no adapter data is read back from the GPU during inference. +Nonzero or shared defaults are preserved and execute the adapter calculation. + +Eliding zero default branches assumes finite activations. IEEE NaN/Inf +propagation through multiplication by zero can differ when that branch is +removed, so this optimization is disabled by default. + +Adapters can be replaced between requests on the same model. Captured execution +must use a distinct capture ID when adapter bindings or the active/off state +change, as with ordinary WebGPU graph capture. GenAI generators own separate +capture IDs; create a new generator when selecting another adapter and use +rewind only while retaining that generator's adapter selection. + +## Tests + +Build the native ORT host and provider test targets, then run from the binary +directory: + +```text +onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsLoRA.* +onnxruntime_test_all.exe --gtest_filter=*MatMulNBitsLoRA* +``` + +The tests cover empty and partially supplied adapters, numerical equivalence +with the original subgraph, adapter replacement, overridable input retention, +invalid shapes, and patterns that must not fuse. Native GenAI evaluation must +also check capture on/off, generator reuse, and the complete model before +interpreting performance results. In a correctness harness, use +`ConfigSetProviderOption` to change `enableGraphCapture`; merely appending a +provider option through a JSON overlay can retain the model's earlier setting. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 19108b848c420..0833372c4e04c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -4,6 +4,8 @@ #param acc_f32 #param block_size #param n_bits +#param has_lora_up +#param lora_transposed #param has_bias #param has_zero_points #param is_qualcomm @@ -12,6 +14,35 @@ #use .getByOffset .setByOffset +fn write_lora_output(offset: u32, value: vec4) { +#if has_lora_up + let row = (offset * 4u) / uniforms.N; + let col = (offset * 4u) % uniforms.N; + var delta = vec4(0); +#if lora_transposed + for (var r4 = 0u; r4 < 8u; r4++) { + let tv = lora_t[row * 8u + r4]; + let b0 = lora_b[col * 8u + r4]; + let b1 = lora_b[(col + 1u) * 8u + r4]; + let b2 = lora_b[(col + 2u) * 8u + r4]; + let b3 = lora_b[(col + 3u) * 8u + r4]; + delta = vec4(b0.x, b1.x, b2.x, b3.x) * tv.x + delta; + delta = vec4(b0.y, b1.y, b2.y, b3.y) * tv.y + delta; + delta = vec4(b0.z, b1.z, b2.z, b3.z) * tv.z + delta; + delta = vec4(b0.w, b1.w, b2.w, b3.w) * tv.w + delta; + } +#else + for (var r = 0u; r < 32u; r++) { + delta = vec4(lora_b[r * uniforms.N + col], lora_b[r * uniforms.N + col + 1u], lora_b[r * uniforms.N + col + 2u], lora_b[r * uniforms.N + col + 3u]) * lora_t[row * 32u + r] + delta; + } +#endif + output.setByOffset(offset, value + delta); +#else + output.setByOffset(offset, value); +#endif +} + + #include "quantization/dp4a_matmul_common.wgsl.template" // This shader implements co-operative matrix multiply. The key idea here is to @@ -467,15 +498,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - output.setByOffset(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); - output.setByOffset(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); - output.setByOffset(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); - output.setByOffset(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); + write_lora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); + write_lora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); + write_lora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); + write_lora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); #else - output.setByOffset(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); - output.setByOffset(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); - output.setByOffset(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); - output.setByOffset(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); + write_lora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); + write_lora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); + write_lora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); + write_lora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); #endif #else #if has_bias @@ -506,15 +537,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - output.setByOffset(output_idx, vec4(lane_output1 + bias_vec1)); - output.setByOffset(output_idx+1, vec4(lane_output2 + bias_vec2)); - output.setByOffset(output_idx+2, vec4(lane_output3 + bias_vec3)); - output.setByOffset(output_idx+3, vec4(lane_output4 + bias_vec4)); + write_lora_output(output_idx, vec4(lane_output1 + bias_vec1)); + write_lora_output(output_idx+1, vec4(lane_output2 + bias_vec2)); + write_lora_output(output_idx+2, vec4(lane_output3 + bias_vec3)); + write_lora_output(output_idx+3, vec4(lane_output4 + bias_vec4)); #else - output.setByOffset(output_idx, vec4(lane_output1)); - output.setByOffset(output_idx+1, vec4(lane_output2)); - output.setByOffset(output_idx+2, vec4(lane_output3)); - output.setByOffset(output_idx+3, vec4(lane_output4)); + write_lora_output(output_idx, vec4(lane_output1)); + write_lora_output(output_idx+1, vec4(lane_output2)); + write_lora_output(output_idx+2, vec4(lane_output3)); + write_lora_output(output_idx+3, vec4(lane_output4)); #endif #endif } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index d9c5f5839a3bf..4438bf556ed44 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -3,6 +3,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" #include "core/providers/webgpu/shader_helper.h" namespace onnxruntime { @@ -33,15 +34,24 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_lora_up_) { + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); + if (has_lora_down_) { + GenerateLoRADown(shader, "uniforms.batch_count * uniforms.num_M_tile * uniforms.num_N_tile"); + } return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(block_size, block_size_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), WGSL_TEMPLATE_VARIABLE(a, a), @@ -66,6 +76,10 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_lora_up_) { + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); @@ -76,9 +90,11 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), @@ -106,7 +122,11 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, - const Tensor* weight_index_indirect) { + const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, + const Tensor* lora_at, Tensor* lora_down, bool lora_transposed) { + const bool has_lora_down = lora_at != nullptr; + ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && M >= min_M_for_tile_optimization), + "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; constexpr uint32_t kU32Components = 4; @@ -126,6 +146,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor .AddUniformVariable({batch_count * M * K / kU32Components}); ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program)); + const bool has_lora_up = lora_t != nullptr; const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; @@ -138,7 +159,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components); const bool broadcast_a = dispatch_M > M; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32}; + DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile); @@ -148,7 +169,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -158,6 +179,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_lora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + } return context.RunProgram(mul_program); } @@ -167,9 +192,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; + DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed}; mul_program.SetWorkgroupSize(256); - mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile); + mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile + (has_lora_down ? CeilDiv(batch_count * M, 32u) : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, @@ -185,7 +210,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); + .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -195,6 +220,15 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_lora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + } + if (has_lora_down) { + mul_program.AddInput({a, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type, 4}); + } return context.RunProgram(mul_program); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 66dc5819e37a8..f5b0004faa980 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -26,15 +26,19 @@ class DP4AMatMulNBitsProgram final : public Program { DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32) : Program{"DP4AMatMulNBits"}, - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32) {} + bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false) : Program{has_lora_down ? "DP4AMatMulNBitsLoRADown" : has_lora_up ? "DP4AMatMulNBitsLoRAUp" + : "DP4AMatMulNBits"}, + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32), + has_lora_up_(has_lora_up), + has_lora_down_(has_lora_down), + lora_transposed_(lora_transposed) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"batch_count", ProgramUniformVariableDataType::Uint32}, @@ -57,6 +61,9 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_weight_idx_indirect_; bool is_qualcomm_; bool acc_f32_; + bool has_lora_up_; + bool has_lora_down_; + bool lora_transposed_; }; class DP4AMatMulNBitsSmallMProgram final : public Program { @@ -65,17 +72,19 @@ class DP4AMatMulNBitsSmallMProgram final : public ProgramShape(), b_shape, false, true)); + const bool has_lora_up = lora_t != nullptr; + const bool has_lora_down = lora_a_transposed != nullptr; + ORT_RETURN_IF_NOT(!has_lora_down || (lora_a_output != nullptr && !has_lora_up && a->Shape().Size() == K_op && + nbits == 4 && block_size_op == 32), + "Unsupported joint LoRA projection."); const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; @@ -226,17 +249,25 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // apple|intel - Experimental dawn support for subgroup matrix matmul. std::optional subgroup_matrix_config; - if (CanApplySubgroupMatrixMatMulNBits(context, - accuracy_level, - block_size, - batch_count, - N, - K, - static_cast(nbits), - y->DataType() == DataTypeImpl::GetType(), - subgroup_matrix_config, - M, - has_weight_idx_indirect)) { + if (!has_lora_down && CanApplySubgroupMatrixMatMulNBits(context, + accuracy_level, + block_size, + batch_count, + N, + K, + static_cast(nbits), + y->DataType() == DataTypeImpl::GetType(), + subgroup_matrix_config, + M, + has_weight_idx_indirect)) { + if (has_lora_up) { + auto base_output = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, + static_cast(nbits), zero_blocks_per_col, + *subgroup_matrix_config, context, &base_output, weight_index, + weight_index_indirect)); + return ApplyLoRAOutput(context, &base_output, lora_t, lora_b, y, lora_transposed); + } return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), zero_blocks_per_col, *subgroup_matrix_config, context, y, weight_index, @@ -245,20 +276,20 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of length of M. // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). - if (CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, - M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect); + if (!has_lora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, + M, has_weight_idx_indirect, y)) { + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed); } // WideTileProgram // This program is optimized for Block32 prefill. - const bool use_wide_tile_program = CanApplyWideTileMatMulNBits(M, - K, - block_size, - nbits, - has_weight_idx_indirect, - components_a, - components_b); + const bool use_wide_tile_program = !has_lora_down && CanApplyWideTileMatMulNBits(M, + K, + block_size, + nbits, + has_weight_idx_indirect, + components_a, + components_b); const bool acc_f32 = context.EnableMatmulFp32Accumulation(); @@ -283,7 +314,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, ? context.AdapterInfo().subgroupMinSize : 0u; - MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, static_cast(nbits), subgroup_min_size, acc_f32}; + MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, static_cast(nbits), subgroup_min_size, acc_f32, has_lora_up, lora_transposed}; program.SetWorkgroupSize(workgroup_size); program.SetDispatchGroupSize(num_N_tile, num_M_tile, batch_count); @@ -324,8 +355,12 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {num_N_tile}, {num_M_tile}, {weight_index}}); - program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32); + program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32, has_lora_up, lora_transposed); + if (has_lora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + } return context.RunProgram(program); } @@ -339,10 +374,10 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, constexpr uint32_t kU32Components = 4; uint32_t components_b_with_u32 = components_b * kU32Components; uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; - MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32}; + MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed}; program.SetWorkgroupSize(workgroup_size); uint32_t num_N_tile = (N + tile_size - 1) / tile_size; - program.SetDispatchGroupSize(num_N_tile, dispatch_M, batch_count); + program.SetDispatchGroupSize(num_N_tile + (has_lora_down ? 8 : 0), dispatch_M, batch_count); program .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_a)}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_b_with_u32)}, @@ -360,7 +395,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {batch_count}, {weight_index}, {dispatch_M}}) - .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32); + .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed); if (has_zero_points) { program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -370,6 +405,15 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, if (has_weight_idx_indirect) { program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_lora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + } + if (has_lora_down) { + ORT_RETURN_IF_NOT(components_a == 4, "Joint LoRA projection requires vec4 X."); + program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, 4}); + program.AddOutput({lora_a_output, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(program); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index 06e43ce949d64..d75d1e7644452 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -9,13 +9,19 @@ namespace onnxruntime { namespace contrib { namespace webgpu { +// Adds an independent FP16 rank-32 down-projection branch to a 256-thread shader. +void GenerateLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups); + +// Adds the same FP32 GEMV reduction to standalone and fused 128-thread shaders. +void GenerateLoRAGemv(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, + const std::string& x_value); using namespace onnxruntime::webgpu; class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32) - : Program{"MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32) {} + MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, bool has_lora_up = false, bool lora_transposed = false) + : Program{"MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32), has_lora_up_(has_lora_up), lora_transposed_(lora_transposed) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"Batch", ProgramUniformVariableDataType::Uint32}, @@ -39,12 +45,28 @@ class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false) - : Program{"MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), has_bias_(has_bias), has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, single_scale_weights_(single_scale_weights), tile_size_k_vec_(tile_size_k_vec), broadcast_a_row_(broadcast_a_row), acc_f32_(acc_f32) {} + MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false) + : Program{has_lora_down ? "MatMulNBitsLoRADown" : has_lora_up ? "MatMulNBitsLoRAUp" + : "MatMulNBits"}, + tile_size_(tile_size), + nbits_(nbits), + has_zero_points_(has_zero_points), + has_bias_(has_bias), + has_weight_idx_{has_weight_idx}, + has_weight_idx_indirect_{has_weight_idx_indirect}, + single_scale_weights_(single_scale_weights), + tile_size_k_vec_(tile_size_k_vec), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32), + has_lora_up_(has_lora_up), + has_lora_down_(has_lora_down), + lora_transposed_(lora_transposed) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"M", ProgramUniformVariableDataType::Uint32}, @@ -71,6 +93,9 @@ class MatMulNBitsProgram final : public Program { uint32_t tile_size_k_vec_; bool broadcast_a_row_; bool acc_f32_; + bool has_lora_up_; + bool has_lora_down_; + bool lora_transposed_; }; class MatMulNBits final : public WebGpuKernel { @@ -99,7 +124,11 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, int64_t bits_op, onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index = 0, const Tensor* weight_index_indirect = nullptr, - uint32_t override_M = 0); + uint32_t override_M = 0, + const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, + const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, bool lora_transposed = false); + +Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed = false); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template index b75bbbebcddc2..e38088548d63e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template @@ -13,6 +13,8 @@ #param tile_size_k_vec #param tile_size_k #param tile_size +#param has_lora_up +#param lora_transposed #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -264,7 +266,26 @@ $MAIN { output_value += acc_element_t(bias[b_global + b_bias_offset]); #endif // Downcast to the output type only at the final write. +#if has_lora_up + var delta = output_element_t(0); +#if lora_transposed + for (var r4 = 0u; r4 < 8u; r4++) { + let bv = lora_b[b_global * 8u + r4]; + let tv = lora_t[(batch * uniforms.M + a_global) * 8u + r4]; + delta = bv.x * tv.x + delta; + delta = bv.y * tv.y + delta; + delta = bv.z * tv.z + delta; + delta = bv.w * tv.w + delta; + } +#else + for (var r = 0u; r < 32u; r++) { + delta = lora_b[r * uniforms.N + b_global] * lora_t[(batch * uniforms.M + a_global) * 32u + r] + delta; + } +#endif + output.setByOffset(output_idx, output_element_t(output_value) + delta); +#else output.setByOffset(output_idx, output_element_t(output_value)); +#endif } } } // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc new file mode 100644 index 0000000000000..40564a57ea205 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc @@ -0,0 +1,303 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/common/safeint.h" +#include "core/providers/webgpu/math/matmul.h" +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "core/providers/webgpu/shader_helper.h" + +namespace onnxruntime::contrib::webgpu { + +void GenerateLoRADown(ShaderHelper& shader, const std::string& base_groups) { + shader.AddInput("lora_x"); + shader.AddInput("lora_at"); + shader.AddOutput("lora_down", ShaderUsage::UseValueTypeAlias); + shader.AdditionalImplementation() << R"( +var lora_x_tile: array, 256>; +var lora_a_tile: array, 256>; +)"; + shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" + << R"( +if (workgroup_idx >= lora_base_groups) { + let group = workgroup_idx - lora_base_groups; + let row_local = local_idx / 8u; + let k4_local = local_idx % 8u; + let row = group * 32u + row_local; + let rank_base = k4_local * 4u; + var acc = vec4(0); + for (var start = 0u; start < uniforms.K / 4u; start += 8u) { + var xv = vec4(0); + if (row < uniforms.M * uniforms.batch_count) { + xv = lora_x[row * (uniforms.K / 4u) + start + k4_local]; + } + lora_x_tile[local_idx] = xv; + lora_a_tile[local_idx] = lora_at[row_local * (uniforms.K / 4u) + start + k4_local]; + workgroupBarrier(); + for (var kk = 0u; kk < 8u; kk++) { + let x = lora_x_tile[row_local * 8u + kk]; + let a0 = lora_a_tile[rank_base * 8u + kk]; + let a1 = lora_a_tile[(rank_base + 1u) * 8u + kk]; + let a2 = lora_a_tile[(rank_base + 2u) * 8u + kk]; + let a3 = lora_a_tile[(rank_base + 3u) * 8u + kk]; + acc = vec4(a0.x, a1.x, a2.x, a3.x) * x.x + acc; + acc = vec4(a0.y, a1.y, a2.y, a3.y) * x.y + acc; + acc = vec4(a0.z, a1.z, a2.z, a3.z) * x.z + acc; + acc = vec4(a0.w, a1.w, a2.w, a3.w) * x.w + acc; + } + workgroupBarrier(); + } + if (row < uniforms.M * uniforms.batch_count) { + lora_down[row * 8u + k4_local] = acc; + } + return; +} +)"; +} + +void GenerateLoRAGemv(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value) { + shader.AdditionalImplementation() << "var lora_partials: array, 128>;\n"; + shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" + << R"( + // Extra workgroups compute each rank output once, alongside independent base + // tiles. The next dispatch supplies the global dependency before consuming T. + if (workgroup_idx >= lora_base_groups) { + let lane = local_idx % 32u; + let rank = (workgroup_idx - lora_base_groups) * 4u + local_idx / 32u; + var acc = vec4(0); + for (var k4 = lane; k4 < uniforms.K / 4u; k4 += 32u) { + acc = fma(vec4()" << x_value + << R"(), vec4(lora_at[rank * (uniforms.K / 4u) + k4]), acc); + } + lora_partials[local_idx] = acc; + workgroupBarrier(); + for (var stride = 16u; stride > 0u; stride /= 2u) { + if (lane < stride) { lora_partials[local_idx] += lora_partials[local_idx + stride]; } + workgroupBarrier(); + } + if (lane == 0u) { + let v = lora_partials[local_idx]; + lora_partials[local_idx] = vec4(v.x + v.z, v.y + v.w, 0, 0); + } + // Materialize the penultimate reduction to prevent reassociation of the + // final four terms, matching the original GEMV's workgroup barriers. + workgroupBarrier(); + if (lane == 0u) { + let v = lora_partials[local_idx]; + lora_down[rank] = lora_down_element_t(v.x + v.y); + } + return; + } +)"; +} + +class LoRAGemvProgram final : public Program { + public: + LoRAGemvProgram() : Program{"LoRADownGemv"} {} + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("lora_x"); + shader.AddInput("lora_at"); + shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); + GenerateLoRAGemv(shader, "0u", "lora_x[k4]"); + return Status::OK(); + } +}; + +class LoRADownProgram final : public Program { + public: + LoRADownProgram() : Program{"LoRADownVec4"} {} + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + GenerateLoRADown(shader, "0u"); + return Status::OK(); + } +}; + +// Fallback epilogue for base kernels without an integrated adapter output path. +class LoRAOutputProgram final : public Program { + public: + explicit LoRAOutputProgram(bool transposed = false) + : Program{transposed ? "LoRAOutputVec4" : "LoRAOutput"}, transposed_(transposed) {} + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("base", ShaderUsage::UseElementTypeAlias); + shader.AddInput("t"); + shader.AddInput("b"); + shader.AddOutput("output"); + shader.MainFunctionBody() + << "if (global_idx >= uniforms.size) { return; }\n" + << "let row = global_idx / uniforms.N; let col = global_idx % uniforms.N;\n" + << "var delta = base_element_t(0);\n"; + if (transposed_) { + shader.MainFunctionBody() + << "for (var r4 = 0u; r4 < 8u; r4++) {\n" + << " let bv = b[col * 8u + r4]; let tv = t[row * 8u + r4];\n" + << " delta = bv.x * tv.x + delta; delta = bv.y * tv.y + delta;\n" + << " delta = bv.z * tv.z + delta; delta = bv.w * tv.w + delta;\n}\n"; + } else { + shader.MainFunctionBody() + << "for (var r = 0u; r < 32u; r++) { delta = b[r * uniforms.N + col] * t[row * 32u + r] + delta; }\n"; + } + shader.MainFunctionBody() << "output[global_idx] = base[global_idx] + delta;\n"; + return Status::OK(); + } + + private: + bool transposed_; +}; + +Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed) { + const auto size = narrow(y->Shape().Size()); + LoRAOutputProgram program(transposed); + program.AddInputs({{base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {t, ProgramTensorMetadataDependency::Type, transposed ? 4 : 1}, + {b, ProgramTensorMetadataDependency::Type, transposed ? 4 : 1}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) + .AddUniformVariables({{size}, {narrow(b->Shape()[transposed ? 0 : 1])}}) + .SetDispatchGroupSize(CeilDiv(size, 64u)); + return context.RunProgram(program); +} + +Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const Tensor* x, const Tensor* weight, + const Tensor* scales, const Tensor* at, const Tensor* bt, Tensor* y, + int64_t k, int64_t n, int64_t accuracy_level, bool fuse_up) { + const auto rows = narrow(x->Shape().Size() / k); + const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); + const bool use_gemv = rows == 1 && k >= 2048 && k <= 8192 && + context.AdapterInfo().vendor == std::string_view{"nvidia"}; + auto down_shape = x->Shape().AsShapeVector(); + down_shape.back() = 32; + auto t = context.CreateGPUTensor(x->DataType(), TensorShape(down_shape)); + if (fuse_up) { + if (use_gemv) { + LoRAGemvProgram down; + down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, + {at, ProgramTensorMetadataDependency::Type, 4}}) + .AddOutput({&t, ProgramTensorMetadataDependency::Type}) + .SetWorkgroupSize(128) + .SetDispatchGroupSize(8) + .AddUniformVariables({{narrow(k)}}); + ORT_RETURN_IF_ERROR(context.RunProgram(down)); + } else { + LoRADownProgram down; + down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, + {at, ProgramTensorMetadataDependency::Type, 4}}) + .AddOutput({&t, ProgramTensorMetadataDependency::Type, 4}) + .SetWorkgroupSize(256) + .SetDispatchGroupSize(CeilDiv(rows, 32u)) + .AddUniformVariables({{rows}, {narrow(k)}, {1u}}); + ORT_RETURN_IF_ERROR(context.RunProgram(down)); + } + return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, + context, y, 0, nullptr, 0, &t, bt, nullptr, nullptr, true); + } + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + if (use_gemv) { + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, + context, &base, 0, nullptr, 0, nullptr, nullptr, at, &t)); + } else if (m >= kMinMForTileOptimization && context.AdapterInfo().vendor == std::string_view{"nvidia"} && + CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, 32, narrow(n), narrow(k), + 4, m, false, y)) { + // Preserve DP4A's activation quantization and compute T from original FP16 X + // in extra workgroups of the base dispatch. + ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits(x, weight, scales, nullptr, nullptr, rows / m, m, m, + narrow(n), narrow(k), 32, 0, + kMinMForTileOptimization, 4, context, &base, 0, + nullptr, nullptr, nullptr, at, &t)); + } else { + LoRADownProgram down; + down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, + {at, ProgramTensorMetadataDependency::Type, 4}}) + .AddOutput({&t, ProgramTensorMetadataDependency::Type, 4}) + .SetWorkgroupSize(256) + .SetDispatchGroupSize(CeilDiv(rows, 32u)) + .AddUniformVariables({{rows}, {narrow(k)}, {1u}}); + ORT_RETURN_IF_ERROR(context.RunProgram(down)); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, + context, &base)); + } + return ApplyLoRAOutput(context, &base, &t, bt, y, true); +} + +class MatMulNBitsLoRA final : public WebGpuKernel { + public: + explicit MatMulNBitsLoRA(const OpKernelInfo& info) + : WebGpuKernel(info), k_(info.GetAttr("K")), n_(info.GetAttr("N")), accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)), transposed_(info.GetAttrOrDefault("adapter_transposed", 0)) { + const auto path = info.GetConfigOptions().GetConfigOrDefault("session.webgpu.lora_fusion_mode", "auto"); + ORT_ENFORCE(path == "auto" || path == "fuse_down" || path == "fuse_up", + "session.webgpu.lora_fusion_mode must be auto, fuse_down or fuse_up."); + fusion_mode_ = path == "auto" ? LoRAFusionMode::Auto : path == "fuse_up" ? LoRAFusionMode::FuseUp + : LoRAFusionMode::FuseDown; + ORT_ENFORCE(transposed_ == 0 || transposed_ == 1, "Invalid adapter_transposed value."); + ORT_ENFORCE(info.GetAttr("bits") == 4 && info.GetAttr("block_size") == 32, + "MatMulNBitsLoRA currently requires INT4 block size 32."); + ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0 && n_ % 4 == 0, "Unsupported projection dimensions."); + } + + Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override { + const auto* x = context.Input(0); + const auto* weight = context.Input(1); + const auto* scales = context.Input(2); + const auto* a = context.Input(3); + const auto* b = context.Input(4); + ORT_RETURN_IF_NOT(x->Shape().NumDimensions() >= 2 && a->Shape().NumDimensions() == 2 && + b->Shape().NumDimensions() == 2, + "Invalid LoRA input ranks."); + ORT_RETURN_IF_NOT(a->Shape()[transposed_] == k_ && b->Shape()[1 - transposed_] == n_, "Invalid LoRA K/N dimensions."); + const int64_t rank_a = a->Shape()[1 - transposed_]; + const int64_t rank_b = b->Shape()[transposed_]; + ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (rank_b == 0 || rank_b == 32), + "MatMulNBitsLoRA currently supports rank 32 or an empty adapter."); + ORT_RETURN_IF_NOT(weight->Shape() == TensorShape({n_, k_ / 32, 16}) && + scales->Shape().Size() == SafeInt(n_) * (k_ / 32), + "Invalid packed weight/scales shape."); + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n_, k_}), false, true)); + auto* y = context.Output(0, helper.OutputShape()); + if (y->Shape().Size() == 0) return Status::OK(); + if (rank_a == 0 || rank_b == 0) { + return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y); + } + if (transposed_) { + // Use the per-matrix row count so batched one-token decode still selects fuse_down. + const auto m = x->Shape()[x->Shape().NumDimensions() - 2]; + const bool fuse_up = fusion_mode_ == LoRAFusionMode::FuseUp || + (fusion_mode_ == LoRAFusionMode::Auto && m > 1); + return ApplyTransposedLoRA(context, x, weight, scales, a, b, y, k_, n_, accuracy_level_, fuse_up); + } + MatMulComputeHelper lora_helper; + ORT_RETURN_IF_ERROR(lora_helper.Compute(x->Shape(), a->Shape())); + auto t = context.CreateGPUTensor(x->DataType(), lora_helper.OutputShape()); + std::vector inputs{x, a}; + ORT_RETURN_IF_ERROR(ComputeMatMul(&context, Activation(), inputs, &t, true, cache_, false)); + return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, + context, y, 0, nullptr, 0, &t, b); + } + + private: + enum class LoRAFusionMode { Auto, + FuseDown, + FuseUp }; + + int64_t k_; + int64_t n_; + int64_t accuracy_level_; + int64_t transposed_; + LoRAFusionMode fusion_mode_; + mutable MatMulOptImplCache cache_; +}; + +ONNX_OPERATOR_KERNEL_EX( + MatMulNBitsLoRA, kMSDomain, 1, kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), + MatMulNBitsLoRA); + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template index a433f8064f5e5..61af4960c7250 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template @@ -2,6 +2,8 @@ // Licensed under the MIT License. #param has_zero_points +#param has_lora_up +#param lora_transposed #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -113,7 +115,26 @@ fn load_scale(row : u32, block_idx : u32) -> output_element_t { fn write_output(batch : u32, row : u32, col : u32, value : output_element_t) { if (batch < uniforms.Batch && row < uniforms.M && col < uniforms.N) { let offset = batch * uniforms.M * uniforms.N + row * uniforms.N + col; +#if has_lora_up + var delta = output_element_t(0); +#if lora_transposed + for (var r4 = 0u; r4 < 8u; r4++) { + let bv = lora_b[col * 8u + r4]; + let tv = lora_t[(batch * uniforms.M + row) * 8u + r4]; + delta = bv.x * tv.x + delta; + delta = bv.y * tv.y + delta; + delta = bv.z * tv.z + delta; + delta = bv.w * tv.w + delta; + } +#else + for (var r = 0u; r < 32u; r++) { + delta = lora_b[r * uniforms.N + col] * lora_t[(batch * uniforms.M + row) * 32u + r] + delta; + } +#endif + output.setByOffset(offset, value + delta); +#else output.setByOffset(offset, value); +#endif } } diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 1574ee7d0bcb3..9fef9b55601ee 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -62,6 +62,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index c3a0f94607870..08d949b83adbd 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3941,6 +3941,56 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by } }); + ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsLoRA) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc( + "Computes MatMulNBits(X, W, scales) + MatMul(MatMul(X, lora_A), lora_B). " + "With adapter_transposed=1, transposes both stored adapter inputs in this formula. " + "The adapter remains replaceable. A zero rank in either adapter input disables the adapter. " + "Intermediate results are rounded to the input type before addition. " + "This prototype supports FP16, rank 32, INT4 block size 32, implicit zero point 8, and no bias.") + .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) + .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) + .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) + .Attr("block_size", "Base weight block size; must be 32.", AttributeProto::INT, int64_t{32}) + .Attr("accuracy_level", "MatMulNBits base accumulation policy.", AttributeProto::INT, int64_t{0}) + .Attr("adapter_transposed", "Whether adapters are stored as [R_A,K] and [N,R_B].", AttributeProto::INT, int64_t{0}) + .Input(0, "X", "Input with rank at least two and last dimension K.", "T") + .Input(1, "W", "INT4 packed weights of shape [N, K/32, 16].", "TQ") + .Input(2, "scales", "FP16 block scales containing N * K/32 elements; may be flattened.", "T") + .Input(3, "lora_A", "Adapter of shape [K, R_A], or [R_A, K] if adapter_transposed=1; R_A is 0 or 32.", "T") + .Input(4, "lora_B", "Adapter of shape [R_B, N], or [N, R_B] if adapter_transposed=1; R_B is 0 or 32.", "T") + .Output(0, "Y", "Input shape with last dimension replaced by N.", "T") + .TypeConstraint("T", {"tensor(float16)"}, "FP16 activations and adapter weights.") + .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized base weights.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const auto k = getAttribute(ctx, "K", -1); + const auto n = getAttribute(ctx, "N", -1); + const auto transposed = getAttribute(ctx, "adapter_transposed", 0); + if ((transposed != 0 && transposed != 1) || k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || + getAttribute(ctx, "bits", 4) != 4 || getAttribute(ctx, "block_size", 32) != 32) { + fail_shape_inference("Unsupported MatMulNBitsLoRA attributes."); + } + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() < 2) { + fail_shape_inference("X must have rank at least two."); + } + MatmulWithQuantWeightShapeInference(ctx, k, n, true); + for (size_t i : {size_t{3}, size_t{4}}) { + if (!hasInputShape(ctx, i)) continue; + const auto& s = getInputShape(ctx, i); + if (s.dim_size() != 2) fail_shape_inference("Adapter inputs must have rank two."); + const int feature_axis = (i == 3 ? 0 : 1) ^ static_cast(transposed); + const auto& feature = s.dim(feature_axis); + const auto& rank = s.dim(1 - feature_axis); + if ((feature.has_dim_value() && feature.dim_value() != (i == 3 ? k : n)) || + (rank.has_dim_value() && rank.dim_value() != 0 && rank.dim_value() != 32)) { + fail_shape_inference("Invalid adapter dimensions."); + } + } + }); + static const char* MatMulNBitsMlp_ver1_doc = R"DOC( MatMulNBitsMlp fuses two MatMulNBits projections that share the same input and computes diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index cc2c49dbc49e3..ac2c1e1ac5d91 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -60,6 +60,7 @@ #include "core/optimizer/group_query_attention_pre_norm_fusion.h" #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/matmul_nbits_mlp_fusion.h" +#include "core/optimizer/matmul_nbits_lora_fusion.h" #include "core/optimizer/matmul_bn_fusion.h" #include "core/optimizer/matmul_integer_to_float.h" #include "core/optimizer/matmul_scale_fusion.h" @@ -465,6 +466,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kCudaExecutionProvider, onnxruntime::kWebGpuExecutionProvider})); + bool has_matmul_nbits_lora_kernel = false; bool has_matmul_nbits_mlp_kernel = false; bool has_matmul_nbits_qkv_kernel = false; if (execution_providers != nullptr) { @@ -482,11 +484,16 @@ InlinedVector> GenerateTransformers( nullptr) .IsOK(); }; + has_matmul_nbits_lora_kernel = has_webgpu_kernel("MatMulNBitsLoRA"); has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp"); has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv"); } } } + if (has_matmul_nbits_lora_kernel && + session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { + transformers.emplace_back(std::make_unique()); + } if (has_matmul_nbits_mlp_kernel) { transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); diff --git a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc new file mode 100644 index 0000000000000..d7f93540b8176 --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc @@ -0,0 +1,188 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/matmul_nbits_lora_fusion.h" + +#include "core/graph/graph_utils.h" +#include "core/graph/node_attr_utils.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/matmul_nbits_fusion_utils.h" +#include "core/optimizer/utils.h" + +namespace onnxruntime { +namespace { + +const Node* Producer(const Graph& graph, const Node& node, int input) { + const auto* edge = graph_utils::GetInputEdge(node, input); + return edge ? graph.GetNode(edge->GetNode().Index()) : nullptr; +} + +bool SingleConsumer(const Graph& graph, const Node& node) { + return !graph.NodeProducesGraphOutput(node) && optimizer_utils::CheckOutputEdges(graph, node, 1); +} + +bool IsMatMul(const Node* node) { + return node && graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMul", {9, 13}) && + node->GetExecutionProviderType() == kWebGpuExecutionProvider; +} + +bool MatrixShape(const NodeArg& arg, int64_t rows, int64_t cols) { + const auto* s = arg.Shape(); + return s && s->dim_size() == 2 && s->dim(0).has_dim_value() && s->dim(1).has_dim_value() && + s->dim(0).dim_value() == rows && s->dim(1).dim_value() == cols; +} + +const Node* AdapterTranspose(const Graph& graph, const Node& matmul) { + const auto* node = Producer(graph, matmul, 1); + if (!node || !graph_utils::IsSupportedOptypeVersionAndDomain(*node, "Transpose", {1, 13, 21, 23, 24}) || + node->GetExecutionProviderType() != kWebGpuExecutionProvider || !SingleConsumer(graph, *node)) return nullptr; + const auto* shape = node->InputDefs()[0]->Shape(); + if (!shape || shape->dim_size() != 2) return nullptr; + auto perm = node->GetAttributes().find("perm"); + if (perm != node->GetAttributes().end() && + (perm->second.ints_size() != 2 || perm->second.ints(0) != 1 || perm->second.ints(1) != 0)) return nullptr; + return node; +} + +bool ZeroDefault(const Graph& graph, const NodeArg& arg) { + if (!graph.IsInputsIncludingInitializers(&arg) || graph.IsOutput(&arg)) { + return false; + } + const ONNX_NAMESPACE::TensorProto* proto = nullptr; + if (!graph.GetInitializedTensor(arg.Name(), proto)) { + return false; + } + // Only rewrite a private adapter parameter: another consumer may depend on + // its original dimensions or on IEEE zero-times-infinity semantics. + const auto consumers = graph.GetConsumerNodes(arg.Name()); + if (consumers.size() != 1) { + return false; + } + Initializer data(graph, *proto, graph.ModelPath()); + auto bytes = data.DataAsByteSpan(); + return std::all_of(bytes.begin(), bytes.end(), [](uint8_t v) { return v == 0; }); +} + +void EmptyDefault(Graph& graph, NodeArg& arg, int rank_axis) { + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(graph.GetInitializedTensor(arg.Name(), original)); + ONNX_NAMESPACE::TensorProto empty; + empty.set_name(original->name()); + empty.set_data_type(original->data_type()); + for (int i = 0; i < original->dims_size(); ++i) { + empty.add_dims(i == rank_axis ? 0 : original->dims(i)); + } + auto shape = *arg.Shape(); + shape.mutable_dim(rank_axis)->Clear(); + arg.SetShape(shape); + graph.RemoveInitializedTensor(arg.Name()); + graph.AddInitializedTensor(empty); +} + +} // namespace + +Status MatMulNBitsLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, + const logging::Logger& logger) const { + GraphViewer viewer(graph); + for (auto index : viewer.GetNodesInTopologicalOrder()) { + auto* add = graph.GetNode(index); + if (!add) continue; + ORT_RETURN_IF_ERROR(Recurse(*add, modified, graph_level, logger)); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(*add, "Add", {7, 13, 14}) || + add->GetExecutionProviderType() != kWebGpuExecutionProvider) continue; + const Node* base = Producer(graph, *add, 0); + const Node* up = Producer(graph, *add, 1); + if (base && base->OpType() != "MatMulNBits") std::swap(base, up); + if (!base || !graph_utils::IsSupportedOptypeVersionAndDomain(*base, "MatMulNBits", {1}, kMSDomain) || + base->GetExecutionProviderType() != kWebGpuExecutionProvider || !IsMatMul(up)) continue; + const Node* down = Producer(graph, *up, 0); + if (!IsMatMul(down) || down->InputDefs()[0] != base->InputDefs()[0] || + !SingleConsumer(graph, *base) || !SingleConsumer(graph, *up) || !SingleConsumer(graph, *down)) continue; + + using namespace matmul_nbits_fusion_utils; + const int64_t k = GetIntAttr(*base, "K", -1, true); + const int64_t n = GetIntAttr(*base, "N", -1, true); + const auto* shape = down->InputDefs()[1]->Shape(); + if (!shape || shape->dim_size() != 2 || !shape->dim(1).has_dim_value()) continue; + const int64_t rank = shape->dim(1).dim_value(); + if (k <= 0 || n <= 0 || rank != 32 || k % 32 != 0 || n % 4 != 0 || + GetIntAttr(*base, "bits", 4) != 4 || GetIntAttr(*base, "block_size", -1, true) != 32 || + GetIntAttr(*base, "weight_prepacked", 0) != 0 || HasInput(*base, 3) || HasInput(*base, 4) || + HasInput(*base, 5) || !MatrixShape(*down->InputDefs()[1], k, rank) || + !MatrixShape(*up->InputDefs()[1], rank, n) || + down->InputDefs()[1]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto::FLOAT16 || + !base->InputDefs()[0]->Shape() || base->InputDefs()[0]->Shape()->dim_size() < 2) continue; + + // MatMulNBits also accepts legacy flattened weights. Leave those graphs + // alone until the fused kernel supports that representation. + const auto* weight_shape = base->InputDefs()[1]->Shape(); + if (!weight_shape || weight_shape->dim_size() != 3 || + !weight_shape->dim(0).has_dim_value() || weight_shape->dim(0).dim_value() != n || + !weight_shape->dim(1).has_dim_value() || weight_shape->dim(1).dim_value() != k / 32 || + !weight_shape->dim(2).has_dim_value() || weight_shape->dim(2).dim_value() != 16) continue; + const auto* scale_shape = base->InputDefs()[2]->Shape(); + const bool flat_scales = scale_shape && scale_shape->dim_size() == 1 && + scale_shape->dim(0).has_dim_value() && + scale_shape->dim(0).dim_value() / n == k / 32 && + scale_shape->dim(0).dim_value() % n == 0; + if (!flat_scales && !MatrixShape(*base->InputDefs()[2], n, k / 32)) continue; + + auto* a = const_cast(down->InputDefs()[1]); + auto* b = const_cast(up->InputDefs()[1]); + const auto* transpose_a = AdapterTranspose(graph, *down); + const auto* transpose_b = AdapterTranspose(graph, *up); + const bool transposed = transpose_a && transpose_b && + MatrixShape(*transpose_a->InputDefs()[0], rank, k) && + MatrixShape(*transpose_b->InputDefs()[0], n, rank); + if (transposed) { + a = const_cast(transpose_a->InputDefs()[0]); + b = const_cast(transpose_b->InputDefs()[0]); + } + const bool empty_defaults = ZeroDefault(graph, *a) && ZeroDefault(graph, *b); + InlinedVector inputs{const_cast(base->InputDefs()[0]), + const_cast(base->InputDefs()[1]), + const_cast(base->InputDefs()[2]), a, b}; + auto attrs = base->GetAttributes(); + attrs.erase("weight_prepacked"); + utils::SetNodeAttribute(utils::MakeAttribute("adapter_transposed", int64_t{transposed}), attrs); + Node& fused = graph.AddNode(graph.GenerateNodeName("MatMulNBitsLoRA"), "MatMulNBitsLoRA", + "fused unmerged rank-32 LoRA projection", inputs, + add->MutableOutputDefs(), &attrs, kMSDomain); + fused.SetExecutionProviderType(kWebGpuExecutionProvider); + // All input producer edges and final consumers must be preserved before + // removal, including weights supplied by upstream nodes. + const Node* sources[] = {base, base, base, transposed ? transpose_a : down, transposed ? transpose_b : up}; + const int slots[] = {0, 1, 2, transposed ? 0 : 1, transposed ? 0 : 1}; + for (int i = 0; i < 5; ++i) { + if (const auto* edge = graph_utils::GetInputEdge(*sources[i], slots[i])) { + graph.AddEdge(edge->GetNode().Index(), fused.Index(), edge->GetSrcArgIndex(), i); + } + } + const auto outputs = graph_utils::GraphEdge::GetNodeOutputEdges(*add); + for (const Node* node : {base, down, up, static_cast(add)}) { + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(node->Index())); + } + const NodeIndex indices[] = {base->Index(), down->Index(), up->Index(), add->Index()}; + if (transposed) { + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(transpose_a->Index())); + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(transpose_b->Index())); + } + for (auto idx : indices) graph.RemoveNode(idx); + if (transposed) { + graph.RemoveNode(transpose_a->Index()); + graph.RemoveNode(transpose_b->Index()); + } + for (const auto& edge : outputs) { + graph.AddEdge(fused.Index(), edge.dst_node, 0, edge.dst_arg_index); + } + if (empty_defaults) { + EmptyDefault(graph, *a, transposed ? 0 : 1); + EmptyDefault(graph, *b, transposed ? 1 : 0); + } + LOGS(logger, VERBOSE) << "MatMulNBitsLoRAFusion: " << fused.Name() << " empty defaults=" << empty_defaults; + modified = true; + } + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h new file mode 100644 index 0000000000000..055f7cb72bcc4 --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +// Opt-in fusion for finite-valued inference. Zero overridable defaults can be +// represented with rank zero, so the kernel can skip inactive adapters. +class MatMulNBitsLoRAFusion final : public GraphTransformer { + public: + MatMulNBitsLoRAFusion() : GraphTransformer("MatMulNBitsLoRAFusion", {kWebGpuExecutionProvider}) {} + + private: + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc new file mode 100644 index 0000000000000..1fd0bda7bb857 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) + +TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { + for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { + SCOPED_TRACE(mode[0] ? mode : "default"); + for (bool transposed : {false, true}) { + for (int64_t m : {1, 37}) { + for (auto ranks : {std::pair{0, 0}, {0, 32}, {32, 0}, {32, 32}}) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 128); + test.AddAttribute("bits", 4); + test.AddAttribute("block_size", 32); + test.AddAttribute("adapter_transposed", transposed); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, MLFloat16(0.5f))); + // Packed nibbles 9 => dequantized weight 1/32, so the base is exactly 1. + test.AddInput("W", {128, 2, 16}, std::vector(128 * 2 * 16, 0x99), true); + test.AddInput("scales", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); + test.AddInput("lora_A", transposed ? std::vector{ranks.first, 64} : std::vector{64, ranks.first}, + std::vector(64 * ranks.first, MLFloat16(0.0625f))); + test.AddInput("lora_B", transposed ? std::vector{128, ranks.second} : std::vector{ranks.second, 128}, + std::vector(128 * ranks.second, MLFloat16(0.03125f))); + const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; + test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, MLFloat16(expected))); + std::vector> eps; + eps.push_back(std::move(ep)); + SessionOptions options; + if (mode[0]) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + } + test.Config(options); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } + } + } +} + +TEST(MatMulNBitsLoRA, RejectInvalidAdapterShape) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 128); + test.AddInput("X", {1, 64}, std::vector(64)); + test.AddInput("W", {128, 2, 16}, std::vector(4096), true); + test.AddInput("scales", {128, 2}, std::vector(256), true); + test.AddInput("lora_A", {63, 32}, std::vector(63 * 32)); + test.AddInput("lora_B", {32, 128}, std::vector(32 * 128)); + test.AddOutput("Y", {1, 128}, std::vector(128)); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid adapter dimensions", {}, nullptr, &eps); +} + +#endif +} // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc new file mode 100644 index 0000000000000..be9a0a1fdf804 --- /dev/null +++ b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc @@ -0,0 +1,294 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/graph/node_attr_utils.h" +#include "core/optimizer/matmul_nbits_lora_fusion.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/optimizer/webgpu_fusion_test_util.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +#ifndef DISABLE_CONTRIB_OPS +namespace { + +void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, + bool expose_intermediate = false, bool different_x = false, int64_t rank = 32, + bool flattened_weights = false, bool transposed = false, int64_t n = 128, + int isolate = 0, int64_t k_override = 0, int64_t batch = 1) { + const int64_t k = k_override ? k_override : transposed ? 2048 + : 128; + auto* x = builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)); + auto weight_shape = flattened_weights ? std::vector{n * k / 2} : std::vector{n, k / 32, 16}; + auto* w = builder.MakeInitializer(weight_shape, isolate ? uint8_t{0x88} : uint8_t{0}, isolate ? uint8_t{0x89} : uint8_t{255}); + auto* scales = builder.MakeInitializer({n, k / 32}, MLFloat16(0.03125f), MLFloat16(0.0625f)); + auto* a = builder.MakeInput(transposed ? std::vector{rank, k} : std::vector{k, rank}, MLFloat16(-0.0625f), MLFloat16(0.0625f)); + auto* b = builder.MakeInput(transposed ? std::vector{n, rank} : std::vector{rank, n}, MLFloat16(-0.0625f), MLFloat16(0.0625f)); + if (isolate == 1 || isolate == 2 || isolate == 4) { + auto* arg = isolate == 2 ? a : b; + auto values = builder.feeds_.at(arg->Name()).GetMutable()->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), MLFloat16(0.0f)); + const int64_t cols = arg->Shape()->dim(1).dim_value(); + for (int64_t r = 0; r < rank; ++r) values[r * cols + r] = MLFloat16(1.0f); + } + if (isolate == 4) { + auto xv = builder.feeds_.at(x->Name()).GetMutable()->MutableDataAsSpan(); + auto av = builder.feeds_.at(a->Name()).GetMutable()->MutableDataAsSpan(); + std::fill(xv.begin(), xv.end(), MLFloat16(0.0f)); + std::fill(av.begin(), av.end(), MLFloat16(0.0f)); + xv[0] = MLFloat16(10000.0f); + xv[1] = MLFloat16(0.001f); + xv[2] = MLFloat16(-10000.0f); + xv[3] = MLFloat16(0.001f); + for (int i = 0; i < 4; ++i) av[i] = MLFloat16(1.0f); + } + if (zeros) { + for (auto* arg : {a, b}) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(arg->Name()); + init.set_data_type(ONNX_NAMESPACE::TensorProto::FLOAT16); + int64_t count = 1; + for (const auto& dim : arg->Shape()->dim()) { + init.add_dims(dim.dim_value()); + count *= dim.dim_value(); + } + init.set_raw_data(std::string(static_cast(count) * sizeof(MLFloat16), '\0')); + builder.graph_.AddInitializedTensor(init); + } + builder.graph_.SetInputs({x, a, b}); + } + auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* t = expose_intermediate ? builder.MakeOutput(std::vector{batch, m, rank}) + : builder.MakeIntermediate(std::vector{batch, m, rank}); + auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* y = builder.MakeOutput(std::vector{batch, m, n}); + NodeAttributes attrs; + for (auto [name, value] : {std::pair{"K", k}, {"N", n}, {"bits", int64_t{4}}, {"block_size", int64_t{32}}}) { + utils::SetNodeAttribute(utils::MakeAttribute(name, value), attrs); + } + auto& base_node = builder.AddNode("MatMulNBits", {x, w, scales}, {base}, kMSDomain); + for (const auto& [name, attr] : attrs) base_node.AddAttributeProto(attr); + base_node.AddAttribute("accuracy_level", int64_t{4}); + auto* lora_x = different_x ? builder.MakeInput({batch, m, k}, MLFloat16(0.0f), MLFloat16(1.0f)) : x; + auto* matmul_a = a; + auto* matmul_b = b; + if (transposed) { + matmul_a = builder.MakeIntermediate(std::vector{k, rank}); + matmul_b = builder.MakeIntermediate(std::vector{rank, n}); + auto& ta = builder.AddNode("Transpose", {a}, {matmul_a}); + auto& tb = builder.AddNode("Transpose", {b}, {matmul_b}); + for (auto* node : {&ta, &tb}) { + node->AddAttribute("perm", std::vector{1, 0}); + node->SetExecutionProviderType(kWebGpuExecutionProvider); + } + } + auto& down = builder.AddNode("MatMul", {lora_x, matmul_a}, {t}); + auto& up = builder.AddNode("MatMul", {t, matmul_b}, {delta}); + auto& add = builder.AddNode("Add", {delta, base}, {y}); + for (auto* node : {&base_node, &down, &up, &add}) node->SetExecutionProviderType(kWebGpuExecutionProvider); +} + +Status CheckCount(const Graph& graph, size_t count) { + size_t found = 0; + for (const auto& node : graph.Nodes()) found += node.OpType() == "MatMulNBitsLoRA"; + ORT_RETURN_IF_NOT(found == count, "Unexpected fused LoRA count: ", found); + return Status::OK(); +} + +} // namespace + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionPreservesOverridableInputs) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { + ORT_RETURN_IF_ERROR(CheckCount(g, 1)); + ORT_RETURN_IF_NOT(g.GetOverridableInitializers().size() == 2, "Adapter inputs were removed."); + for (const auto* arg : g.GetOverridableInitializers()) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(g.GetInitializedTensor(arg->Name(), init), "Missing adapter default."); + ORT_RETURN_IF_NOT(init->dims(0) == 0 || init->dims(1) == 0, "Default was not emptied."); + } + return Status::OK(); + })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedAdapters) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, true, false, false, 32, false, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { + ORT_RETURN_IF_ERROR(CheckCount(g, 1)); + for (const auto& node : g.Nodes()) { + ORT_RETURN_IF_NOT(node.OpType() != "Transpose", "Adapter transpose was not absorbed."); + if (node.OpType() == "MatMulNBitsLoRA") { + ORT_RETURN_IF_NOT(node.GetAttributes().at("adapter_transposed").i() == 1, "Missing layout flag."); + } + } + ORT_RETURN_IF_NOT(g.GetOverridableInitializers().size() == 2, "Adapter inputs were removed."); + return Status::OK(); + })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedNumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { + SCOPED_TRACE(mode[0] ? mode : "default"); + for (int64_t m : {1, 2, 37}) { + for (int64_t n : {128, 132}) { + SCOPED_TRACE(m); + SCOPED_TRACE(n); + RunWebGpuFusionTransformerTest( + [m, n](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, n); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + if (mode[0]) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + } + }); + } + } + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedProjectionNumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + // An identity up/down adapter isolates each projection. Zero base weights + // expose low-magnitude errors; the last case stresses FP32 cancellation. + for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { + SCOPED_TRACE(mode[0] ? mode : "default"); + for (int isolate : {1, 2, 3, 4}) { + for (int64_t k : {2048, 3072, 8192}) { + SCOPED_TRACE(isolate); + SCOPED_TRACE(k); + RunWebGpuFusionTransformerTest( + [isolate, k](ModelTestBuilder& b) { BuildLoRA(b, 1, true, false, false, 32, false, true, 2048, isolate, k); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + if (mode[0]) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + } + }); + } + } + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedBatchedNumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { + SCOPED_TRACE(mode[0] ? mode : "default"); + for (int64_t batch : {2, 4}) { + for (int64_t m : {1, 3, 37}) { + SCOPED_TRACE(batch); + SCOPED_TRACE(m); + RunWebGpuFusionTransformerTest( + [batch, m](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, 128, 0, 0, batch); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + if (mode[0]) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + } + }); + } + } + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsObservableIntermediate) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckCount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsDifferentActivation) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckCount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsUnsupportedRank) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, false, 16); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckCount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionLeavesLegacyPackedShapesAlone) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, false, 32, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckCount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionMatchesWebGpuDecodeAndPrefill) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (int64_t m : {1, 37}) { + SCOPED_TRACE(m); + RunWebGpuFusionTransformerTest( + [m](ModelTestBuilder& b) { BuildLoRA(b, m, true); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + Model model("LoRA switching", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + auto& graph = model.MainGraph(); + ModelTestBuilder helper(graph); + BuildLoRA(helper, 1, true); + helper.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + const auto adapters = graph.GetOverridableInitializers(); + ASSERT_EQ(adapters.size(), 2u); + std::string data; + model.ToProto().SerializeToString(&data); + std::vector expected; + for (bool fused : {false, true}) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + if (fused) { + ASSERT_STATUS_OK(session.RegisterGraphTransformer(std::make_unique(), + TransformerLevel::Level2)); + } + ASSERT_STATUS_OK(session.Initialize()); + // Off -> adapter A -> adapter B -> partially supplied -> off. + for (int state = 0; state < 5; ++state) { + auto feeds = helper.feeds_; + for (size_t i = 0; i < adapters.size(); ++i) { + const auto& name = adapters[i]->Name(); + if (state == 0 || state == 4 || (state == 3 && i == 1)) { + feeds.erase(name); + } else { + auto* tensor = feeds.at(name).GetMutable(); + auto values = tensor->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), MLFloat16(state == 2 ? -0.03125f : 0.0625f)); + } + } + std::vector outputs; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, helper.output_names_, &outputs)); + ASSERT_EQ(outputs.size(), 1u); + if (fused) { + auto result = CompareOrtValue(outputs[0], expected[state], 0.0, 0.0, false); + EXPECT_EQ(result.first, COMPARE_RESULT::SUCCESS) << "state=" << state << " " << result.second; + } else { + expected.push_back(std::move(outputs[0])); + } + } + } +} +#endif +} // namespace onnxruntime::test From 5a5d325833a29afcc376f4d2a1038054128a00c9 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:30:22 +0800 Subject: [PATCH 04/35] [WebGPU] Fuse rank-32 INT8 adapters with MatMulNBitsQLoRA Add opt-in graph fusion and WebGPU kernels for INT4 base projections with block-32 INT8 adapters. Preserve replaceable packed weights, original INT8 arithmetic, partial overrides, and base-only execution for empty adapters. Use fuse_up for prefill and fuse_down for one-row NVIDIA decode, with packed vec4 loads. Add strict numerical and graph-matching tests and document native GenAI usage. --- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 108 ++++++++ .../quantization/dp4a_matmul.wgsl.template | 5 +- .../webgpu/quantization/dp4a_matmul_nbits.cc | 30 ++- .../webgpu/quantization/dp4a_matmul_nbits.h | 58 ++-- .../dp4a_matmul_small_m.wgsl.template | 5 +- .../webgpu/quantization/matmul_nbits.cc | 51 +++- .../webgpu/quantization/matmul_nbits.h | 20 +- .../quantization/matmul_nbits.wgsl.template | 5 +- .../webgpu/quantization/matmul_nbits_lora.cc | 4 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 255 ++++++++++++++++++ .../webgpu/quantization/matmul_nbits_qlora.h | 16 ++ .../matmul_nbits_wide_tile.wgsl.template | 5 +- .../webgpu/webgpu_contrib_kernels.cc | 1 + .../core/graph/contrib_ops/contrib_defs.cc | 56 ++++ .../core/optimizer/graph_transformer_utils.cc | 7 + .../optimizer/matmul_nbits_qlora_fusion.cc | 156 +++++++++++ .../optimizer/matmul_nbits_qlora_fusion.h | 19 ++ .../contrib_ops/matmul_nbits_qlora_test.cc | 65 +++++ .../matmul_nbits_qlora_fusion_test.cc | 246 +++++++++++++++++ 19 files changed, 1052 insertions(+), 60 deletions(-) create mode 100644 docs/WebGPU_MatMulNBitsQLoRA_Prototype.md create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h create mode 100644 onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc create mode 100644 onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc create mode 100644 onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md new file mode 100644 index 0000000000000..903573f3bbc00 --- /dev/null +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -0,0 +1,108 @@ +# WebGPU MatMulNBitsQLoRA prototype + +`com.microsoft.MatMulNBitsQLoRA` fuses an INT4 base projection, two INT8 adapter +projections, and their residual Add. Adapter weights remain replaceable. This +is an inference operator for quantized adapters; it does not merge adapter +weights into the base or perform adapter training. + +Enable it with `session.enable_matmul_nbits_lora_fusion=1` in the native GenAI +decoder session options, as described in [the FP16 prototype](WebGPU_MatMulNBitsLoRA_Prototype.md). +The same option selects the appropriate transformer for FP16 or INT8 adapters. +Both built-in and plugin WebGPU are supported; plugin use requires the host +and provider library from this branch. Disable the option for a comparison +with the original graph. + +## Weight format and arithmetic + +The operator takes `X, W, scales, lora_A, lora_A_scales, lora_B, lora_B_scales`. +For rank 32, positive K divisible by 32, and positive N divisible by 4: + +| Tensor | Type | Shape | +|---|---|---| +| X | FP16 | [..., M, K] | +| W | Packed INT4 in UINT8 | [N, K/32, 16] | +| scales | FP16 | [N, K/32] or flattened | +| lora_A | Packed INT8 in UINT8 | [32, K/32, 32] | +| lora_A_scales | FP16 | [32, K/32] | +| lora_B | Packed INT8 in UINT8 | [N, 1, 32] | +| lora_B_scales | FP16 | [N, 1] | + +Packed adapter bytes encode signed integers with an implicit zero point of +128. For down weight index `(r,k)` and up weight index `(n,r)`: + +```text +At[r,k] = FP16((FP16(Aq[r,k]) - 128) * As[r,floor(k/32)]) +Bt[n,r] = FP16((FP16(Bq[n,r]) - 128) * Bs[n,0]) +T = FP16(X * At^T) +delta = FP16(T * Bt^T) +Y = FP16(MatMulNBits(X,W,scales) + delta) +``` + +Base quantization stays INT4 with block size 32 and implicit zero point 8. +The adapter projections retain accuracy level 0 and their original dot-product +grouping, accumulation policy, and FP16 intermediate rounding. The provider's +`enableMatmulFp32Accumulation` option applies to adapter accumulation as it does +to the original INT8 MatMulNBits kernels. + +The input layout is already transposed and block packed. The fusion does not +insert a transpose or materialize a complete dequantized adapter. + +## Dispatch + +`session.webgpu.lora_fusion_mode=auto` is the default. M is the second-to-last +dimension of X; a one-token prompt follows the decode selection. + +| Phase | Selected mode | NVIDIA implementation | +|---|---|---| +| Prefill, M > 1 | `fuse_up` | Original INT8 down projection, then base with INT8 up projection and addition | +| Decode, M = 1 | `fuse_down` | Base and INT8 down projection in one dispatch, then INT8 up projection and addition | + +The joint decode kernel supports one total input row and K from 2048 through +8192. It retains 128 threads per workgroup and adds four workgroups for the +32 down-projection outputs. Its reduction matches the original generic INT8 +kernel. `QLoRAOutputVec4` uses 64 threads and one output per thread, loading +packed weights as `vec4` and T as `vec4` along rank. + +DP4A prefill retains the base activation quantization dispatch, making three +dispatches per adapted projection. Explicit `fuse_down` and `fuse_up` remain +available. Shapes outside the joint decode specialization calculate T +separately. Devices that select subgroup matrix arithmetic for the original +INT8 up projection retain that kernel and a separate Add to preserve arithmetic. + +## Adapter off and graph matching + +The matcher requires two INT8 MatMulNBits nodes with block size 32 and accuracy +level 0, an INT4 base MatMulNBits node, and a final Add. All three projections +must run on WebGPU; the down projection must consume the base's X. Intermediate +outputs must be private. Bias, explicit zero points, group indices, prepacked +weights, and incompatible shapes prevent this fusion. + +Private overridable defaults containing only byte value 128, with finite +default scales, are replaced by empty packed adapter tensors in the session +graph. Scale defaults remain unchanged, preserving partial override semantics. +Empty A rank or empty B rank selects the ordinary base kernel before adapter +allocation or dispatch. Nonzero and shared packed defaults retain their data. +Input names and the adapter file stay unchanged. + +Like the FP16 prototype, eliding a zero branch assumes finite activations, +adapter scales, and adapter intermediates. NaN/Inf propagation through zero multiplication can differ; +the transformer remains opt-in. Use separate GenAI generators when switching +adapters or their on/off state with graph capture, and rewind only while +retaining the same generator's adapter selection. + +## Validation + +Run from the native build output directory: + +```text +onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.*:MatMulNBitsLoRA.* +onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBits*LoRA* +``` + +The quantized tests check original-versus-fused results with zero tolerance, +FP32 accumulation, batched and partial-tile shapes, empty adapters, partial +overrides, nonzero/shared defaults, unsupported shapes and attributes, and +patterns that must remain unfused. Full-model evaluation must separately +compare fusion equivalence against the original INT8 model and quantization +differences against FP16. Performance runs should use no profiling and record +the actual runtime libraries, provider, capture setting, and adapter state. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 0833372c4e04c..790c0d629a43c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -5,6 +5,7 @@ #param block_size #param n_bits #param has_lora_up +#param lora_quantized #param lora_transposed #param has_bias #param has_zero_points @@ -19,7 +20,9 @@ fn write_lora_output(offset: u32, value: vec4) { let row = (offset * 4u) / uniforms.N; let col = (offset * 4u) % uniforms.N; var delta = vec4(0); -#if lora_transposed +#if lora_quantized + delta = vec4(qlora_delta(row, col), qlora_delta(row, col + 1u), qlora_delta(row, col + 2u), qlora_delta(row, col + 3u)); +#elif lora_transposed for (var r4 = 0u; r4 < 8u; r4++) { let tv = lora_t[row * 8u + r4]; let b0 = lora_b[col * 8u + r4]; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 4438bf556ed44..be61d2f356606 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -4,6 +4,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "core/providers/webgpu/shader_helper.h" namespace onnxruntime { @@ -37,6 +38,10 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_lora_up_) { shader.AddInput("lora_t"); shader.AddInput("lora_b"); + if (lora_quantized_) { + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); + } } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); if (has_lora_down_) { @@ -51,6 +56,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), @@ -79,6 +85,10 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co if (has_lora_up_) { shader.AddInput("lora_t"); shader.AddInput("lora_b"); + if (lora_quantized_) { + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); + } } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); @@ -94,6 +104,7 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), @@ -123,8 +134,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, - const Tensor* lora_at, Tensor* lora_down, bool lora_transposed) { + const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales) { const bool has_lora_down = lora_at != nullptr; + const bool lora_quantized = lora_b_scales != nullptr; ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && M >= min_M_for_tile_optimization), "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; @@ -159,7 +171,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components); const bool broadcast_a = dispatch_M > M; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed}; + DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile); @@ -169,7 +181,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -181,7 +193,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor } if (has_lora_up) { mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 + : 1}); + if (lora_quantized) mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(mul_program); } @@ -192,7 +206,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed}; + DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; mul_program.SetWorkgroupSize(256); mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile + (has_lora_down ? CeilDiv(batch_count * M, 32u) : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, @@ -210,7 +224,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed); + .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -222,7 +236,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor } if (has_lora_up) { mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 + : 1}); + if (lora_quantized) mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } if (has_lora_down) { mul_program.AddInput({a, ProgramTensorMetadataDependency::Type, 4}); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index f5b0004faa980..4d30a9f2544f7 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -26,19 +26,20 @@ class DP4AMatMulNBitsProgram final : public Program { DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false) : Program{has_lora_down ? "DP4AMatMulNBitsLoRADown" : has_lora_up ? "DP4AMatMulNBitsLoRAUp" - : "DP4AMatMulNBits"}, - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32), - has_lora_up_(has_lora_up), - has_lora_down_(has_lora_down), - lora_transposed_(lora_transposed) {} + bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false) : Program{has_lora_down ? (lora_quantized ? "DP4AMatMulNBitsQLoRADown" : "DP4AMatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBitsLoRAUp") + : "DP4AMatMulNBits"}, + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32), + has_lora_up_(has_lora_up), + has_lora_down_(has_lora_down), + lora_transposed_(lora_transposed), + lora_quantized_(lora_quantized) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"batch_count", ProgramUniformVariableDataType::Uint32}, @@ -64,6 +65,7 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_lora_up_; bool has_lora_down_; bool lora_transposed_; + bool lora_quantized_; }; class DP4AMatMulNBitsSmallMProgram final : public Program { @@ -72,19 +74,20 @@ class DP4AMatMulNBitsSmallMProgram final : public Program #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" @@ -46,6 +47,10 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons if (has_lora_up_) { shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + if (lora_quantized_) { + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); + } } const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -61,6 +66,7 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(nbits, nbits_), WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_), @@ -89,12 +95,23 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_lora_up_) { shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + if (lora_quantized_) { + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); + } + } + if (has_lora_down_) { + shader.AddInput("lora_at"); + if (lora_quantized_) shader.AddInput("lora_a_scales"); } - if (has_lora_down_) shader.AddInput("lora_at"); const auto& output = shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); if (has_lora_down_) { shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - GenerateLoRAGemv(shader, "uniforms.num_N_tile", a.GetByOffset("k4")); + if (lora_quantized_) { + GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_); + } else { + GenerateLoRAGemv(shader, "uniforms.num_N_tile", a.GetByOffset("k4")); + } } const uint32_t components_a = a.NumComponents(); @@ -117,6 +134,7 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, false), @@ -204,13 +222,15 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const uint32_t weight_index, const Tensor* weight_index_indirect, uint32_t override_M, const Tensor* lora_t, const Tensor* lora_b, - const Tensor* lora_a_transposed, Tensor* lora_a_output, bool lora_transposed) { + const Tensor* lora_a_transposed, Tensor* lora_a_output, bool lora_transposed, + const Tensor* lora_b_scales, const Tensor* lora_a_scales) { TensorShape b_shape({N_op, K_op}); MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); const bool has_lora_up = lora_t != nullptr; const bool has_lora_down = lora_a_transposed != nullptr; + const bool lora_quantized = lora_b_scales != nullptr || lora_a_scales != nullptr; ORT_RETURN_IF_NOT(!has_lora_down || (lora_a_output != nullptr && !has_lora_up && a->Shape().Size() == K_op && nbits == 4 && block_size_op == 32), "Unsupported joint LoRA projection."); @@ -266,7 +286,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, static_cast(nbits), zero_blocks_per_col, *subgroup_matrix_config, context, &base_output, weight_index, weight_index_indirect)); - return ApplyLoRAOutput(context, &base_output, lora_t, lora_b, y, lora_transposed); + return ApplyLoRAOutput(context, &base_output, lora_t, lora_b, y, lora_transposed, lora_b_scales); } return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), zero_blocks_per_col, @@ -278,7 +298,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). if (!has_lora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed); + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed, lora_b_scales); } // WideTileProgram @@ -314,7 +334,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, ? context.AdapterInfo().subgroupMinSize : 0u; - MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, static_cast(nbits), subgroup_min_size, acc_f32, has_lora_up, lora_transposed}; + MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, static_cast(nbits), subgroup_min_size, acc_f32, has_lora_up, lora_transposed, lora_quantized}; program.SetWorkgroupSize(workgroup_size); program.SetDispatchGroupSize(num_N_tile, num_M_tile, batch_count); @@ -355,11 +375,13 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {num_N_tile}, {num_M_tile}, {weight_index}}); - program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32, has_lora_up, lora_transposed); + program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32, has_lora_up, lora_transposed, lora_quantized); if (has_lora_up) { program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 + : 1}); + if (lora_quantized) program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(program); } @@ -374,10 +396,10 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, constexpr uint32_t kU32Components = 4; uint32_t components_b_with_u32 = components_b * kU32Components; uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; - MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed}; + MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; program.SetWorkgroupSize(workgroup_size); uint32_t num_N_tile = (N + tile_size - 1) / tile_size; - program.SetDispatchGroupSize(num_N_tile + (has_lora_down ? 8 : 0), dispatch_M, batch_count); + program.SetDispatchGroupSize(num_N_tile + (has_lora_down ? (lora_quantized ? 4 : 8) : 0), dispatch_M, batch_count); program .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_a)}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_b_with_u32)}, @@ -395,7 +417,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {batch_count}, {weight_index}, {dispatch_M}}) - .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed); + .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); if (has_zero_points) { program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -407,11 +429,14 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, } if (has_lora_up) { program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 + : 1}); + if (lora_quantized) program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } if (has_lora_down) { ORT_RETURN_IF_NOT(components_a == 4, "Joint LoRA projection requires vec4 X."); - program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, 4}); + program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : 4}); + if (lora_quantized) program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); program.AddOutput({lora_a_output, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(program); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index d75d1e7644452..513bbf0b8fe50 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -20,8 +20,8 @@ using namespace onnxruntime::webgpu; class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, bool has_lora_up = false, bool lora_transposed = false) - : Program{"MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32), has_lora_up_(has_lora_up), lora_transposed_(lora_transposed) {} + MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, bool has_lora_up = false, bool lora_transposed = false, bool lora_quantized = false) + : Program{lora_quantized ? "MatMulNBitsQLoRAUpWideTile" : "MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32), has_lora_up_(has_lora_up), lora_transposed_(lora_transposed), lora_quantized_(lora_quantized) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"Batch", ProgramUniformVariableDataType::Uint32}, @@ -47,13 +47,14 @@ class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false) - : Program{has_lora_down ? "MatMulNBitsLoRADown" : has_lora_up ? "MatMulNBitsLoRAUp" - : "MatMulNBits"}, + MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false) + : Program{has_lora_down ? (lora_quantized ? "MatMulNBitsQLoRADown" : "MatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "MatMulNBitsQLoRAUp" : "MatMulNBitsLoRAUp") + : "MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), @@ -66,7 +67,8 @@ class MatMulNBitsProgram final : public Program { acc_f32_(acc_f32), has_lora_up_(has_lora_up), has_lora_down_(has_lora_down), - lora_transposed_(lora_transposed) {} + lora_transposed_(lora_transposed), + lora_quantized_(lora_quantized) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"M", ProgramUniformVariableDataType::Uint32}, @@ -96,6 +98,7 @@ class MatMulNBitsProgram final : public Program { bool has_lora_up_; bool has_lora_down_; bool lora_transposed_; + bool lora_quantized_; }; class MatMulNBits final : public WebGpuKernel { @@ -126,9 +129,10 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* weight_index_indirect = nullptr, uint32_t override_M = 0, const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, - const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, bool lora_transposed = false); + const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, bool lora_transposed = false, + const Tensor* lora_b_scales = nullptr, const Tensor* lora_a_scales = nullptr); -Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed = false); +Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed = false, const Tensor* quantized_scales = nullptr); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template index e38088548d63e..4a5b40d3cd4de 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template @@ -14,6 +14,7 @@ #param tile_size_k #param tile_size #param has_lora_up +#param lora_quantized #param lora_transposed #param has_bias #param has_weight_idx @@ -268,7 +269,9 @@ $MAIN { // Downcast to the output type only at the final write. #if has_lora_up var delta = output_element_t(0); -#if lora_transposed +#if lora_quantized + delta = qlora_delta(batch * uniforms.M + a_global, b_global); +#elif lora_transposed for (var r4 = 0u; r4 < 8u; r4++) { let bv = lora_b[b_global * 8u + r4]; let tv = lora_t[(batch * uniforms.M + a_global) * 8u + r4]; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc index 40564a57ea205..55ec6ab229d2d 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/common/safeint.h" #include "core/providers/webgpu/math/matmul.h" @@ -152,7 +153,8 @@ class LoRAOutputProgram final : public Program { bool transposed_; }; -Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed) { +Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed, const Tensor* quantized_scales) { + if (quantized_scales) return ApplyQLoRAOutput(context, base, t, b, quantized_scales, y); const auto size = narrow(y->Shape().Size()); LoRAOutputProgram program(transposed); program.AddInputs({{base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc new file mode 100644 index 0000000000000..efee298dcebdc --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -0,0 +1,255 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/common/safeint.h" +#include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/shader_helper.h" + +namespace onnxruntime::contrib::webgpu { + +void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32) { + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" + << R"( +var qlora_x_tile: array, 128>; +var qlora_partials: array, 8>; +)"; + shader.MainFunctionBody() << "let qlora_base_groups = " << base_groups << ";\n" + << R"( +if (workgroup_idx >= qlora_base_groups) { + let rank_base = (workgroup_idx - qlora_base_groups) * 8u; + let lane = local_idx % 32u; + let rank_local = local_idx / 32u; + // Match the generic INT8 MatMulNBits tile, dot grouping and reduction order. + for (var start = 0u; start < uniforms.K; start += 512u) { + let k4 = start / 4u + local_idx; + var xv = vec4(0); + if (k4 < uniforms.K / 4u) { xv = )" + << x_value << R"(; } + qlora_x_tile[local_idx] = xv; + workgroupBarrier(); + for (var rank_offset = 0u; rank_offset < 8u; rank_offset += 4u) { + let rank = rank_base + rank_offset + rank_local; + let k = start + lane * 16u; + if (k < uniforms.K) { + let packed = lora_at[rank * (uniforms.K / 16u) + k / 16u]; + let scale = lora_a_scales[rank * (uniforms.K / 32u) + k / 32u]; + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; + sum += dot(vec4(qlora_x_tile[lane * 4u + i]), vec4(bv)); + } + qlora_partials[rank_offset + rank_local][lane] += sum; + } + } + workgroupBarrier(); + } + if (local_idx < 8u) { + var value = qlora_acc_t(0); + for (var i = 0u; i < 32u; i++) { value += qlora_partials[local_idx][i]; } + lora_down[rank_base + local_idx] = f16(value); + } + return; +} +)"; +} + +void GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32) { + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" + << R"( +fn qlora_weight(col: u32, r4: u32) -> vec4 { + let packed = lora_b[col * 2u + r4 / 4u][r4 % 4u]; + return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; +} +fn qlora_delta(row: u32, col: u32) -> f16 { + var delta = qlora_acc_t(0); + // The original K=32 INT8 up projection uses wide tiles at M>=4. + if (uniforms.M >= 4u) { + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = lora_t[row * 8u + r4]; + let a1 = lora_t[row * 8u + r4 + 1u]; + delta += dot(vec4(a0), vec4(qlora_weight(col, r4))) + + dot(vec4(a1), vec4(qlora_weight(col, r4 + 1u))); + } + } else { + for (var block = 0u; block < 2u; block++) { + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let r4 = block * 4u + i; + sum += dot(vec4(lora_t[row * 8u + r4]), vec4(qlora_weight(col, r4))); + } + delta += sum; + } + // The remaining generic-kernel partials are zero, including their sign. + delta += qlora_acc_t(0); + } + return f16(delta); +} +)"; +} + +namespace { + +class QLoRAOutputProgram final : public Program { + public: + explicit QLoRAOutputProgram(bool acc_f32) : Program{"QLoRAOutputVec4"}, acc_f32_(acc_f32) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(QLoRAOutputProgram); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("base"); + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + shader.AddInput("lora_b_scales"); + shader.AddOutput("output"); + GenerateQLoRAUp(shader, acc_f32_); + shader.MainFunctionBody() + << "if (global_idx >= uniforms.size) { return; }\n" + << "output[global_idx] = base[global_idx] + qlora_delta(global_idx / uniforms.N, global_idx % uniforms.N);\n"; + return Status::OK(); + } + + private: + bool acc_f32_; +}; + +class QLoRAAddProgram final : public Program { + public: + QLoRAAddProgram() : Program{"QLoRAAdd"} {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(QLoRAAddProgram); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("base"); + shader.AddInput("delta"); + shader.AddOutput("output"); + shader.MainFunctionBody() + << "if (global_idx < uniforms.size) { output[global_idx] = base[global_idx] + delta[global_idx]; }\n"; + return Status::OK(); + } +}; + +} // namespace + +Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, + const Tensor* b, const Tensor* scales, Tensor* y) { + const auto size = narrow(y->Shape().Size()); + QLoRAOutputProgram program(context.EnableMatmulFp32Accumulation()); + program + .AddInputs({{base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {t, ProgramTensorMetadataDependency::Type, 4}, + {b, ProgramTensorMetadataDependency::Type, 16}, + {scales, ProgramTensorMetadataDependency::Type}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) + .AddUniformVariables( + {{size}, {narrow(b->Shape()[0])}, {narrow(y->Shape()[y->Shape().NumDimensions() - 2])}}) + .SetWorkgroupSize(64) + .SetDispatchGroupSize(CeilDiv(size, 64u)) + .CacheHint(context.EnableMatmulFp32Accumulation()); + return context.RunProgram(program); +} + +class MatMulNBitsQLoRA final : public WebGpuKernel { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulNBitsQLoRA); + explicit MatMulNBitsQLoRA(const OpKernelInfo& info) + : WebGpuKernel(info), + k_(info.GetAttr("K")), + n_(info.GetAttr("N")), + accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)) { + const auto mode = info.GetConfigOptions().GetConfigOrDefault("session.webgpu.lora_fusion_mode", "auto"); + ORT_ENFORCE(mode == "auto" || mode == "fuse_up" || mode == "fuse_down", + "session.webgpu.lora_fusion_mode must be auto, fuse_down or fuse_up."); + automatic_ = mode == "auto"; + fuse_up_ = mode == "fuse_up"; + ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0 && n_ % 4 == 0 && info.GetAttr("bits") == 4 && + info.GetAttr("block_size") == 32 && info.GetAttr("lora_bits") == 8 && + info.GetAttr("lora_block_size") == 32, + "MatMulNBitsQLoRA requires INT4 base and INT8 adapters with block size 32."); + } + + Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override { + const auto* x = context.Input(0); + const auto* w = context.Input(1); + const auto* s = context.Input(2); + const auto* aq = context.Input(3); + const auto* as = context.Input(4); + const auto* bq = context.Input(5); + const auto* bs = context.Input(6); + ORT_RETURN_IF_NOT( + x->Shape().NumDimensions() >= 2 && aq->Shape().NumDimensions() == 3 && bq->Shape().NumDimensions() == 3, + "Invalid quantized LoRA input ranks."); + ORT_RETURN_IF_NOT( + w->Shape() == TensorShape({n_, k_ / 32, 16}) && s->Shape().Size() == SafeInt(n_) * (k_ / 32), + "Invalid base weight/scales shape."); + const auto rank_a = aq->Shape()[0]; + const auto blocks_b = bq->Shape()[1]; + ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (blocks_b == 0 || blocks_b == 1) && aq->Shape()[1] == k_ / 32 && + aq->Shape()[2] == 32 && bq->Shape()[0] == n_ && bq->Shape()[2] == 32, + "Invalid packed adapter dimensions."); + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n_, k_}), false, true)); + auto* y = context.Output(0, helper.OutputShape()); + if (y->Shape().Size() == 0) return Status::OK(); + if (rank_a == 0 || blocks_b == 0) { + return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y); + } + ORT_RETURN_IF_NOT(as->Shape() == TensorShape({32, k_ / 32}) && bs->Shape() == TensorShape({n_, 1}), + "Invalid adapter scales shape."); + auto t_shape = x->Shape().AsShapeVector(); + t_shape.back() = 32; + auto t = context.CreateGPUTensor(x->DataType(), TensorShape(t_shape)); + const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); + const bool fuse_up = fuse_up_ || (automatic_ && m > 1); + const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && + context.AdapterInfo().vendor == std::string_view{"nvidia"}; + if (joint_down) { + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, + nullptr, 0, nullptr, nullptr, aq, &t, true, nullptr, as)); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + } + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, 0, 8, context, &t)); + std::optional up_config; + const bool subgroup_up = + CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), + narrow(n_), 32, 8, true, up_config, m); + if (fuse_up && !subgroup_up) { + return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, &t, + bq, nullptr, nullptr, true, bs); + } + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base)); + if (!subgroup_up) return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + // Keep the original up kernel on devices with a different matrix-accumulation path. + auto delta = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(&t, bq, bs, nullptr, nullptr, 32, n_, 32, 0, 8, context, &delta)); + const auto size = narrow(y->Shape().Size()); + QLoRAAddProgram add; + add.AddInputs({{&base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {&delta, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) + .AddUniformVariables({{size}}) + .SetWorkgroupSize(64) + .SetDispatchGroupSize(CeilDiv(size, 64u)); + return context.RunProgram(add); + } + + private: + int64_t k_; + int64_t n_; + int64_t accuracy_level_; + bool automatic_; + bool fuse_up_; +}; + +ONNX_OPERATOR_KERNEL_EX(MatMulNBitsQLoRA, kMSDomain, 1, kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), + MatMulNBitsQLoRA); + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h new file mode 100644 index 0000000000000..befb368db025b --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" + +namespace onnxruntime::contrib::webgpu { + +void GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, + const std::string& x_value, bool acc_f32); +void GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32); +Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, + const Tensor* b, const Tensor* scales, Tensor* y); + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template index 61af4960c7250..25379afea40e0 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template @@ -3,6 +3,7 @@ #param has_zero_points #param has_lora_up +#param lora_quantized #param lora_transposed #param has_bias #param has_weight_idx @@ -117,7 +118,9 @@ fn write_output(batch : u32, row : u32, col : u32, value : output_element_t) { let offset = batch * uniforms.M * uniforms.N + row * uniforms.N + col; #if has_lora_up var delta = output_element_t(0); -#if lora_transposed +#if lora_quantized + delta = qlora_delta(batch * uniforms.M + row, col); +#elif lora_transposed for (var r4 = 0u; r4 < 8u; r4++) { let bv = lora_b[col * 8u + r4]; let tv = lora_t[(batch * uniforms.M + row) * 8u + r4]; diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 9fef9b55601ee..b098284c9bea3 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -63,6 +63,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 08d949b83adbd..fe9cafbc8bdff 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3991,6 +3991,62 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by } }); + ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsQLoRA) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc( + "Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). " + "The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights " + "stored as unsigned bytes with implicit zero point 128 and FP16 block scales. Adapter projections use " + "accuracy level 0. Rank is 32; an empty rank dimension in either packed adapter disables the adapter. " + "Intermediate projections are rounded to FP16 before the final addition. Adapters remain replaceable.") + .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) + .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) + .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) + .Attr("block_size", "Base weight block size; must be 32.", AttributeProto::INT, int64_t{32}) + .Attr("accuracy_level", "Base MatMulNBits accumulation policy.", AttributeProto::INT, int64_t{0}) + .Attr("lora_bits", "Adapter weight bit width; must be 8.", AttributeProto::INT, int64_t{8}) + .Attr("lora_block_size", "Adapter weight block size; must be 32.", AttributeProto::INT, int64_t{32}) + .Input(0, "X", "Input of rank at least two, with last dimension K.", "T") + .Input(1, "W", "Packed INT4 base weights [N,K/32,16].", "TQ") + .Input(2, "scales", "Base block scales with N*K/32 elements; may be flattened.", "T") + .Input(3, "lora_A", "Packed INT8 down weights [R_A,K/32,32], R_A=0 or 32.", "TQ") + .Input(4, "lora_A_scales", "Down block scales [32,K/32]; retained even when the adapter is empty.", "T") + .Input(5, "lora_B", "Packed INT8 up weights [N,R_B/32,32], R_B=0 or 32.", "TQ") + .Input(6, "lora_B_scales", "Up block scales [N,1]; retained even when the adapter is empty.", "T") + .Output(0, "Y", "Input shape with the last dimension replaced by N.", "T") + .TypeConstraint("T", {"tensor(float16)"}, "FP16 activations and block scales.") + .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized weights.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const auto k = getAttribute(ctx, "K", -1); + const auto n = getAttribute(ctx, "N", -1); + if (k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || + getAttribute(ctx, "bits", 4) != 4 || getAttribute(ctx, "block_size", 32) != 32 || + getAttribute(ctx, "lora_bits", 8) != 8 || getAttribute(ctx, "lora_block_size", 32) != 32) { + fail_shape_inference("Unsupported MatMulNBitsQLoRA attributes."); + } + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() < 2) { + fail_shape_inference("X must have rank at least two."); + } + MatmulWithQuantWeightShapeInference(ctx, k, n, true); + for (size_t index : {size_t{3}, size_t{4}, size_t{5}, size_t{6}}) { + if (!hasInputShape(ctx, index)) continue; + const auto& shape = getInputShape(ctx, index); + const bool packed = index == 3 || index == 5; + if (shape.dim_size() != (packed ? 3 : 2)) fail_shape_inference("Invalid quantized adapter rank."); + const int64_t dims[4][3] = {{32, k / 32, 32}, {32, k / 32, 0}, {n, 1, 32}, {n, 1, 0}}; + for (int axis = 0; axis < shape.dim_size(); ++axis) { + const auto& dim = shape.dim(axis); + const bool empty_axis = (index == 3 && axis == 0) || (index == 5 && axis == 1); + if (dim.has_dim_value() && dim.dim_value() != dims[index - 3][axis] && + !(empty_axis && dim.dim_value() == 0)) { + fail_shape_inference("Invalid quantized adapter dimensions."); + } + } + } + }); + static const char* MatMulNBitsMlp_ver1_doc = R"DOC( MatMulNBitsMlp fuses two MatMulNBits projections that share the same input and computes diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index ac2c1e1ac5d91..800b0340e912c 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -61,6 +61,7 @@ #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/matmul_nbits_mlp_fusion.h" #include "core/optimizer/matmul_nbits_lora_fusion.h" +#include "core/optimizer/matmul_nbits_qlora_fusion.h" #include "core/optimizer/matmul_bn_fusion.h" #include "core/optimizer/matmul_integer_to_float.h" #include "core/optimizer/matmul_scale_fusion.h" @@ -467,6 +468,7 @@ InlinedVector> GenerateTransformers( InlinedHashSet{onnxruntime::kCudaExecutionProvider, onnxruntime::kWebGpuExecutionProvider})); bool has_matmul_nbits_lora_kernel = false; + bool has_matmul_nbits_qlora_kernel = false; bool has_matmul_nbits_mlp_kernel = false; bool has_matmul_nbits_qkv_kernel = false; if (execution_providers != nullptr) { @@ -485,6 +487,7 @@ InlinedVector> GenerateTransformers( .IsOK(); }; has_matmul_nbits_lora_kernel = has_webgpu_kernel("MatMulNBitsLoRA"); + has_matmul_nbits_qlora_kernel = has_webgpu_kernel("MatMulNBitsQLoRA"); has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp"); has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv"); } @@ -494,6 +497,10 @@ InlinedVector> GenerateTransformers( session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { transformers.emplace_back(std::make_unique()); } + if (has_matmul_nbits_qlora_kernel && + session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { + transformers.emplace_back(std::make_unique()); + } if (has_matmul_nbits_mlp_kernel) { transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc new file mode 100644 index 0000000000000..62395631a14e1 --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -0,0 +1,156 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/matmul_nbits_qlora_fusion.h" + +#include + +#include "core/graph/graph_utils.h" +#include "core/graph/node_attr_utils.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/matmul_nbits_fusion_utils.h" +#include "core/optimizer/utils.h" + +namespace onnxruntime { +namespace { + +const Node* Producer(const Graph& graph, const Node& node, int input) { + const auto* edge = graph_utils::GetInputEdge(node, input); + return edge ? graph.GetNode(edge->GetNode().Index()) : nullptr; +} + +bool Private(const Graph& graph, const Node& node) { + return !graph.NodeProducesGraphOutput(node) && optimizer_utils::CheckOutputEdges(graph, node, 1); +} + +bool Shape(const NodeArg& arg, std::initializer_list expected) { + const auto* shape = arg.Shape(); + if (!shape || shape->dim_size() != static_cast(expected.size())) return false; + int index = 0; + for (auto value : expected) { + const auto& dim = shape->dim(index++); + if (!dim.has_dim_value() || dim.dim_value() != value) return false; + } + return true; +} + +bool Projection(const Node* node, int64_t bits) { + using namespace matmul_nbits_fusion_utils; + return node && graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMulNBits", {1}, kMSDomain) && + node->GetExecutionProviderType() == kWebGpuExecutionProvider && GetIntAttr(*node, "bits", 4) == bits && + GetIntAttr(*node, "block_size", -1, true) == 32 && GetIntAttr(*node, "weight_prepacked", 0) == 0 && + !HasInput(*node, 3) && !HasInput(*node, 4) && !HasInput(*node, 5) && + (bits != 8 || GetIntAttr(*node, "accuracy_level", 0) == 0); +} + +bool ZeroDefault(const Graph& graph, const NodeArg& q, const NodeArg& scales) { + if (!graph.IsInputsIncludingInitializers(&q) || graph.IsOutput(&q) || graph.GetConsumerNodes(q.Name()).size() != 1) + return false; + const ONNX_NAMESPACE::TensorProto* qp = nullptr; + const ONNX_NAMESPACE::TensorProto* sp = nullptr; + if (!graph.GetInitializedTensor(q.Name(), qp) || !graph.GetInitializedTensor(scales.Name(), sp)) return false; + Initializer qdata(graph, *qp, graph.ModelPath()); + Initializer sdata(graph, *sp, graph.ModelPath()); + const auto bytes = qdata.DataAsByteSpan(); + const auto values = sdata.DataAsSpan(); + return std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 128; }) && + std::all_of(values.begin(), values.end(), [](MLFloat16 value) { return std::isfinite(value.ToFloat()); }); +} + +void EmptyDefault(Graph& graph, NodeArg& arg, int axis) { + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(graph.GetInitializedTensor(arg.Name(), original)); + ONNX_NAMESPACE::TensorProto empty; + empty.set_name(arg.Name()); + empty.set_data_type(original->data_type()); + for (int i = 0; i < original->dims_size(); ++i) empty.add_dims(i == axis ? 0 : original->dims(i)); + auto shape = *arg.Shape(); + shape.mutable_dim(axis)->Clear(); + arg.SetShape(shape); + graph.RemoveInitializedTensor(arg.Name()); + graph.AddInitializedTensor(empty); +} + +} // namespace + +Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, + const logging::Logger& logger) const { + using namespace matmul_nbits_fusion_utils; + GraphViewer viewer(graph); + for (auto index : viewer.GetNodesInTopologicalOrder()) { + auto* add = graph.GetNode(index); + if (!add) continue; + ORT_RETURN_IF_ERROR(Recurse(*add, modified, graph_level, logger)); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(*add, "Add", {7, 13, 14}) || + add->GetExecutionProviderType() != kWebGpuExecutionProvider) + continue; + auto* base = Producer(graph, *add, 0); + auto* up = Producer(graph, *add, 1); + if (!Projection(base, 4)) std::swap(base, up); + if (!Projection(base, 4) || !Projection(up, 8)) continue; + auto* down = Producer(graph, *up, 0); + if (!Projection(down, 8) || down->InputDefs()[0] != base->InputDefs()[0] || !Private(graph, *base) || + !Private(graph, *down) || !Private(graph, *up)) + continue; + const int64_t k = GetIntAttr(*base, "K", -1, true); + const int64_t n = GetIntAttr(*base, "N", -1, true); + if (k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || GetIntAttr(*down, "K", -1, true) != k || + GetIntAttr(*down, "N", -1, true) != 32 || GetIntAttr(*up, "K", -1, true) != 32 || + GetIntAttr(*up, "N", -1, true) != n) + continue; + InlinedVector inputs{ + const_cast(base->InputDefs()[0]), const_cast(base->InputDefs()[1]), + const_cast(base->InputDefs()[2]), const_cast(down->InputDefs()[1]), + const_cast(down->InputDefs()[2]), const_cast(up->InputDefs()[1]), + const_cast(up->InputDefs()[2])}; + bool types_ok = true; + for (size_t i = 0; i < inputs.size(); ++i) { + const auto* type = inputs[i]->TypeAsProto(); + const auto expected = + (i == 1 || i == 3 || i == 5) ? ONNX_NAMESPACE::TensorProto::UINT8 : ONNX_NAMESPACE::TensorProto::FLOAT16; + types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == expected; + } + if (!types_ok || !inputs[0]->Shape() || inputs[0]->Shape()->dim_size() < 2 || !Shape(*inputs[1], {n, k / 32, 16}) || + !Shape(*inputs[3], {32, k / 32, 32}) || !Shape(*inputs[4], {32, k / 32}) || !Shape(*inputs[5], {n, 1, 32}) || + !Shape(*inputs[6], {n, 1})) + continue; + const auto* s = inputs[2]->Shape(); + const bool flat_scales = s && s->dim_size() == 1 && s->dim(0).has_dim_value() && + s->dim(0).dim_value() / n == k / 32 && s->dim(0).dim_value() % n == 0; + if (!flat_scales && !Shape(*inputs[2], {n, k / 32})) continue; + const bool empty_defaults = + ZeroDefault(graph, *inputs[3], *inputs[4]) && ZeroDefault(graph, *inputs[5], *inputs[6]); + auto attrs = base->GetAttributes(); + attrs.erase("weight_prepacked"); + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", int64_t{8}), attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", int64_t{32}), attrs); + auto& fused = graph.AddNode(graph.GenerateNodeName("MatMulNBitsQLoRA"), "MatMulNBitsQLoRA", + "fused unmerged rank-32 INT8 LoRA projection", inputs, add->MutableOutputDefs(), &attrs, + kMSDomain); + fused.SetExecutionProviderType(kWebGpuExecutionProvider); + const Node* sources[] = {base, base, base, down, down, up, up}; + const int slots[] = {0, 1, 2, 1, 2, 1, 2}; + for (int i = 0; i < 7; ++i) { + if (const auto* edge = graph_utils::GetInputEdge(*sources[i], slots[i])) { + graph.AddEdge(edge->GetNode().Index(), fused.Index(), edge->GetSrcArgIndex(), i); + } + } + const auto outputs = graph_utils::GraphEdge::GetNodeOutputEdges(*add); + for (const Node* node : {base, down, up, static_cast(add)}) { + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(node->Index())); + } + const NodeIndex removed[] = {base->Index(), down->Index(), up->Index(), add->Index()}; + for (auto node : removed) graph.RemoveNode(node); + for (const auto& edge : outputs) graph.AddEdge(fused.Index(), edge.dst_node, 0, edge.dst_arg_index); + if (empty_defaults) { + // Keep scale defaults: overriding only a packed weight must still use its original scale. + EmptyDefault(graph, *inputs[3], 0); + EmptyDefault(graph, *inputs[5], 1); + } + modified = true; + LOGS(logger, VERBOSE) << "MatMulNBitsQLoRAFusion: " << fused.Name() << " empty defaults=" << empty_defaults; + } + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h new file mode 100644 index 0000000000000..943fe431c61e3 --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +class MatMulNBitsQLoRAFusion final : public GraphTransformer { + public: + MatMulNBitsQLoRAFusion() : GraphTransformer("MatMulNBitsQLoRAFusion", {kWebGpuExecutionProvider}) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulNBitsQLoRAFusion); + + private: + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc new file mode 100644 index 0000000000000..a862ef0279410 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) + +TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { + for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { + for (int64_t m : {1, 37}) { + for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 128); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, MLFloat16(0.5f))); + test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); + test.AddInput("S", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); + test.AddInput("AS", {32, 2}, std::vector(64, MLFloat16(0.03125f))); + test.AddInput("BQ", {128, ranks.second, 32}, std::vector(128 * ranks.second * 32, 129)); + test.AddInput("BS", {128, 1}, std::vector(128, MLFloat16(0.0625f))); + const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; + test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, MLFloat16(expected))); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + test.Config(options); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } + } +} + +TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { + for (bool wrong_bits : {false, true}) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 128); + if (wrong_bits) test.AddAttribute("lora_bits", 4); + test.AddInput("X", {1, 1, 64}, std::vector(64, MLFloat16(0.5f))); + test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); + test.AddInput("S", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); + const int64_t rank = wrong_bits ? 32 : 16; + test.AddInput("AQ", {rank, 2, 32}, std::vector(rank * 64, 129)); + test.AddInput("AS", {32, 2}, std::vector(64, MLFloat16(0.03125f))); + test.AddInput("BQ", {128, 1, 32}, std::vector(4096, 129)); + test.AddInput("BS", {128, 1}, std::vector(128, MLFloat16(0.0625f))); + test.AddOutput("Y", {1, 1, 128}, std::vector(128, MLFloat16(3.0f))); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectFailure, + wrong_bits ? "Unsupported MatMulNBitsQLoRA attributes" : "Invalid quantized adapter dimensions", {}, + nullptr, &eps); + } +} + +#endif +} // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc new file mode 100644 index 0000000000000..598c1947982c6 --- /dev/null +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -0,0 +1,246 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/matmul_nbits_qlora_fusion.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/optimizer/webgpu_fusion_test_util.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +#ifndef DISABLE_CONTRIB_OPS +namespace { + +void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int64_t n = 128, int64_t batch = 1, + bool expose = false, bool different_x = false, bool zero_base = false, int64_t adapter_accuracy = 0) { + auto* x = builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)); + auto* w = builder.MakeInitializer({n, k / 32, 16}, zero_base ? uint8_t{0x88} : uint8_t{0}, + zero_base ? uint8_t{0x88} : uint8_t{255}); + auto* s = builder.MakeInitializer({n, k / 32}, MLFloat16(0.03125f), MLFloat16(0.0625f)); + auto* aq = builder.MakeInput({32, k / 32, 32}, uint8_t{100}, uint8_t{156}); + auto* as = builder.MakeInput({32, k / 32}, MLFloat16(0.001f), MLFloat16(0.002f)); + auto* bq = builder.MakeInput({n, 1, 32}, uint8_t{100}, uint8_t{156}); + auto* bs = builder.MakeInput({n, 1}, MLFloat16(0.001f), MLFloat16(0.002f)); + for (auto* arg : {aq, as, bq, bs}) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(arg->Name()); + const bool packed = arg == aq || arg == bq; + init.set_data_type(packed ? ONNX_NAMESPACE::TensorProto::UINT8 : ONNX_NAMESPACE::TensorProto::FLOAT16); + int64_t count = 1; + for (const auto& dim : arg->Shape()->dim()) { + init.add_dims(dim.dim_value()); + count *= dim.dim_value(); + } + if (packed) { + init.set_raw_data(std::string(static_cast(count), static_cast(128))); + } else { + const std::vector ones(static_cast(count), MLFloat16(1.0f)); + init.set_raw_data(ones.data(), ones.size() * sizeof(MLFloat16)); + } + builder.graph_.AddInitializedTensor(init); + } + auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* t = expose ? builder.MakeOutput(std::vector{batch, m, 32}) + : builder.MakeIntermediate(std::vector{batch, m, 32}); + auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* y = builder.MakeOutput(std::vector{batch, m, n}); + auto* adapter_x = different_x ? builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)) : x; + if (different_x) { + builder.graph_.SetInputs({x, aq, as, bq, bs, adapter_x}); + } else { + builder.graph_.SetInputs({x, aq, as, bq, bs}); + } + auto& base_node = builder.AddNode("MatMulNBits", {x, w, s}, {base}, kMSDomain); + auto& down = builder.AddNode("MatMulNBits", {adapter_x, aq, as}, {t}, kMSDomain); + auto& up = builder.AddNode("MatMulNBits", {t, bq, bs}, {delta}, kMSDomain); + for (auto* node : {&base_node, &down, &up}) { + node->AddAttribute("K", node == &up ? int64_t{32} : k); + node->AddAttribute("N", node == &down ? int64_t{32} : n); + node->AddAttribute("bits", node == &base_node ? int64_t{4} : int64_t{8}); + node->AddAttribute("block_size", int64_t{32}); + node->AddAttribute("accuracy_level", node == &base_node ? int64_t{4} : adapter_accuracy); + node->SetExecutionProviderType(kWebGpuExecutionProvider); + } + builder.AddNode("Add", {base, delta}, {y}).SetExecutionProviderType(kWebGpuExecutionProvider); +} + +Status CheckQLoRACount(const Graph& graph, size_t expected) { + size_t count = 0; + for (const auto& node : graph.Nodes()) count += node.OpType() == "MatMulNBitsQLoRA"; + ORT_RETURN_IF_NOT(count == expected, "Unexpected quantized LoRA fusion count."); + return Status::OK(); +} + +} // namespace + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPreservesDefaults) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { BuildQLoRA(b); }, 21, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, nullptr, + [](const Graph& graph) { + ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); + const auto inputs = graph.GetOverridableInitializers(); + ORT_RETURN_IF_NOT(inputs.size() == 4, "Adapter parameters were removed."); + for (const auto* input : inputs) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(graph.GetInitializedTensor(input->Name(), init), "Missing default."); + if (init->data_type() == ONNX_NAMESPACE::TensorProto::UINT8) { + ORT_RETURN_IF_NOT(init->dims(0) == 0 || init->dims(1) == 0, "Packed default was not emptied."); + } else { + ORT_RETURN_IF_NOT(init->dims(0) > 0 && init->dims(1) > 0, "Scale default was emptied."); + } + } + return Status::OK(); + })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + struct ShapeCase { + int64_t m, k, n, batch; + bool zero_base; + }; + const ShapeCase shapes[] = {{1, 64, 128, 1, false}, {1, 2048, 128, 1, false}, {1, 3072, 132, 1, true}, {1, 8192, 128, 1, true}, {2, 2048, 132, 1, true}, {37, 2048, 128, 1, false}, {37, 3072, 132, 1, true}, {1, 2048, 128, 2, true}, {37, 2048, 132, 2, true}}; + for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { + for (const auto& shape : shapes) { + SCOPED_TRACE(mode[0] ? mode : "default"); + SCOPED_TRACE(shape.m); + SCOPED_TRACE(shape.k); + SCOPED_TRACE(shape.n); + SCOPED_TRACE(shape.batch); + RunWebGpuFusionTransformerTest( + [shape](ModelTestBuilder& b) { + BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); + }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, std::make_unique(), + [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + if (mode[0]) + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + }); + } + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsObservableIntermediate) { + ASSERT_STATUS_OK(TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, + nullptr, [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Accumulation) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (int64_t m : {1, 2, 37}) { + SCOPED_TRACE(m); + RunWebGpuFusionTransformerTest( + [m](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, 132, 1, false, false, true); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, std::make_unique(), + [] { + ConfigOptions options; + ORT_THROW_IF_ERROR(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + return WebGpuExecutionProviderWithOptions(options); + }); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionKeepsNonzeroOrSharedDefaults) { + for (bool shared : {false, true}) { + ASSERT_STATUS_OK(TestGraphTransformer( + [shared](ModelTestBuilder& b) { + BuildQLoRA(b); + const NodeArg* input = nullptr; + for (const auto& [name, tensor] : b.graph_.GetAllInitializedTensors()) { + if (tensor->data_type() == ONNX_NAMESPACE::TensorProto::UINT8 && tensor->dims(0) == 32) { + input = b.graph_.GetNodeArg(name); + break; + } + } + ORT_ENFORCE(input != nullptr); + if (shared) { + auto* observable = b.MakeOutput(std::vector{32, 64, 32}); + b.AddNode("Identity", {const_cast(input)}, {observable}); + } else { + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(b.graph_.GetInitializedTensor(input->Name(), original)); + auto replacement = *original; + replacement.mutable_raw_data()->at(0) = static_cast(129); + b.graph_.RemoveInitializedTensor(input->Name()); + b.graph_.AddInitializedTensor(replacement); + } + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& graph) { + ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); + for (const auto* input : graph.GetOverridableInitializers()) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(graph.GetInitializedTensor(input->Name(), init), "Missing default."); + for (const auto dim : init->dims()) + ORT_RETURN_IF_NOT(dim > 0, "Observable or nonzero default was emptied."); + } + return Status::OK(); + })); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentInput) { + ASSERT_STATUS_OK(TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, false, true); }, 21, + *logger_, std::make_unique(), TransformerLevel::Level2, + 1, nullptr, [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentAdapterAccuracy) { + ASSERT_STATUS_OK( + TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, false, false, false, 4); }, 21, + *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + Model model("quantized LoRA switching", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + BuildQLoRA(builder, 1, 64); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + const auto adapters = graph.GetOverridableInitializers(); + ASSERT_EQ(adapters.size(), 4u); + std::string data; + model.ToProto().SerializeToString(&data); + std::vector expected; + for (bool fused : {false, true}) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + if (fused) + ASSERT_STATUS_OK( + session.RegisterGraphTransformer(std::make_unique(), TransformerLevel::Level2)); + ASSERT_STATUS_OK(session.Initialize()); + // Off, packed weights only, scales only, all parameters, then off again. + for (int state = 0; state < 5; ++state) { + auto feeds = builder.feeds_; + for (const auto* input : adapters) { + const bool packed = input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto::UINT8; + if (state == 0 || state == 4 || (state == 1 && !packed) || (state == 2 && packed)) feeds.erase(input->Name()); + } + std::vector output; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, builder.output_names_, &output)); + ASSERT_EQ(output.size(), 1u); + if (!fused) { + expected.push_back(std::move(output[0])); + } else { + const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << "state=" << state << " " << comparison.second; + } + } + } +} + +#endif +} // namespace onnxruntime::test From 2638ae21c0aef94249e1ec626c19d93970ba3517 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:21:03 +0800 Subject: [PATCH 05/35] Optimize WebGPU QLoRA prefill with shared DP4A input and narrow tiles --- .../quantization/dp4a_matmul.wgsl.template | 420 +++++++++--------- .../webgpu/quantization/dp4a_matmul_nbits.cc | 59 ++- .../webgpu/quantization/dp4a_matmul_nbits.h | 18 +- .../webgpu/quantization/matmul_nbits.cc | 5 +- .../webgpu/quantization/matmul_nbits.h | 4 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 37 +- .../webgpu/quantization/matmul_nbits_qlora.h | 2 +- .../test/contrib_ops/matmul_2bits_test.cc | 20 + .../test/contrib_ops/matmul_4bits_test.cc | 21 + .../test/contrib_ops/matmul_8bits_test.cc | 22 + .../contrib_ops/matmul_nbits_qlora_test.cc | 32 ++ 11 files changed, 389 insertions(+), 251 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 790c0d629a43c..eee5453d68959 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -12,6 +12,10 @@ #param is_qualcomm #param has_weight_idx #param has_weight_idx_indirect +#param tile_size_k_vec +#param tile_size_m +#param tile_size_n +#param workgroup_size #use .getByOffset .setByOffset @@ -21,7 +25,7 @@ fn write_lora_output(offset: u32, value: vec4) { let col = (offset * 4u) % uniforms.N; var delta = vec4(0); #if lora_quantized - delta = vec4(qlora_delta(row, col), qlora_delta(row, col + 1u), qlora_delta(row, col + 2u), qlora_delta(row, col + 3u)); + delta = qlora_delta4(row, col); #elif lora_transposed for (var r4 = 0u; r4 < 8u; r4++) { let tv = lora_t[row * 8u + r4]; @@ -48,46 +52,25 @@ fn write_lora_output(offset: u32, value: vec4) { #include "quantization/dp4a_matmul_common.wgsl.template" -// This shader implements co-operative matrix multiply. The key idea here is to -// assume there is a primitive for medium size matrix multiply a subgroup can perform, -// using all its lanes and pooling all its registers to keep the values in registry. -// -// The entire workgroup which has N subgroups first loads a tile into shared memory, -// Then each subgroup loads a subtile from shared memory into registers and uses -// the medium size matrix multiply primitive to perform the math. -// The values for tile/subtile size are chosen to conform to the resource limits -// of an alderlake/tiger lake gpu. A tile is 64x64, workgroup is 256 threads - -// therefore there are 16 subgroups and 16 lanes in each subgroup. -// K the hidden dimension is paged in from RAM at k tile size which is 64. -// All this puts the shared memory requirement slightly above 16KB. -// WebGPU limit is 16KB, output is moved to registers instead of SHM to make -// everything fit in shared memory. -// -// Each subgroup performs a 16 x 64 x 16 multiply which is implemented with -// subgroup shuffle as a placeholder for the day the medium matrix mul primitive -// becomes available in WGSL. The registry requirements is ~2KB per subgroup, on -// Alderlake/Tigerlake subgroup has 8KB of registry space pooling the -// 512B of registry from each lane. -// -// The medium size matmul is implemented using dot4I8Packed, so the inputs for -// this shader require A to be int8 quantized with block size 64. B is regular -// matmulnbits input with block size 32. - -const tile_size = 64; +// Each 16-thread group computes a 16x16 output subtile. Rectangular workgroup +// tiles avoid wasting half the columns on narrow projections. Wider K loads +// amortize barriers while retaining the original 32-element dot/scale order. +// A is dynamically INT8-quantized per 128 elements; B uses its model block size. + +const subtiles_n = tile_size_n / 16; const subtile_size = 16; -const tile_size_k = 32; const vec_factor = 4; const u32_factor = 4; -const tile_size_k_vec = 2; + // Shared memory -var tile_A : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_A : array; // 64 x 1 -var tile_B : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_B : array; // 64 x 1 +var tile_A : array, tile_size_m>, tile_size_k_vec>; +var scale_A : array; +var tile_B : array, tile_size_n>, tile_size_k_vec>; +var scale_B : array; #if has_zero_points - var zeroes : array; + var zeroes : array; #endif fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) @@ -125,14 +108,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const actual_weight_idx : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16+kidx_v + col); #endif - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom4BitsTo8Bits(b_value, zero); - if (col == 0) + if (col % 2u == 0u) { // kidx_v - each kidx_v covers 16 values of k let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } } #endif @@ -160,14 +143,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif tile_B[col][row] = AlignWithZeroPoint(b_value); - if (col == 0) + if (col % 2u == 0u) { // kidx_v - each kidx_v covers 16 values of k - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); #if has_zero_points - zeroes[row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); + zeroes[(col / 2u) * tile_size_n + row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); #endif } } @@ -194,15 +177,17 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const b_weight_offset : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); #if has_zero_points let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value, zero); #else tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value); #endif - let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + if (col % 2u == 0u) { + let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + } } #endif @@ -217,31 +202,27 @@ alias acc_element_t = output_element_t; $MAIN { #if n_bits == 2 - LoadDequantizationTable(local_idx); #if has_zero_points - LoadDequantizationTable(local_idx + 256); - LoadDequantizationTable(local_idx + 512); - LoadDequantizationTable(local_idx + 768); + const table_size = 1024u; +#else + const table_size = 256u; #endif + for (var i = local_idx; i < table_size; i += workgroup_size) { + LoadDequantizationTable(i); + } workgroupBarrier(); #endif - // During the load phase we use all 256 threads to load 64 rows of A/B. - // For each row we load tile_size_k_vec (2) vectorized elements, which are 32 elements of K. let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; } - let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size; - let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size; - let load_AorB = u32(local_idx/128); - let load_row = u32((local_idx%128)/2); - let load_col = u32(local_idx%2); - - // During the compute phase, we have the 64x64 tile split into - // subtiles of 16x16. We have a grid of 4x4 subtiles. + let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size_m; + let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size_n; + + // Each group of 16 threads computes one 16x16 output subtile. var subtile_id = u32(local_idx / subtile_size); - var subtile_idx = u32(subtile_id / 4); - var subtile_idy = u32(subtile_id % 4); + var subtile_idx = u32(subtile_id / subtiles_n); + var subtile_idy = u32(subtile_id % subtiles_n); var base_A = subtile_idx * 16; var base_B = subtile_idy * 16; // For each subtile we have 16 threads assigned. @@ -255,8 +236,8 @@ $MAIN { // we would need var lane_outputs: array; if (sg_size == 64) { subtile_id = u32(local_idx / sg_size); - subtile_idx = u32(subtile_id / 4); - subtile_idy = u32(subtile_id % 4); + subtile_idx = u32(subtile_id / subtiles_n); + subtile_idy = u32(subtile_id % subtiles_n); base_A = subtile_idx * sg_size; base_B = subtile_idy * 16; a_idx = sg_id; @@ -269,189 +250,190 @@ $MAIN { var lane_output4: vec4; #endif // K's vectorization is 16 items per index. See input_a/input_b. - // tile_size_k_vec - is the k tile size in vectorized space (1/16). That is - // k tile size is 32. In vectorized space that is 32/16 = 2. + // Load 32 or 128 K elements, represented by 2 or 8 vec4 values. for (var kidx_v:u32 = 0; kidx_v < uniforms.K16; kidx_v+=tile_size_k_vec) { // Load Phase: Populate shared memory for the workgroup. - if (load_AorB == 0) - { - loadSHMA(batch, a_global_base, kidx_v, load_row, load_col); - } - else - { - loadSHMB(b_global_base, kidx_v, load_row, load_col); + for (var i = local_idx; i < (tile_size_m + tile_size_n) * tile_size_k_vec; i += workgroup_size) { + if (i < tile_size_m * tile_size_k_vec) { + loadSHMA(batch, a_global_base, kidx_v, i / tile_size_k_vec, i % tile_size_k_vec); + } else { + let j = i - tile_size_m * tile_size_k_vec; + loadSHMB(b_global_base, kidx_v, j / tile_size_k_vec, j % tile_size_k_vec); + } } workgroupBarrier(); - // Compute phase: Perform matmul for this subtile 16 x 32 x 16. - // Step 1: Load from shared memory into registers across entire subgroup. - var own_a0: vec4 = tile_A[0][base_A + a_idx]; - var own_a1: vec4 = tile_A[1][base_A + a_idx]; - var own_scale_a: output_element_t = scale_A[base_A + a_idx]; + // Keep the original per-block scaling and accumulation order. + for (var k_sub = 0u; k_sub < tile_size_k_vec / 2u; k_sub++) { + // Step 1: Load from shared memory into registers across entire subgroup. + var own_a0: vec4 = tile_A[k_sub * 2u][base_A + a_idx]; + var own_a1: vec4 = tile_A[k_sub * 2u + 1u][base_A + a_idx]; + var own_scale_a: output_element_t = scale_A[base_A + a_idx]; #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == 64) - { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - var zero: i32; - if (sg_id < 16) + if (sg_size == 64) { - own_b0 = tile_B[0][base_B + sg_id]; - own_b1 = tile_B[1][base_B + sg_id]; - own_scale_b = scale_B[base_B + sg_id]; - zero = zeroes[base_B + sg_id]; + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + var zero: i32; + if (sg_id < 16) + { + own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; + own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; + zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; + } + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); + } } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) + else { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i], zeroes[k_sub * tile_size_n + base_B + i]); + } } - } - else - { - for (var i = 0u; i < 16u; i++) + #else + if (sg_size == 16) { - lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i], zeroes[base_B + i]); + var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; + var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; + var zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); + + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); + + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); + + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0], zeroes[k_sub * tile_size_n + base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1], zeroes[k_sub * tile_size_n + base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2], zeroes[k_sub * tile_size_n + base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3], zeroes[k_sub * tile_size_n + base_B + 3]); + + lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4], zeroes[k_sub * tile_size_n + base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5], zeroes[k_sub * tile_size_n + base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6], zeroes[k_sub * tile_size_n + base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7], zeroes[k_sub * tile_size_n + base_B + 7]); + + lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8], zeroes[k_sub * tile_size_n + base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9], zeroes[k_sub * tile_size_n + base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10], zeroes[k_sub * tile_size_n + base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11], zeroes[k_sub * tile_size_n + base_B + 11]); + + lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12], zeroes[k_sub * tile_size_n + base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13], zeroes[k_sub * tile_size_n + base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14], zeroes[k_sub * tile_size_n + base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15], zeroes[k_sub * tile_size_n + base_B + 15]); } - } - #else - if (sg_size == 16) - { - var own_b0: vec4 = tile_B[0][base_B + sg_id]; - var own_b1: vec4 = tile_B[1][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[base_B + sg_id]; - var zero = zeroes[base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); - - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); - - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); - - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0], zeroes[base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1], zeroes[base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2], zeroes[base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3], zeroes[base_B + 3]); - - lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4], zeroes[base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5], zeroes[base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6], zeroes[base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7], zeroes[base_B + 7]); - - lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8], zeroes[base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9], zeroes[base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10], zeroes[base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11], zeroes[base_B + 11]); - - lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12], zeroes[base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13], zeroes[base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14], zeroes[base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15], zeroes[base_B + 15]); - } #endif #else #if is_qualcomm - if (sg_size == 64) - { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - if (sg_id < 16) + if (sg_size == 64) { - own_b0 = tile_B[0][base_B + sg_id]; - own_b1 = tile_B[1][base_B + sg_id]; - own_scale_b = scale_B[base_B + sg_id]; + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + if (sg_id < 16) + { + own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; + own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; + } + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); + } } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) + else { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i]); + } } - } - else - { - for (var i = 0u; i < 16u; i++) + #else + if (sg_size == 16) { - lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i]); + var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; + var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); + + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); + + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); + + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3]); + + lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7]); + + lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11]); + + lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15]); } - } - #else - if (sg_size == 16) - { - var own_b0: vec4 = tile_B[0][base_B + sg_id]; - var own_b1: vec4 = tile_B[1][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); - - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); - - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); - - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3]); - - lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7]); - - lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11]); - - lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15]); - } #endif #endif + } workgroupBarrier(); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index be61d2f356606..0815265090513 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -40,7 +40,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("lora_b"); if (lora_quantized_) { shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); + GenerateQLoRAUp(shader, acc_f32_, true); } } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); @@ -60,6 +60,10 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), + WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), + WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), + WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), @@ -134,7 +138,8 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, - const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales) { + const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales, + DP4AQuantizedInput* quantized_input) { const bool has_lora_down = lora_at != nullptr; const bool lora_quantized = lora_b_scales != nullptr; ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && M >= min_M_for_tile_optimization), @@ -143,20 +148,28 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor constexpr uint32_t kVec2Components = 2; constexpr uint32_t kU32Components = 4; - constexpr uint32_t kBlockSizeA = 128; - DP4AMatMulQuantizeProgram quantize_program; - quantize_program.SetWorkgroupSize(64); - uint32_t tile_size = 64 * kVec4Components; - quantize_program.SetDispatchGroupSize((batch_count * M * K + tile_size - 1) / tile_size, 1, 1); - TensorShape a_quant_shape{batch_count, M, K / kU32Components}; - Tensor a_quant = context.CreateGPUTensor(DataTypeImpl::GetType(), a_quant_shape); - TensorShapeVector a_scales_dims({batch_count, 1, M, K / kBlockSizeA}); - Tensor a_scale = context.CreateGPUTensor(a->DataType(), a_scales_dims); - quantize_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}}) - .AddOutputs({{&a_quant, ProgramTensorMetadataDependency::Rank, a_quant.Shape(), 1}, - {&a_scale, ProgramTensorMetadataDependency::Rank, 1}}) - .AddUniformVariable({batch_count * M * K / kU32Components}); - ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program)); + DP4AQuantizedInput local_input; + auto& input = quantized_input ? *quantized_input : local_input; + if (!input.values) { + constexpr uint32_t kBlockSizeA = 128; + DP4AMatMulQuantizeProgram quantize_program; + quantize_program.SetWorkgroupSize(64); + uint32_t tile_size = 64 * kVec4Components; + quantize_program.SetDispatchGroupSize((batch_count * M * K + tile_size - 1) / tile_size, 1, 1); + TensorShape a_quant_shape{batch_count, M, K / kU32Components}; + Tensor a_quant = context.CreateGPUTensor(DataTypeImpl::GetType(), a_quant_shape); + TensorShapeVector a_scales_dims({batch_count, 1, M, K / kBlockSizeA}); + Tensor a_scale = context.CreateGPUTensor(a->DataType(), a_scales_dims); + quantize_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}}) + .AddOutputs({{&a_quant, ProgramTensorMetadataDependency::Rank, a_quant.Shape(), 1}, + {&a_scale, ProgramTensorMetadataDependency::Rank, 1}}) + .AddUniformVariable({batch_count * M * K / kU32Components}); + ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program)); + input.values.emplace(std::move(a_quant)); + input.scales.emplace(std::move(a_scale)); + } + const auto& a_quant = *input.values; + const auto& a_scale = *input.scales; const bool has_lora_up = lora_t != nullptr; const bool has_zero_points = zero_points != nullptr; @@ -200,14 +213,16 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } - constexpr uint32_t kTileSize = 64; + const uint32_t tile_size_m = N < 64 && context.AdapterInfo().vendor == std::string_view{"nvidia"} ? 32 : 64; + const uint32_t tile_size_n = N < 64 ? 32 : 64; + const uint32_t tile_size_k_vec = N < 64 && nbits == 8 && block_size == 32 ? 8 : 2; TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components}; - uint32_t num_M_tile = (M + kTileSize - 1) / kTileSize; - uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize; + uint32_t num_M_tile = CeilDiv(M, tile_size_m); + uint32_t num_N_tile = CeilDiv(N, tile_size_n); bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; - mul_program.SetWorkgroupSize(256); + DP4AMatMulNBitsProgram mul_program{tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; + mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile + (has_lora_down ? CeilDiv(batch_count * M, 32u) : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, @@ -224,7 +239,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); + .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 4d30a9f2544f7..3505033aad1f5 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/webgpu_kernel.h" @@ -14,6 +15,12 @@ namespace webgpu { using namespace onnxruntime::webgpu; +// Invocation-local workspace, reusable only by projections of the same input tensor. +struct DP4AQuantizedInput { + std::optional values; + std::optional scales; +}; + class DP4AMatMulQuantizeProgram final : public Program { public: DP4AMatMulQuantizeProgram() : Program{"DP4AMatMulQuantize"} {} @@ -23,11 +30,14 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, + DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false) : Program{has_lora_down ? (lora_quantized ? "DP4AMatMulNBitsQLoRADown" : "DP4AMatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBitsLoRAUp") : "DP4AMatMulNBits"}, + tile_size_m_(tile_size_m), + tile_size_n_(tile_size_n), + tile_size_k_vec_(tile_size_k_vec), block_size_(block_size), nbits_(nbits), has_bias_(has_bias), @@ -54,6 +64,9 @@ class DP4AMatMulNBitsProgram final : public Program { {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: + uint32_t tile_size_m_; + uint32_t tile_size_n_; + uint32_t tile_size_k_vec_; uint32_t block_size_; uint32_t nbits_; bool has_bias_; @@ -134,7 +147,8 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const uint32_t weight_index, const Tensor* weight_index_indirect = nullptr, const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, - const Tensor* lora_at = nullptr, Tensor* lora_down = nullptr, bool lora_transposed = false, const Tensor* lora_b_scales = nullptr); + const Tensor* lora_at = nullptr, Tensor* lora_down = nullptr, bool lora_transposed = false, const Tensor* lora_b_scales = nullptr, + DP4AQuantizedInput* quantized_input = nullptr); // The optional M / has_weight_idx_indirect / y arguments fold the original // dispatch-precondition (DP4A is preferred when M is large enough, or diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 0c080104535fd..ee40f49dc2ffc 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -223,7 +223,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* weight_index_indirect, uint32_t override_M, const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_a_transposed, Tensor* lora_a_output, bool lora_transposed, - const Tensor* lora_b_scales, const Tensor* lora_a_scales) { + const Tensor* lora_b_scales, const Tensor* lora_a_scales, + DP4AQuantizedInput* quantized_input) { TensorShape b_shape({N_op, K_op}); MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); @@ -298,7 +299,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). if (!has_lora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed, lora_b_scales); + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed, lora_b_scales, quantized_input); } // WideTileProgram diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index 513bbf0b8fe50..91802885fd847 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -9,6 +9,7 @@ namespace onnxruntime { namespace contrib { namespace webgpu { +struct DP4AQuantizedInput; // Adds an independent FP16 rank-32 down-projection branch to a 256-thread shader. void GenerateLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups); @@ -130,7 +131,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, uint32_t override_M = 0, const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, bool lora_transposed = false, - const Tensor* lora_b_scales = nullptr, const Tensor* lora_a_scales = nullptr); + const Tensor* lora_b_scales = nullptr, const Tensor* lora_a_scales = nullptr, + DP4AQuantizedInput* quantized_input = nullptr); Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed = false, const Tensor* quantized_scales = nullptr); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index efee298dcebdc..56a7222041f63 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -2,7 +2,9 @@ // Licensed under the MIT License. #include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/common/safeint.h" @@ -57,7 +59,7 @@ if (workgroup_idx >= qlora_base_groups) { )"; } -void GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32) { +void GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32, bool vector_output) { shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" << R"( fn qlora_weight(col: u32, r4: u32) -> vec4 { @@ -89,6 +91,28 @@ fn qlora_delta(row: u32, col: u32) -> f16 { return f16(delta); } )"; + if (vector_output) { + shader.AdditionalImplementation() << R"( +fn qlora_delta4(row: u32, col: u32) -> vec4 { + var delta = vec4(0); + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = vec4(lora_t[row * 8u + r4]); + let a1 = vec4(lora_t[row * 8u + r4 + 1u]); + delta += vec4( + dot(a0, vec4(qlora_weight(col, r4))), + dot(a0, vec4(qlora_weight(col + 1u, r4))), + dot(a0, vec4(qlora_weight(col + 2u, r4))), + dot(a0, vec4(qlora_weight(col + 3u, r4)))) + + vec4( + dot(a1, vec4(qlora_weight(col, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 1u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 2u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 3u, r4 + 1u)))); + } + return vec4(delta); +} +)"; + } } namespace { @@ -212,17 +236,22 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { nullptr, 0, nullptr, nullptr, aq, &t, true, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, 0, 8, context, &t)); + // Allow activation quantization for prefill only when the model opts into it. + const auto down_accuracy = fuse_up && m >= kMinMForTileOptimization ? accuracy_level_ : 0; + DP4AQuantizedInput quantized_x; + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, + nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); std::optional up_config; const bool subgroup_up = CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), narrow(n_), 32, 8, true, up_config, m); if (fuse_up && !subgroup_up) { return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, &t, - bq, nullptr, nullptr, true, bs); + bq, nullptr, nullptr, true, bs, nullptr, &quantized_x); } auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base)); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, + nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); if (!subgroup_up) return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); // Keep the original up kernel on devices with a different matrix-accumulation path. auto delta = context.CreateGPUTensor(y->DataType(), y->Shape()); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h index befb368db025b..91e2242737262 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -9,7 +9,7 @@ namespace onnxruntime::contrib::webgpu { void GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32); -void GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32); +void GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32, bool vector_output = false); Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, const Tensor* scales, Tensor* y); diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 34544efa0ac2e..26106177956d1 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1520,6 +1520,26 @@ TEST(MatMul2Bits, Float32_2b_BlkLen32_Accuracy0) { TestMatMul2BitsTyped(); } +#ifdef USE_WEBGPU +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int2) { + for (int64_t n : {16, 32, 48, 64, 80}) { + for (bool has_zero_point : {false, true}) { + TestOptions2Bits opts{}; + opts.M = 65; + opts.N = n; + opts.K = 256; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.05f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + RunTest2Bits(opts); + } + } +} +#endif + #if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) namespace { diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index e7ffb62fbcb4d..0fc9214b6c88f 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -452,6 +452,27 @@ TEST(MatMulNBits, Float32_4b_Accuracy0) { TestMatMulNBitsTyped(); } +#ifdef USE_WEBGPU +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int4) { + for (int64_t n : {16, 32, 48, 64, 80}) { + for (bool has_zero_point : {false, true}) { + TestOptions opts{}; + opts.batch_count = 2; + opts.M = 65; + opts.N = n; + opts.K = 256; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.05f; + opts.output_rel_error = 0.02f; + RunTest(opts); + RunTest(opts); + } + } +} +#endif + TEST(MatMulNBits, Float32_4b_Accuracy1) { TestMatMulNBitsTyped(); TestMatMulNBitsTyped(); diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 65ea7f5bf2ecb..d03639e7a5e6f 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1057,6 +1057,28 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { } #endif +#if defined(USE_WEBGPU) +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { + for (int64_t m : {4, 65}) { + for (int64_t n : {16, 32, 48, 64, 80}) { + for (bool has_zero_point : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m; + opts.N = n; + opts.K = 256; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.05f; + opts.output_rel_error = 0.02f; + RunTest8Bits(opts); + RunTest8Bits(opts); + } + } + } +} +#endif + #if !defined(USE_CUDA) && !defined(USE_WEBGPU) #ifndef ENABLE_TRAINING // Pre-packing (and therefore cross-session sharing of pre-packed weights) is disabled in a full diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index a862ef0279410..89cdeca89ff5f 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -61,5 +61,37 @@ TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { } } +TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { + for (int64_t m : {1, 3, 4, 65}) { + for (int64_t accuracy : {0, 4}) { + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + constexpr int64_t k = 256; + constexpr int64_t n = 128; + constexpr int64_t batches = 2; + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", accuracy); + std::vector x(batches * m * k); + std::vector expected(batches * m * n); + for (int64_t row = 0; row < batches * m; ++row) { + const float value = row % 2 == 0 ? 0.5f : -0.25f; + std::fill_n(x.begin() + row * k, k, MLFloat16(value)); + std::fill_n(expected.begin() + row * n, n, MLFloat16(value * 24.0f)); + } + test.AddInput("X", {batches, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {32, k / 32, 32}, std::vector(32 * k, 129)); + test.AddInput("AS", {32, k / 32}, std::vector(k, MLFloat16(0.03125f))); + test.AddInput("BQ", {n, 1, 32}, std::vector(n * 32, 129)); + test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(0.0625f))); + test.AddOutput("Y", {batches, m, n}, expected); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + #endif } // namespace onnxruntime::test From 8eab6582927d643d50d822381e1406e5d92ec7a1 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:21:04 +0800 Subject: [PATCH 06/35] Batch WebGPU input uploads to amortize adapter transfer overhead --- .../core/providers/webgpu/buffer_manager.cc | 56 +++++++++++++++++++ .../core/providers/webgpu/buffer_manager.h | 6 ++ .../core/providers/webgpu/data_transfer.cc | 19 +++++++ .../core/providers/webgpu/data_transfer.h | 5 +- .../providers/webgpu/data_transfer_test.cc | 51 +++++++++++++++++ 5 files changed, 136 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 228a118f1ce96..9ea2f30781c67 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -533,6 +533,62 @@ void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuff ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); } +void BufferManager::UploadBatch(CommandRecordingState& recording, gsl::span uploads) const { + // Bound staging memory while amortizing buffer creation and queue submission across small inputs. + constexpr size_t kBatchBytes = 16 * 1024 * 1024; + size_t begin = 0; + while (begin < uploads.size()) { + size_t end = begin; + size_t bytes = 0; + while (end < uploads.size()) { + const auto aligned_size = NormalizeCopySize(uploads[end].size); + if (aligned_size > kBatchBytes - bytes) break; + bytes += aligned_size; + ++end; + } + if (end == begin) { + const auto& upload = uploads[begin++]; + Upload(recording, const_cast(upload.source), upload.target, upload.size); + continue; + } + if (bytes != 0) { + wgpu::BufferDescriptor desc{}; + desc.size = bytes; + desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite; + desc.mappedAtCreation = true; + auto staging_buffer = context_.Device().CreateBuffer(&desc); + auto* mapped = static_cast(staging_buffer.GetMappedRange()); + size_t offset = 0; + for (size_t i = begin; i < end; ++i) { + const auto& upload = uploads[i]; + const auto aligned_size = NormalizeCopySize(upload.size); + if (upload.size != 0) { + memcpy(mapped + offset, upload.source, upload.size); + memset(mapped + offset + upload.size, 0, aligned_size - upload.size); + } + offset += aligned_size; + } + staging_buffer.Unmap(); + ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); + context_.EndComputePass(recording); + auto& encoder = context_.GetCommandEncoder(recording); + offset = 0; + for (size_t i = begin; i < end; ++i) { + const auto& upload = uploads[i]; + const auto aligned_size = NormalizeCopySize(upload.size); + if (aligned_size != 0) { + ORT_ENFORCE(aligned_size <= wgpuBufferGetSize(upload.target)); + EnforceBufferUnmapped(context_, upload.target); + encoder.CopyBufferToBuffer(staging_buffer, offset, upload.target, 0, aligned_size); + } + offset += aligned_size; + } + ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); + } + begin = end; + } +} + void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const { ORT_ENFORCE(src != dst, "Source and destination buffers must be different."); EnforceBufferUnmapped(context_, src); diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index f7e1fb033c816..acb3f0f262402 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -94,6 +94,12 @@ class BufferManager { public: BufferManager(WebGpuContext& context, BufferCacheMode storage_buffer_cache_mode, BufferCacheMode uniform_buffer_cache_mode, BufferCacheMode query_resolve_buffer_cache_mode, BufferCacheMode default_buffer_cache_mode); void Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const; + struct UploadInfo { + const void* source; + WGPUBuffer target; + size_t size; + }; + void UploadBatch(CommandRecordingState& recording, gsl::span uploads) const; void MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const; WGPUBuffer Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage, bool initialize_to_zero = false, diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 4411c922a6fc3..0bf7b0319a7ea 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -71,5 +71,24 @@ common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { src.SizeInBytes()); } +common::Status DataTransfer::CopyTensors(const std::vector& src_dst_pairs) const { + if (src_dst_pairs.size() < 2 || buffer_manager_.SupportsUMA()) { + return IDataTransfer::CopyTensors(src_dst_pairs); + } + InlinedVector uploads; + uploads.reserve(src_dst_pairs.size()); + for (const auto& pair : src_dst_pairs) { + const auto& src = pair.src.get(); + auto& dst = pair.dst.get(); + if (pair.src_stream != nullptr || src.Location().device.Type() != OrtDevice::CPU || + dst.Location().device.Type() != OrtDevice::GPU) { + return IDataTransfer::CopyTensors(src_dst_pairs); + } + uploads.push_back({src.DataRaw(), static_cast(dst.MutableDataRaw()), src.SizeInBytes()}); + } + buffer_manager_.UploadBatch(recording_, uploads); + return Status::OK(); +} + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/data_transfer.h b/onnxruntime/core/providers/webgpu/data_transfer.h index 1ebc869f66d39..7951ff032ddd1 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.h +++ b/onnxruntime/core/providers/webgpu/data_transfer.h @@ -73,7 +73,7 @@ class DataTransferImpl { class DataTransfer : public IDataTransfer { public: DataTransfer(const BufferManager& buffer_manager, CommandRecordingState& recording, int context_id) - : impl_{buffer_manager, recording}, context_id_{context_id} {} + : impl_{buffer_manager, recording}, context_id_{context_id}, buffer_manager_{buffer_manager}, recording_{recording} {} ~DataTransfer() {}; // Device-compatibility half of CanCopy, split out because it needs no BufferManager and so can @@ -83,10 +83,13 @@ class DataTransfer : public IDataTransfer { bool CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const override; common::Status CopyTensor(const Tensor& src, Tensor& dst) const override; + common::Status CopyTensors(const std::vector& src_dst_pairs) const override; private: DataTransferImpl impl_; int context_id_; + const BufferManager& buffer_manager_; + CommandRecordingState& recording_; }; } // namespace webgpu diff --git a/onnxruntime/test/providers/webgpu/data_transfer_test.cc b/onnxruntime/test/providers/webgpu/data_transfer_test.cc index 23ac26fbfe227..050ca1206634f 100644 --- a/onnxruntime/test/providers/webgpu/data_transfer_test.cc +++ b/onnxruntime/test/providers/webgpu/data_transfer_test.cc @@ -5,6 +5,12 @@ #include "core/framework/ortdevice.h" #include "core/providers/webgpu/data_transfer.h" +#include "core/providers/webgpu/allocator.h" +#include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_execution_provider.h" +#include "core/providers/webgpu/webgpu_provider_factory_creator.h" +#include "test/util/include/asserts.h" namespace onnxruntime { namespace test { @@ -71,5 +77,50 @@ TEST(WebGpuDataTransferTest, RejectsForeignGpuVendors) { } } +TEST(WebGpuDataTransferTest, BatchedUploadsPreserveDataAcrossChunksAndReuse) { + ConfigOptions options; + auto provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& ep = static_cast(*provider); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + auto& recording = ep.Recording(); + CPUAllocator cpu_allocator; + const OrtMemoryInfo gpu_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault); + webgpu::DataTransfer transfer(ep.BufferManager(), recording, 0); + + // Include empty/unaligned tensors, a batch boundary, and an upload larger than the batch limit. + const std::vector sizes{0, 1, 3, 4, 31, 65, 4097, 9 * 1024 * 1024 + 1, + 9 * 1024 * 1024 + 3, 17 * 1024 * 1024 + 1}; + std::vector> data; + std::vector buffers; + std::vector sources, targets; + sources.reserve(sizes.size()); + targets.reserve(sizes.size()); + for (const auto size : sizes) { + data.emplace_back(static_cast(size)); + wgpu::BufferDescriptor desc{}; + // Match BufferManager's allocation alignment, which Download also relies on. + desc.size = std::max(16, (static_cast(size) + 15) / 16 * 16); + desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + buffers.push_back(context.Device().CreateBuffer(&desc)); + sources.emplace_back(DataTypeImpl::GetType(), TensorShape({size}), data.back().data(), cpu_allocator.Info()); + targets.emplace_back(DataTypeImpl::GetType(), TensorShape({size}), buffers.back().Get(), gpu_info); + } + std::vector uploads; + for (size_t i = 0; i < sizes.size(); ++i) uploads.push_back({sources[i], targets[i], nullptr}); + for (size_t pass = 0; pass < 2; ++pass) { + for (size_t i = 0; i < sizes.size(); ++i) { + for (size_t j = 0; j < data[i].size(); ++j) data[i][j] = static_cast(i * 17 + j * 7 + pass); + } + ASSERT_STATUS_OK(transfer.CopyTensors(uploads)); + for (size_t i = 0; i < sizes.size(); ++i) { + std::vector actual(static_cast(sizes[i])); + Tensor result(DataTypeImpl::GetType(), TensorShape({sizes[i]}), actual.data(), cpu_allocator.Info()); + ASSERT_STATUS_OK(transfer.CopyTensor(targets[i], result)); + EXPECT_EQ(actual, data[i]) << "tensor " << i << " pass " << pass; + } + } + ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), recording)); +} + } // namespace test } // namespace onnxruntime From 210584c743785f949cfc417a185366fa0343b932 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 16:59:11 +0800 Subject: [PATCH 07/35] Optimize narrow WebGPU QLoRA projections with small-M DP4A --- .../webgpu/quantization/dp4a_matmul_nbits.cc | 14 +++- .../webgpu/quantization/dp4a_matmul_nbits.h | 4 +- .../dp4a_matmul_small_m.wgsl.template | 2 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 16 +++- .../contrib_ops/matmul_nbits_qlora_test.cc | 80 +++++++++++++++++++ 5 files changed, 108 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 0815265090513..6e750083db2e6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -139,9 +139,13 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales, - DP4AQuantizedInput* quantized_input) { + DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n) { const bool has_lora_down = lora_at != nullptr; const bool lora_quantized = lora_b_scales != nullptr; + ORT_RETURN_IF_NOT(small_m_tile_n == 0 || + (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && + !zero_points && !bias && !weight_index && !weight_index_indirect && !lora_t && !lora_at), + "Small-M override requires a plain rank-32 INT8 projection."); ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && M >= min_M_for_tile_optimization), "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; @@ -177,13 +181,15 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index != 0 || has_weight_idx_indirect; const bool single_scale_weights = (block_size == K * N); - if (has_weight_idx_indirect || M < min_M_for_tile_optimization) { + if (small_m_tile_n || has_weight_idx_indirect || M < min_M_for_tile_optimization) { uint32_t tile_size_k_vec = 32; - uint32_t tile_size_n = 4; + uint32_t tile_size_n = small_m_tile_n ? small_m_tile_n : 4; const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components); const bool broadcast_a = dispatch_M > M; - const bool acc_f32 = context.EnableMatmulFp32Accumulation(); + // The narrow projection changes the K reduction order. Keep partial sums + // and the final reduction in FP32, then cast T once when storing it. + const bool acc_f32 = small_m_tile_n != 0 || context.EnableMatmulFp32Accumulation(); DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 3505033aad1f5..84ddf55f62cfb 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -148,7 +148,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const Tensor* weight_index_indirect = nullptr, const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, const Tensor* lora_at = nullptr, Tensor* lora_down = nullptr, bool lora_transposed = false, const Tensor* lora_b_scales = nullptr, - DP4AQuantizedInput* quantized_input = nullptr); + DP4AQuantizedInput* quantized_input = nullptr, + // Nonzero selects the K-parallel kernel for a rank-32 INT8 projection. + uint32_t small_m_tile_n = 0); // The optional M / has_weight_idx_indirect / y arguments fold the original // dispatch-precondition (DP4A is preferred when M is large enough, or diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template index 41380419e25ff..645178a10501e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template @@ -95,7 +95,7 @@ fn loadSHMA(batch: u32, a_global: u32, kidx_v: u32, col: u32) } tile_A[col] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+k_offset); - if (col < scale_a_size_in_tile_a) + if (col < scale_a_size_in_tile_a && kidx_v/8 + col < uniforms.K16/8) { // kidx_v - covers 16 values of k in input_a scale_A[col] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8 + col); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 56a7222041f63..3099da06cdcda 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -239,8 +239,20 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // Allow activation quantization for prefill only when the model opts into it. const auto down_accuracy = fuse_up && m >= kMinMForTileOptimization ? accuracy_level_ : 0; DP4AQuantizedInput quantized_x; - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, - nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); + // Rank 32 exposes too few output tiles for the tiled DP4A kernel at these sizes. + const bool small_m_down = fuse_up && m <= 4096 && k_ >= 2048 && + context.AdapterInfo().vendor == std::string_view{"nvidia"} && + CanApplyDP4AMatrixMatMulNBits(context, down_accuracy, 32, 32, + narrow(k_), 4, m, false, &t); + if (small_m_down) { + ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( + x, aq, as, nullptr, nullptr, narrow(x->Shape().Size() / k_ / m), m, m, 32, + narrow(k_), 32, narrow(k_ / 32), kMinMForTileOptimization, 8, context, &t, 0, + nullptr, nullptr, nullptr, nullptr, nullptr, false, nullptr, &quantized_x, 32)); + } else { + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, + nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); + } std::optional up_config; const bool subgroup_up = CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 89cdeca89ff5f..09916e7e369f3 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -4,6 +4,11 @@ #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" +#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) +#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "core/providers/webgpu/webgpu_context.h" +#endif + namespace onnxruntime::test { #if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) @@ -93,5 +98,80 @@ TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { } } +TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId()) || + gpu.AdapterInfo().vendor != std::string_view{"nvidia"}) { + GTEST_SKIP() << "The narrow small-M projection is selected on NVIDIA DP4A devices."; + } + struct Shape { + int64_t m; + int64_t k; + int64_t batch; + }; + for (const auto shape : {Shape{4, 2048, 1}, Shape{37, 2176, 2}, Shape{128, 3072, 1}, Shape{65, 8192, 1}}) { + const auto [m, k, batch] = shape; + SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", batch=", batch)); + constexpr int64_t n = 64; + std::vector x(batch * m * k); + std::vector aq(32 * k); + std::vector as(k); + std::vector bq(n * 32, 128); + std::vector bs(n, MLFloat16(1.0f)); + std::vector expected(batch * m * n); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // The maximum in every 128-element block is 127 * scale, so DP4A's + // input quantization preserves these signed integer values exactly. + const float scale = 1.0f / static_cast(1 << (7 + (row + col / 128) % 4)); + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + x[row * k + col] = MLFloat16(row % 7 == 0 ? 0.0f : value * scale); + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) { + aq[r * k + col] = static_cast((r * 17 + col * 13) % 256); + } + for (int64_t block = 0; block < k / 32; ++block) { + as[r * (k / 32) + block] = MLFloat16(1.0f / static_cast(1 << (8 + (r + block) % 4))); + } + } + // A zero base and an identity up projection expose every element of T. + // Compute the reference as a dense dequantized dot product in double. + for (int64_t col = 0; col < n; ++col) { + bq[col * 32 + col % 32] = 129; + } + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + sum += static_cast(x[row * k + col].ToFloat()) * + (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32].ToFloat(); + } + expected[row * n + r] = expected[row * n + r + 32] = MLFloat16(static_cast(sum)); + } + } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x88), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(1.0f)), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, bs); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.002f); + test.SetOutputRelErr("Y", 0.001f); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } +} + #endif } // namespace onnxruntime::test From b3e1633c9567459041474114256d389bc824e569 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:05:41 +0800 Subject: [PATCH 08/35] Optimize WebGPU QLoRA decode with register reuse and subgroup reduction --- .../webgpu/quantization/matmul_nbits.cc | 12 +++- .../webgpu/quantization/matmul_nbits.h | 10 +-- .../webgpu/quantization/matmul_nbits_qlora.cc | 47 ++++++++++++- .../webgpu/quantization/matmul_nbits_qlora.h | 2 +- .../contrib_ops/matmul_nbits_qlora_test.cc | 70 +++++++++++++++++++ .../matmul_nbits_qlora_fusion_test.cc | 13 +++- 6 files changed, 143 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index ee40f49dc2ffc..0b5f7284b7b26 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -108,7 +108,7 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_lora_down_) { shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); if (lora_quantized_) { - GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_); + GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_, qlora_down_subgroups_); } else { GenerateLoRAGemv(shader, "uniforms.num_N_tile", a.GetByOffset("k4")); } @@ -397,7 +397,13 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, constexpr uint32_t kU32Components = 4; uint32_t components_b_with_u32 = components_b * kU32Components; uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; - MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; + // The decode reduction assigns one rank to each 32-lane subgroup. + const auto& adapter = context.AdapterInfo(); + const bool qlora_down_subgroups = has_lora_down && lora_quantized && + adapter.vendor == std::string_view{"nvidia"} && + context.HasFeature(wgpu::FeatureName::Subgroups) && + adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; + MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, qlora_down_subgroups}; program.SetWorkgroupSize(workgroup_size); uint32_t num_N_tile = (N + tile_size - 1) / tile_size; program.SetDispatchGroupSize(num_N_tile + (has_lora_down ? (lora_quantized ? 4 : 8) : 0), dispatch_M, batch_count); @@ -418,7 +424,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {batch_count}, {weight_index}, {dispatch_M}}) - .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); + .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, qlora_down_subgroups); if (has_zero_points) { program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index 91802885fd847..0b2d7b8201b79 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -53,9 +53,9 @@ class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false) - : Program{has_lora_down ? (lora_quantized ? "MatMulNBitsQLoRADown" : "MatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "MatMulNBitsQLoRAUp" : "MatMulNBitsLoRAUp") - : "MatMulNBits"}, + MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false, bool qlora_down_subgroups = false) + : Program{has_lora_down ? (lora_quantized ? (qlora_down_subgroups ? "MatMulNBitsQLoRADownSubgroup" : "MatMulNBitsQLoRADown") : "MatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "MatMulNBitsQLoRAUp" : "MatMulNBitsLoRAUp") + : "MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), @@ -69,7 +69,8 @@ class MatMulNBitsProgram final : public Program { has_lora_up_(has_lora_up), has_lora_down_(has_lora_down), lora_transposed_(lora_transposed), - lora_quantized_(lora_quantized) {} + lora_quantized_(lora_quantized), + qlora_down_subgroups_(qlora_down_subgroups) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"M", ProgramUniformVariableDataType::Uint32}, @@ -100,6 +101,7 @@ class MatMulNBitsProgram final : public Program { bool has_lora_down_; bool lora_transposed_; bool lora_quantized_; + bool qlora_down_subgroups_; }; class MatMulNBits final : public WebGpuKernel { diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 3099da06cdcda..89149c96f82df 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -13,7 +13,52 @@ namespace onnxruntime::contrib::webgpu { -void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32) { +void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, + bool acc_f32, bool use_subgroups) { + if (use_subgroups) { + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n"; + auto& body = shader.MainFunctionBody(); + body << "let qlora_base_groups = " << base_groups << ";\n" + << R"( +if (workgroup_idx >= qlora_base_groups) { + let rank_base = (workgroup_idx - qlora_base_groups) * 8u; + let lane = local_idx % 32u; + let rank_local = local_idx / 32u; + var partial0 = qlora_acc_t(0); + var partial1 = qlora_acc_t(0); + for (var k = lane * 16u; k < uniforms.K; k += 512u) { + var k4 = k / 4u; +)"; + // Keep each X vector in a register for both ranks assigned to this lane. + for (uint32_t i = 0; i < 4; ++i) { + body << "let x" << i << " = vec4(" << x_value << ");\n"; + if (i != 3) body << "k4 += 1u;\n"; + } + for (uint32_t r = 0; r < 2; ++r) { + body << "let rank" << r << " = rank_base + rank_local + " << r * 4 << "u;\n" + << "let packed" << r << " = lora_at[rank" << r << " * (uniforms.K / 16u) + k / 16u];\n" + << "let scale" << r << " = lora_a_scales[rank" << r << " * (uniforms.K / 32u) + k / 32u];\n" + << "var sum" << r << " = qlora_acc_t(0);\n"; + for (uint32_t i = 0; i < 4; ++i) { + body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" + << r << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; + } + body << "partial" << r << " += sum" << r << ";\n"; + } + // FP32 avoids half-precision rounding at each step of the subgroup tree. + body << R"( + } + let result0 = subgroupAdd(f32(partial0)); + let result1 = subgroupAdd(f32(partial1)); + if (lane == 0u) { + lora_down[rank_base + rank_local] = f16(result0); + lora_down[rank_base + rank_local + 4u] = f16(result1); + } + return; +} +)"; + return; + } shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" << R"( var qlora_x_tile: array, 128>; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h index 91e2242737262..66148cdebf8ee 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -8,7 +8,7 @@ namespace onnxruntime::contrib::webgpu { void GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, - const std::string& x_value, bool acc_f32); + const std::string& x_value, bool acc_f32, bool use_subgroups = false); void GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32, bool vector_output = false); Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, const Tensor* scales, Tensor* y); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 09916e7e369f3..1e7460e7d5c3a 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -7,6 +7,7 @@ #if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_provider_options.h" #endif namespace onnxruntime::test { @@ -173,5 +174,74 @@ TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { } } +TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); + const auto& adapter = gpu.AdapterInfo(); + if (!gpu.DeviceHasFeature(wgpu::FeatureName::Subgroups) || + adapter.vendor != std::string_view{"nvidia"} || + adapter.subgroupMinSize != 32 || adapter.subgroupMaxSize != 32) { + GTEST_SKIP() << "The decode specialization requires NVIDIA 32-lane subgroups."; + } + for (bool acc_f32 : {false, true}) { + for (const auto [k, n] : {std::pair{2048, 36}, {4096, 132}, {8192, 64}}) { + SCOPED_TRACE(MakeString("K=", k, ", N=", n, ", acc_f32=", acc_f32)); + std::vector x(k); + std::vector aq(32 * k); + std::vector as(k); + std::vector bq(n * 32, 128); + std::vector expected(n); + double base = 0; + for (int64_t col = 0; col < k; ++col) { + x[col] = MLFloat16((static_cast(col % 17) - 8) / 64.0f); + base += x[col].ToFloat() * 0.03125; + } + std::vector t(32); + for (int64_t rank = 0; rank < 32; ++rank) { + for (int64_t block = 0; block < k / 32; ++block) { + as[rank * (k / 32) + block] = MLFloat16(0.001f * (1 + (rank + block) % 4)); + } + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + const int q = rank == 0 ? 0 : static_cast((rank * 19 + col * 13) % 65) - 32; + aq[rank * k + col] = static_cast(q + 128); + const auto weight = MLFloat16(q * as[rank * (k / 32) + col / 32].ToFloat()); + sum += static_cast(x[col].ToFloat()) * weight.ToFloat(); + } + t[rank] = MLFloat16(static_cast(sum)); + } + // Identity rows expose all 32 T values; the nonzero base also exercises + // ordinary workgroups in the same dispatch and partial output tiles. + const auto rounded_base = MLFloat16(static_cast(base)); + for (int64_t col = 0; col < n; ++col) { + bq[col * 32 + col % 32] = 129; + expected[col] = MLFloat16(rounded_base.ToFloat() + t[col % 32].ToFloat()); + } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddInput("X", {1, 1, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(1.0f))); + test.AddOutput("Y", {1, 1, n}, expected); + test.SetOutputAbsErr("Y", 0.002f); + test.SetOutputRelErr("Y", 0.001f); + ConfigOptions options; + if (acc_f32) { + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + } + std::vector> eps; + eps.push_back(WebGpuExecutionProviderWithOptions(options)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + #endif } // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 598c1947982c6..631a24f0aa436 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -108,12 +108,19 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { SCOPED_TRACE(shape.k); SCOPED_TRACE(shape.n); SCOPED_TRACE(shape.batch); + // DP4A prefill quantizes X; subgroup decode changes the K reduction + // order. Keep exact comparisons for the other dispatch paths. + const bool approximate = + (shape.m >= 4 && std::string_view{mode} != "fuse_down") || + (shape.m == 1 && shape.batch == 1 && shape.k >= 2048 && shape.k <= 8192 && + std::string_view{mode} != "fuse_up"); RunWebGpuFusionTransformerTest( [shape](ModelTestBuilder& b) { BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, std::make_unique(), + TransformerLevel::Level1, TransformerLevel::Level2, 21, + approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) @@ -133,10 +140,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Accumulation) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 2, 37}) { SCOPED_TRACE(m); + const bool approximate = m == 1 || m >= 4; RunWebGpuFusionTransformerTest( [m](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, 132, 1, false, false, true); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, std::make_unique(), + TransformerLevel::Level1, TransformerLevel::Level2, 21, + approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), [] { ConfigOptions options; ORT_THROW_IF_ERROR(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, From 6a1cf7d38da330edab6272fba843353972cc6981 Mon Sep 17 00:00:00 2001 From: qjia7 <4221210+qjia7@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:15:43 +0800 Subject: [PATCH 09/35] Document WebGPU LoRA performance and numerical limits --- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 116 ++++++++++++++++++ docs/WebGPU_MatMulNBitsLoRA_Prototype.md | 4 +- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 55 ++++++--- .../core/graph/contrib_ops/contrib_defs.cc | 8 +- 4 files changed, 163 insertions(+), 20 deletions(-) create mode 100644 docs/WebGPU_MatMulNBitsLoRA_Performance.md diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md new file mode 100644 index 0000000000000..15a57a22dc5af --- /dev/null +++ b/docs/WebGPU_MatMulNBitsLoRA_Performance.md @@ -0,0 +1,116 @@ +# WebGPU LoRA and QLoRA evaluation + +This report compares upstream main with the optimized QLoRA implementation +using one evaluation model and two adapters, labeled A and B. The model uses +FP16 activations, INT4 base weights, and rank-32 INT8 adapters. Model names, +prompt contents, local paths, raw logs, and model assets are omitted. + +## Main versus optimized + +Measurements were collected on 2026-10-14 with native GenAI C++, built-in +WebGPU, Windows D3D12, and an NVIDIA RTX 5080 (driver 610.88). + +**The unchanged original optional-adapter graph fails with the adapter off on +main**, at both prompt lengths: empty adapter defaults reach DequantizeLinear, +which rejects dispatch size `(0, 1, 1)`. Main's **Off*** values below therefore +use a separately pruned base-only graph. Optimized Off uses the actual +optional-adapter graph with empty adapters. These are different off baselines. + +| Prompt tokens | Adapter | Main TTFT ms | Optimized TTFT ms | TTFT reduction | Main decode tok/s | Optimized decode tok/s | Decode gain | +|---:|---|---:|---:|---:|---:|---:|---:| +| 128 | Off* | 29.86 | 31.18 | -4.4% | 261.49 | 236.06 | -9.7% | +| 128 | A | 97.44 | 40.46 | 58.5% | 115.73 | 205.66 | 77.7% | +| 128 | B | 92.48 | 40.34 | 56.4% | 115.11 | 205.36 | 78.4% | +| 1024 | Off* | 162.67 | 162.98 | -0.2% | 258.23 | 232.99 | -9.8% | +| 1024 | A | 242.25 | 172.87 | 28.6% | 113.52 | 201.19 | 77.2% | +| 1024 | B | 246.09 | 172.97 | 29.7% | 113.64 | 199.22 | 75.3% | + +TTFT includes prompt processing and first-token sampling/readback. Lower TTFT +and higher throughput are better. Negative gains indicate regressions against +the listed reference. Pure prompt-processing throughput is reported separately: + +| Prompt tokens | Adapter | Main prefill tok/s | Optimized prefill tok/s | Gain | +|---:|---|---:|---:|---:| +| 128 | Off* | 4296.0 | 4114.5 | -4.2% | +| 128 | A | 1314.8 | 3170.1 | 141.1% | +| 128 | B | 1385.3 | 3179.5 | 129.5% | +| 1024 | Off* | 6299.4 | 6289.8 | -0.2% | +| 1024 | A | 4228.5 | 5926.8 | 40.2% | +| 1024 | B | 4162.7 | 5925.0 | 42.3% | + +Within the optimized optional-adapter graph, enabling an adapter adds +29.4-29.8% TTFT at 128 prompt tokens and 6.1% at 1024. Decode latency overhead +is 14.8-14.9% and 15.8-16.9%, respectively, computed as +`off tokens/s / on tokens/s - 1`. The 10% decode-overhead target is not met. + +Main's base-only graph receives MatMulNBitsMlpFusion and MatMulNBitsQkvFusion. +The optimized contrib nodes prevent those separate base fusions. Skipping +adapter arithmetic therefore does not give parity with the standalone base +graph: optimized off decode is approximately 10% slower than that reference. + +## Measurement method + +- Main revision: `9c87bdf2e6676eee0073c711a32e75183dc41199`. +- Optimized source tree: `b3e1633c9567459041474114256d389bc824e569`. + Publication changed commit metadata only; the tested source tree is identical. +- Both configurations use the same native benchmark executable, GenAI DLL, and + DXC. The benchmark includes adapter selection and fixed-token inputs. Only + the ORT runtime and model representation differ. +- Batch size 1, identical fixed prompt IDs, greedy decoding, 128 generated + tokens, maximum length 4096, graph capture, generator reuse, and token readback. +- Two rounds per configuration, each with three warmups and five measured + requests. The second round reverses configuration and prompt-length order. + Tables use the median of all ten individual requests. All 120 measured + requests completed with the expected token counts and stable token sequences + within each configuration. +- Measurements run serially after builds finish, with profiling, verbose + logging, shader dumping, and other evaluation work disabled. Each cell checks + for competing build/test/benchmark work and waits for three idle-GPU samples. + Separate diagnostic runs confirm all nodes execute on WebGPU and graph + capture replays; diagnostic timings are excluded. +- Main prefill varies more: adapter-on TTFT ranges are 91.69-107.07 ms (A) and + 91.12-106.52 ms (B) at 128 tokens; 234.83-252.03 ms (A) and 236.42-250.94 ms + (B) at 1024. Optimized ranges are 40.06-41.18 ms and 40.01-41.08 ms at 128; + 172.42-173.09 ms and 172.34-173.23 ms at 1024. + +The original adapters use signed INT8 weights and FP32 scales. Offline +conversion transposes and packs the same integer weights and rounds scales to +FP16; the base external weights remain byte-identical. DP4A prefill and decode +reduction changes also affect arithmetic. This compares the same source +workload after conversion, with the numerical limits below. + +Main includes newer unrelated upstream changes than the optimization branch's +common ancestor. This is an end-to-end configuration comparison, not an +identical-base ablation. Results cover one GPU and model; final performance +measurements cover built-in WebGPU only. Earlier built-in/plugin compatibility +checks do not constitute a plugin rerun of the final optimizations. + +## Correctness and tests + +The model-level check uses six prompts across both adapters, 72 full-vocabulary +logit rows evaluated at identical forced prefixes, adapter on/off, capture +on/off, and repeated generation after rewind. + +| Comparison | Adapter | Identical responses | Matching top-1 rows | Logit RMSE | Maximum absolute difference | +|---|---|---:|---:|---:|---:| +| Final vs previous QLoRA | Off | 6/6 | 72/72 | 0 | 0 | +| Final vs previous QLoRA | On | 6/6 | 71/72 | 0.008171 | 0.0625 | +| Final vs original source model | On | 4/6 | 70/72 | 0.075461 | 1.03125 | +| Capture on vs off | Off | 6/6 | 72/72 | 0 | 0 | +| Capture on vs off | On | 6/6 | 72/72 | 0 | 0 | + +The previous QLoRA implementation is `210584c743785f949cfc417a185366fa0343b932`, +before the final decode register-reuse/subgroup-reduction change. Adapter-on +arithmetic is not bitwise equivalent to the original model. Six prompts are +limited quality coverage, and larger quality evaluation remains necessary. + +Final validation passed seven operator tests, eight graph-fusion tests, and +two WGSL smoke tests. The decode regression checks against a dense +double-precision reference at K=2048/4096/8192, with mixed signs, varying INT8 +scales, a zero rank, tail columns, and both accumulation options. Quantized +prefill and reordered decode comparisons use absolute tolerance 0.002 and +relative tolerance 0.001; unaffected fusion paths retain zero tolerance. +The input-upload change also has a regression test for chunking and reuse. + +See [FP16 LoRA usage](WebGPU_MatMulNBitsLoRA_Prototype.md) and +[QLoRA format, dispatch, and test commands](WebGPU_MatMulNBitsQLoRA_Prototype.md). diff --git a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md index e802852b6079d..576b36697a90a 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md @@ -8,8 +8,8 @@ model file. The current matcher supports FP16 rank-32 adapters, INT4 base weights with block size 32, an implicit zero point, no base bias, and no separate adapter scaling -node. Both MatMuls must use the same activation as the base projection, and -their intermediate outputs must have no other consumers or graph outputs. +node. The down projection must use the same activation as the base projection, +and intermediate outputs must have no other consumers or graph outputs. ## Enable in native GenAI diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 903573f3bbc00..c135542a7e07a 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -12,6 +12,12 @@ Both built-in and plugin WebGPU are supported; plugin use requires the host and provider library from this branch. Disable the option for a comparison with the original graph. +The transformer matches packed INT8 MatMulNBits adapter projections, not +arbitrary DequantizeLinear/Cast/MatMul graphs. Those graphs require an offline +conversion to the weight format below. A model exported with explicit +MatMulNBitsQLoRA nodes does not need the transformer enabled, and disabling +the transformer does not undo those explicit nodes. + ## Weight format and arithmetic The operator takes `X, W, scales, lora_A, lora_A_scales, lora_B, lora_B_scales`. @@ -39,10 +45,14 @@ Y = FP16(MatMulNBits(X,W,scales) + delta) ``` Base quantization stays INT4 with block size 32 and implicit zero point 8. -The adapter projections retain accuracy level 0 and their original dot-product -grouping, accumulation policy, and FP16 intermediate rounding. The provider's -`enableMatmulFp32Accumulation` option applies to adapter accumulation as it does -to the original INT8 MatMulNBits kernels. +The equations describe the weight interpretation and mathematical projections; +the optimized kernels can change reduction order and activation precision. +Eligible prefill down projections use DP4A activation quantization when the +model's `accuracy_level` permits it. Decode and up projections do not add +activation quantization. The provider's `enableMatmulFp32Accumulation` option +controls floating-point lane accumulation. Intermediate T remains FP16; +the specialized decode kernel performs its final subgroup reduction in FP32. +These paths are not claimed to be bitwise equivalent to the unfused graph. The input layout is already transposed and block packed. The fusion does not insert a transpose or materialize a complete dequantized adapter. @@ -54,18 +64,25 @@ dimension of X; a one-token prompt follows the decode selection. | Phase | Selected mode | NVIDIA implementation | |---|---|---| -| Prefill, M > 1 | `fuse_up` | Original INT8 down projection, then base with INT8 up projection and addition | +| Prefill, M > 1 | `fuse_up` | INT8 down projection, then base with INT8 up projection and addition; eligible DP4A paths share quantized X | | Decode, M = 1 | `fuse_down` | Base and INT8 down projection in one dispatch, then INT8 up projection and addition | The joint decode kernel supports one total input row and K from 2048 through 8192. It retains 128 threads per workgroup and adds four workgroups for the -32 down-projection outputs. Its reduction matches the original generic INT8 -kernel. `QLoRAOutputVec4` uses 64 threads and one output per thread, loading -packed weights as `vec4` and T as `vec4` along rank. - -DP4A prefill retains the base activation quantization dispatch, making three -dispatches per adapted projection. Explicit `fuse_down` and `fuse_up` remain -available. Shapes outside the joint decode specialization calculate T +32 down-projection outputs. On devices with fixed 32-lane subgroups, each lane +loads X into registers and reuses it for two rank outputs, then reduces lane +partials with FP32 subgroup sums. The previous workgroup-memory reduction +remains the fallback. `QLoRAOutputVec4` uses 64 threads and one output per +thread, loading packed weights as `vec4` and T as `vec4` along rank. + +DP4A prefill quantizes X once and shares the quantized tensor and scales +between the down projection and base projection, making three dispatches per +adapted projection. For eligible NVIDIA inputs with K >= 2048 and M <= 4096, +the rank-32 down projection uses the small-M DP4A reduction with a 32-column +tile. This narrow output remains underutilized by the general tiled kernel +even at larger M. The general DP4A kernel also uses 32-column tiles when N < 64. +Explicit `fuse_down` and `fuse_up` remain available. Shapes outside the joint +decode specialization calculate T separately. Devices that select subgroup matrix arithmetic for the original INT8 up projection retain that kernel and a separate Add to preserve arithmetic. @@ -99,10 +116,18 @@ onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.*:MatMulNBitsLoRA. onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBits*LoRA* ``` -The quantized tests check original-versus-fused results with zero tolerance, -FP32 accumulation, batched and partial-tile shapes, empty adapters, partial -overrides, nonzero/shared defaults, unsupported shapes and attributes, and +The quantized tests check original-versus-fused results, FP32 accumulation, +batched and partial-tile shapes, empty adapters, partial overrides, +nonzero/shared defaults, unsupported shapes and attributes, and patterns that must remain unfused. Full-model evaluation must separately compare fusion equivalence against the original INT8 model and quantization differences against FP16. Performance runs should use no profiling and record the actual runtime libraries, provider, capture setting, and adapter state. + +Quantized-prefill and reordered-decode fusion comparisons use 0.002 absolute +and 0.001 relative tolerances; unaffected paths retain zero tolerance. A dense +reference test separately checks the specialized decode reduction. +See [the performance evaluation](WebGPU_MatMulNBitsLoRA_Performance.md) for model-level +accuracy limits and the main-versus-optimized performance comparison. WebGPU +also batches CPU-to-GPU input uploads to amortize per-tensor adapter transfer +cost; this is independent of the shader fusion. diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index fe9cafbc8bdff..29f353f4547a3 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3997,14 +3997,16 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .SetDoc( "Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). " "The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights " - "stored as unsigned bytes with implicit zero point 128 and FP16 block scales. Adapter projections use " - "accuracy level 0. Rank is 32; an empty rank dimension in either packed adapter disables the adapter. " + "stored as unsigned bytes with implicit zero point 128 and FP16 block scales. Eligible adapter down " + "projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension " + "in either packed adapter disables the adapter. " "Intermediate projections are rounded to FP16 before the final addition. Adapters remain replaceable.") .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) .Attr("block_size", "Base weight block size; must be 32.", AttributeProto::INT, int64_t{32}) - .Attr("accuracy_level", "Base MatMulNBits accumulation policy.", AttributeProto::INT, int64_t{0}) + .Attr("accuracy_level", "Accuracy policy for the base and eligible adapter down projections.", AttributeProto::INT, + int64_t{0}) .Attr("lora_bits", "Adapter weight bit width; must be 8.", AttributeProto::INT, int64_t{8}) .Attr("lora_block_size", "Adapter weight block size; must be 32.", AttributeProto::INT, int64_t{32}) .Input(0, "X", "Input of rank at least two, with last dimension K.", "T") From 788de1b1f323ad4bf2a6f1b018d4167f3d5f76f5 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 15:05:43 +0800 Subject: [PATCH 10/35] Validate compact QLoRA capture and scope FP16 test tolerances --- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 85 ++++++++++ docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 79 ++++++++- .../contrib_ops/matmul_nbits_qlora_test.cc | 5 +- .../matmul_nbits_lora_fusion_test.cc | 17 +- .../matmul_nbits_qlora_fusion_test.cc | 152 ++++++++++++++++++ 5 files changed, 330 insertions(+), 8 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md index bb29a1e3bf1ee..06424e51a7316 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Performance.md @@ -16,6 +16,10 @@ which rejects dispatch size `(0, 1, 1)`. Main's **Off*** values below therefore use a separately pruned base-only graph. Optimized Off uses the actual optional-adapter graph with empty adapters. These are different off baselines. +This failure describes the recorded main revision below. Newer main includes +the empty-rank MatMul/DequantizeLinear fix (#33171), now merged into this branch. +The historical table has not been relabeled as measurements of that newer main. + | Prompt tokens | Adapter | Main TTFT ms | Optimized TTFT ms | TTFT reduction | Main decode tok/s | Optimized decode tok/s | Decode gain | |---:|---|---:|---:|---:|---:|---:|---:| | 128 | Off* | 29.86 | 31.18 | -4.4% | 261.49 | 236.06 | -9.7% | @@ -183,3 +187,84 @@ QLoRA matrix. They do not establish zero change or results on other devices. runtime's outputs; it does not change the original-model numerical limits above. - Formatting and whitespace checks pass. Generated contrib documentation now includes both new operator schemas. Remote CI must validate other build targets. + +## Review-feedback verification (2026-10-15) + +Merged main at `9cdf0f11fde12c171eceb4e4ca3b31be8777e7f0`, including the +empty-rank MatMul/DequantizeLinear fix (#33171) and graph-capture buffer-copy +fix (#33362). The LoRA review-feedback patch changes tests and documentation; +it does not change LoRA shaders or workgroup sizes. + +The following comparison uses the previously validated PR runtime snapshot +(`7e7079934a`) as Before. Updated uses the merged ORT build and GenAI +`b4023152` (adapter-aware capture identity, based on GenAI main `bdfadcf4`). +These are update-regression measurements, not a new main-versus-fusion study. + +Native C++ benchmark, built-in WebGPU / Windows D3D12 / NVIDIA RTX 5080: +batch 1, fixed prompt IDs, 128 generated tokens, maximum length 4096, greedy +decoding, capture, generator reuse, and token readback. Profiling and competing +build/test work were stopped throughout measurement. Each cell has two rounds +of three warmups and five measured requests, with reversed order in round two. +The table reports medians over ten individual requests per runtime/cell. +All 160 measured requests completed and generated identical sequences across +the two runtimes for each case. + +| Model layout | Prompt tokens | Adapter | Before TTFT ms | Updated TTFT ms | TTFT change | Before decode tok/s | Updated decode tok/s | Decode change | +|---|---:|---|---:|---:|---:|---:|---:|---:| +| Packed INT8 | 128 | Off | 31.48 | 31.43 | -0.17% | 228.69 | 233.12 | +1.94% | +| Packed INT8 | 128 | On | 40.17 | 40.31 | +0.34% | 205.68 | 205.60 | -0.04% | +| Packed INT8 | 1024 | Off | 162.62 | 162.52 | -0.06% | 233.16 | 233.74 | +0.25% | +| Packed INT8 | 1024 | On | 172.67 | 174.30 | +0.95% | 204.11 | 197.54 | -3.22% | +| FP16 original | 128 | Off | 31.95 | 31.90 | -0.17% | 229.53 | 229.59 | +0.03% | +| FP16 original | 128 | On | 97.34 | 96.19 | -1.18% | 140.95 | 141.15 | +0.14% | +| FP16 original | 1024 | Off | 163.49 | 163.97 | +0.29% | 226.22 | 219.26 | -3.07% | +| FP16 original | 1024 | On | 232.54 | 233.87 | +0.57% | 137.72 | 138.31 | +0.43% | + +Positive TTFT change means slower; positive decode change means faster. +The two approximately 3% decode slowdowns prompted a longer, independent +confirmation separating the ORT merge from the GenAI update. This used two +rounds with reversed order, six warmups and eight measured requests each: +16 measurements per runtime/case, 96 measured requests total. All sequences +matched across the three runtimes. + +| 1024-token case | Runtime | TTFT ms | Decode tok/s | Decode range tok/s | +|---|---|---:|---:|---:| +| FP16, off | Previous | 164.59 | 219.34 | 217.18-221.37 | +| FP16, off | Merged ORT / previous GenAI | 163.76 | 218.64 | 214.75-221.01 | +| FP16, off | Merged ORT / fixed GenAI | 163.91 | 219.66 | 216.30-224.18 | +| Packed INT8, on | Previous | 175.74 | 198.31 | 197.07-199.59 | +| Packed INT8, on | Merged ORT / previous GenAI | 176.12 | 198.22 | 195.54-199.43 | +| Packed INT8, on | Merged ORT / fixed GenAI | 176.00 | 197.80 | 195.55-198.86 | + +In this confirmation, final versus previous decode changes are +0.14% for +FP16 adapter-off and -0.25% for packed INT8 adapter-on. TTFT changes are -0.41% +and +0.15%, respectively. The earlier decode slowdowns did not repeat. These +results show no repeatable regression on the tested GPU; they do not establish +zero change or performance on other devices. The initial measurements remain +above rather than being discarded. This does not resolve the pre-existing +approximately 10% adapter-off gap versus a graph with QKV/MLP base fusion. + +### Correctness and capture + +- Built-in and plugin Release builds pass. Built-in provider/empty-tensor tests: + 23 passed; fusion tests: 20 passed; graph-capture buffer-copy tests: 6 passed. +- Plugin provider tests: 5 passed, with 2 expected hardware-inspection skips; + fusion tests: 20 passed. Both skipped specializations pass in built-in mode. +- Empty-default serialization/reload tests cover decode and prefill, full and + partial overrides, and separate on/off capture IDs. Both actual captures are + asserted; alternating replay overwrites sentinel outputs correctly with CPU + fallback disabled. Empty packed defaults have no serialized payload. +- Reported RTX 4070 tolerance failures are addressed with narrowly scoped + tolerances: K=8192 FP16 lane accumulation allows 0.003 absolute error; FP16 + fusion comparisons allow 0.0001 absolute plus two FP16 epsilons relative. + Shorter K, FP32 accumulation, and deliberately exact switching/cancellation + checks keep their previous bounds. RTX 4070 was not available locally. +- GenAI previously reused a base capture after late adapter activation. In a + full-model forced-prefix reproduction, capture-off results matched a fresh + active generator exactly, while captured decode differed by up to 2.82764. + With [GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740), + all six full-vocabulary logit rows match exactly for built-in capture off/on + and plugin capture on. The added Python regression passes both WebGPU + parameterizations; its CUDA variants were not run locally. +- Formatting and whitespace checks pass. Wider device/build validation remains + the responsibility of fresh CI and hardware testing. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index c135542a7e07a..332291591d966 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -103,9 +103,72 @@ Input names and the adapter file stay unchanged. Like the FP16 prototype, eliding a zero branch assumes finite activations, adapter scales, and adapter intermediates. NaN/Inf propagation through zero multiplication can differ; -the transformer remains opt-in. Use separate GenAI generators when switching -adapters or their on/off state with graph capture, and rewind only while -retaining the same generator's adapter selection. +the transformer remains opt-in. GenAI builds without adapter-aware capture +identity require separate generators when changing the adapter selection. +[GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740) assigns +separate capture IDs after additive adapter activation. Activation does not +recompute an existing KV cache: rewind to zero and append the prompt again +if the adapter should apply to the earlier tokens. This does not add an +adapter-deactivation API; use a new base generator to return to adapter-off. + +Graph replay does not re-evaluate the empty-rank branch. Native callers must +use distinct `gpu_graph_id` run options for adapter-off and adapter-on captures, +with stable GPU bindings and buffer lifetimes for each ID. Captures also depend +on tensor shapes and execution mode; prefill and decode cannot share an ID +when their dispatches differ. The provider does not automatically derive an +ID from adapter state. Release and recapture an ID before changing the execution +path or bindings associated with it. Replacing active adapter values in place +is a separate case: the captured path can remain valid only when its shapes, +layout, and bound buffers stay compatible. + +### Compact model packaging + +An explicitly exported `MatMulNBitsQLoRA` node accepts empty packed defaults +directly in the serialized model. Full-size buffers filled with 128 are only +needed by the current three-MatMulNBits fusion pattern, not by the fused operator. +For an adapter-off default, export overridable initializers as follows: + +| Initializer | Default shape | Default payload | +|---|---|---| +| lora_A | [0, K/32, 32] | Empty | +| lora_B | [N, 0, 32] | Empty | +| lora_A_scales | [32, K/32] | Retain the scale defaults | +| lora_B_scales | [N, 1] | Retain the scale defaults | + +Keep these tensors in the graph inputs. Declare the rank dimension of each +packed graph input as dynamic, so a rank-32 adapter can override its empty +initializer. Either empty packed weight disables the adapter. Scale-only and +single-weight overrides therefore continue to produce the base result; supplying +both packed weights uses the supplied scales or their retained defaults. +Saving a graph after offline LoRA fusion also produces this representation. + +This compact export requires a runtime with the custom operator. Disabling +fusion does not convert it to a portable fallback graph. Do not place empty +packed weights into the existing fixed-rank MatMulNBits chain: its N/K +attributes still require nonempty rank-32 projections. For an unfused fallback, +retain the original empty-default DequantizeLinear/MatMul graph. WebGPU's +empty-tensor handling supports that graph, but this transformer does not yet +match it. Empty DequantizeLinear outputs avoid dequantization work; a zero-K +MatMul with a nonempty output still initializes that output, and the residual +Add remains. + +### FP32 packages and composed base fusions + +The current schema and kernel require FP16 activations, base scales, adapter +scales, and outputs. FP32 accumulation is an arithmetic option, not FP32 input +support. Sharing one adapter between FP16 and FP32 model packages will require +FP32 kernel support and a scale-type contract that permits the same adapter +scale tensors in both packages. Casting only initializer defaults at fusion +time would not handle later adapter overrides. + +Adapter-off execution skips the adapter projections but calls a standalone +base MatMulNBits. The current QKV/MLP transformers do not compose with the +LoRA operators, so this can lose the throughput benefit of those base fusions. +Empty defaults alone do not resolve that difference. A follow-up composed +QKV/MLP+LoRA implementation should dispatch the existing fused base kernels +when every relevant adapter is empty, and preserve per-projection adapter +semantics when any adapter is active. Merely reordering the transformers does +not provide both behaviors. ## Validation @@ -127,6 +190,16 @@ the actual runtime libraries, provider, capture setting, and adapter state. Quantized-prefill and reordered-decode fusion comparisons use 0.002 absolute and 0.001 relative tolerances; unaffected paths retain zero tolerance. A dense reference test separately checks the specialized decode reduction. +That reference test allows 0.003 absolute error for K=8192 with FP16 lane +accumulation; shorter K and FP32 accumulation retain 0.002. Its relative +tolerance stays 0.001. FP16 LoRA fusion comparisons use a 0.0001 absolute +floor plus two FP16 epsilons of relative tolerance; deliberately exact +switching and cancellation checks remain exact. A serialization regression +checks that empty packed defaults have no payload and that a reloaded explicit +operator accepts full and partial overrides without a load-time LoRA transform. +It also captures adapter-off and adapter-on under separate IDs and alternates +replay with sentinel-filled outputs, checking that both captures execute and +produce the correct results with CPU fallback disabled. See [the performance evaluation](WebGPU_MatMulNBitsLoRA_Performance.md) for model-level accuracy limits and the main-versus-optimized performance comparison. WebGPU also batches CPU-to-GPU input uploads to amortize per-tensor adapter transfer diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 0a632c981309c..a447cbd7936de 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -239,7 +239,10 @@ TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { test.AddInput("BQ", {n, 1, 32}, bq); test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(1.0f))); test.AddOutput("Y", {1, 1, n}, expected); - test.SetOutputAbsErr("Y", 0.002f); + // The long-K half-precision lane accumulation can drift before the + // FP32 subgroup reduction (up to 0.00238 on RTX 4070 at K=8192). + // Keep the tighter bound for shorter reductions and FP32 accumulation. + test.SetOutputAbsErr("Y", !acc_f32 && k == 8192 ? 0.003f : 0.002f); test.SetOutputRelErr("Y", 0.001f); ConfigOptions options; if (acc_f32) { diff --git a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc index be9a0a1fdf804..9ed1690ee5cef 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc @@ -11,6 +11,11 @@ namespace onnxruntime::test { #ifndef DISABLE_CONTRIB_OPS namespace { +// Fused FP16 reductions can differ by a few ULPs across GPU architectures. +// Use two FP16 epsilons plus a small floor for values near cancellation. +constexpr double kFp16FusionAbsTolerance = 1e-4; +constexpr double kFp16FusionRelTolerance = 2.0 / 1024.0; + void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, bool expose_intermediate = false, bool different_x = false, int64_t rank = 32, bool flattened_weights = false, bool transposed = false, int64_t n = 128, @@ -140,7 +145,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedNumerics) { RunWebGpuFusionTransformerTest( [m, n](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, n); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + kFp16FusionAbsTolerance, kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) { @@ -165,7 +171,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedProjectionNumeri RunWebGpuFusionTransformerTest( [isolate, k](ModelTestBuilder& b) { BuildLoRA(b, 1, true, false, false, 32, false, true, 2048, isolate, k); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + isolate == 4 ? 0.0 : kFp16FusionAbsTolerance, isolate == 4 ? 0.0 : kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) { @@ -188,7 +195,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedBatchedNumerics) RunWebGpuFusionTransformerTest( [batch, m](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, 128, 0, 0, batch); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + kFp16FusionAbsTolerance, kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) { @@ -235,7 +243,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionMatchesWebGpuDecodeAndPref RunWebGpuFusionTransformerTest( [m](ModelTestBuilder& b) { BuildLoRA(b, m, true); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 0.0, 0.0, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + kFp16FusionAbsTolerance, kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); } } diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 631a24f0aa436..c84ff23f99545 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -3,6 +3,8 @@ #include "core/optimizer/matmul_nbits_qlora_fusion.h" #include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/IOBinding.h" +#include "core/session/onnxruntime_run_options_config_keys.h" #include "test/optimizer/graph_transform_test_fixture.h" #include "test/optimizer/webgpu_fusion_test_util.h" #include "test/util/include/default_providers.h" @@ -251,5 +253,155 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { } } +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (int64_t m : {1, 37}) { + SCOPED_TRACE(m); + Model model("quantized LoRA empty defaults", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + BuildQLoRA(builder, m, 64); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + const auto adapters = graph.GetOverridableInitializers(); + ASSERT_EQ(adapters.size(), 4u); + std::string a_name, b_name; + // Dyadic inputs keep the switching checks exact on all reduction paths. + for (auto& [name, value] : builder.feeds_) { + auto* tensor = value.GetMutable(); + if (tensor->IsDataType()) { + auto bytes = tensor->MutableDataAsSpan(); + std::fill(bytes.begin(), bytes.end(), uint8_t{129}); + (tensor->Shape()[0] == 32 ? a_name : b_name) = name; + } else { + auto values = tensor->MutableDataAsSpan(); + const float fill = tensor->Shape().NumDimensions() == 3 ? 0.25f : 0.03125f; + std::fill(values.begin(), values.end(), MLFloat16(fill)); + } + } + ASSERT_FALSE(a_name.empty()); + ASSERT_FALSE(b_name.empty()); + std::vector expected; + for (bool fused : {false, true}) { + if (fused) { + bool modified = false; + ASSERT_STATUS_OK(MatMulNBitsQLoRAFusion().Apply(graph, modified, *logger_)); + ASSERT_TRUE(modified); + ASSERT_STATUS_OK(graph.Resolve()); + } + const auto proto = model.ToProto(); + if (fused) { + ASSERT_EQ(proto.graph().node_size(), 1); + ASSERT_EQ(proto.graph().node(0).op_type(), "MatMulNBitsQLoRA"); + int empty_count = 0; + for (const auto& init : proto.graph().initializer()) { + if (init.name() != a_name && init.name() != b_name) continue; + ++empty_count; + ASSERT_EQ(init.dims_size(), 3); + EXPECT_EQ(init.dims(init.name() == a_name ? 0 : 1), 0); + EXPECT_TRUE(init.raw_data().empty()); + EXPECT_EQ(init.int32_data_size(), 0); + EXPECT_EQ(init.external_data_size(), 0); + } + ASSERT_EQ(empty_count, 2); + } + std::string data; + ASSERT_TRUE(proto.SerializeToString(&data)); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + // No LoRA transformer is registered at load time: the serialized model + // itself must retain empty defaults and accept nonempty rank-32 overrides. + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); + // Off, A only, B only, scales only, packed weights only, all, off again. + for (int state = 0; state < 7; ++state) { + SCOPED_TRACE(state); + auto feeds = builder.feeds_; + for (const auto* input : adapters) { + const auto& name = input->Name(); + const bool packed = name == a_name || name == b_name; + const bool supply = (state == 1 && name == a_name) || (state == 2 && name == b_name) || + (state == 3 && !packed) || (state == 4 && packed) || state == 5; + if (!supply) feeds.erase(name); + } + std::vector output; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, builder.output_names_, &output)); + ASSERT_EQ(output.size(), 1u); + if (!fused) { + expected.push_back(std::move(output[0])); + } else { + const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + } + } + EXPECT_NE(CompareOrtValue(expected[0], expected[5], 0.0, 0.0, false).first, COMPARE_RESULT::SUCCESS); + + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto ep = WebGpuExecutionProviderWithOptions(config); + ASSERT_NE(ep, nullptr); + auto* capture_provider = ep.get(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + const auto data = model.ToProto().SerializeAsString(); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + const auto allocators = capture_provider->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto allocator = allocators.back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr bindings[2]; + OrtValue outputs[2]; + for (int state : {0, 1}) { + ASSERT_STATUS_OK(session.NewIOBinding(&bindings[state])); + for (const auto& [name, value] : builder.feeds_) { + const bool adapter = std::any_of(adapters.begin(), adapters.end(), + [&name](const NodeArg* arg) { return arg->Name() == name; }); + if (state == 0 && adapter) continue; + const auto& source = value.Get(); + OrtValue input; + Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); + ASSERT_STATUS_OK(bindings[state]->BindInput(name, input)); + } + const auto& reference = expected[state == 0 ? 0 : 5].Get(); + Tensor::InitOrtValue(reference.DataType(), reference.Shape(), allocator, outputs[state]); + ASSERT_STATUS_OK(bindings[state]->BindOutput(builder.output_names_[0], outputs[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeInputs()); + } + OrtValue actual; + const auto& reference = expected[0].Get(); + Tensor::InitOrtValue(reference.DataType(), reference.Shape(), CPUAllocator::DefaultInstance(), actual); + // Each adapter state owns a capture ID and stable GPU bindings. Replay + // skips branch selection, so one ID cannot be reused across on/off states. + for (int state : {0, 0, 0, 1, 1, 1, 0, 1, 0, 1}) { + SCOPED_TRACE(MakeString("capture state=", state)); + auto* cpu_output = actual.GetMutable(); + auto values = cpu_output->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), MLFloat16(-99.0f)); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(*cpu_output, *outputs[state].GetMutable())); + RunOptions run_options; + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, + state == 0 ? "1" : "2")); + ASSERT_STATUS_OK(session.Run(run_options, *bindings[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeOutputs()); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(outputs[state].Get(), *cpu_output)); + const auto comparison = CompareOrtValue(actual, expected[state == 0 ? 0 : 5], 0.0, 0.0, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + EXPECT_TRUE(capture_provider->IsGraphCaptured(1)); + EXPECT_TRUE(capture_provider->IsGraphCaptured(2)); + } +} + #endif } // namespace onnxruntime::test From 8df7833750ce2739547f786c2445d1bc9c5ecc89 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 15:51:59 +0800 Subject: [PATCH 11/35] Fix GCC range-loop warning in QLoRA decode test --- onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index a447cbd7936de..ded2b7108ff7b 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -195,7 +195,7 @@ TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { } #endif for (bool acc_f32 : {false, true}) { - for (const auto [k, n] : {std::pair{2048, 36}, {4096, 132}, {8192, 64}}) { + for (const auto& [k, n] : {std::pair{2048, 36}, {4096, 132}, {8192, 64}}) { SCOPED_TRACE(MakeString("K=", k, ", N=", n, ", acc_f32=", acc_f32)); std::vector x(k); std::vector aq(32 * k); From c9d3692113304a99832f4f854616da4a8e314a74 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 16:56:22 +0800 Subject: [PATCH 12/35] [WebGPU] Support FP16 and FP32 models in LoRA and QLoRA --- docs/ContribOperators.md | 14 +-- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 116 +++++++++++++++++ docs/WebGPU_MatMulNBitsLoRA_Prototype.md | 23 ++-- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 69 ++++++---- .../webgpu/quantization/matmul_nbits_lora.cc | 45 +++---- .../webgpu/quantization/matmul_nbits_qlora.cc | 47 +++---- .../core/graph/contrib_ops/contrib_defs.cc | 12 +- .../optimizer/matmul_nbits_lora_fusion.cc | 13 +- .../optimizer/matmul_nbits_qlora_fusion.cc | 15 ++- .../contrib_ops/matmul_nbits_lora_test.cc | 105 ++++++++++++++-- .../contrib_ops/matmul_nbits_qlora_test.cc | 119 +++++++++++++++--- .../matmul_nbits_lora_fusion_test.cc | 94 +++++++++----- .../matmul_nbits_qlora_fusion_test.cc | 71 +++++++---- 13 files changed, 574 insertions(+), 169 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index b2991ffe40aba..9eee6d78725fc 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4108,7 +4108,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulNBitsLoRA** - Computes MatMulNBits(X, W, scales) + MatMul(MatMul(X, lora_A), lora_B). With adapter_transposed=1, transposes both stored adapter inputs in this formula. The adapter remains replaceable. A zero rank in either adapter input disables the adapter. Intermediate results are rounded to the input type before addition. This prototype supports FP16, rank 32, INT4 block size 32, implicit zero point 8, and no bias. + Computes MatMulNBits(X, W, scales) + MatMul(MatMul(X, lora_A), lora_B). With adapter_transposed=1, transposes both stored adapter inputs in this formula. The adapter remains replaceable. A zero rank in either adapter input disables the adapter. Intermediate results are rounded to the input type before addition. Supports FP16 or FP32, rank 32, INT4 block size 32, implicit zero point 8, and no bias. #### Version @@ -4139,7 +4139,7 @@ This version of the operator has been available since version 1 of the 'com.micr
W : TQ
INT4 packed weights of shape [N, K/32, 16].
scales : T
-
FP16 block scales containing N * K/32 elements; may be flattened.
+
Block scales containing N * K/32 elements; may be flattened.
lora_A : T
Adapter of shape [K, R_A], or [R_A, K] if adapter_transposed=1; R_A is 0 or 32.
lora_B : T
@@ -4156,8 +4156,8 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
T : tensor(float16)
-
FP16 activations and adapter weights.
+
T : tensor(float16), tensor(float)
+
FP16 or FP32 activations, block scales, and adapter weights.
TQ : tensor(uint8)
Packed quantized base weights.
@@ -4264,7 +4264,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulNBitsQLoRA** - Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights stored as unsigned bytes with implicit zero point 128 and FP16 block scales. Eligible adapter down projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension in either packed adapter disables the adapter. Intermediate projections are rounded to FP16 before the final addition. Adapters remain replaceable. + Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights stored as unsigned bytes with implicit zero point 128 and block scales matching the activation type. Eligible adapter down projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension in either packed adapter disables the adapter. Intermediate projections use the activation type (FP16 or FP32) before the final addition. Adapters remain replaceable. #### Version @@ -4318,8 +4318,8 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
T : tensor(float16)
-
FP16 activations and block scales.
+
T : tensor(float16), tensor(float)
+
FP16 or FP32 activations and block scales.
TQ : tensor(uint8)
Packed quantized weights.
diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md index 06424e51a7316..40fa01afa7928 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Performance.md @@ -268,3 +268,119 @@ approximately 10% adapter-off gap versus a graph with QKV/MLP base fusion. parameterizations; its CUDA variants were not run locally. - Formatting and whitespace checks pass. Wider device/build validation remains the responsibility of fresh CI and hardware testing. + +## FP16 and FP32 support verification (2026-10-15) + +Both `MatMulNBitsLoRA` and `MatMulNBitsQLoRA` now accept FP16 and FP32 models. +All floating inputs to an operator share one type, including base scales, +adapter scales or floating adapter weights, intermediate T, and output. +LoRA stores floating adapter weights; QLoRA stores packed INT8 adapter weights. +Cross-type adapter loading is not implicit. Rank 32 and INT4 block-32 base +restrictions remain. FP32 decode preserves the ordinary base DP4A policy even +when the adapter is active. + +### Validation + +- Built-in Release: 28 provider/empty-tensor and 24 fusion tests pass. +- Plugin Release: 10 provider and 24 fusion tests pass; 2 hardware-inspection + tests skip in plugin mode and pass with the built-in provider. +- FP32 dense FP64-reference tests cover prefill/decode, batching, partial N, + K through 8192, accuracy levels 0/4, and all fusion modes, with 1e-5 + absolute/relative bounds on inputs exactly representable by DP4A. +- Both original and transposed LoRA layouts, empty/full/partial adapters, + QLoRA serialization/reload and distinct on/off capture replay are tested. +- The two changed kernel files pass wasm32 syntax checks. Formatting and + whitespace checks pass. This is not a full linked WebAssembly build. + +### FP32 full-model numerical checks + +Two adapters, three prompts each, up to 64 greedy tokens, and 12 forced-prefix +full-vocabulary logit rows per prompt were compared against the original +unfused package. The converted base data is unchanged; both conversions +preserve the original effective FP32 adapter weights exactly. + +| FP32 path | Adapter | Max logit difference | Max per-case RMSE | Top-1 rows match | Greedy outputs match | +| --- | --- | --- | --- | --- | --- | +| LORA | Off | 0 | 0 | 72/72 | 6/6 | +| LORA | On | 0.68059158 | 0.079206155 | 70/72 | 2/6 | +| QLORA | Off | 0 | 0 | 72/72 | 6/6 | +| QLORA | On | 1.3800392 | 0.12060427 | 71/72 | 3/6 | + +These are the default `accuracy_level=4` results. Adapter-off logits match +exactly. Adapter-on differences include reduction-order changes and DP4A +activation quantization; they are not a bitwise-equivalence claim. QLoRA +`fuse_down`, which avoids quantizing the adapter down input, still differs by +up to 0.918821, with 71/72 top-1 rows and 5/6 greedy outputs matching, because +the base path still uses DP4A. The difference must not be attributed solely +to the new adapter down projection. + +To isolate floating arithmetic, diagnostic copies set `accuracy_level=0` +for both reference and fused models: + +| FP32 path, accuracy=0 | Max logit difference | Max per-case RMSE | Top-1 rows match | Greedy outputs match | +| --- | --- | --- | --- | --- | +| LORA | 3.0517578e-05 | 3.3701491e-06 | 72/72 | 6/6 | +| QLORA | 2.3409724e-05 | 2.7866782e-06 | 72/72 | 6/6 | + +Built-in capture off/on and plugin capture on produce exactly equal logits +for each fused representation. Rewind/repeat generation is identical. +These checks use the GenAI adapter capture-identity fix from #2740. +This small prompt set is a numerical regression check, not a broad quality +evaluation. + +### FP32 throughput + +Native GenAI C++, built-in WebGPU, Windows D3D12, NVIDIA RTX 5080. Batch 1, +fixed identical prompt IDs, 128 generated tokens, maximum length 4096, +graph capture, generator reuse and token readback. Each cell uses three +warmups and five measured requests; values are medians. No builds, tests, +profiling or other GPU workloads run during timing. The unfused baseline +uses the same updated runtime; this is not a new upstream-main comparison. + +| Prompt tokens | Adapter | FP32 path | TTFT ms | Decode tok/s | +| --- | --- | --- | --- | --- | +| 128 | Off | Unfused | 31.48 | 174.42 | +| 128 | Off | LoRA | 29.56 | 201.51 | +| 128 | Off | QLoRA | 29.67 | 200.70 | +| 128 | A | Unfused | 77.99 | 37.62 | +| 128 | A | LoRA | 123.84 | 142.93 | +| 128 | A | QLoRA | 43.46 | 133.93 | +| 128 | B | Unfused | 77.59 | 37.68 | +| 128 | B | LoRA | 123.56 | 143.19 | +| 128 | B | QLoRA | 43.57 | 133.50 | +| 1024 | Off | Unfused | 174.11 | 171.00 | +| 1024 | Off | LoRA | 162.04 | 198.04 | +| 1024 | Off | QLoRA | 162.34 | 199.26 | +| 1024 | A | Unfused | 227.59 | 37.52 | +| 1024 | A | LoRA | 270.12 | 142.24 | +| 1024 | A | QLoRA | 192.25 | 131.47 | +| 1024 | B | Unfused | 227.28 | 37.69 | +| 1024 | B | LoRA | 270.76 | 141.47 | +| 1024 | B | QLoRA | 192.22 | 132.22 | + +FP32 QLoRA improves both prefill and decode over this unfused package. FP32 +LoRA improves decode but its current floating down-projection shader makes +prefill slower than the original package; this support change does not solve +that performance limitation. The two representations use different shaders +and need not have identical speed or floating-point reduction order. + +### FP16 before/after regression check + +The previous validated PR runtime and this update use the same fixed GenAI +runtime and the same converted FP16 models. Token IDs match between revisions +in every measured cell. + +| FP16 adapter | Prompt tokens | Adapter | TTFT ms before → after | Change | Decode tok/s before → after | Change | +| --- | --- | --- | --- | --- | --- | --- | +| LORA | 128 | Off | 31.24 → 31.22 | -0.06% | 237.29 → 235.36 | -0.82% | +| LORA | 128 | A | 81.33 → 81.42 | +0.12% | 205.54 → 204.45 | -0.53% | +| LORA | 1024 | Off | 162.74 → 162.58 | -0.10% | 234.15 → 233.27 | -0.38% | +| LORA | 1024 | A | 223.72 → 223.39 | -0.15% | 204.52 → 205.28 | +0.37% | +| QLORA | 128 | Off | 31.33 → 31.32 | -0.03% | 235.81 → 236.37 | +0.24% | +| QLORA | 128 | A | 40.86 → 40.35 | -1.26% | 204.94 → 202.06 | -1.41% | +| QLORA | 1024 | Off | 162.93 → 162.60 | -0.21% | 233.84 → 234.40 | +0.24% | +| QLORA | 1024 | A | 172.54 → 172.80 | +0.15% | 204.50 → 204.34 | -0.08% | + +Across these measured cells, the largest TTFT increase is 0.15% and the +largest decode decrease is 1.41%. No large FP16 regression appears in this +single-round check; smaller changes remain within typical run variation. diff --git a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md index 576b36697a90a..de536e884803f 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md @@ -2,15 +2,21 @@ This opt-in prototype replaces `MatMulNBits(X, W) + MatMul(MatMul(X, A), B)` with one `com.microsoft.MatMulNBitsLoRA` node during session optimization. It retains -the packed base weights and the replaceable FP16 adapter inputs. It does not +the packed base weights and the replaceable FP16 or FP32 adapter inputs. It does not merge or requantize the adapter into the base weights, or change the source model file. -The current matcher supports FP16 rank-32 adapters, INT4 base weights with block +The current matcher supports FP16 and FP32 rank-32 adapters, INT4 base weights with block size 32, an implicit zero point, no base bias, and no separate adapter scaling node. The down projection must use the same activation as the base projection, and intermediate outputs must have no other consumers or graph outputs. +All floating tensors in one operator share a type: FP16 models use FP16 adapter +weights and base scales; FP32 models use FP32 adapter weights and base scales. +There are no implicit casts between those types. `MatMulNBitsQLoRA` has the same +floating-type contract but stores adapter weights as INT8 with scales matching +the model type. Both operators support adapter on/off and prefill/decode. + ## Enable in native GenAI Add this entry to the decoder's `session_options`, or supply it in the native @@ -38,8 +44,8 @@ graph for a matched comparison. With an adapter, the kernel computes `T = X * A`, then incorporates `T * B` and the addition into the base kernel's output computation. The small -intermediate `T` remains FP16. The base result and adapter result are rounded -separately before addition. Decode uses two GPU dispatches per projection. +intermediate `T` uses the model type. The base result and adapter result are rounded +separately before addition. FP16 NVIDIA decode uses two GPU dispatches per projection. DP4A prefill retains the existing activation quantization dispatch, so it uses three. Subgroup matrix base kernels use a separate output epilogue fallback. @@ -61,14 +67,17 @@ chooses between the two existing implementations; it adds no shader or GPU dispatch. Explicit `fuse_down` and `fuse_up` settings force that path for all input sizes and remain available for comparisons. -On NVIDIA decode (one input row, K from 2048 through 8192), `fuse_down` uses +On NVIDIA decode (one input row, K from 2048 through 8192), when the base uses +floating arithmetic, `fuse_down` uses `MatMulNBitsLoRADown` with eight additional workgroups that calculate T once, followed by `LoRAOutputVec4`. `fuse_up` uses `LoRADownGemv`, followed by -`MatMulNBitsLoRAUp`. Both use the same FP32 GEMV reduction and FP16 T, with +`MatMulNBitsLoRAUp`. Both use the same FP32 GEMV reduction and model-type T, with 128 threads per producer workgroup. The standalone output shader retains 64 threads and one output per thread. Both output implementations load Bt and T as vec4 along rank, preserving scalar rank accumulation order and separate -FP16 rounding before the base/adapter addition. +model-type rounding before the base/adapter addition. FP32 decode can select +DP4A for the base even at M=1. That case keeps the ordinary base DP4A dispatch +and computes T separately, preserving the base activation-quantization policy. DP4A prefill retains its activation quantization dispatch. `fuse_down` adds the tiled down projection to the base dispatch; `fuse_up` runs that same down diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 332291591d966..79d01aafd293b 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -6,8 +6,8 @@ is an inference operator for quantized adapters; it does not merge adapter weights into the base or perform adapter training. Enable it with `session.enable_matmul_nbits_lora_fusion=1` in the native GenAI -decoder session options, as described in [the FP16 prototype](WebGPU_MatMulNBitsLoRA_Prototype.md). -The same option selects the appropriate transformer for FP16 or INT8 adapters. +decoder session options, as described in [the LoRA prototype](WebGPU_MatMulNBitsLoRA_Prototype.md). +The same option selects the appropriate transformer for floating-point or INT8 adapters. Both built-in and plugin WebGPU are supported; plugin use requires the host and provider library from this branch. Disable the option for a comparison with the original graph. @@ -25,23 +25,24 @@ For rank 32, positive K divisible by 32, and positive N divisible by 4: | Tensor | Type | Shape | |---|---|---| -| X | FP16 | [..., M, K] | +| X | FP16 or FP32 | [..., M, K] | | W | Packed INT4 in UINT8 | [N, K/32, 16] | -| scales | FP16 | [N, K/32] or flattened | +| scales | Same as X | [N, K/32] or flattened | | lora_A | Packed INT8 in UINT8 | [32, K/32, 32] | -| lora_A_scales | FP16 | [32, K/32] | +| lora_A_scales | Same as X | [32, K/32] | | lora_B | Packed INT8 in UINT8 | [N, 1, 32] | -| lora_B_scales | FP16 | [N, 1] | +| lora_B_scales | Same as X | [N, 1] | Packed adapter bytes encode signed integers with an implicit zero point of -128. For down weight index `(r,k)` and up weight index `(n,r)`: +128. For down weight index `(r,k)` and up weight index `(n,r)`, with `F` +denoting the activation type (FP16 or FP32): ```text -At[r,k] = FP16((FP16(Aq[r,k]) - 128) * As[r,floor(k/32)]) -Bt[n,r] = FP16((FP16(Bq[n,r]) - 128) * Bs[n,0]) -T = FP16(X * At^T) -delta = FP16(T * Bt^T) -Y = FP16(MatMulNBits(X,W,scales) + delta) +At[r,k] = F((F(Aq[r,k]) - 128) * As[r,floor(k/32)]) +Bt[n,r] = F((F(Bq[n,r]) - 128) * Bs[n,0]) +T = F(X * At^T) +delta = F(T * Bt^T) +Y = F(MatMulNBits(X,W,scales) + delta) ``` Base quantization stays INT4 with block size 32 and implicit zero point 8. @@ -50,7 +51,8 @@ the optimized kernels can change reduction order and activation precision. Eligible prefill down projections use DP4A activation quantization when the model's `accuracy_level` permits it. Decode and up projections do not add activation quantization. The provider's `enableMatmulFp32Accumulation` option -controls floating-point lane accumulation. Intermediate T remains FP16; +controls floating-point lane accumulation for FP16 inputs. FP32 inputs always +use FP32 floating-point accumulation. Intermediate T uses the activation type; the specialized decode kernel performs its final subgroup reduction in FP32. These paths are not claimed to be bitwise equivalent to the unfused graph. @@ -68,12 +70,15 @@ dimension of X; a one-token prompt follows the decode selection. | Decode, M = 1 | `fuse_down` | Base and INT8 down projection in one dispatch, then INT8 up projection and addition | The joint decode kernel supports one total input row and K from 2048 through -8192. It retains 128 threads per workgroup and adds four workgroups for the +8192, when the base projection does not select DP4A. FP32 base projections +can select DP4A even at M=1; those retain the original base kernel and compute +the adapter down projection separately. The joint kernel retains 128 threads +per workgroup and adds four workgroups for the 32 down-projection outputs. On devices with fixed 32-lane subgroups, each lane loads X into registers and reuses it for two rank outputs, then reduces lane partials with FP32 subgroup sums. The previous workgroup-memory reduction remains the fallback. `QLoRAOutputVec4` uses 64 threads and one output per -thread, loading packed weights as `vec4` and T as `vec4` along rank. +thread, loading packed weights as `vec4` and T as `vec4` along rank. DP4A prefill quantizes X once and shares the quantized tensor and scales between the down projection and base projection, making three dispatches per @@ -101,7 +106,7 @@ Empty A rank or empty B rank selects the ordinary base kernel before adapter allocation or dispatch. Nonzero and shared packed defaults retain their data. Input names and the adapter file stay unchanged. -Like the FP16 prototype, eliding a zero branch assumes finite activations, +Like the LoRA prototype, eliding a zero branch assumes finite activations, adapter scales, and adapter intermediates. NaN/Inf propagation through zero multiplication can differ; the transformer remains opt-in. GenAI builds without adapter-aware capture identity require separate generators when changing the adapter selection. @@ -154,12 +159,18 @@ Add remains. ### FP32 packages and composed base fusions -The current schema and kernel require FP16 activations, base scales, adapter -scales, and outputs. FP32 accumulation is an arithmetic option, not FP32 input -support. Sharing one adapter between FP16 and FP32 model packages will require -FP32 kernel support and a scale-type contract that permits the same adapter -scale tensors in both packages. Casting only initializer defaults at fusion -time would not handle later adapter overrides. +The schema, kernel, and transformer support both FP16 and FP32 activations, +base scales, adapter scales, and outputs. All floating tensors in a node must +have the same type. FP32 packages retain FP32 intermediate projections and +scales; they do not convert the model to FP16. Shader cache keys include the +input/output types, so FP16 and FP32 programs have distinct cache entries. + +The packed INT8 weights use the same layout for both types, but an adapter +file with FP16 scales cannot directly override FP32 scale inputs (or vice +versa). Export scales in the package's activation type. Sharing one adapter +file across different activation types still requires scale conversion at +the graph boundary or a separate mixed-scale type contract. Casting only +initializer defaults would not handle later adapter overrides. Adapter-off execution skips the adapter projections but calls a standalone base MatMulNBits. The current QKV/MLP transformers do not compose with the @@ -179,7 +190,7 @@ onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.*:MatMulNBitsLoRA. onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBits*LoRA* ``` -The quantized tests check original-versus-fused results, FP32 accumulation, +The quantized tests check original-versus-fused FP16 and FP32 results, FP32 accumulation, batched and partial-tile shapes, empty adapters, partial overrides, nonzero/shared defaults, unsupported shapes and attributes, and patterns that must remain unfused. Full-model evaluation must separately @@ -187,9 +198,15 @@ compare fusion equivalence against the original INT8 model and quantization differences against FP16. Performance runs should use no profiling and record the actual runtime libraries, provider, capture setting, and adapter state. -Quantized-prefill and reordered-decode fusion comparisons use 0.002 absolute -and 0.001 relative tolerances; unaffected paths retain zero tolerance. A dense -reference test separately checks the specialized decode reduction. +FP32 tests include dense FP64 references with non-half-representable scales, +both fusion modes, DP4A and floating-point down projections, serialized empty +defaults, partial overrides, and separate adapter-on/off graph captures. + +Quantized-prefill and reordered-decode FP16 fusion comparisons use 0.002 absolute +and 0.001 relative tolerances; unaffected FP16 paths retain zero tolerance. +FP32 quantized-prefill comparisons use the same approximation bounds; other +FP32 comparisons use 1e-5 absolute/relative bounds. A dense reference test +separately checks the specialized decode reduction. That reference test allows 0.003 absolute error for K=8192 with FP16 lane accumulation; shorter K and FP32 accumulation retain 0.002. Its relative tolerance stays 0.001. FP16 LoRA fusion comparisons use a 0.0001 absolute diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc index 80807e942d58e..26829944487d7 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc @@ -9,6 +9,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" namespace onnxruntime::contrib::webgpu { @@ -17,8 +18,8 @@ void GenerateLoRADown(ShaderHelper& shader, const std::string& base_groups) { shader.AddInput("lora_at"); shader.AddOutput("lora_down", ShaderUsage::UseValueTypeAlias); shader.AdditionalImplementation() << R"( -var lora_x_tile: array, 256>; -var lora_a_tile: array, 256>; +var lora_x_tile: array; +var lora_a_tile: array; )"; shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" << R"( @@ -28,9 +29,9 @@ if (workgroup_idx >= lora_base_groups) { let k4_local = local_idx % 8u; let row = group * 32u + row_local; let rank_base = k4_local * 4u; - var acc = vec4(0); + var acc = lora_down_value_t(0); for (var start = 0u; start < uniforms.K / 4u; start += 8u) { - var xv = vec4(0); + var xv = lora_down_value_t(0); if (row < uniforms.M * uniforms.batch_count) { xv = lora_x[row * (uniforms.K / 4u) + start + k4_local]; } @@ -43,10 +44,10 @@ if (workgroup_idx >= lora_base_groups) { let a1 = lora_a_tile[(rank_base + 1u) * 8u + kk]; let a2 = lora_a_tile[(rank_base + 2u) * 8u + kk]; let a3 = lora_a_tile[(rank_base + 3u) * 8u + kk]; - acc = vec4(a0.x, a1.x, a2.x, a3.x) * x.x + acc; - acc = vec4(a0.y, a1.y, a2.y, a3.y) * x.y + acc; - acc = vec4(a0.z, a1.z, a2.z, a3.z) * x.z + acc; - acc = vec4(a0.w, a1.w, a2.w, a3.w) * x.w + acc; + acc = lora_down_value_t(a0.x, a1.x, a2.x, a3.x) * x.x + acc; + acc = lora_down_value_t(a0.y, a1.y, a2.y, a3.y) * x.y + acc; + acc = lora_down_value_t(a0.z, a1.z, a2.z, a3.z) * x.z + acc; + acc = lora_down_value_t(a0.w, a1.w, a2.w, a3.w) * x.w + acc; } workgroupBarrier(); } @@ -176,7 +177,7 @@ Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const T auto down_shape = x->Shape().AsShapeVector(); down_shape.back() = 32; auto t = context.CreateGPUTensor(x->DataType(), TensorShape(down_shape)); - if (fuse_up) { + const auto compute_down = [&]() -> Status { if (use_gemv) { LoRAGemvProgram down; down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, @@ -196,31 +197,31 @@ Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const T .AddUniformVariables({{rows}, {narrow(k)}, {1u}}); ORT_RETURN_IF_ERROR(context.RunProgram(down)); } + return Status::OK(); + }; + if (fuse_up) { + ORT_RETURN_IF_ERROR(compute_down()); return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, context, y, 0, nullptr, 0, &t, bt, nullptr, nullptr, true); } auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); - if (use_gemv) { + // FP32 can select DP4A even at M=1. Keep that base policy when the + // floating joint down kernel would otherwise replace it. + const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits( + context, accuracy_level, 32, narrow(n), narrow(k), 4, m, false, y); + if (use_gemv && !base_dp4a) { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, context, &base, 0, nullptr, 0, nullptr, nullptr, at, &t)); } else if (m >= kMinMForTileOptimization && context.AdapterInfo().vendor == std::string_view{"nvidia"} && - CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, 32, narrow(n), narrow(k), - 4, m, false, y)) { - // Preserve DP4A's activation quantization and compute T from original FP16 X + base_dp4a) { + // Preserve DP4A's activation quantization and compute T from original X // in extra workgroups of the base dispatch. ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits(x, weight, scales, nullptr, nullptr, rows / m, m, m, narrow(n), narrow(k), 32, 0, kMinMForTileOptimization, 4, context, &base, 0, nullptr, nullptr, nullptr, at, &t)); } else { - LoRADownProgram down; - down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, - {at, ProgramTensorMetadataDependency::Type, 4}}) - .AddOutput({&t, ProgramTensorMetadataDependency::Type, 4}) - .SetWorkgroupSize(256) - .SetDispatchGroupSize(CeilDiv(rows, 32u)) - .AddUniformVariables({{rows}, {narrow(k)}, {1u}}); - ORT_RETURN_IF_ERROR(context.RunProgram(down)); + ORT_RETURN_IF_ERROR(compute_down()); ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, context, &base)); } @@ -300,7 +301,7 @@ class MatMulNBitsLoRA final : public WebGpuKernel { ONNX_OPERATOR_KERNEL_EX( MatMulNBitsLoRA, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()) .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), MatMulNBitsLoRA); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 89149c96f82df..0a959afd46695 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -10,13 +10,14 @@ #include "core/common/safeint.h" #include "core/providers/cpu/math/matmul_helper.h" #include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" namespace onnxruntime::contrib::webgpu { void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32, bool use_subgroups) { if (use_subgroups) { - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n"; + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n"; auto& body = shader.MainFunctionBody(); body << "let qlora_base_groups = " << base_groups << ";\n" << R"( @@ -40,8 +41,8 @@ if (workgroup_idx >= qlora_base_groups) { << "let scale" << r << " = lora_a_scales[rank" << r << " * (uniforms.K / 32u) + k / 32u];\n" << "var sum" << r << " = qlora_acc_t(0);\n"; for (uint32_t i = 0; i < 4; ++i) { - body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" - << r << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; + body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" + << r << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; } body << "partial" << r << " += sum" << r << ";\n"; } @@ -51,17 +52,17 @@ if (workgroup_idx >= qlora_base_groups) { let result0 = subgroupAdd(f32(partial0)); let result1 = subgroupAdd(f32(partial1)); if (lane == 0u) { - lora_down[rank_base + rank_local] = f16(result0); - lora_down[rank_base + rank_local + 4u] = f16(result1); + lora_down[rank_base + rank_local] = output_element_t(result0); + lora_down[rank_base + rank_local + 4u] = output_element_t(result1); } return; } )"; return; } - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n" << R"( -var qlora_x_tile: array, 128>; +var qlora_x_tile: array, 128>; var qlora_partials: array, 8>; )"; shader.MainFunctionBody() << "let qlora_base_groups = " << base_groups << ";\n" @@ -73,7 +74,7 @@ if (workgroup_idx >= qlora_base_groups) { // Match the generic INT8 MatMulNBits tile, dot grouping and reduction order. for (var start = 0u; start < uniforms.K; start += 512u) { let k4 = start / 4u + local_idx; - var xv = vec4(0); + var xv = vec4(0); if (k4 < uniforms.K / 4u) { xv = )" << x_value << R"(; } qlora_x_tile[local_idx] = xv; @@ -86,7 +87,7 @@ if (workgroup_idx >= qlora_base_groups) { let scale = lora_a_scales[rank * (uniforms.K / 32u) + k / 32u]; var sum = qlora_acc_t(0); for (var i = 0u; i < 4u; i++) { - let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; + let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; sum += dot(vec4(qlora_x_tile[lane * 4u + i]), vec4(bv)); } qlora_partials[rank_offset + rank_local][lane] += sum; @@ -97,7 +98,7 @@ if (workgroup_idx >= qlora_base_groups) { if (local_idx < 8u) { var value = qlora_acc_t(0); for (var i = 0u; i < 32u; i++) { value += qlora_partials[local_idx][i]; } - lora_down[rank_base + local_idx] = f16(value); + lora_down[rank_base + local_idx] = output_element_t(value); } return; } @@ -105,13 +106,13 @@ if (workgroup_idx >= qlora_base_groups) { } void GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32, bool vector_output) { - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "f16") << ";\n" + shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n" << R"( -fn qlora_weight(col: u32, r4: u32) -> vec4 { +fn qlora_weight(col: u32, r4: u32) -> vec4 { let packed = lora_b[col * 2u + r4 / 4u][r4 % 4u]; - return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; + return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; } -fn qlora_delta(row: u32, col: u32) -> f16 { +fn qlora_delta(row: u32, col: u32) -> output_element_t { var delta = qlora_acc_t(0); // The original K=32 INT8 up projection uses wide tiles at M>=4. if (uniforms.M >= 4u) { @@ -133,12 +134,12 @@ fn qlora_delta(row: u32, col: u32) -> f16 { // The remaining generic-kernel partials are zero, including their sign. delta += qlora_acc_t(0); } - return f16(delta); + return output_element_t(delta); } )"; if (vector_output) { shader.AdditionalImplementation() << R"( -fn qlora_delta4(row: u32, col: u32) -> vec4 { +fn qlora_delta4(row: u32, col: u32) -> vec4 { var delta = vec4(0); for (var r4 = 0u; r4 < 8u; r4 += 2u) { let a0 = vec4(lora_t[row * 8u + r4]); @@ -154,7 +155,7 @@ fn qlora_delta4(row: u32, col: u32) -> vec4 { dot(a1, vec4(qlora_weight(col + 2u, r4 + 1u))), dot(a1, vec4(qlora_weight(col + 3u, r4 + 1u)))); } - return vec4(delta); + return vec4(delta); } )"; } @@ -174,7 +175,7 @@ class QLoRAOutputProgram final : public Program { shader.AddInput("lora_t"); shader.AddInput("lora_b"); shader.AddInput("lora_b_scales"); - shader.AddOutput("output"); + shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); GenerateQLoRAUp(shader, acc_f32_); shader.MainFunctionBody() << "if (global_idx >= uniforms.size) { return; }\n" @@ -273,8 +274,12 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { auto t = context.CreateGPUTensor(x->DataType(), TensorShape(t_shape)); const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); const bool fuse_up = fuse_up_ || (automatic_ && m > 1); + // FP32 may select DP4A even for M=1. The joint floating-point down + // shader must not replace that base projection's activation quantization. + const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), + narrow(k_), 4, m, false, y); const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && - context.AdapterInfo().vendor == std::string_view{"nvidia"}; + context.AdapterInfo().vendor == std::string_view{"nvidia"} && !base_dp4a; if (joint_down) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, @@ -301,7 +306,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { std::optional up_config; const bool subgroup_up = CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), - narrow(n_), 32, 8, true, up_config, m); + narrow(n_), 32, 8, x->IsDataType(), up_config, m); if (fuse_up && !subgroup_up) { return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, &t, bq, nullptr, nullptr, true, bs, nullptr, &quantized_x); @@ -334,7 +339,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { ONNX_OPERATOR_KERNEL_EX(MatMulNBitsQLoRA, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()) .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), MatMulNBitsQLoRA); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 12bc4b97a97c7..1f3fcc967144a 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3949,7 +3949,7 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by "With adapter_transposed=1, transposes both stored adapter inputs in this formula. " "The adapter remains replaceable. A zero rank in either adapter input disables the adapter. " "Intermediate results are rounded to the input type before addition. " - "This prototype supports FP16, rank 32, INT4 block size 32, implicit zero point 8, and no bias.") + "Supports FP16 or FP32, rank 32, INT4 block size 32, implicit zero point 8, and no bias.") .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) @@ -3958,11 +3958,11 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .Attr("adapter_transposed", "Whether adapters are stored as [R_A,K] and [N,R_B].", AttributeProto::INT, int64_t{0}) .Input(0, "X", "Input with rank at least two and last dimension K.", "T") .Input(1, "W", "INT4 packed weights of shape [N, K/32, 16].", "TQ") - .Input(2, "scales", "FP16 block scales containing N * K/32 elements; may be flattened.", "T") + .Input(2, "scales", "Block scales containing N * K/32 elements; may be flattened.", "T") .Input(3, "lora_A", "Adapter of shape [K, R_A], or [R_A, K] if adapter_transposed=1; R_A is 0 or 32.", "T") .Input(4, "lora_B", "Adapter of shape [R_B, N], or [N, R_B] if adapter_transposed=1; R_B is 0 or 32.", "T") .Output(0, "Y", "Input shape with last dimension replaced by N.", "T") - .TypeConstraint("T", {"tensor(float16)"}, "FP16 activations and adapter weights.") + .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, "FP16 or FP32 activations, block scales, and adapter weights.") .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized base weights.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { const auto k = getAttribute(ctx, "K", -1); @@ -3997,10 +3997,10 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .SetDoc( "Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). " "The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights " - "stored as unsigned bytes with implicit zero point 128 and FP16 block scales. Eligible adapter down " + "stored as unsigned bytes with implicit zero point 128 and block scales matching the activation type. Eligible adapter down " "projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension " "in either packed adapter disables the adapter. " - "Intermediate projections are rounded to FP16 before the final addition. Adapters remain replaceable.") + "Intermediate projections use the activation type (FP16 or FP32) before the final addition. Adapters remain replaceable.") .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) @@ -4017,7 +4017,7 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .Input(5, "lora_B", "Packed INT8 up weights [N,R_B/32,32], R_B=0 or 32.", "TQ") .Input(6, "lora_B_scales", "Up block scales [N,1]; retained even when the adapter is empty.", "T") .Output(0, "Y", "Input shape with the last dimension replaced by N.", "T") - .TypeConstraint("T", {"tensor(float16)"}, "FP16 activations and block scales.") + .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, "FP16 or FP32 activations and block scales.") .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized weights.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { const auto k = getAttribute(ctx, "K", -1); diff --git a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc index d7f93540b8176..8796ee372c6c0 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc @@ -110,9 +110,20 @@ Status MatMulNBitsLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_ GetIntAttr(*base, "weight_prepacked", 0) != 0 || HasInput(*base, 3) || HasInput(*base, 4) || HasInput(*base, 5) || !MatrixShape(*down->InputDefs()[1], k, rank) || !MatrixShape(*up->InputDefs()[1], rank, n) || - down->InputDefs()[1]->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto::FLOAT16 || !base->InputDefs()[0]->Shape() || base->InputDefs()[0]->Shape()->dim_size() < 2) continue; + const auto* x_type = base->InputDefs()[0]->TypeAsProto(); + if (!x_type || !x_type->has_tensor_type()) continue; + const auto float_type = x_type->tensor_type().elem_type(); + if (float_type != ONNX_NAMESPACE::TensorProto::FLOAT16 && float_type != ONNX_NAMESPACE::TensorProto::FLOAT) + continue; + bool types_ok = true; + for (const auto* arg : {base->InputDefs()[2], down->InputDefs()[1], up->InputDefs()[1]}) { + const auto* type = arg->TypeAsProto(); + types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == float_type; + } + if (!types_ok) continue; + // MatMulNBits also accepts legacy flattened weights. Leave those graphs // alone until the fused kernel supports that representation. const auto* weight_shape = base->InputDefs()[1]->Shape(); diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc index 62395631a14e1..aa7e6ec9ecb5d 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -52,9 +52,13 @@ bool ZeroDefault(const Graph& graph, const NodeArg& q, const NodeArg& scales) { Initializer qdata(graph, *qp, graph.ModelPath()); Initializer sdata(graph, *sp, graph.ModelPath()); const auto bytes = qdata.DataAsByteSpan(); + if (!std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 128; })) return false; + if (sp->data_type() == ONNX_NAMESPACE::TensorProto::FLOAT) { + const auto values = sdata.DataAsSpan(); + return std::all_of(values.begin(), values.end(), [](float value) { return std::isfinite(value); }); + } const auto values = sdata.DataAsSpan(); - return std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 128; }) && - std::all_of(values.begin(), values.end(), [](MLFloat16 value) { return std::isfinite(value.ToFloat()); }); + return std::all_of(values.begin(), values.end(), [](MLFloat16 value) { return std::isfinite(value.ToFloat()); }); } void EmptyDefault(Graph& graph, NodeArg& arg, int axis) { @@ -103,11 +107,16 @@ Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph const_cast(base->InputDefs()[2]), const_cast(down->InputDefs()[1]), const_cast(down->InputDefs()[2]), const_cast(up->InputDefs()[1]), const_cast(up->InputDefs()[2])}; + const auto* x_type = inputs[0]->TypeAsProto(); + if (!x_type || !x_type->has_tensor_type()) continue; + const auto float_type = x_type->tensor_type().elem_type(); + if (float_type != ONNX_NAMESPACE::TensorProto::FLOAT16 && float_type != ONNX_NAMESPACE::TensorProto::FLOAT) + continue; bool types_ok = true; for (size_t i = 0; i < inputs.size(); ++i) { const auto* type = inputs[i]->TypeAsProto(); const auto expected = - (i == 1 || i == 3 || i == 5) ? ONNX_NAMESPACE::TensorProto::UINT8 : ONNX_NAMESPACE::TensorProto::FLOAT16; + (i == 1 || i == 3 || i == 5) ? ONNX_NAMESPACE::TensorProto::UINT8 : float_type; types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == expected; } if (!types_ok || !inputs[0]->Shape() || inputs[0]->Shape()->dim_size() < 2 || !Shape(*inputs[1], {n, k / 32, 16}) || diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc index 1fd0bda7bb857..da84940078d7a 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc @@ -7,7 +7,8 @@ namespace onnxruntime::test { #if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) -TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { +template +void RunEmptyAdapterSkipsProjection() { for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { SCOPED_TRACE(mode[0] ? mode : "default"); for (bool transposed : {false, true}) { @@ -21,16 +22,16 @@ TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { test.AddAttribute("bits", 4); test.AddAttribute("block_size", 32); test.AddAttribute("adapter_transposed", transposed); - test.AddInput("X", {1, m, 64}, std::vector(m * 64, MLFloat16(0.5f))); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); // Packed nibbles 9 => dequantized weight 1/32, so the base is exactly 1. test.AddInput("W", {128, 2, 16}, std::vector(128 * 2 * 16, 0x99), true); - test.AddInput("scales", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); - test.AddInput("lora_A", transposed ? std::vector{ranks.first, 64} : std::vector{64, ranks.first}, - std::vector(64 * ranks.first, MLFloat16(0.0625f))); - test.AddInput("lora_B", transposed ? std::vector{128, ranks.second} : std::vector{ranks.second, 128}, - std::vector(128 * ranks.second, MLFloat16(0.03125f))); + test.AddInput("scales", {128, 2}, std::vector(256, T(0.03125f)), true); + test.AddInput("lora_A", transposed ? std::vector{ranks.first, 64} : std::vector{64, ranks.first}, + std::vector(64 * ranks.first, T(0.0625f))); + test.AddInput("lora_B", transposed ? std::vector{128, ranks.second} : std::vector{ranks.second, 128}, + std::vector(128 * ranks.second, T(0.03125f))); const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; - test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, MLFloat16(expected))); + test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); std::vector> eps; eps.push_back(std::move(ep)); SessionOptions options; @@ -45,6 +46,94 @@ TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { } } +TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { RunEmptyAdapterSkipsProjection(); } +TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjectionFp32) { RunEmptyAdapterSkipsProjection(); } + +TEST(MatMulNBitsLoRA, Fp32Numerics) { + struct Shape { + int64_t m, k, n, batch; + }; + for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, + Shape{2, 2048, 36, 2}, Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { + const auto [m, k, n, batch] = shape; + SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); + std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); + std::vector aq(32 * k), bq(n * 32); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // Exactly representable by DP4A quantization, including zero rows. + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + x[row * k + col] = row % 7 == 6 ? 0.0f : value / 128.0f; + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + aq[r * k + col] = static_cast(96 + (r * 17 + col * 13) % 65); + for (int64_t block = 0; block < k / 32; ++block) + as[r * (k / 32) + block] = 0.001234567f * (1 + (r + block) % 4); + } + for (int64_t col = 0; col < n; ++col) { + bs[col] = 0.002345678f * (1 + col % 3); + for (int64_t r = 0; r < 32; ++r) + bq[col * 32 + r] = static_cast(112 + (col * 7 + r * 3) % 33); + } + std::vector at(32 * k), bt(n * 32); + for (int64_t r = 0; r < 32; ++r) + for (int64_t col = 0; col < k; ++col) + at[r * k + col] = (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; + for (int64_t col = 0; col < n; ++col) + for (int64_t r = 0; r < 32; ++r) + bt[col * 32 + r] = (static_cast(bq[col * 32 + r]) - 128) * bs[col]; + // Independent dense FP64 reference, rounding intermediate T + // to FP32. Non-half-representable weights detect accidental FP16 arithmetic. + for (int64_t row = 0; row < batch * m; ++row) { + double base = 0; + for (int64_t col = 0; col < k; ++col) base += x[row * k + col] * 0.03125; + std::vector t(32); + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + const float weight = at[r * k + col]; + sum += static_cast(x[row * k + col]) * weight; + } + t[r] = static_cast(sum); + } + for (int64_t col = 0; col < n; ++col) { + double delta = 0; + for (int64_t r = 0; r < 32; ++r) { + const float weight = bt[col * 32 + r]; + delta += static_cast(t[r]) * weight; + } + expected[row * n + col] = static_cast(base) + static_cast(delta); + } + } + for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { + for (int64_t accuracy : {0, 4}) { + SCOPED_TRACE(MakeString(mode, ", accuracy=", accuracy)); + OpTester test("MatMulNBitsLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", accuracy); + test.AddAttribute("adapter_transposed", 1); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("A", {32, k}, at); + test.AddInput("B", {n, 32}, bt); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + test.Config(options); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } + } +} + TEST(MatMulNBitsLoRA, RejectInvalidAdapterShape) { auto ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(ep, nullptr); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index ded2b7108ff7b..a492a0c0a1a6e 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -15,7 +15,8 @@ namespace onnxruntime::test { #if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) -TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { +template +void RunEmptyAndActiveAdapters() { for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { for (int64_t m : {1, 37}) { for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { @@ -24,15 +25,15 @@ TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); test.AddAttribute("K", 64); test.AddAttribute("N", 128); - test.AddInput("X", {1, m, 64}, std::vector(m * 64, MLFloat16(0.5f))); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); - test.AddInput("S", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); + test.AddInput("S", {128, 2}, std::vector(256, T(0.03125f)), true); test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); - test.AddInput("AS", {32, 2}, std::vector(64, MLFloat16(0.03125f))); + test.AddInput("AS", {32, 2}, std::vector(64, T(0.03125f))); test.AddInput("BQ", {128, ranks.second, 32}, std::vector(128 * ranks.second * 32, 129)); - test.AddInput("BS", {128, 1}, std::vector(128, MLFloat16(0.0625f))); + test.AddInput("BS", {128, 1}, std::vector(128, T(0.0625f))); const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; - test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, MLFloat16(expected))); + test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); SessionOptions options; ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); test.Config(options); @@ -44,6 +45,9 @@ TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { } } +TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { RunEmptyAndActiveAdapters(); } +TEST(MatMulNBitsQLoRA, EmptyAndActiveAdaptersFp32) { RunEmptyAndActiveAdapters(); } + TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { for (bool wrong_bits : {false, true}) { auto ep = DefaultWebGpuExecutionProvider(); @@ -69,7 +73,8 @@ TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { } } -TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { +template +void RunQuantizedPrefillAndReusedInput() { for (int64_t m : {1, 3, 4, 65}) { for (int64_t accuracy : {0, 4}) { OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); @@ -79,21 +84,21 @@ TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { test.AddAttribute("K", k); test.AddAttribute("N", n); test.AddAttribute("accuracy_level", accuracy); - std::vector x(batches * m * k); - std::vector expected(batches * m * n); + std::vector x(batches * m * k); + std::vector expected(batches * m * n); for (int64_t row = 0; row < batches * m; ++row) { const float value = row % 2 == 0 ? 0.5f : -0.25f; - std::fill_n(x.begin() + row * k, k, MLFloat16(value)); - std::fill_n(expected.begin() + row * n, n, MLFloat16(value * 24.0f)); + std::fill_n(x.begin() + row * k, k, T(value)); + std::fill_n(expected.begin() + row * n, n, T(value * 24.0f)); } - test.AddInput("X", {batches, m, k}, x); + test.AddInput("X", {batches, m, k}, x); test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, T(0.03125f)), true); test.AddInput("AQ", {32, k / 32, 32}, std::vector(32 * k, 129)); - test.AddInput("AS", {32, k / 32}, std::vector(k, MLFloat16(0.03125f))); + test.AddInput("AS", {32, k / 32}, std::vector(k, T(0.03125f))); test.AddInput("BQ", {n, 1, 32}, std::vector(n * 32, 129)); - test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(0.0625f))); - test.AddOutput("Y", {batches, m, n}, expected); + test.AddInput("BS", {n, 1}, std::vector(n, T(0.0625f))); + test.AddOutput("Y", {batches, m, n}, expected); std::vector> eps; eps.push_back(DefaultWebGpuExecutionProvider()); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); @@ -101,6 +106,88 @@ TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { } } +TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { RunQuantizedPrefillAndReusedInput(); } +TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInputFp32) { RunQuantizedPrefillAndReusedInput(); } + +TEST(MatMulNBitsQLoRA, Fp32Numerics) { + struct Shape { + int64_t m, k, n, batch; + }; + for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, + Shape{2, 2048, 36, 2}, Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { + const auto [m, k, n, batch] = shape; + SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); + std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); + std::vector aq(32 * k), bq(n * 32); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // Exactly representable by DP4A quantization, including zero rows. + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + x[row * k + col] = row % 7 == 6 ? 0.0f : value / 128.0f; + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + aq[r * k + col] = static_cast(96 + (r * 17 + col * 13) % 65); + for (int64_t block = 0; block < k / 32; ++block) + as[r * (k / 32) + block] = 0.001234567f * (1 + (r + block) % 4); + } + for (int64_t col = 0; col < n; ++col) { + bs[col] = 0.002345678f * (1 + col % 3); + for (int64_t r = 0; r < 32; ++r) + bq[col * 32 + r] = static_cast(112 + (col * 7 + r * 3) % 33); + } + // Independent dense FP64 reference, rounding weights and intermediate T + // to FP32. Non-half-representable scales detect accidental FP16 arithmetic. + for (int64_t row = 0; row < batch * m; ++row) { + double base = 0; + for (int64_t col = 0; col < k; ++col) base += x[row * k + col] * 0.03125; + std::vector t(32); + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + const float weight = (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; + sum += static_cast(x[row * k + col]) * weight; + } + t[r] = static_cast(sum); + } + for (int64_t col = 0; col < n; ++col) { + double delta = 0; + for (int64_t r = 0; r < 32; ++r) { + const float weight = (static_cast(bq[col * 32 + r]) - 128) * bs[col]; + delta += static_cast(t[r]) * weight; + } + expected[row * n + col] = static_cast(base) + static_cast(delta); + } + } + for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { + for (int64_t accuracy : {0, 4}) { + SCOPED_TRACE(MakeString(mode, ", accuracy=", accuracy)); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", accuracy); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, bs); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + test.Config(options); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } + } +} + TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; diff --git a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc index 9ed1690ee5cef..9aaf6602bbbdb 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc @@ -16,56 +16,57 @@ namespace { constexpr double kFp16FusionAbsTolerance = 1e-4; constexpr double kFp16FusionRelTolerance = 2.0 / 1024.0; +template void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, bool expose_intermediate = false, bool different_x = false, int64_t rank = 32, bool flattened_weights = false, bool transposed = false, int64_t n = 128, int isolate = 0, int64_t k_override = 0, int64_t batch = 1) { const int64_t k = k_override ? k_override : transposed ? 2048 : 128; - auto* x = builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)); + auto* x = builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)); auto weight_shape = flattened_weights ? std::vector{n * k / 2} : std::vector{n, k / 32, 16}; auto* w = builder.MakeInitializer(weight_shape, isolate ? uint8_t{0x88} : uint8_t{0}, isolate ? uint8_t{0x89} : uint8_t{255}); - auto* scales = builder.MakeInitializer({n, k / 32}, MLFloat16(0.03125f), MLFloat16(0.0625f)); - auto* a = builder.MakeInput(transposed ? std::vector{rank, k} : std::vector{k, rank}, MLFloat16(-0.0625f), MLFloat16(0.0625f)); - auto* b = builder.MakeInput(transposed ? std::vector{n, rank} : std::vector{rank, n}, MLFloat16(-0.0625f), MLFloat16(0.0625f)); + auto* scales = builder.MakeInitializer({n, k / 32}, T(0.03125f), T(0.0625f)); + auto* a = builder.MakeInput(transposed ? std::vector{rank, k} : std::vector{k, rank}, T(-0.0625f), T(0.0625f)); + auto* b = builder.MakeInput(transposed ? std::vector{n, rank} : std::vector{rank, n}, T(-0.0625f), T(0.0625f)); if (isolate == 1 || isolate == 2 || isolate == 4) { auto* arg = isolate == 2 ? a : b; - auto values = builder.feeds_.at(arg->Name()).GetMutable()->MutableDataAsSpan(); - std::fill(values.begin(), values.end(), MLFloat16(0.0f)); + auto values = builder.feeds_.at(arg->Name()).GetMutable()->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(0.0f)); const int64_t cols = arg->Shape()->dim(1).dim_value(); - for (int64_t r = 0; r < rank; ++r) values[r * cols + r] = MLFloat16(1.0f); + for (int64_t r = 0; r < rank; ++r) values[r * cols + r] = T(1.0f); } if (isolate == 4) { - auto xv = builder.feeds_.at(x->Name()).GetMutable()->MutableDataAsSpan(); - auto av = builder.feeds_.at(a->Name()).GetMutable()->MutableDataAsSpan(); - std::fill(xv.begin(), xv.end(), MLFloat16(0.0f)); - std::fill(av.begin(), av.end(), MLFloat16(0.0f)); - xv[0] = MLFloat16(10000.0f); - xv[1] = MLFloat16(0.001f); - xv[2] = MLFloat16(-10000.0f); - xv[3] = MLFloat16(0.001f); - for (int i = 0; i < 4; ++i) av[i] = MLFloat16(1.0f); + auto xv = builder.feeds_.at(x->Name()).GetMutable()->MutableDataAsSpan(); + auto av = builder.feeds_.at(a->Name()).GetMutable()->MutableDataAsSpan(); + std::fill(xv.begin(), xv.end(), T(0.0f)); + std::fill(av.begin(), av.end(), T(0.0f)); + xv[0] = T(10000.0f); + xv[1] = T(0.001f); + xv[2] = T(-10000.0f); + xv[3] = T(0.001f); + for (int i = 0; i < 4; ++i) av[i] = T(1.0f); } if (zeros) { for (auto* arg : {a, b}) { ONNX_NAMESPACE::TensorProto init; init.set_name(arg->Name()); - init.set_data_type(ONNX_NAMESPACE::TensorProto::FLOAT16); + init.set_data_type(utils::ToTensorProtoElementType()); int64_t count = 1; for (const auto& dim : arg->Shape()->dim()) { init.add_dims(dim.dim_value()); count *= dim.dim_value(); } - init.set_raw_data(std::string(static_cast(count) * sizeof(MLFloat16), '\0')); + init.set_raw_data(std::string(static_cast(count) * sizeof(T), '\0')); builder.graph_.AddInitializedTensor(init); } builder.graph_.SetInputs({x, a, b}); } - auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* t = expose_intermediate ? builder.MakeOutput(std::vector{batch, m, rank}) - : builder.MakeIntermediate(std::vector{batch, m, rank}); - auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* y = builder.MakeOutput(std::vector{batch, m, n}); + auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* t = expose_intermediate ? builder.MakeOutput(std::vector{batch, m, rank}) + : builder.MakeIntermediate(std::vector{batch, m, rank}); + auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* y = builder.MakeOutput(std::vector{batch, m, n}); NodeAttributes attrs; for (auto [name, value] : {std::pair{"K", k}, {"N", n}, {"bits", int64_t{4}}, {"block_size", int64_t{32}}}) { utils::SetNodeAttribute(utils::MakeAttribute(name, value), attrs); @@ -73,12 +74,12 @@ void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, auto& base_node = builder.AddNode("MatMulNBits", {x, w, scales}, {base}, kMSDomain); for (const auto& [name, attr] : attrs) base_node.AddAttributeProto(attr); base_node.AddAttribute("accuracy_level", int64_t{4}); - auto* lora_x = different_x ? builder.MakeInput({batch, m, k}, MLFloat16(0.0f), MLFloat16(1.0f)) : x; + auto* lora_x = different_x ? builder.MakeInput({batch, m, k}, T(0.0f), T(1.0f)) : x; auto* matmul_a = a; auto* matmul_b = b; if (transposed) { - matmul_a = builder.MakeIntermediate(std::vector{k, rank}); - matmul_b = builder.MakeIntermediate(std::vector{rank, n}); + matmul_a = builder.MakeIntermediate(std::vector{k, rank}); + matmul_b = builder.MakeIntermediate(std::vector{rank, n}); auto& ta = builder.AddNode("Transpose", {a}, {matmul_a}); auto& tb = builder.AddNode("Transpose", {b}, {matmul_b}); for (auto* node : {&ta, &tb}) { @@ -249,13 +250,14 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionMatchesWebGpuDecodeAndPref } } -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) { +template +void RunLoRASwitchesFeedsAndDefaults(const logging::Logger& logger) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; Model model("LoRA switching", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), - {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); auto& graph = model.MainGraph(); ModelTestBuilder helper(graph); - BuildLoRA(helper, 1, true); + BuildLoRA(helper, 1, true); helper.SetGraphOutputs(); ASSERT_STATUS_OK(graph.Resolve()); const auto adapters = graph.GetOverridableInitializers(); @@ -283,8 +285,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) feeds.erase(name); } else { auto* tensor = feeds.at(name).GetMutable(); - auto values = tensor->MutableDataAsSpan(); - std::fill(values.begin(), values.end(), MLFloat16(state == 2 ? -0.03125f : 0.0625f)); + auto values = tensor->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(state == 2 ? -0.03125f : 0.0625f)); } } std::vector outputs; @@ -299,5 +301,35 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) } } } +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) { + RunLoRASwitchesFeedsAndDefaults(*logger_); +} +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaultsFp32) { + RunLoRASwitchesFeedsAndDefaults(*logger_); +} + +TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionFp32Numerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (bool transposed : {false, true}) { + for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { + for (int64_t m : {1, 2, 37}) { + for (int64_t n : {128, 132}) { + for (int64_t batch : {1, 2}) { + SCOPED_TRACE(MakeString("transposed=", transposed, ", mode=", mode, ", M=", m, + ", N=", n, ", batch=", batch)); + RunWebGpuFusionTransformerTest( + [=](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, transposed, n, 0, 2048, batch); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 1e-5, 1e-5, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + }); + } + } + } + } + } +} #endif } // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index c84ff23f99545..00fd195b7928e 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -13,21 +13,22 @@ namespace onnxruntime::test { #ifndef DISABLE_CONTRIB_OPS namespace { +template void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int64_t n = 128, int64_t batch = 1, bool expose = false, bool different_x = false, bool zero_base = false, int64_t adapter_accuracy = 0) { - auto* x = builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)); + auto* x = builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)); auto* w = builder.MakeInitializer({n, k / 32, 16}, zero_base ? uint8_t{0x88} : uint8_t{0}, zero_base ? uint8_t{0x88} : uint8_t{255}); - auto* s = builder.MakeInitializer({n, k / 32}, MLFloat16(0.03125f), MLFloat16(0.0625f)); + auto* s = builder.MakeInitializer({n, k / 32}, T(0.03125f), T(0.0625f)); auto* aq = builder.MakeInput({32, k / 32, 32}, uint8_t{100}, uint8_t{156}); - auto* as = builder.MakeInput({32, k / 32}, MLFloat16(0.001f), MLFloat16(0.002f)); + auto* as = builder.MakeInput({32, k / 32}, T(0.001f), T(0.002f)); auto* bq = builder.MakeInput({n, 1, 32}, uint8_t{100}, uint8_t{156}); - auto* bs = builder.MakeInput({n, 1}, MLFloat16(0.001f), MLFloat16(0.002f)); + auto* bs = builder.MakeInput({n, 1}, T(0.001f), T(0.002f)); for (auto* arg : {aq, as, bq, bs}) { ONNX_NAMESPACE::TensorProto init; init.set_name(arg->Name()); const bool packed = arg == aq || arg == bq; - init.set_data_type(packed ? ONNX_NAMESPACE::TensorProto::UINT8 : ONNX_NAMESPACE::TensorProto::FLOAT16); + init.set_data_type(packed ? ONNX_NAMESPACE::TensorProto::UINT8 : utils::ToTensorProtoElementType()); int64_t count = 1; for (const auto& dim : arg->Shape()->dim()) { init.add_dims(dim.dim_value()); @@ -36,17 +37,17 @@ void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int6 if (packed) { init.set_raw_data(std::string(static_cast(count), static_cast(128))); } else { - const std::vector ones(static_cast(count), MLFloat16(1.0f)); - init.set_raw_data(ones.data(), ones.size() * sizeof(MLFloat16)); + const std::vector ones(static_cast(count), T(1.0f)); + init.set_raw_data(ones.data(), ones.size() * sizeof(T)); } builder.graph_.AddInitializedTensor(init); } - auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* t = expose ? builder.MakeOutput(std::vector{batch, m, 32}) - : builder.MakeIntermediate(std::vector{batch, m, 32}); - auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* y = builder.MakeOutput(std::vector{batch, m, n}); - auto* adapter_x = different_x ? builder.MakeInput({batch, m, k}, MLFloat16(-0.5f), MLFloat16(0.5f)) : x; + auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* t = expose ? builder.MakeOutput(std::vector{batch, m, 32}) + : builder.MakeIntermediate(std::vector{batch, m, 32}); + auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* y = builder.MakeOutput(std::vector{batch, m, n}); + auto* adapter_x = different_x ? builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)) : x; if (different_x) { builder.graph_.SetInputs({x, aq, as, bq, bs, adapter_x}); } else { @@ -132,6 +133,25 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { } } +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Numerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { + for (const auto& [m, n] : {std::pair{1, 128}, {1, 132}, {2, 128}, {37, 132}}) { + SCOPED_TRACE(MakeString(mode, ", M=", m, ", N=", n)); + const bool quantized_down = m >= 4 && std::string_view{mode} != "fuse_down"; + RunWebGpuFusionTransformerTest( + [m, n](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, n); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + quantized_down ? 0.002 : 0.00001, quantized_down ? 0.001 : 0.00001, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, + [mode](SessionOptions& options) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + }); + } + } +} + TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsObservableIntermediate) { ASSERT_STATUS_OK(TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, true); }, 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, @@ -253,15 +273,16 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { } } -TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) { +template +void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 37}) { SCOPED_TRACE(m); Model model("quantized LoRA empty defaults", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); auto& graph = model.MainGraph(); ModelTestBuilder builder(graph); - BuildQLoRA(builder, m, 64); + BuildQLoRA(builder, m, 64); builder.SetGraphOutputs(); ASSERT_STATUS_OK(graph.Resolve()); const auto adapters = graph.GetOverridableInitializers(); @@ -275,9 +296,9 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) std::fill(bytes.begin(), bytes.end(), uint8_t{129}); (tensor->Shape()[0] == 32 ? a_name : b_name) = name; } else { - auto values = tensor->MutableDataAsSpan(); + auto values = tensor->MutableDataAsSpan(); const float fill = tensor->Shape().NumDimensions() == 3 ? 0.25f : 0.03125f; - std::fill(values.begin(), values.end(), MLFloat16(fill)); + std::fill(values.begin(), values.end(), T(fill)); } } ASSERT_FALSE(a_name.empty()); @@ -286,7 +307,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) for (bool fused : {false, true}) { if (fused) { bool modified = false; - ASSERT_STATUS_OK(MatMulNBitsQLoRAFusion().Apply(graph, modified, *logger_)); + ASSERT_STATUS_OK(MatMulNBitsQLoRAFusion().Apply(graph, modified, logger)); ASSERT_TRUE(modified); ASSERT_STATUS_OK(graph.Resolve()); } @@ -386,8 +407,8 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) for (int state : {0, 0, 0, 1, 1, 1, 0, 1, 0, 1}) { SCOPED_TRACE(MakeString("capture state=", state)); auto* cpu_output = actual.GetMutable(); - auto values = cpu_output->MutableDataAsSpan(); - std::fill(values.begin(), values.end(), MLFloat16(-99.0f)); + auto values = cpu_output->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(-99.0f)); ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(*cpu_output, *outputs[state].GetMutable())); RunOptions run_options; ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, @@ -403,5 +424,13 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) } } +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) { + RunQLoRASerializedEmptyDefaults(*logger_); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaultsFp32) { + RunQLoRASerializedEmptyDefaults(*logger_); +} + #endif } // namespace onnxruntime::test From 4e8a53738de04af23f72967aeb326ebb43ec8634 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 17:44:32 +0800 Subject: [PATCH 13/35] [WebGPU] Fuse FP32 QLoRA decode base and down projection with DP4A --- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 45 ++++++++++++ docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 19 +++-- .../webgpu/quantization/dp4a_matmul_nbits.cc | 29 ++++++-- .../webgpu/quantization/dp4a_matmul_nbits.h | 35 +++++---- .../dp4a_matmul_small_m.wgsl.template | 41 +++++++++++ .../webgpu/quantization/matmul_nbits_qlora.cc | 10 +++ .../contrib_ops/matmul_nbits_qlora_test.cc | 71 +++++++++++++++++++ .../matmul_nbits_qlora_fusion_test.cc | 12 +++- 8 files changed, 235 insertions(+), 27 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md index 40fa01afa7928..d442f06df53db 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Performance.md @@ -384,3 +384,48 @@ in every measured cell. Across these measured cells, the largest TTFT increase is 0.15% and the largest decode decrease is 1.41%. No large FP16 regression appears in this single-round check; smaller changes remain within typical run variation. + +## Joint DP4A FP32 QLoRA decode (2026-10-15) + +Eligible NVIDIA FP32 decode now uses three dispatches: quantize X once, +jointly compute the INT4 base and INT8 down projections, then apply the +vectorized INT8 up projection and addition. This replaces the separate +floating-input down projection in the previous four-dispatch implementation. +The base workgroup size remains 128. Eight extra workgroups per input row +compute rank 32 from the shared quantized activation and scale buffers. + +The baseline is the preceding FP16/FP32 support implementation, not upstream +main. These native C++ measurements use built-in WebGPU/D3D12 on RTX 5080, +batch 1, accuracy level 4, capture and generator reuse, 128 generated tokens, +three warmups, and five measured requests. Values are medians. Timing runs +were isolated from builds, tests, profiling, and verbose diagnostics. + +| Prompt | Adapter | TTFT before → after (ms) | Decode before → after (tok/s) | Decode change | +|---:|---|---:|---:|---:| +| 128 | Off | 29.61 → 29.60 | 203.20 → 203.10 | -0.05% | +| 128 | A | 43.64 → 43.42 | 134.55 → 180.44 | +34.11% | +| 1024 | Off | 161.88 → 162.15 | 199.03 → 199.29 | +0.13% | +| 1024 | A | 192.12 → 195.07 | 129.82 → 174.29 | +34.26% | + +The adapter-on throughput deficit relative to adapter-off is now 11.2% and +12.5%, respectively. Control measurements of FP16 QLoRA, FP16 LoRA, and +FP32 LoRA at both prompt lengths show absolute decode changes below 0.9% +and TTFT changes below 0.9%. Across all measured cells the largest TTFT +increase is 1.53%; TTFT includes the first generated token. + +Both adapters were checked on six full-model cases against the previous +FP32 QLoRA runtime. All six generated sequences and 72 forced-prefix next-token +choices match. Adapter-on decode logits are approximate: maximum absolute +difference 0.822086 and largest per-case RMSE 0.087801. This follows from +using quantized X in the down projection; it is not a bitwise-equivalence +claim or a broad quality evaluation. Adapter-off logits, prefill logits, +and capture OFF/ON logits are identical. All 12 before/after benchmark +output sequences also match. + +Validation includes the provider/empty-tensor suite, expanded FP32 numerical +coverage, an independent dense reference on a known activation quantization +grid, and all 12 QLoRA fusion tests. The new cases cover batched decode, +zero rows, K tails, the full INT8 byte range, and distinct adapter-on/off +capture IDs with replay. Separate dispatch logging confirms quantize → +joint base/down → up/add. Both changed kernel translation units pass WASM32 +syntax checks. Only the built-in EP was tested for this experiment. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 79d01aafd293b..6b36666a87cfd 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -49,8 +49,10 @@ Base quantization stays INT4 with block size 32 and implicit zero point 8. The equations describe the weight interpretation and mathematical projections; the optimized kernels can change reduction order and activation precision. Eligible prefill down projections use DP4A activation quantization when the -model's `accuracy_level` permits it. Decode and up projections do not add -activation quantization. The provider's `enableMatmulFp32Accumulation` option +model's `accuracy_level` permits it. Eligible FP32 DP4A decode also shares the +base projection's quantized X with the down projection. Other decode down +projections and all up projections retain floating-point activations. +The provider's `enableMatmulFp32Accumulation` option controls floating-point lane accumulation for FP16 inputs. FP32 inputs always use FP32 floating-point accumulation. Intermediate T uses the activation type; the specialized decode kernel performs its final subgroup reduction in FP32. @@ -71,8 +73,17 @@ dimension of X; a one-token prompt follows the decode selection. The joint decode kernel supports one total input row and K from 2048 through 8192, when the base projection does not select DP4A. FP32 base projections -can select DP4A even at M=1; those retain the original base kernel and compute -the adapter down projection separately. The joint kernel retains 128 threads +can select DP4A even at M=1. On NVIDIA, eligible FP32 inputs with K from 2048 +through 8192 use three dispatches: quantize X once, jointly compute the INT4 +base and INT8 down projections, then perform the INT8 up projection and addition. +The joint DP4A shader retains the base's 128-thread workgroups and adds eight +workgroups per input row for rank 32, with four ranks per workgroup. Both +projections read the same quantized activation buffer and scales; each writes +its own output. The next dispatch provides synchronization. The existing DP4A +eligibility checks, including `accuracy_level=4`, still apply. Adapter-off, +FP16 decode, and prefill selection are unchanged. + +The floating-input joint kernel retains 128 threads per workgroup and adds four workgroups for the 32 down-projection outputs. On devices with fixed 32-lane subgroups, each lane loads X into registers and reuses it for two rank outputs, then reduces lane diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 6e750083db2e6..18cb40d3bde9e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -100,11 +100,18 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_; ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count"); + if (has_qlora_down_) { + shader.AddInput("lora_a"); + shader.AddInput("lora_a_scales"); + shader.AddOutput("lora_down"); + } + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_small_m.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), + WGSL_TEMPLATE_PARAMETER(has_qlora_down, has_qlora_down_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), @@ -139,14 +146,21 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales, - DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n) { + DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n, const Tensor* lora_a_scales) { const bool has_lora_down = lora_at != nullptr; + const bool has_qlora_down = lora_a_scales != nullptr; const bool lora_quantized = lora_b_scales != nullptr; ORT_RETURN_IF_NOT(small_m_tile_n == 0 || (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && !zero_points && !bias && !weight_index && !weight_index_indirect && !lora_t && !lora_at), "Small-M override requires a plain rank-32 INT8 projection."); - ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && M >= min_M_for_tile_optimization), + ORT_RETURN_IF_NOT(!has_qlora_down || + (has_lora_down && lora_down && a->IsDataType() && nbits == 4 && block_size == 32 && + M == 1 && dispatch_M == 1 && !small_m_tile_n && !weight_index && !weight_index_indirect && + !zero_points && !bias && !lora_t), + "Unsupported joint DP4A QLoRA decode projection."); + ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && + (has_qlora_down || M >= min_M_for_tile_optimization)), "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; @@ -190,17 +204,17 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor // The narrow projection changes the K reduction order. Keep partial sums // and the final reduction in FP32, then cast T once when storing it. const bool acc_f32 = small_m_tile_n != 0 || context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized}; + DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); - mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile); + mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -216,6 +230,11 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor : 1}); if (lora_quantized) mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } + if (has_qlora_down) { + mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); + mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(mul_program); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 84ddf55f62cfb..017036ae6b7a6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -87,20 +87,23 @@ class DP4AMatMulNBitsSmallMProgram final : public Program(0x80808080u); + let w1 = lora_a[offset * 2u + 1u] ^ vec4(0x80808080u); + partial += SDP8AI(tile_A[lane * 2u], w0, tile_A[lane * 2u + 1u], w1, + scale_A[lane / 4u] * lora_a_scales[offset]); + } + workgroupBarrier(); + } + inter_results[rank][lane] = partial; + workgroupBarrier(); + if (local_idx < 4u) { + var sum = f32(0); + for (var i = 0u; i < 32u; i++) { sum += inter_results[local_idx][i]; } + lora_down[batch * 32u + rank_base + local_idx] = sum; + } +} +#endif + $MAIN { +#if has_qlora_down + let base_groups = uniforms.batch_count * uniforms.dispatch_M * uniforms.num_N_tile; + if (workgroup_idx >= base_groups) { + qlora_down(workgroup_idx - base_groups, local_idx); + return; + } +#endif let batch = workgroup_idx / (uniforms.dispatch_M * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 0a959afd46695..d6e34ad20007a 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -278,6 +278,16 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // shader must not replace that base projection's activation quantization. const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), narrow(k_), 4, m, false, y); + if (!fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192 && + context.AdapterInfo().vendor == std::string_view{"nvidia"}) { + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( + x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, + narrow(n_), narrow(k_), 32, narrow(k_ / 32), + kMinMForTileOptimization, 4, context, &base, 0, + nullptr, nullptr, nullptr, aq, &t, true, nullptr, nullptr, 0, as)); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + } const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && context.AdapterInfo().vendor == std::string_view{"nvidia"} && !base_dp4a; if (joint_down) { diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index a492a0c0a1a6e..b5680c4b5c981 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -114,6 +114,7 @@ TEST(MatMulNBitsQLoRA, Fp32Numerics) { int64_t m, k, n, batch; }; for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, + Shape{1, 2176, 64, 7}, Shape{1, 8192, 64, 1}, Shape{2, 2048, 36, 2}, Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { const auto [m, k, n, batch] = shape; SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); @@ -188,6 +189,76 @@ TEST(MatMulNBitsQLoRA, Fp32Numerics) { } } +TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; +#else + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId()) || + gpu.AdapterInfo().vendor != std::string_view{"nvidia"}) { + GTEST_SKIP() << "The joint DP4A decode projection is selected on NVIDIA DP4A devices."; + } +#endif + constexpr int64_t k = 2176, n = 64, batch = 3; + std::vector x(batch * k), quantized_x(batch * k), as(k), expected(batch * n); + std::vector aq(32 * k), bq(n * 32, 128); + for (int64_t row = 0; row < batch; ++row) { + for (int64_t col = 0; col < k; ++col) { + const float scale = 1.0f / static_cast(1 << (8 + (row + col / 128) % 3)); + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + quantized_x[row * k + col] = row == 1 ? 0.0f : value * scale; + // Stay inside the rounding interval of the known quantization grid. + const float offset = value == 127 || value == -127 ? 0.0f : 0.2f; + x[row * k + col] = row == 1 ? 0.0f : (value + offset) * scale; + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + aq[r * k + col] = static_cast((r * 17 + col * 13) % 256); + for (int64_t block = 0; block < k / 32; ++block) + as[r * (k / 32) + block] = (1 + (r + block) % 4) / 2048.0f; + } + for (int64_t col = 0; col < n; ++col) bq[col * 32 + col % 32] = 129; + // An identity up projection exposes every T value. The reference uses + // dense FP64 products on the known grid, including a zero row and a K tail. + for (int64_t row = 0; row < batch; ++row) { + double base = 0; + for (int64_t col = 0; col < k; ++col) base += quantized_x[row * k + col] * 0.03125; + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) + sum += static_cast(quantized_x[row * k + col]) * + (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; + expected[row * n + r] = expected[row * n + r + 32] = static_cast(base) + static_cast(sum); + } + } + for (const char* mode : {"auto", "fuse_down"}) { + SCOPED_TRACE(mode); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {batch, 1, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, std::vector(n, 1.0f)); + test.AddOutput("Y", {batch, 1, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); + test.Config(options); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } +} + TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 00fd195b7928e..614cba24d6cd8 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -138,7 +138,9 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Numerics) { for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { for (const auto& [m, n] : {std::pair{1, 128}, {1, 132}, {2, 128}, {37, 132}}) { SCOPED_TRACE(MakeString(mode, ", M=", m, ", N=", n)); - const bool quantized_down = m >= 4 && std::string_view{mode} != "fuse_down"; + // Eligible FP32 decode shares the base projection's quantized X. + const bool quantized_down = (m >= 4 && std::string_view{mode} != "fuse_down") || + (m == 1 && n % 16 == 0 && std::string_view{mode} != "fuse_up"); RunWebGpuFusionTransformerTest( [m, n](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, n); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, @@ -274,7 +276,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { } template -void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger) { +void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 37}) { SCOPED_TRACE(m); @@ -282,7 +284,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger) { IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); auto& graph = model.MainGraph(); ModelTestBuilder builder(graph); - BuildQLoRA(builder, m, 64); + BuildQLoRA(builder, m, k); builder.SetGraphOutputs(); ASSERT_STATUS_OK(graph.Resolve()); const auto adapters = graph.GetOverridableInitializers(); @@ -432,5 +434,9 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaultsFp RunQLoRASerializedEmptyDefaults(*logger_); } +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionJointDP4ACapture) { + RunQLoRASerializedEmptyDefaults(*logger_, 2048); +} + #endif } // namespace onnxruntime::test From 0243e40267de450b79166afe60f9e10c833edabe Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 18:23:59 +0800 Subject: [PATCH 14/35] [WebGPU] Parallelize floating LoRA projections for prefill and decode --- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 69 ++++++++++++ .../quantization/dp4a_matmul.wgsl.template | 20 ++++ .../webgpu/quantization/dp4a_matmul_nbits.cc | 29 ++++- .../webgpu/quantization/dp4a_matmul_nbits.h | 71 ++++++------ .../webgpu/quantization/matmul_nbits.h | 1 + .../webgpu/quantization/matmul_nbits_lora.cc | 106 +++++++++++++++++- .../contrib_ops/matmul_nbits_lora_test.cc | 62 ++++++++++ .../matmul_nbits_lora_fusion_test.cc | 10 +- 8 files changed, 325 insertions(+), 43 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md index d442f06df53db..c02fc3dde4f25 100644 --- a/docs/WebGPU_MatMulNBitsLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsLoRA_Performance.md @@ -429,3 +429,72 @@ zero rows, K tails, the full INT8 byte range, and distinct adapter-on/off capture IDs with replay. Separate dispatch logging confirms quantize → joint base/down → up/add. Both changed kernel translation units pass WASM32 syntax checks. Only the built-in EP was tested for this experiment. + +## Floating LoRA prefill and decode follow-up + +The baseline for this experiment includes the three-dispatch QLoRA change above. +The same optional graphs and floating adapters were used before and after. The +native built-in WebGPU measurements use batch 1, 128 generated tokens, fixed +prompt IDs, capacity 4096, capture/reuse and token readback, three warmups and +five measured requests. Values are medians on RTX 5080/D3D12. Timing ran alone; +profiling and correctness checks ran separately. A reversed runtime ordering +confirmed the baseline after the candidate measurements. + +| Type | Prompt | Off TTFT (ms) | LoRA TTFT before → after (ms) | QLoRA TTFT (ms) | LoRA decode before → after (tok/s) | QLoRA decode (tok/s) | +|---|---:|---:|---:|---:|---:|---:| +| FP16 | 128 | 31.30 | 83.72 → 46.74 | 40.56 | 200.04 → 197.53 | 195.13 | +| FP16 | 1024 | 162.84 | 225.95 → 190.35 | 172.53 | 197.64 → 195.02 | 195.08 | +| FP32 | 128 | 29.77 | 123.51 → 59.69 | 45.18 | 138.98 → 165.35 | 170.26 | +| FP32 | 1024 | 162.09 | 270.70 → 220.47 | 192.53 | 136.07 → 162.33 | 167.47 | + +FP16 adapter-off decode measured 225.01/223.75 tok/s, and FP32 measured +195.75/192.21 tok/s. The floating LoRA decode deficit is now 12–16%, close +to QLoRA's 11–13% on its separately measured adapter-off controls. FP16 +batch-one decode uses its existing algorithm; its measured change was -1.3%. +Adapter-off controls changed by at most 1.3% in decode and 0.6% in TTFT. +QLoRA controls showed no regression above 1.7% in TTFT or 1% in decode. +These control differences are timing variation, not claimed improvements. + +The retained changes are: + +- A K-parallel floating down projection for NVIDIA devices with fixed + 32-lane subgroups. A 128-thread workgroup computes four ranks for four + rows, reusing each vec4 weight across rows. FP32 partials are reduced + with subgroup operations and cast once to T's model dtype. +- Shared transposed B tiles in the FP16 prefill DP4A base/up epilogue. + This preserves the existing workgroup size and arithmetic order. + The FP32 shared-tile candidate did not help and is not enabled. +- FP32 decode uses quantize X → joint base/down → up/add. Each floating + rank gets one 128-thread workgroup, four subgroup reductions, and one + workgroup barrier. The down projection reads original floating X; + only the base projection uses quantized X. Eligibility retains ordinary + DP4A checks and additionally requires NVIDIA, fixed 32-lane subgroups, + M=1, and K=2048–8192. Explicit `fuse_up` retains its existing routing. + +Prefill down-projection GPU time fell from 36.97/37.39 ms to 4.63/12.01 ms +for FP16 at 128/1024 tokens, and from 68.64/71.11 ms to 4.39/17.17 ms for +FP32. These are hardware timestamp sums from separate uncaptured profiles, +not end-to-end latencies. Full floating prefill still trails QLoRA. A +private diagnostic that moved identical adapter values into initializers +saved about 11 ms (FP16) and 21 ms (FP32) on short prompts, showing substantial +per-request input/upload cost outside the shader timestamps. Persistent +device adapter storage remains separate work requiring correct ownership, +device identity, unload, and capture invalidation handling. + +Validation passed 15 provider tests and 25 fusion tests, including new dense +FP64 references for both floating types, off-grid X, zero rows, batched +decode, K tails, and 128/1024-row prefill. Two affected FP16 fusion +comparisons use a scoped 4/1024 absolute bound for changing a sequential +FP16 reduction to FP32 partial sums; independent dense-reference tests +retain tighter bounds. Native build, lint, and WASM32 syntax checks passed. + +With adapters A/B on six full-model cases per dtype, five of six generated +answers were unchanged for each dtype; the changed answers retained their +facts with different wording. Forced-prefix next-token agreement was +72/72 for FP16 and 70/72 for FP32. Maximum logit differences were 0.689453 +and 0.911110, with largest per-case RMSE 0.071453 and 0.111359. Adapter-off +logits and capture OFF/ON logits were exactly equal, and rewind checks passed. +In an additional FP32 diagnostic with base activation quantization disabled, +the maximum logit difference dropped to 0.0000253 and all 72/72 choices +matched. This supports DP4A threshold amplification of small upstream +rounding changes. The checks are not a broad model-quality evaluation. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index eee5453d68959..eb78cd0dac5d4 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -7,6 +7,7 @@ #param has_lora_up #param lora_quantized #param lora_transposed +#param lora_shared_b #param has_bias #param has_zero_points #param is_qualcomm @@ -19,6 +20,10 @@ #use .getByOffset .setByOffset +#if lora_shared_b +var lora_b_tile: array, tile_size_n * 8>; +#endif + fn write_lora_output(offset: u32, value: vec4) { #if has_lora_up let row = (offset * 4u) / uniforms.N; @@ -29,10 +34,18 @@ fn write_lora_output(offset: u32, value: vec4) { #elif lora_transposed for (var r4 = 0u; r4 < 8u; r4++) { let tv = lora_t[row * 8u + r4]; +#if lora_shared_b + let local_col = col % tile_size_n; + let b0 = lora_b_tile[local_col * 8u + r4]; + let b1 = lora_b_tile[(local_col + 1u) * 8u + r4]; + let b2 = lora_b_tile[(local_col + 2u) * 8u + r4]; + let b3 = lora_b_tile[(local_col + 3u) * 8u + r4]; +#else let b0 = lora_b[col * 8u + r4]; let b1 = lora_b[(col + 1u) * 8u + r4]; let b2 = lora_b[(col + 2u) * 8u + r4]; let b3 = lora_b[(col + 3u) * 8u + r4]; +#endif delta = vec4(b0.x, b1.x, b2.x, b3.x) * tv.x + delta; delta = vec4(b0.y, b1.y, b2.y, b3.y) * tv.y + delta; delta = vec4(b0.z, b1.z, b2.z, b3.z) * tv.z + delta; @@ -437,6 +450,13 @@ $MAIN { workgroupBarrier(); } +#if lora_shared_b + for (var i = local_idx; i < tile_size_n * 8u; i += workgroup_size) { + let col = b_global_base + i / 8u; + if (col < uniforms.N) { lora_b_tile[i] = lora_b[col * 8u + i % 8u]; } + } + workgroupBarrier(); +#endif let a_global = a_global_base + base_A + a_idx; let b_global = b_global_base + base_B; let output_idx = (batch * uniforms.M * uniforms.N + a_global * uniforms.N + b_global)/4; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 18cb40d3bde9e..19e8fbf094ae6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -57,6 +57,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), + WGSL_TEMPLATE_PARAMETER(lora_shared_b, lora_shared_b_), WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), @@ -106,6 +107,10 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co shader.AddOutput("lora_down"); } + if (has_lora_down_) { + GenerateLoRADecode(shader, "uniforms.batch_count * uniforms.dispatch_M * uniforms.num_N_tile"); + } + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_small_m.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), @@ -149,6 +154,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n, const Tensor* lora_a_scales) { const bool has_lora_down = lora_at != nullptr; const bool has_qlora_down = lora_a_scales != nullptr; + const bool has_float_decode = has_lora_down && !has_qlora_down && M == 1 && dispatch_M == 1 && + a->IsDataType() && context.AdapterInfo().vendor == std::string_view{"nvidia"} && + context.HasFeature(wgpu::FeatureName::Subgroups) && + context.AdapterInfo().subgroupMinSize == 32 && context.AdapterInfo().subgroupMaxSize == 32; const bool lora_quantized = lora_b_scales != nullptr; ORT_RETURN_IF_NOT(small_m_tile_n == 0 || (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && @@ -160,7 +169,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor !zero_points && !bias && !lora_t), "Unsupported joint DP4A QLoRA decode projection."); ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && - (has_qlora_down || M >= min_M_for_tile_optimization)), + (has_qlora_down || has_float_decode || M >= min_M_for_tile_optimization)), "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; @@ -204,17 +213,18 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor // The narrow projection changes the K reduction order. Keep partial sums // and the final reduction in FP32, then cast T once when storing it. const bool acc_f32 = small_m_tile_n != 0 || context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down}; + DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down, has_float_decode}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); - mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : 0)); + mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : has_float_decode ? batch_count * 32 + : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down, has_float_decode); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -235,6 +245,11 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor mul_program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); } + if (has_float_decode) { + mul_program.AddInput({a, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(mul_program); } @@ -246,7 +261,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor uint32_t num_N_tile = CeilDiv(N, tile_size_n); bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized}; + const bool lora_shared_b = has_lora_up && lora_transposed && !lora_quantized && + y->IsDataType() && context.AdapterInfo().vendor == std::string_view{"nvidia"}; + DP4AMatMulNBitsProgram mul_program{tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, lora_shared_b}; mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile + (has_lora_down ? CeilDiv(batch_count * M, 32u) : 0)); mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, @@ -264,7 +281,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized); + .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, lora_shared_b); if (has_zero_points) { mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 017036ae6b7a6..962295e1671eb 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -33,23 +33,24 @@ class DP4AMatMulNBitsProgram final : public Program { DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false) : Program{has_lora_down ? (lora_quantized ? "DP4AMatMulNBitsQLoRADown" : "DP4AMatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBitsLoRAUp") - : "DP4AMatMulNBits"}, - tile_size_m_(tile_size_m), - tile_size_n_(tile_size_n), - tile_size_k_vec_(tile_size_k_vec), - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32), - has_lora_up_(has_lora_up), - has_lora_down_(has_lora_down), - lora_transposed_(lora_transposed), - lora_quantized_(lora_quantized) {} + bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false, bool lora_shared_b = false) : Program{has_lora_down ? (lora_quantized ? "DP4AMatMulNBitsQLoRADown" : "DP4AMatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBitsLoRAUp") + : "DP4AMatMulNBits"}, + tile_size_m_(tile_size_m), + tile_size_n_(tile_size_n), + tile_size_k_vec_(tile_size_k_vec), + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32), + has_lora_up_(has_lora_up), + has_lora_down_(has_lora_down), + lora_transposed_(lora_transposed), + lora_quantized_(lora_quantized), + lora_shared_b_(lora_shared_b) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"batch_count", ProgramUniformVariableDataType::Uint32}, @@ -79,6 +80,7 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_lora_down_; bool lora_transposed_; bool lora_quantized_; + bool lora_shared_b_; }; class DP4AMatMulNBitsSmallMProgram final : public Program { @@ -88,22 +90,24 @@ class DP4AMatMulNBitsSmallMProgram final : public Program { } }; +static void GenerateLoRAParallel(ShaderHelper& shader, const std::string& base_groups, uint32_t tile_rows) { + shader.AddInput("lora_x"); + shader.AddInput("lora_at"); + shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); + auto& body = shader.MainFunctionBody(); + body << "let lora_base_groups = " << base_groups << ";\n" + << "if (workgroup_idx >= lora_base_groups) {\n" + << "let group = workgroup_idx - lora_base_groups;\n" + << "let row = (group / 8u) * " << tile_rows << "u;\n" + << R"( + let rank = (group % 8u) * 4u + local_idx / 32u; + let lane = local_idx % 32u; +)"; + // Reuse each contiguous weight vector across prompt rows. FP32 partials + // avoid accumulating FP16 rounding error over the long K dimension. + for (uint32_t row = 0; row < tile_rows; ++row) body << "var acc" << row << " = vec4(0);\n"; + body << "for (var k4 = lane; k4 < uniforms.K / 4u; k4 += 32u) {\n" + << "let av = vec4(lora_at[rank * (uniforms.K / 4u) + k4]);\n"; + for (uint32_t row = 0; row < tile_rows; ++row) { + body << "if (row + " << row << "u < uniforms.M * uniforms.batch_count) {\n" + << "acc" << row << " = fma(vec4(lora_x[(row + " << row + << "u) * (uniforms.K / 4u) + k4]), av, acc" << row << ");\n}\n"; + } + body << "}\n"; + for (uint32_t row = 0; row < tile_rows; ++row) { + body << "let sum" << row << " = subgroupAdd((acc" << row << ".x + acc" << row << ".z) + (acc" + << row << ".y + acc" << row << ".w));\n" + << "if (lane == 0u && row + " << row << "u < uniforms.M * uniforms.batch_count) {\n" + << "lora_down[(row + " << row << "u) * 32u + rank] = lora_down_element_t(sum" << row << ");\n}\n"; + } + body << "return;\n}\n"; +} + +void GenerateLoRADecode(ShaderHelper& shader, const std::string& base_groups) { + // One workgroup per rank exposes more parallelism than four ranks per + // workgroup; the four subgroup partials need only one workgroup barrier. + shader.AddInput("lora_x"); + shader.AddInput("lora_at"); + shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); + shader.AdditionalImplementation() << "var lora_sums: array;\n"; + shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" + << R"( +if (workgroup_idx >= lora_base_groups) { + let group = workgroup_idx - lora_base_groups; + let row = group / 32u; + let rank = group % 32u; + var acc = vec4(0); + for (var k4 = local_idx; k4 < uniforms.K / 4u; k4 += 128u) { + acc = fma(vec4(lora_x[row * (uniforms.K / 4u) + k4]), + vec4(lora_at[rank * (uniforms.K / 4u) + k4]), acc); + } + let sum = subgroupAdd((acc.x + acc.z) + (acc.y + acc.w)); + if (local_idx % 32u == 0u) { lora_sums[local_idx / 32u] = sum; } + workgroupBarrier(); + if (local_idx == 0u) { + lora_down[row * 32u + rank] = lora_down_element_t((lora_sums[0] + lora_sums[2]) + (lora_sums[1] + lora_sums[3])); + } + return; +} +)"; +} + +class LoRADownParallelProgram final : public Program { + public: + explicit LoRADownParallelProgram(uint32_t tile_rows) : Program{"LoRADownParallel"}, tile_rows_(tile_rows) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(LoRADownParallelProgram); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + GenerateLoRAParallel(shader, "0u", tile_rows_); + return Status::OK(); + } + + private: + uint32_t tile_rows_; +}; + // Fallback epilogue for base kernels without an integrated adapter output path. class LoRAOutputProgram final : public Program { public: @@ -178,7 +256,21 @@ Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const T down_shape.back() = 32; auto t = context.CreateGPUTensor(x->DataType(), TensorShape(down_shape)); const auto compute_down = [&]() -> Status { - if (use_gemv) { + const auto& adapter = context.AdapterInfo(); + if (rows > 1 && adapter.vendor == std::string_view{"nvidia"} && + context.HasFeature(wgpu::FeatureName::Subgroups) && + adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) { + constexpr uint32_t tile_rows = 4; + LoRADownParallelProgram down(tile_rows); + down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, + {at, ProgramTensorMetadataDependency::Type, 4}}) + .AddOutput({&t, ProgramTensorMetadataDependency::Type}) + .SetWorkgroupSize(128) + .SetDispatchGroupSize(CeilDiv(rows, tile_rows) * 8) + .AddUniformVariables({{narrow(k)}, {rows}, {1u}}) + .CacheHint(tile_rows); + ORT_RETURN_IF_ERROR(context.RunProgram(down)); + } else if (use_gemv) { LoRAGemvProgram down; down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, {at, ProgramTensorMetadataDependency::Type, 4}}) @@ -209,7 +301,17 @@ Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const T // floating joint down kernel would otherwise replace it. const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits( context, accuracy_level, 32, narrow(n), narrow(k), 4, m, false, y); - if (use_gemv && !base_dp4a) { + const auto& adapter = context.AdapterInfo(); + if (m == 1 && k >= 2048 && k <= 8192 && base_dp4a && x->IsDataType() && + adapter.vendor == std::string_view{"nvidia"} && context.HasFeature(wgpu::FeatureName::Subgroups) && + adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) { + // Quantize only the base's X. The appended floating down projection + // reads original X and writes T for the subsequent up/add dispatch. + ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits(x, weight, scales, nullptr, nullptr, rows, m, m, + narrow(n), narrow(k), 32, 0, + kMinMForTileOptimization, 4, context, &base, 0, + nullptr, nullptr, nullptr, at, &t)); + } else if (use_gemv && !base_dp4a) { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, context, &base, 0, nullptr, 0, nullptr, nullptr, at, &t)); } else if (m >= kMinMForTileOptimization && context.AdapterInfo().vendor == std::string_view{"nvidia"} && diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc index da84940078d7a..f09f780aee3c8 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc @@ -134,6 +134,68 @@ TEST(MatMulNBitsLoRA, Fp32Numerics) { } } +template +void RunParallelDownNumerics() { + struct Shape { + int64_t batch, m, k, n; + }; + for (const auto shape : {Shape{1, 1, 2048, 64}, Shape{2, 1, 8192, 64}, Shape{2, 33, 2176, 64}, + Shape{1, 128, 4096, 64}, Shape{1, 1024, 2048, 16}}) { + const auto [batch, m, k, n] = shape; + SCOPED_TRACE(MakeString("batch=", batch, ", M=", m, ", K=", k, ", N=", n)); + std::vector x(batch * m * k), a(32 * k), b(n * 32), expected(batch * m * n); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // Off-grid floating X detects accidental reuse of the base's INT8 X. + x[row * k + col] = T(row % 7 == 6 ? 0.0f : (static_cast((row * 19 + col * 11) % 255) - 127) * 0.00712345f); + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + a[r * k + col] = T((static_cast((r * 17 + col * 13) % 65) - 32) * 0.001234567f); + } + for (int64_t col = 0; col < n; ++col) + for (int64_t r = 0; r < 32; ++r) + b[col * 32 + r] = T((static_cast((col * 7 + r * 3) % 33) - 16) * 0.002345678f); + // Dense FP64 reference with exactly one cast at each MatMul boundary. + // A zero base isolates adapter arithmetic from activation quantization. + for (int64_t row = 0; row < batch * m; ++row) { + std::vector t(32); + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) + sum += static_cast(static_cast(x[row * k + col])) * static_cast(a[r * k + col]); + t[r] = T(static_cast(sum)); + } + for (int64_t col = 0; col < n; ++col) { + double sum = 0; + for (int64_t r = 0; r < 32; ++r) + sum += static_cast(static_cast(t[r])) * static_cast(b[col * 32 + r]); + expected[row * n + col] = T(static_cast(sum)); + } + } + OpTester test("MatMulNBitsLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddAttribute("adapter_transposed", 1); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x88), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, T(0.03125f)), true); + test.AddInput("A", {32, k}, a); + test.AddInput("B", {n, 32}, b); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", std::is_same_v ? 2e-6f : 1e-4f); + test.SetOutputRelErr("Y", std::is_same_v ? 2e-5f : 0.003f); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } +} + +TEST(MatMulNBitsLoRA, ParallelDownNumericsFp16) { RunParallelDownNumerics(); } +TEST(MatMulNBitsLoRA, ParallelDownNumericsFp32) { RunParallelDownNumerics(); } + TEST(MatMulNBitsLoRA, RejectInvalidAdapterShape) { auto ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(ep, nullptr); diff --git a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc index 9aaf6602bbbdb..4aa38acf3bc17 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc @@ -15,6 +15,10 @@ namespace { // Use two FP16 epsilons plus a small floor for values near cancellation. constexpr double kFp16FusionAbsTolerance = 1e-4; constexpr double kFp16FusionRelTolerance = 2.0 / 1024.0; +// The K-parallel down projection uses FP32 partials instead of the unfused +// sequential FP16 sum over K=2048. Allow its accumulated rounding difference +// only for affected cases; ParallelDownNumerics checks a dense FP64 reference. +constexpr double kFp16ParallelDownAbsTolerance = 4.0 / 1024.0; template void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, @@ -143,11 +147,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedNumerics) { for (int64_t n : {128, 132}) { SCOPED_TRACE(m); SCOPED_TRACE(n); + const bool parallel_down = m > 1 && (std::string_view(mode) != "fuse_down" || m < 4 || n % 16 != 0); RunWebGpuFusionTransformerTest( [m, n](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, n); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, TransformerLevel::Level1, TransformerLevel::Level2, 21, - kFp16FusionAbsTolerance, kFp16FusionRelTolerance, + parallel_down ? kFp16ParallelDownAbsTolerance : kFp16FusionAbsTolerance, kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) { @@ -193,11 +198,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedBatchedNumerics) for (int64_t m : {1, 3, 37}) { SCOPED_TRACE(batch); SCOPED_TRACE(m); + const bool parallel_down = std::string_view(mode) != "fuse_down" || m < 4; RunWebGpuFusionTransformerTest( [batch, m](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, 128, 0, 0, batch); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, TransformerLevel::Level1, TransformerLevel::Level2, 21, - kFp16FusionAbsTolerance, kFp16FusionRelTolerance, + parallel_down ? kFp16ParallelDownAbsTolerance : kFp16FusionAbsTolerance, kFp16FusionRelTolerance, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, [mode](SessionOptions& options) { if (mode[0]) { From 139ad24cd663a22dc8a071429359b9f9d238ba29 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 15 Oct 2026 22:37:51 +0800 Subject: [PATCH 15/35] WebGPU: separate floating LoRA from QLoRA support Retain FP16/FP32 MatMulNBitsQLoRA and remove floating-only schema, transformer, kernels and shader branches. Preserve QLoRA dispatch policy and packed adapter arithmetic. Validation: 38 native tests, 32 exact full-model comparisons, WASM syntax checks and lint. Isolated built-in WebGPU before/after timings show no repeatable material regression; record sanitized results. --- docs/ContribOperators.md | 58 -- docs/WebGPU_MatMulNBitsLoRA_Performance.md | 500 ------------------ docs/WebGPU_MatMulNBitsLoRA_Prototype.md | 125 ----- docs/WebGPU_MatMulNBitsQLoRA_Performance.md | 116 ++++ docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 37 +- .../quantization/dp4a_matmul.wgsl.template | 80 +-- .../webgpu/quantization/dp4a_matmul_nbits.cc | 257 ++++----- .../webgpu/quantization/dp4a_matmul_nbits.h | 156 ++---- .../dp4a_matmul_small_m.wgsl.template | 30 +- .../webgpu/quantization/matmul_nbits.cc | 297 +++++------ .../webgpu/quantization/matmul_nbits.h | 97 ++-- .../quantization/matmul_nbits.wgsl.template | 24 +- .../webgpu/quantization/matmul_nbits_lora.cc | 410 -------------- .../webgpu/quantization/matmul_nbits_qlora.cc | 38 +- .../matmul_nbits_wide_tile.wgsl.template | 24 +- .../webgpu/webgpu_contrib_kernels.cc | 1 - .../core/graph/contrib_ops/contrib_defs.cc | 50 -- .../core/optimizer/graph_transformer_utils.cc | 7 - .../optimizer/matmul_nbits_lora_fusion.cc | 199 ------- .../core/optimizer/matmul_nbits_lora_fusion.h | 20 - .../contrib_ops/matmul_nbits_lora_test.cc | 217 -------- .../matmul_nbits_lora_fusion_test.cc | 341 ------------ 22 files changed, 523 insertions(+), 2561 deletions(-) delete mode 100644 docs/WebGPU_MatMulNBitsLoRA_Performance.md delete mode 100644 docs/WebGPU_MatMulNBitsLoRA_Prototype.md create mode 100644 docs/WebGPU_MatMulNBitsQLoRA_Performance.md delete mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc delete mode 100644 onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc delete mode 100644 onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h delete mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc delete mode 100644 onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 9eee6d78725fc..a97ca9dd6b0f1 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -69,7 +69,6 @@ Do not modify directly.* * com.microsoft.MatMulInteger16 * com.microsoft.MatMulIntegerToFloat * com.microsoft.MatMulNBits - * com.microsoft.MatMulNBitsLoRA * com.microsoft.MatMulNBitsMlp * com.microsoft.MatMulNBitsQLoRA * com.microsoft.MatMulNBitsQkv @@ -4106,63 +4105,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.MatMulNBitsLoRA** - - Computes MatMulNBits(X, W, scales) + MatMul(MatMul(X, lora_A), lora_B). With adapter_transposed=1, transposes both stored adapter inputs in this formula. The adapter remains replaceable. A zero rank in either adapter input disables the adapter. Intermediate results are rounded to the input type before addition. Supports FP16 or FP32, rank 32, INT4 block size 32, implicit zero point 8, and no bias. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
K : int (required)
-
Input feature dimension; a positive multiple of 32.
-
N : int (required)
-
Output feature dimension; a positive multiple of 4.
-
accuracy_level : int
-
MatMulNBits base accumulation policy.
-
adapter_transposed : int
-
Whether adapters are stored as [R_A,K] and [N,R_B].
-
bits : int
-
Base weight bit width; must be 4.
-
block_size : int
-
Base weight block size; must be 32.
-
- -#### Inputs - -
-
X : T
-
Input with rank at least two and last dimension K.
-
W : TQ
-
INT4 packed weights of shape [N, K/32, 16].
-
scales : T
-
Block scales containing N * K/32 elements; may be flattened.
-
lora_A : T
-
Adapter of shape [K, R_A], or [R_A, K] if adapter_transposed=1; R_A is 0 or 32.
-
lora_B : T
-
Adapter of shape [R_B, N], or [N, R_B] if adapter_transposed=1; R_B is 0 or 32.
-
- -#### Outputs - -
-
Y : T
-
Input shape with last dimension replaced by N.
-
- -#### Type Constraints - -
-
T : tensor(float16), tensor(float)
-
FP16 or FP32 activations, block scales, and adapter weights.
-
TQ : tensor(uint8)
-
Packed quantized base weights.
-
- - ### **com.microsoft.MatMulNBitsMlp** MatMulNBitsMlp fuses two MatMulNBits projections that share the same input and computes diff --git a/docs/WebGPU_MatMulNBitsLoRA_Performance.md b/docs/WebGPU_MatMulNBitsLoRA_Performance.md deleted file mode 100644 index c02fc3dde4f25..0000000000000 --- a/docs/WebGPU_MatMulNBitsLoRA_Performance.md +++ /dev/null @@ -1,500 +0,0 @@ -# WebGPU LoRA and QLoRA evaluation - -This report compares upstream main with the optimized QLoRA implementation -using one evaluation model and two adapters, labeled A and B. The model uses -FP16 activations, INT4 base weights, and rank-32 INT8 adapters. Model names, -prompt contents, local paths, raw logs, and model assets are omitted. - -## Main versus optimized - -Measurements were collected on 2026-10-14 with native GenAI C++, built-in -WebGPU, Windows D3D12, and an NVIDIA RTX 5080 (driver 610.88). - -**The unchanged original optional-adapter graph fails with the adapter off on -main**, at both prompt lengths: empty adapter defaults reach DequantizeLinear, -which rejects dispatch size `(0, 1, 1)`. Main's **Off*** values below therefore -use a separately pruned base-only graph. Optimized Off uses the actual -optional-adapter graph with empty adapters. These are different off baselines. - -This failure describes the recorded main revision below. Newer main includes -the empty-rank MatMul/DequantizeLinear fix (#33171), now merged into this branch. -The historical table has not been relabeled as measurements of that newer main. - -| Prompt tokens | Adapter | Main TTFT ms | Optimized TTFT ms | TTFT reduction | Main decode tok/s | Optimized decode tok/s | Decode gain | -|---:|---|---:|---:|---:|---:|---:|---:| -| 128 | Off* | 29.86 | 31.18 | -4.4% | 261.49 | 236.06 | -9.7% | -| 128 | A | 97.44 | 40.46 | 58.5% | 115.73 | 205.66 | 77.7% | -| 128 | B | 92.48 | 40.34 | 56.4% | 115.11 | 205.36 | 78.4% | -| 1024 | Off* | 162.67 | 162.98 | -0.2% | 258.23 | 232.99 | -9.8% | -| 1024 | A | 242.25 | 172.87 | 28.6% | 113.52 | 201.19 | 77.2% | -| 1024 | B | 246.09 | 172.97 | 29.7% | 113.64 | 199.22 | 75.3% | - -TTFT includes prompt processing and first-token sampling/readback. Lower TTFT -and higher throughput are better. Negative gains indicate regressions against -the listed reference. Pure prompt-processing throughput is reported separately: - -| Prompt tokens | Adapter | Main prefill tok/s | Optimized prefill tok/s | Gain | -|---:|---|---:|---:|---:| -| 128 | Off* | 4296.0 | 4114.5 | -4.2% | -| 128 | A | 1314.8 | 3170.1 | 141.1% | -| 128 | B | 1385.3 | 3179.5 | 129.5% | -| 1024 | Off* | 6299.4 | 6289.8 | -0.2% | -| 1024 | A | 4228.5 | 5926.8 | 40.2% | -| 1024 | B | 4162.7 | 5925.0 | 42.3% | - -Within the optimized optional-adapter graph, enabling an adapter adds -29.4-29.8% TTFT at 128 prompt tokens and 6.1% at 1024. Decode latency overhead -is 14.8-14.9% and 15.8-16.9%, respectively, computed as -`off tokens/s / on tokens/s - 1`. The 10% decode-overhead target is not met. - -Main's base-only graph receives MatMulNBitsMlpFusion and MatMulNBitsQkvFusion. -The optimized contrib nodes prevent those separate base fusions. Skipping -adapter arithmetic therefore does not give parity with the standalone base -graph: optimized off decode is approximately 10% slower than that reference. - -## Measurement method - -- Main revision: `9c87bdf2e6676eee0073c711a32e75183dc41199`. -- Optimized source tree: `b3e1633c9567459041474114256d389bc824e569`. - Publication changed commit metadata only; the tested source tree is identical. -- Both configurations use the same native benchmark executable, GenAI DLL, and - DXC. The benchmark includes adapter selection and fixed-token inputs. Only - the ORT runtime and model representation differ. -- Batch size 1, identical fixed prompt IDs, greedy decoding, 128 generated - tokens, maximum length 4096, graph capture, generator reuse, and token readback. -- Two rounds per configuration, each with three warmups and five measured - requests. The second round reverses configuration and prompt-length order. - Tables use the median of all ten individual requests. All 120 measured - requests completed with the expected token counts and stable token sequences - within each configuration. -- Measurements run serially after builds finish, with profiling, verbose - logging, shader dumping, and other evaluation work disabled. Each cell checks - for competing build/test/benchmark work and waits for three idle-GPU samples. - Separate diagnostic runs confirm all nodes execute on WebGPU and graph - capture replays; diagnostic timings are excluded. -- Main prefill varies more: adapter-on TTFT ranges are 91.69-107.07 ms (A) and - 91.12-106.52 ms (B) at 128 tokens; 234.83-252.03 ms (A) and 236.42-250.94 ms - (B) at 1024. Optimized ranges are 40.06-41.18 ms and 40.01-41.08 ms at 128; - 172.42-173.09 ms and 172.34-173.23 ms at 1024. - -The original adapters use signed INT8 weights and FP32 scales. Offline -conversion transposes and packs the same integer weights and rounds scales to -FP16; the base external weights remain byte-identical. DP4A prefill and decode -reduction changes also affect arithmetic. This compares the same source -workload after conversion, with the numerical limits below. - -Main includes newer unrelated upstream changes than the optimization branch's -common ancestor. This is an end-to-end configuration comparison, not an -identical-base ablation. Results cover one GPU and model; final performance -measurements cover built-in WebGPU only. Earlier built-in/plugin compatibility -checks do not constitute a plugin rerun of the final optimizations. - -## Correctness and tests - -The model-level check uses six prompts across both adapters, 72 full-vocabulary -logit rows evaluated at identical forced prefixes, adapter on/off, capture -on/off, and repeated generation after rewind. - -| Comparison | Adapter | Identical responses | Matching top-1 rows | Logit RMSE | Maximum absolute difference | -|---|---|---:|---:|---:|---:| -| Final vs previous QLoRA | Off | 6/6 | 72/72 | 0 | 0 | -| Final vs previous QLoRA | On | 6/6 | 71/72 | 0.008171 | 0.0625 | -| Final vs original source model | On | 4/6 | 70/72 | 0.075461 | 1.03125 | -| Capture on vs off | Off | 6/6 | 72/72 | 0 | 0 | -| Capture on vs off | On | 6/6 | 72/72 | 0 | 0 | - -The previous QLoRA implementation is `210584c743785f949cfc417a185366fa0343b932`, -before the final decode register-reuse/subgroup-reduction change. Adapter-on -arithmetic is not bitwise equivalent to the original model. Six prompts are -limited quality coverage, and larger quality evaluation remains necessary. - -Initial optimization validation passed seven operator tests, eight graph-fusion tests, and -two WGSL smoke tests. The decode regression checks against a dense -double-precision reference at K=2048/4096/8192, with mixed signs, varying INT8 -scales, a zero rank, tail columns, and both accumulation options. Quantized -prefill and reordered decode comparisons use absolute tolerance 0.002 and -relative tolerance 0.001; unaffected fusion paths retain zero tolerance. -The input-upload change also has a regression test for chunking and reuse. - -See [FP16 LoRA usage](WebGPU_MatMulNBitsLoRA_Prototype.md) and -[QLoRA format, dispatch, and test commands](WebGPU_MatMulNBitsQLoRA_Prototype.md). - -## CI integration verification - -The CI integration merges main at `a5e205be24` and updates the FP16 original-layout -fallback to the current `MatMulComputeDispatcher` API. Adapter weights remain -nonconstant so that replacement cannot reuse a stale transformed weight. It also -uses explicit `size_t` shape indices for wasm32 and removes plugin-test linkage -to built-in-only capability-inspection symbols. The optimized LoRA/QLoRA shader -code and workgroup sizes are unchanged by these fixes. - -### Performance against the pre-fix optimization - -The following measurements compare the previously tested optimization with the -merged/fixed build, using the same models, native benchmark, GenAI DLL, and DXC. -This includes upstream main changes, not just the compatibility edit. Hardware -and settings match the measurement method above. All other build, test, and -profiling work was stopped during measurement, with idle checks before each cell. -There were two rounds in reversed order, three warmups and five measured -requests per round: 200 measured requests in total. Values are medians of ten -requests. Every before/after generated token sequence matches. - -| Model layout | Prompt tokens | Adapter | Before TTFT ms | Fixed TTFT ms | TTFT change | Before decode tok/s | Fixed decode tok/s | Decode change | -|---|---:|---|---:|---:|---:|---:|---:|---:| -| Packed INT8 | 128 | Off | 31.43 | 31.39 | -0.13% | 228.07 | 226.45 | -0.71% | -| Packed INT8 | 128 | A | 40.63 | 40.28 | -0.88% | 197.45 | 198.80 | +0.68% | -| Packed INT8 | 128 | B | 40.36 | 40.54 | +0.44% | 197.33 | 196.49 | -0.43% | -| Packed INT8 | 1024 | Off | 163.28 | 163.32 | +0.03% | 224.57 | 223.77 | -0.36% | -| Packed INT8 | 1024 | A | 173.95 | 173.51 | -0.25% | 196.29 | 196.00 | -0.15% | -| Packed INT8 | 1024 | B | 174.23 | 173.57 | -0.38% | 195.98 | 194.35 | -0.83% | -| FP16 original | 128 | Off | 31.94 | 31.95 | +0.01% | 225.84 | 221.95 | -1.73% | -| FP16 original | 128 | A | 94.74 | 96.33 | +1.68% | 138.70 | 137.97 | -0.53% | -| FP16 original | 1024 | Off | 163.72 | 163.79 | +0.04% | 222.65 | 222.52 | -0.06% | -| FP16 original | 1024 | A | 230.97 | 230.88 | -0.04% | 137.42 | 138.36 | +0.69% | - -Positive TTFT change is slower; positive decode change is faster. QLoRA changes -are all within 0.9%. The initial FP16 128-token results warranted an independent, -longer confirmation: four rounds with alternating order, three warmups and eight -measured requests per cell (32 requests per variant/adapter; 128 requests total). -These additional requests use the same settings and also match token-for-token. - -| FP16, 128 tokens | Before TTFT ms | Fixed TTFT ms | TTFT change | Before decode tok/s | Fixed decode tok/s | Decode change | -|---|---:|---:|---:|---:|---:|---:| -| Off | 31.82 | 31.82 | -0.00% | 222.76 | 223.70 | +0.42% | -| A | 94.78 | 95.26 | +0.51% | 138.78 | 138.44 | -0.25% | - -In this confirmation, adapter-on TTFT ranges overlap: 94.10-100.15 ms before and -93.68-101.22 ms after. Adapter-off decode ranges also overlap: 218.36-228.44 and -219.64-228.64 tok/s. These measurements show no material performance regression -on this GPU, with residual sub-1% median differences in the confirmation and -QLoRA matrix. They do not establish zero change or results on other devices. - -### Build and correctness coverage - -- Built-in and dynamic-plugin Release builds and test binaries compile and link. -- Built-in: 8 operator/upload tests and 19 graph-fusion tests pass, with no skips. -- Plugin: 5 operator tests and 19 graph-fusion tests pass. Two hardware-specific - specialization checks explicitly skip because the plugin test binary cannot - inspect internal device capabilities. Both checks run and pass in built-in - mode; general plugin operator/fusion tests remain active. -- Two WGSL smoke tests pass. Six changed WebGPU translation units also pass - wasm32 syntax compilation with Emscripten 4.0.23 and narrowing warnings treated - as errors. This is not a full linked WebAssembly build. -- Full-model QLoRA checks against the pre-fix runtime are bitwise exact for - adapter off and on: 6/6 responses and all 72 forced-prefix logit rows per state. - Built-in/plugin, capture on/off, and generation after rewind also match exactly - in the 12-case matrix. This verifies the CI fixes preserve the optimized - runtime's outputs; it does not change the original-model numerical limits above. -- Formatting and whitespace checks pass. Generated contrib documentation now - includes both new operator schemas. Remote CI must validate other build targets. - -## Review-feedback verification (2026-10-15) - -Merged main at `9cdf0f11fde12c171eceb4e4ca3b31be8777e7f0`, including the -empty-rank MatMul/DequantizeLinear fix (#33171) and graph-capture buffer-copy -fix (#33362). The LoRA review-feedback patch changes tests and documentation; -it does not change LoRA shaders or workgroup sizes. - -The following comparison uses the previously validated PR runtime snapshot -(`7e7079934a`) as Before. Updated uses the merged ORT build and GenAI -`b4023152` (adapter-aware capture identity, based on GenAI main `bdfadcf4`). -These are update-regression measurements, not a new main-versus-fusion study. - -Native C++ benchmark, built-in WebGPU / Windows D3D12 / NVIDIA RTX 5080: -batch 1, fixed prompt IDs, 128 generated tokens, maximum length 4096, greedy -decoding, capture, generator reuse, and token readback. Profiling and competing -build/test work were stopped throughout measurement. Each cell has two rounds -of three warmups and five measured requests, with reversed order in round two. -The table reports medians over ten individual requests per runtime/cell. -All 160 measured requests completed and generated identical sequences across -the two runtimes for each case. - -| Model layout | Prompt tokens | Adapter | Before TTFT ms | Updated TTFT ms | TTFT change | Before decode tok/s | Updated decode tok/s | Decode change | -|---|---:|---|---:|---:|---:|---:|---:|---:| -| Packed INT8 | 128 | Off | 31.48 | 31.43 | -0.17% | 228.69 | 233.12 | +1.94% | -| Packed INT8 | 128 | On | 40.17 | 40.31 | +0.34% | 205.68 | 205.60 | -0.04% | -| Packed INT8 | 1024 | Off | 162.62 | 162.52 | -0.06% | 233.16 | 233.74 | +0.25% | -| Packed INT8 | 1024 | On | 172.67 | 174.30 | +0.95% | 204.11 | 197.54 | -3.22% | -| FP16 original | 128 | Off | 31.95 | 31.90 | -0.17% | 229.53 | 229.59 | +0.03% | -| FP16 original | 128 | On | 97.34 | 96.19 | -1.18% | 140.95 | 141.15 | +0.14% | -| FP16 original | 1024 | Off | 163.49 | 163.97 | +0.29% | 226.22 | 219.26 | -3.07% | -| FP16 original | 1024 | On | 232.54 | 233.87 | +0.57% | 137.72 | 138.31 | +0.43% | - -Positive TTFT change means slower; positive decode change means faster. -The two approximately 3% decode slowdowns prompted a longer, independent -confirmation separating the ORT merge from the GenAI update. This used two -rounds with reversed order, six warmups and eight measured requests each: -16 measurements per runtime/case, 96 measured requests total. All sequences -matched across the three runtimes. - -| 1024-token case | Runtime | TTFT ms | Decode tok/s | Decode range tok/s | -|---|---|---:|---:|---:| -| FP16, off | Previous | 164.59 | 219.34 | 217.18-221.37 | -| FP16, off | Merged ORT / previous GenAI | 163.76 | 218.64 | 214.75-221.01 | -| FP16, off | Merged ORT / fixed GenAI | 163.91 | 219.66 | 216.30-224.18 | -| Packed INT8, on | Previous | 175.74 | 198.31 | 197.07-199.59 | -| Packed INT8, on | Merged ORT / previous GenAI | 176.12 | 198.22 | 195.54-199.43 | -| Packed INT8, on | Merged ORT / fixed GenAI | 176.00 | 197.80 | 195.55-198.86 | - -In this confirmation, final versus previous decode changes are +0.14% for -FP16 adapter-off and -0.25% for packed INT8 adapter-on. TTFT changes are -0.41% -and +0.15%, respectively. The earlier decode slowdowns did not repeat. These -results show no repeatable regression on the tested GPU; they do not establish -zero change or performance on other devices. The initial measurements remain -above rather than being discarded. This does not resolve the pre-existing -approximately 10% adapter-off gap versus a graph with QKV/MLP base fusion. - -### Correctness and capture - -- Built-in and plugin Release builds pass. Built-in provider/empty-tensor tests: - 23 passed; fusion tests: 20 passed; graph-capture buffer-copy tests: 6 passed. -- Plugin provider tests: 5 passed, with 2 expected hardware-inspection skips; - fusion tests: 20 passed. Both skipped specializations pass in built-in mode. -- Empty-default serialization/reload tests cover decode and prefill, full and - partial overrides, and separate on/off capture IDs. Both actual captures are - asserted; alternating replay overwrites sentinel outputs correctly with CPU - fallback disabled. Empty packed defaults have no serialized payload. -- Reported RTX 4070 tolerance failures are addressed with narrowly scoped - tolerances: K=8192 FP16 lane accumulation allows 0.003 absolute error; FP16 - fusion comparisons allow 0.0001 absolute plus two FP16 epsilons relative. - Shorter K, FP32 accumulation, and deliberately exact switching/cancellation - checks keep their previous bounds. RTX 4070 was not available locally. -- GenAI previously reused a base capture after late adapter activation. In a - full-model forced-prefix reproduction, capture-off results matched a fresh - active generator exactly, while captured decode differed by up to 2.82764. - With [GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740), - all six full-vocabulary logit rows match exactly for built-in capture off/on - and plugin capture on. The added Python regression passes both WebGPU - parameterizations; its CUDA variants were not run locally. -- Formatting and whitespace checks pass. Wider device/build validation remains - the responsibility of fresh CI and hardware testing. - -## FP16 and FP32 support verification (2026-10-15) - -Both `MatMulNBitsLoRA` and `MatMulNBitsQLoRA` now accept FP16 and FP32 models. -All floating inputs to an operator share one type, including base scales, -adapter scales or floating adapter weights, intermediate T, and output. -LoRA stores floating adapter weights; QLoRA stores packed INT8 adapter weights. -Cross-type adapter loading is not implicit. Rank 32 and INT4 block-32 base -restrictions remain. FP32 decode preserves the ordinary base DP4A policy even -when the adapter is active. - -### Validation - -- Built-in Release: 28 provider/empty-tensor and 24 fusion tests pass. -- Plugin Release: 10 provider and 24 fusion tests pass; 2 hardware-inspection - tests skip in plugin mode and pass with the built-in provider. -- FP32 dense FP64-reference tests cover prefill/decode, batching, partial N, - K through 8192, accuracy levels 0/4, and all fusion modes, with 1e-5 - absolute/relative bounds on inputs exactly representable by DP4A. -- Both original and transposed LoRA layouts, empty/full/partial adapters, - QLoRA serialization/reload and distinct on/off capture replay are tested. -- The two changed kernel files pass wasm32 syntax checks. Formatting and - whitespace checks pass. This is not a full linked WebAssembly build. - -### FP32 full-model numerical checks - -Two adapters, three prompts each, up to 64 greedy tokens, and 12 forced-prefix -full-vocabulary logit rows per prompt were compared against the original -unfused package. The converted base data is unchanged; both conversions -preserve the original effective FP32 adapter weights exactly. - -| FP32 path | Adapter | Max logit difference | Max per-case RMSE | Top-1 rows match | Greedy outputs match | -| --- | --- | --- | --- | --- | --- | -| LORA | Off | 0 | 0 | 72/72 | 6/6 | -| LORA | On | 0.68059158 | 0.079206155 | 70/72 | 2/6 | -| QLORA | Off | 0 | 0 | 72/72 | 6/6 | -| QLORA | On | 1.3800392 | 0.12060427 | 71/72 | 3/6 | - -These are the default `accuracy_level=4` results. Adapter-off logits match -exactly. Adapter-on differences include reduction-order changes and DP4A -activation quantization; they are not a bitwise-equivalence claim. QLoRA -`fuse_down`, which avoids quantizing the adapter down input, still differs by -up to 0.918821, with 71/72 top-1 rows and 5/6 greedy outputs matching, because -the base path still uses DP4A. The difference must not be attributed solely -to the new adapter down projection. - -To isolate floating arithmetic, diagnostic copies set `accuracy_level=0` -for both reference and fused models: - -| FP32 path, accuracy=0 | Max logit difference | Max per-case RMSE | Top-1 rows match | Greedy outputs match | -| --- | --- | --- | --- | --- | -| LORA | 3.0517578e-05 | 3.3701491e-06 | 72/72 | 6/6 | -| QLORA | 2.3409724e-05 | 2.7866782e-06 | 72/72 | 6/6 | - -Built-in capture off/on and plugin capture on produce exactly equal logits -for each fused representation. Rewind/repeat generation is identical. -These checks use the GenAI adapter capture-identity fix from #2740. -This small prompt set is a numerical regression check, not a broad quality -evaluation. - -### FP32 throughput - -Native GenAI C++, built-in WebGPU, Windows D3D12, NVIDIA RTX 5080. Batch 1, -fixed identical prompt IDs, 128 generated tokens, maximum length 4096, -graph capture, generator reuse and token readback. Each cell uses three -warmups and five measured requests; values are medians. No builds, tests, -profiling or other GPU workloads run during timing. The unfused baseline -uses the same updated runtime; this is not a new upstream-main comparison. - -| Prompt tokens | Adapter | FP32 path | TTFT ms | Decode tok/s | -| --- | --- | --- | --- | --- | -| 128 | Off | Unfused | 31.48 | 174.42 | -| 128 | Off | LoRA | 29.56 | 201.51 | -| 128 | Off | QLoRA | 29.67 | 200.70 | -| 128 | A | Unfused | 77.99 | 37.62 | -| 128 | A | LoRA | 123.84 | 142.93 | -| 128 | A | QLoRA | 43.46 | 133.93 | -| 128 | B | Unfused | 77.59 | 37.68 | -| 128 | B | LoRA | 123.56 | 143.19 | -| 128 | B | QLoRA | 43.57 | 133.50 | -| 1024 | Off | Unfused | 174.11 | 171.00 | -| 1024 | Off | LoRA | 162.04 | 198.04 | -| 1024 | Off | QLoRA | 162.34 | 199.26 | -| 1024 | A | Unfused | 227.59 | 37.52 | -| 1024 | A | LoRA | 270.12 | 142.24 | -| 1024 | A | QLoRA | 192.25 | 131.47 | -| 1024 | B | Unfused | 227.28 | 37.69 | -| 1024 | B | LoRA | 270.76 | 141.47 | -| 1024 | B | QLoRA | 192.22 | 132.22 | - -FP32 QLoRA improves both prefill and decode over this unfused package. FP32 -LoRA improves decode but its current floating down-projection shader makes -prefill slower than the original package; this support change does not solve -that performance limitation. The two representations use different shaders -and need not have identical speed or floating-point reduction order. - -### FP16 before/after regression check - -The previous validated PR runtime and this update use the same fixed GenAI -runtime and the same converted FP16 models. Token IDs match between revisions -in every measured cell. - -| FP16 adapter | Prompt tokens | Adapter | TTFT ms before → after | Change | Decode tok/s before → after | Change | -| --- | --- | --- | --- | --- | --- | --- | -| LORA | 128 | Off | 31.24 → 31.22 | -0.06% | 237.29 → 235.36 | -0.82% | -| LORA | 128 | A | 81.33 → 81.42 | +0.12% | 205.54 → 204.45 | -0.53% | -| LORA | 1024 | Off | 162.74 → 162.58 | -0.10% | 234.15 → 233.27 | -0.38% | -| LORA | 1024 | A | 223.72 → 223.39 | -0.15% | 204.52 → 205.28 | +0.37% | -| QLORA | 128 | Off | 31.33 → 31.32 | -0.03% | 235.81 → 236.37 | +0.24% | -| QLORA | 128 | A | 40.86 → 40.35 | -1.26% | 204.94 → 202.06 | -1.41% | -| QLORA | 1024 | Off | 162.93 → 162.60 | -0.21% | 233.84 → 234.40 | +0.24% | -| QLORA | 1024 | A | 172.54 → 172.80 | +0.15% | 204.50 → 204.34 | -0.08% | - -Across these measured cells, the largest TTFT increase is 0.15% and the -largest decode decrease is 1.41%. No large FP16 regression appears in this -single-round check; smaller changes remain within typical run variation. - -## Joint DP4A FP32 QLoRA decode (2026-10-15) - -Eligible NVIDIA FP32 decode now uses three dispatches: quantize X once, -jointly compute the INT4 base and INT8 down projections, then apply the -vectorized INT8 up projection and addition. This replaces the separate -floating-input down projection in the previous four-dispatch implementation. -The base workgroup size remains 128. Eight extra workgroups per input row -compute rank 32 from the shared quantized activation and scale buffers. - -The baseline is the preceding FP16/FP32 support implementation, not upstream -main. These native C++ measurements use built-in WebGPU/D3D12 on RTX 5080, -batch 1, accuracy level 4, capture and generator reuse, 128 generated tokens, -three warmups, and five measured requests. Values are medians. Timing runs -were isolated from builds, tests, profiling, and verbose diagnostics. - -| Prompt | Adapter | TTFT before → after (ms) | Decode before → after (tok/s) | Decode change | -|---:|---|---:|---:|---:| -| 128 | Off | 29.61 → 29.60 | 203.20 → 203.10 | -0.05% | -| 128 | A | 43.64 → 43.42 | 134.55 → 180.44 | +34.11% | -| 1024 | Off | 161.88 → 162.15 | 199.03 → 199.29 | +0.13% | -| 1024 | A | 192.12 → 195.07 | 129.82 → 174.29 | +34.26% | - -The adapter-on throughput deficit relative to adapter-off is now 11.2% and -12.5%, respectively. Control measurements of FP16 QLoRA, FP16 LoRA, and -FP32 LoRA at both prompt lengths show absolute decode changes below 0.9% -and TTFT changes below 0.9%. Across all measured cells the largest TTFT -increase is 1.53%; TTFT includes the first generated token. - -Both adapters were checked on six full-model cases against the previous -FP32 QLoRA runtime. All six generated sequences and 72 forced-prefix next-token -choices match. Adapter-on decode logits are approximate: maximum absolute -difference 0.822086 and largest per-case RMSE 0.087801. This follows from -using quantized X in the down projection; it is not a bitwise-equivalence -claim or a broad quality evaluation. Adapter-off logits, prefill logits, -and capture OFF/ON logits are identical. All 12 before/after benchmark -output sequences also match. - -Validation includes the provider/empty-tensor suite, expanded FP32 numerical -coverage, an independent dense reference on a known activation quantization -grid, and all 12 QLoRA fusion tests. The new cases cover batched decode, -zero rows, K tails, the full INT8 byte range, and distinct adapter-on/off -capture IDs with replay. Separate dispatch logging confirms quantize → -joint base/down → up/add. Both changed kernel translation units pass WASM32 -syntax checks. Only the built-in EP was tested for this experiment. - -## Floating LoRA prefill and decode follow-up - -The baseline for this experiment includes the three-dispatch QLoRA change above. -The same optional graphs and floating adapters were used before and after. The -native built-in WebGPU measurements use batch 1, 128 generated tokens, fixed -prompt IDs, capacity 4096, capture/reuse and token readback, three warmups and -five measured requests. Values are medians on RTX 5080/D3D12. Timing ran alone; -profiling and correctness checks ran separately. A reversed runtime ordering -confirmed the baseline after the candidate measurements. - -| Type | Prompt | Off TTFT (ms) | LoRA TTFT before → after (ms) | QLoRA TTFT (ms) | LoRA decode before → after (tok/s) | QLoRA decode (tok/s) | -|---|---:|---:|---:|---:|---:|---:| -| FP16 | 128 | 31.30 | 83.72 → 46.74 | 40.56 | 200.04 → 197.53 | 195.13 | -| FP16 | 1024 | 162.84 | 225.95 → 190.35 | 172.53 | 197.64 → 195.02 | 195.08 | -| FP32 | 128 | 29.77 | 123.51 → 59.69 | 45.18 | 138.98 → 165.35 | 170.26 | -| FP32 | 1024 | 162.09 | 270.70 → 220.47 | 192.53 | 136.07 → 162.33 | 167.47 | - -FP16 adapter-off decode measured 225.01/223.75 tok/s, and FP32 measured -195.75/192.21 tok/s. The floating LoRA decode deficit is now 12–16%, close -to QLoRA's 11–13% on its separately measured adapter-off controls. FP16 -batch-one decode uses its existing algorithm; its measured change was -1.3%. -Adapter-off controls changed by at most 1.3% in decode and 0.6% in TTFT. -QLoRA controls showed no regression above 1.7% in TTFT or 1% in decode. -These control differences are timing variation, not claimed improvements. - -The retained changes are: - -- A K-parallel floating down projection for NVIDIA devices with fixed - 32-lane subgroups. A 128-thread workgroup computes four ranks for four - rows, reusing each vec4 weight across rows. FP32 partials are reduced - with subgroup operations and cast once to T's model dtype. -- Shared transposed B tiles in the FP16 prefill DP4A base/up epilogue. - This preserves the existing workgroup size and arithmetic order. - The FP32 shared-tile candidate did not help and is not enabled. -- FP32 decode uses quantize X → joint base/down → up/add. Each floating - rank gets one 128-thread workgroup, four subgroup reductions, and one - workgroup barrier. The down projection reads original floating X; - only the base projection uses quantized X. Eligibility retains ordinary - DP4A checks and additionally requires NVIDIA, fixed 32-lane subgroups, - M=1, and K=2048–8192. Explicit `fuse_up` retains its existing routing. - -Prefill down-projection GPU time fell from 36.97/37.39 ms to 4.63/12.01 ms -for FP16 at 128/1024 tokens, and from 68.64/71.11 ms to 4.39/17.17 ms for -FP32. These are hardware timestamp sums from separate uncaptured profiles, -not end-to-end latencies. Full floating prefill still trails QLoRA. A -private diagnostic that moved identical adapter values into initializers -saved about 11 ms (FP16) and 21 ms (FP32) on short prompts, showing substantial -per-request input/upload cost outside the shader timestamps. Persistent -device adapter storage remains separate work requiring correct ownership, -device identity, unload, and capture invalidation handling. - -Validation passed 15 provider tests and 25 fusion tests, including new dense -FP64 references for both floating types, off-grid X, zero rows, batched -decode, K tails, and 128/1024-row prefill. Two affected FP16 fusion -comparisons use a scoped 4/1024 absolute bound for changing a sequential -FP16 reduction to FP32 partial sums; independent dense-reference tests -retain tighter bounds. Native build, lint, and WASM32 syntax checks passed. - -With adapters A/B on six full-model cases per dtype, five of six generated -answers were unchanged for each dtype; the changed answers retained their -facts with different wording. Forced-prefix next-token agreement was -72/72 for FP16 and 70/72 for FP32. Maximum logit differences were 0.689453 -and 0.911110, with largest per-case RMSE 0.071453 and 0.111359. Adapter-off -logits and capture OFF/ON logits were exactly equal, and rewind checks passed. -In an additional FP32 diagnostic with base activation quantization disabled, -the maximum logit difference dropped to 0.0000253 and all 72/72 choices -matched. This supports DP4A threshold amplification of small upstream -rounding changes. The checks are not a broad model-quality evaluation. diff --git a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsLoRA_Prototype.md deleted file mode 100644 index de536e884803f..0000000000000 --- a/docs/WebGPU_MatMulNBitsLoRA_Prototype.md +++ /dev/null @@ -1,125 +0,0 @@ -# WebGPU MatMulNBitsLoRA prototype - -This opt-in prototype replaces `MatMulNBits(X, W) + MatMul(MatMul(X, A), B)` with -one `com.microsoft.MatMulNBitsLoRA` node during session optimization. It retains -the packed base weights and the replaceable FP16 or FP32 adapter inputs. It does not -merge or requantize the adapter into the base weights, or change the source -model file. - -The current matcher supports FP16 and FP32 rank-32 adapters, INT4 base weights with block -size 32, an implicit zero point, no base bias, and no separate adapter scaling -node. The down projection must use the same activation as the base projection, -and intermediate outputs must have no other consumers or graph outputs. - -All floating tensors in one operator share a type: FP16 models use FP16 adapter -weights and base scales; FP32 models use FP32 adapter weights and base scales. -There are no implicit casts between those types. `MatMulNBitsQLoRA` has the same -floating-type contract but stores adapter weights as INT8 with scales matching -the model type. Both operators support adapter on/off and prefill/decode. - -## Enable in native GenAI - -Add this entry to the decoder's `session_options`, or supply it in the native -benchmark's configuration overlay: - -```json -{ - "model": { - "decoder": { - "session_options": { - "session.enable_matmul_nbits_lora_fusion": "1" - } - } - } -} -``` - -Use the original model directory and original `.onnx_adapter` file. The -transformer is registered only when the selected WebGPU provider advertises -the new kernel. Both the ORT host and the WebGPU plugin must come from this -branch for dynamic plugin use. Setting the option to `0` restores the original -graph for a matched comparison. - -## Execution - -With an adapter, the kernel computes `T = X * A`, then incorporates `T * B` -and the addition into the base kernel's output computation. The small -intermediate `T` uses the model type. The base result and adapter result are rounded -separately before addition. FP16 NVIDIA decode uses two GPU dispatches per projection. -DP4A prefill retains the existing activation quantization dispatch, so it uses -three. Subgroup matrix base kernels use a separate output epilogue fallback. - -For offline-transposed adapters, both fusion modes consume the same `At[32,K]` -and `Bt[N,32]` tensors directly. The optimizer absorbs the adapter transpose -wrappers; neither mode dispatches an adapter transpose. Select the mode with -`session.webgpu.lora_fusion_mode` in the same decoder session options: - -| Mode | First dispatch | Second dispatch | -|---|---|---| -| `fuse_down` | Base + `T = X * At^T` | `Y = base + T * Bt^T` | -| `fuse_up` | `T = X * At^T` | Base + `Y = base + T * Bt^T` | - -The default setting, `auto`, selects `fuse_up` for prefill (`M > 1`) and -`fuse_down` for decode (`M == 1`) on each invocation. Here `M` is the -second-to-last dimension of X, so batched one-token decode also selects -`fuse_down`. A one-token prompt uses the same path as decode. This setting -chooses between the two existing implementations; it adds no shader or GPU -dispatch. Explicit `fuse_down` and `fuse_up` settings force that path for all -input sizes and remain available for comparisons. - -On NVIDIA decode (one input row, K from 2048 through 8192), when the base uses -floating arithmetic, `fuse_down` uses -`MatMulNBitsLoRADown` with eight additional workgroups that calculate T once, -followed by `LoRAOutputVec4`. `fuse_up` uses `LoRADownGemv`, followed by -`MatMulNBitsLoRAUp`. Both use the same FP32 GEMV reduction and model-type T, with -128 threads per producer workgroup. The standalone output shader retains -64 threads and one output per thread. Both output implementations load Bt and -T as vec4 along rank, preserving scalar rank accumulation order and separate -model-type rounding before the base/adapter addition. FP32 decode can select -DP4A for the base even at M=1. That case keeps the ordinary base DP4A dispatch -and computes T separately, preserving the base activation-quantization policy. - -DP4A prefill retains its activation quantization dispatch. `fuse_down` adds -the tiled down projection to the base dispatch; `fuse_up` runs that same down -projection separately and integrates the vectorized up projection into DP4A. -Other shapes use the existing base kernels, with a separate output fallback -when the selected kernel has no integrated up projection. The mode selection -applies to transposed adapters; original-layout models retain their compatible -down-then-base/up implementation. Empty adapters execute only the base kernel -in either mode. Unsupported mode strings produce an error at initialization. - -For an overridable pair of private, all-zero default initializers, the -transformer replaces the defaults in the session graph with empty rank -dimensions and relaxes those graph-input dimensions. Input names, element -types, and non-rank dimensions remain intact. A zero rank in either adapter -input selects the ordinary base kernel directly. The zero check happens once -during optimization; no adapter data is read back from the GPU during inference. -Nonzero or shared defaults are preserved and execute the adapter calculation. - -Eliding zero default branches assumes finite activations. IEEE NaN/Inf -propagation through multiplication by zero can differ when that branch is -removed, so this optimization is disabled by default. - -Adapters can be replaced between requests on the same model. Captured execution -must use a distinct capture ID when adapter bindings or the active/off state -change, as with ordinary WebGPU graph capture. GenAI generators own separate -capture IDs; create a new generator when selecting another adapter and use -rewind only while retaining that generator's adapter selection. - -## Tests - -Build the native ORT host and provider test targets, then run from the binary -directory: - -```text -onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsLoRA.* -onnxruntime_test_all.exe --gtest_filter=*MatMulNBitsLoRA* -``` - -The tests cover empty and partially supplied adapters, numerical equivalence -with the original subgraph, adapter replacement, overridable input retention, -invalid shapes, and patterns that must not fuse. Native GenAI evaluation must -also check capture on/off, generator reuse, and the complete model before -interpreting performance results. In a correctness harness, use -`ConfigSetProviderOption` to change `enableGraphCapture`; merely appending a -provider option through a JSON overlay can retain the model's earlier setting. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md new file mode 100644 index 0000000000000..60bc4593a9cc9 --- /dev/null +++ b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md @@ -0,0 +1,116 @@ +# WebGPU quantized LoRA evaluation + +`MatMulNBitsQLoRA` keeps rank-32 adapters replaceable while combining an INT4 +base projection with packed INT8 adapter projections. Both FP16 and FP32 +models are supported, with scales and intermediate tensors matching the +model's floating type. See the [operator guide](WebGPU_MatMulNBitsQLoRA_Prototype.md) +for layouts, opt-in graph matching, explicit export, and capture requirements. + +## QLoRA-only split regression (2026-10-15) + +A fresh paired comparison against the saved pre-split runtime (`0243e40267`) +uses the benchmark protocol below. Only the ORT DLL changes; GenAI, DXC, models +and adapters are identical. Builds, tests and profiling finished before timing. + +| Type | Prompt | Adapter | TTFT before / after (ms) | Change | Decode before / after (tok/s) | Change | +|---|---:|---|---:|---:|---:|---:| +| FP16 | 128 | off | 31.30 / 31.32 | +0.04% | 224.34 / 227.71 | +1.50% | +| FP16 | 128 | on | 40.48 / 40.20 | -0.70% | 205.08 / 206.54 | +0.71% | +| FP16 | 1024 | off | 163.33 / 163.04 | -0.18% | 222.47 / 225.62 | +1.42% | +| FP16 | 1024 | on | 173.43 / 173.15 | -0.16% | 194.33 / 193.15 | -0.61% | +| FP32 | 128 | off | 29.83 / 29.73 | -0.33% | 195.47 / 192.94 | -1.29% | +| FP32 | 128 | on | 43.41 / 43.51 | +0.25% | 180.45 / 179.47 | -0.54% | +| FP32 | 1024 | off | 162.45 / 162.55 | +0.06% | 192.53 / 192.49 | -0.02% | +| FP32 | 1024 | on | 195.10 / 192.97 | -1.09% | 166.45 / 169.01 | +1.54% | + +No repeatable material regression was observed: TTFT changed by -1.09% to ++0.25%, and decode throughput by -1.29% to +1.54% in these cells. Adapter-on +checks at 512/2048 tokens remained within 1.42%. A noisy 32-token first pass +showed a 2.48% FP16 decode loss and an 8.33% FP32 gain; a reverse-order repeat +measured -0.62%/-1.35% decode changes and -0.19%/+1.55% TTFT changes for +FP16/FP32. The experiment supports preserving performance, not a claimed speedup. +Run-to-run absolute rates varied, so compare paired runtimes rather than +historical absolute timings. + +All 38 native QLoRA/fusion/base regression tests passed with existing tolerances. +All 32 full-model before/after comparisons had exact logits and greedy tokens, +covering both types, two adapters, adapter on/off, and capture on/off. Native +build, Emscripten syntax checks and lint passed. These exact comparisons validate +the split; they do not remove the original-versus-optimized accuracy caveats +below. Original model/adapter hashes are unchanged. + +## Performance reference + +Native GenAI C++, built-in WebGPU, Windows D3D12, RTX 5080. Batch 1, fixed +prompt IDs, 128 generated tokens, capacity 4096, accuracy level 4, graph +capture, generator reuse, and token readback. Three warmups and five measured +requests per cell; medians. Timings were isolated from builds, tests, profiling, +and competing GPU evaluation. These are saved measurements of the QLoRA +implementation before floating-adapter support was separated into a follow-up. + +| Model type | Prompt tokens | TTFT off / on (ms) | Decode off / on (tokens/s) | +|---|---:|---:|---:| +| FP16 | 128 | 31.30 / 40.28 | 235.74 / 204.78 | +| FP16 | 1024 | 163.10 / 172.80 | 232.17 / 204.25 | +| FP32 | 128 | 29.62 / 43.47 | 202.64 / 180.06 | +| FP32 | 1024 | 162.50 / 192.40 | 198.13 / 176.73 | + +TTFT includes prompt processing and first-token sampling/readback. Adapter-off +uses the same optional graph's empty-adapter bypass. The measured TTFT overhead +is 28.7%/5.9% for FP16 and 46.7%/18.4% for FP32 at 128/1024 tokens. Decode +throughput loss is 10.8-13.1%; per-token latency overhead is 12.1-15.1%. +These are local device measurements, not a guarantee for other GPUs. + +An earlier paired comparison against upstream main measured FP16 adapter-on +TTFT reductions of 56-59% at 128 tokens and 29-30% at 1024 tokens, with decode +throughput increasing 75-78%. Those runs used two adapters and ten samples per +configuration on earlier runtime revisions. The original adapter-off graph +failed on that main revision because empty DequantizeLinear was unsupported; +its published base reference was a separately pruned graph. Later main fixes +addressed that empty-tensor issue. The optimized optional graph's adapter-off +decode remained about 10% below the separate base graph with QKV/MLP fusion. +Skipping adapter arithmetic does not recover those composed base fusions. + +## Retained execution paths + +- Prefill defaults to `fuse_up`: compute T, then combine the base projection + with the INT8 up projection and Add. Eligible DP4A projections share one + quantized X allocation. The narrow down projection reuses the small-M DP4A + kernel with 32 output columns. +- Eligible FP16 decode defaults to `fuse_down`: one dispatch computes the base + and down projections in independent workgroups, followed by up/Add. Packed + adapter weights are dequantized as needed; X stays floating point. +- Eligible FP32 decode uses quantize X, joint DP4A base/down, then up/Add. + Quantized X and its scales are reused by the two projections. This increased + decode throughput by approximately 34% against the preceding four-dispatch + FP32 QLoRA implementation in the paired local experiment. +- Empty adapters return through the ordinary base implementation before + adapter allocation or dispatch. Unsupported specializations retain existing + kernels, including subgroup-matrix up projections where selected. + +Workgroup sizes, reduction order, quantization, and dispatch selection are +preserved when removing floating-adapter-only branches. The deferred row-tile +experiment is not included. + +## Numerical interpretation and coverage + +Accuracy level 4 permits DP4A activation quantization. Quantization and changed +reduction order can alter logits and generated tokens relative to the original +unfused package; these performance results do not assert bitwise equivalence. +In the FP32 support evaluation, adapter-off logits were exact. Adapter-on +comparison with the original package reached a maximum logit difference of +1.3800392, 71/72 forced-prefix next-token matches, and 3/6 identical generated +sequences. With activation quantization disabled in both configurations, the +maximum difference was 0.0000234097 and all tested choices and sequences matched. + +The subsequent joint DP4A decode change matched all six generated sequences +and all 72 forced-prefix choices against its preceding QLoRA implementation; +the maximum decode-logit difference was 0.822086. This is a small numerical +regression set, not a broad model-quality evaluation. Adapter-off and prefill +logits were exact across that decode change. + +Provider and fusion coverage includes FP16/FP32, dense FP64 references, K tails, +batching, prefill/decode, empty and partial adapters, compact serialization, +unsupported patterns, and separate adapter-on/off capture IDs. Full-model +capture OFF/ON and generator rewind checks use the GenAI adapter capture-identity +fix from [GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740). diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 6b36666a87cfd..3f22bd2cb51ae 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -5,9 +5,19 @@ projections, and their residual Add. Adapter weights remain replaceable. This is an inference operator for quantized adapters; it does not merge adapter weights into the base or perform adapter training. -Enable it with `session.enable_matmul_nbits_lora_fusion=1` in the native GenAI -decoder session options, as described in [the LoRA prototype](WebGPU_MatMulNBitsLoRA_Prototype.md). -The same option selects the appropriate transformer for floating-point or INT8 adapters. +Enable graph matching with `session.enable_matmul_nbits_lora_fusion=1` in +the decoder session options. For native GenAI, merge these entries into +the model's existing `model.decoder.session_options` object: + +```json +{ + "session.enable_matmul_nbits_lora_fusion": "1", + "session.webgpu.lora_fusion_mode": "auto" +} +``` + +This transformer handles packed INT8 adapters. An explicitly exported +`MatMulNBitsQLoRA` node does not require the graph-matching option. Both built-in and plugin WebGPU are supported; plugin use requires the host and provider library from this branch. Disable the option for a comparison with the original graph. @@ -117,7 +127,7 @@ Empty A rank or empty B rank selects the ordinary base kernel before adapter allocation or dispatch. Nonzero and shared packed defaults retain their data. Input names and the adapter file stay unchanged. -Like the LoRA prototype, eliding a zero branch assumes finite activations, +Eliding a zero branch assumes finite activations, adapter scales, and adapter intermediates. NaN/Inf propagation through zero multiplication can differ; the transformer remains opt-in. GenAI builds without adapter-aware capture identity require separate generators when changing the adapter selection. @@ -156,7 +166,7 @@ packed graph input as dynamic, so a rank-32 adapter can override its empty initializer. Either empty packed weight disables the adapter. Scale-only and single-weight overrides therefore continue to produce the base result; supplying both packed weights uses the supplied scales or their retained defaults. -Saving a graph after offline LoRA fusion also produces this representation. +Saving a graph after offline QLoRA fusion also produces this representation. This compact export requires a runtime with the custom operator. Disabling fusion does not convert it to a portable fallback graph. Do not place empty @@ -185,9 +195,9 @@ initializer defaults would not handle later adapter overrides. Adapter-off execution skips the adapter projections but calls a standalone base MatMulNBits. The current QKV/MLP transformers do not compose with the -LoRA operators, so this can lose the throughput benefit of those base fusions. +QLoRA operator, so this can lose the throughput benefit of those base fusions. Empty defaults alone do not resolve that difference. A follow-up composed -QKV/MLP+LoRA implementation should dispatch the existing fused base kernels +QKV/MLP+QLoRA implementation should dispatch the existing fused base kernels when every relevant adapter is empty, and preserve per-projection adapter semantics when any adapter is active. Merely reordering the transformers does not provide both behaviors. @@ -197,8 +207,8 @@ not provide both behaviors. Run from the native build output directory: ```text -onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.*:MatMulNBitsLoRA.* -onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBits*LoRA* +onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.* +onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBitsQLoRAFusion* ``` The quantized tests check original-versus-fused FP16 and FP32 results, FP32 accumulation, @@ -220,15 +230,14 @@ FP32 comparisons use 1e-5 absolute/relative bounds. A dense reference test separately checks the specialized decode reduction. That reference test allows 0.003 absolute error for K=8192 with FP16 lane accumulation; shorter K and FP32 accumulation retain 0.002. Its relative -tolerance stays 0.001. FP16 LoRA fusion comparisons use a 0.0001 absolute -floor plus two FP16 epsilons of relative tolerance; deliberately exact -switching and cancellation checks remain exact. A serialization regression +tolerance stays 0.001. Deliberately exact switching checks remain exact. +A serialization regression checks that empty packed defaults have no payload and that a reloaded explicit -operator accepts full and partial overrides without a load-time LoRA transform. +operator accepts full and partial overrides without a load-time QLoRA transform. It also captures adapter-off and adapter-on under separate IDs and alternates replay with sentinel-filled outputs, checking that both captures execute and produce the correct results with CPU fallback disabled. -See [the performance evaluation](WebGPU_MatMulNBitsLoRA_Performance.md) for model-level +See [the performance evaluation](WebGPU_MatMulNBitsQLoRA_Performance.md) for model-level accuracy limits and the main-versus-optimized performance comparison. WebGPU also batches CPU-to-GPU input uploads to amortize per-tensor adapter transfer cost; this is independent of the shader fusion. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index eb78cd0dac5d4..e6ecbf96e8407 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -4,10 +4,7 @@ #param acc_f32 #param block_size #param n_bits -#param has_lora_up -#param lora_quantized -#param lora_transposed -#param lora_shared_b +#param has_qlora_up #param has_bias #param has_zero_points #param is_qualcomm @@ -20,42 +17,12 @@ #use .getByOffset .setByOffset -#if lora_shared_b -var lora_b_tile: array, tile_size_n * 8>; -#endif -fn write_lora_output(offset: u32, value: vec4) { -#if has_lora_up +fn write_qlora_output(offset: u32, value: vec4) { +#if has_qlora_up let row = (offset * 4u) / uniforms.N; let col = (offset * 4u) % uniforms.N; - var delta = vec4(0); -#if lora_quantized - delta = qlora_delta4(row, col); -#elif lora_transposed - for (var r4 = 0u; r4 < 8u; r4++) { - let tv = lora_t[row * 8u + r4]; -#if lora_shared_b - let local_col = col % tile_size_n; - let b0 = lora_b_tile[local_col * 8u + r4]; - let b1 = lora_b_tile[(local_col + 1u) * 8u + r4]; - let b2 = lora_b_tile[(local_col + 2u) * 8u + r4]; - let b3 = lora_b_tile[(local_col + 3u) * 8u + r4]; -#else - let b0 = lora_b[col * 8u + r4]; - let b1 = lora_b[(col + 1u) * 8u + r4]; - let b2 = lora_b[(col + 2u) * 8u + r4]; - let b3 = lora_b[(col + 3u) * 8u + r4]; -#endif - delta = vec4(b0.x, b1.x, b2.x, b3.x) * tv.x + delta; - delta = vec4(b0.y, b1.y, b2.y, b3.y) * tv.y + delta; - delta = vec4(b0.z, b1.z, b2.z, b3.z) * tv.z + delta; - delta = vec4(b0.w, b1.w, b2.w, b3.w) * tv.w + delta; - } -#else - for (var r = 0u; r < 32u; r++) { - delta = vec4(lora_b[r * uniforms.N + col], lora_b[r * uniforms.N + col + 1u], lora_b[r * uniforms.N + col + 2u], lora_b[r * uniforms.N + col + 3u]) * lora_t[row * 32u + r] + delta; - } -#endif + let delta = qlora_delta4(row, col); output.setByOffset(offset, value + delta); #else output.setByOffset(offset, value); @@ -450,13 +417,6 @@ $MAIN { workgroupBarrier(); } -#if lora_shared_b - for (var i = local_idx; i < tile_size_n * 8u; i += workgroup_size) { - let col = b_global_base + i / 8u; - if (col < uniforms.N) { lora_b_tile[i] = lora_b[col * 8u + i % 8u]; } - } - workgroupBarrier(); -#endif let a_global = a_global_base + base_A + a_idx; let b_global = b_global_base + base_B; let output_idx = (batch * uniforms.M * uniforms.N + a_global * uniforms.N + b_global)/4; @@ -503,15 +463,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - write_lora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); - write_lora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); - write_lora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); - write_lora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); + write_qlora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); + write_qlora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); + write_qlora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); + write_qlora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); #else - write_lora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); - write_lora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); - write_lora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); - write_lora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); + write_qlora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); + write_qlora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); + write_qlora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); + write_qlora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); #endif #else #if has_bias @@ -542,15 +502,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - write_lora_output(output_idx, vec4(lane_output1 + bias_vec1)); - write_lora_output(output_idx+1, vec4(lane_output2 + bias_vec2)); - write_lora_output(output_idx+2, vec4(lane_output3 + bias_vec3)); - write_lora_output(output_idx+3, vec4(lane_output4 + bias_vec4)); + write_qlora_output(output_idx, vec4(lane_output1 + bias_vec1)); + write_qlora_output(output_idx+1, vec4(lane_output2 + bias_vec2)); + write_qlora_output(output_idx+2, vec4(lane_output3 + bias_vec3)); + write_qlora_output(output_idx+3, vec4(lane_output4 + bias_vec4)); #else - write_lora_output(output_idx, vec4(lane_output1)); - write_lora_output(output_idx+1, vec4(lane_output2)); - write_lora_output(output_idx+2, vec4(lane_output3)); - write_lora_output(output_idx+3, vec4(lane_output4)); + write_qlora_output(output_idx, vec4(lane_output1)); + write_qlora_output(output_idx+1, vec4(lane_output2)); + write_qlora_output(output_idx+2, vec4(lane_output3)); + write_qlora_output(output_idx+3, vec4(lane_output4)); #endif #endif } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 19e8fbf094ae6..98a2bd2ca58dd 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -3,7 +3,6 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" -#include "contrib_ops/webgpu/quantization/matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "core/providers/webgpu/shader_helper.h" @@ -12,17 +11,17 @@ namespace contrib { namespace webgpu { Status DP4AMatMulQuantizeProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | + ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform); const auto& scales = shader.AddOutput("scales", ShaderUsage::UseUniform); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template", - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales, scales)); + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template", WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales)); } Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); + const auto& a = shader.AddInput( + "input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); const auto& scales_b = shader.AddInput("scales_b", ShaderUsage::UseUniform); @@ -35,41 +34,24 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } - if (has_lora_up_) { + if (has_qlora_up_) { shader.AddInput("lora_t"); shader.AddInput("lora_b"); - if (lora_quantized_) { - shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_, true); - } + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_, true); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - if (has_lora_down_) { - GenerateLoRADown(shader, "uniforms.batch_count * uniforms.num_M_tile * uniforms.num_N_tile"); - } - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(block_size, block_size_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), - WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), - WGSL_TEMPLATE_PARAMETER(lora_shared_b, lora_shared_b_), - WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), - WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), - WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), - WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(block_size, block_size_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), + WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), + WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } // scale_A components = 1, b components = 4, output components = 1 @@ -87,17 +69,16 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } - if (has_lora_up_) { + if (has_qlora_up_) { shader.AddInput("lora_t"); shader.AddInput("lora_b"); - if (lora_quantized_) { - shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); - } + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); + ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, + "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_; ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count"); @@ -107,70 +88,39 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co shader.AddOutput("lora_down"); } - if (has_lora_down_) { - GenerateLoRADecode(shader, "uniforms.batch_count * uniforms.dispatch_M * uniforms.num_N_tile"); - } - - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_small_m.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), - WGSL_TEMPLATE_PARAMETER(has_qlora_down, has_qlora_down_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), - WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul_small_m.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_down, has_qlora_down_), WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } -Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, - const Tensor* zero_points, const Tensor* bias, - uint32_t batch_count, - uint32_t M, - uint32_t dispatch_M, - uint32_t N, - uint32_t K, - uint32_t block_size, - uint32_t zero_blocks_per_col, - uint32_t min_M_for_tile_optimization, - uint32_t nbits, - onnxruntime::webgpu::ComputeContext& context, - Tensor* y, - const uint32_t weight_index, +Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, uint32_t batch_count, uint32_t M, uint32_t dispatch_M, uint32_t N, + uint32_t K, uint32_t block_size, uint32_t zero_blocks_per_col, + uint32_t min_M_for_tile_optimization, uint32_t nbits, + onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, - const Tensor* lora_at, Tensor* lora_down, bool lora_transposed, const Tensor* lora_b_scales, - DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n, const Tensor* lora_a_scales) { - const bool has_lora_down = lora_at != nullptr; - const bool has_qlora_down = lora_a_scales != nullptr; - const bool has_float_decode = has_lora_down && !has_qlora_down && M == 1 && dispatch_M == 1 && - a->IsDataType() && context.AdapterInfo().vendor == std::string_view{"nvidia"} && - context.HasFeature(wgpu::FeatureName::Subgroups) && - context.AdapterInfo().subgroupMinSize == 32 && context.AdapterInfo().subgroupMaxSize == 32; - const bool lora_quantized = lora_b_scales != nullptr; - ORT_RETURN_IF_NOT(small_m_tile_n == 0 || - (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && - !zero_points && !bias && !weight_index && !weight_index_indirect && !lora_t && !lora_at), - "Small-M override requires a plain rank-32 INT8 projection."); - ORT_RETURN_IF_NOT(!has_qlora_down || - (has_lora_down && lora_down && a->IsDataType() && nbits == 4 && block_size == 32 && - M == 1 && dispatch_M == 1 && !small_m_tile_n && !weight_index && !weight_index_indirect && - !zero_points && !bias && !lora_t), + const Tensor* lora_at, Tensor* lora_down, const Tensor* lora_b_scales, + DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n, + const Tensor* lora_a_scales) { + const bool has_qlora_down = lora_at != nullptr; + ORT_RETURN_IF_NOT( + small_m_tile_n == 0 || (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && + !zero_points && !bias && !weight_index && !weight_index_indirect && !lora_t && !lora_at), + "Small-M override requires a plain rank-32 INT8 projection."); + ORT_RETURN_IF_NOT(!has_qlora_down || (lora_a_scales && lora_down && a->IsDataType() && nbits == 4 && + block_size == 32 && M == 1 && dispatch_M == 1 && !small_m_tile_n && + !weight_index && !weight_index_indirect && !zero_points && !bias && !lora_t), "Unsupported joint DP4A QLoRA decode projection."); - ORT_RETURN_IF_NOT(!has_lora_down || (lora_down && !lora_t && !weight_index_indirect && - (has_qlora_down || has_float_decode || M >= min_M_for_tile_optimization)), - "Unsupported joint DP4A LoRA projection."); constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; constexpr uint32_t kU32Components = 4; @@ -198,7 +148,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const auto& a_quant = *input.values; const auto& a_scale = *input.scales; - const bool has_lora_up = lora_t != nullptr; + const bool has_qlora_up = lora_t != nullptr; const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; @@ -213,20 +163,25 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor // The narrow projection changes the K reduction order. Keep partial sums // and the final reduction in FP32, then cast T once when storing it. const bool acc_f32 = small_m_tile_n != 0 || context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down, has_float_decode}; + DP4AMatMulNBitsSmallMProgram mul_program{ + tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, + has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_qlora_up, has_qlora_down}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); - mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : has_float_decode ? batch_count * 32 - : 0)); - mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) + mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : 0)); + mul_program + .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, + weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32, has_lora_up, lora_transposed, lora_quantized, has_qlora_down, has_float_decode); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, + has_weight_idx_indirect, broadcast_a, acc_f32, has_qlora_up, has_qlora_down); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput( + {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -234,22 +189,16 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } - if (has_lora_up) { - mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 - : 1}); - if (lora_quantized) mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + if (has_qlora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } if (has_qlora_down) { mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 16}); mul_program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); } - if (has_float_decode) { - mul_program.AddInput({a, ProgramTensorMetadataDependency::Type, 4}); - mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 4}); - mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); - } return context.RunProgram(mul_program); } @@ -261,15 +210,16 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor uint32_t num_N_tile = CeilDiv(N, tile_size_n); bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - const bool lora_shared_b = has_lora_up && lora_transposed && !lora_quantized && - y->IsDataType() && context.AdapterInfo().vendor == std::string_view{"nvidia"}; - DP4AMatMulNBitsProgram mul_program{tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, lora_shared_b}; + DP4AMatMulNBitsProgram mul_program{ + tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, + has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_qlora_up}; mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); - mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile + (has_lora_down ? CeilDiv(batch_count * M, 32u) : 0)); - mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile); + mul_program + .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({{static_cast(batch_count)}, {static_cast(M)}, {static_cast(N)}, @@ -281,9 +231,12 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, lora_shared_b); + .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, + has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, + has_qlora_up); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput( + {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -291,37 +244,23 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } - if (has_lora_up) { - mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 - : 1}); - if (lora_quantized) mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); - } - if (has_lora_down) { - mul_program.AddInput({a, ProgramTensorMetadataDependency::Type, 4}); - mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 4}); - mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type, 4}); + if (has_qlora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(mul_program); } -bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, - uint64_t accuracy_level, - uint32_t block_size, - uint32_t N, - uint32_t K, - uint32_t components_k, - uint32_t M, - bool has_weight_idx_indirect, - const Tensor* y) { +bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, + uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k, uint32_t M, + bool has_weight_idx_indirect, const Tensor* y) { // macOS - Avoid using dp4a on Metal, as it does not appear to have native dp4a support. // https://github.com/gpuweb/gpuweb/issues/2677#issuecomment-1713292226 // Use 'vendor' to check for metal; 'backend' is always WEBGPU when running under wasm. - bool use_dp4a = context.HasFeature(wgpu::FeatureName::Subgroups) && - context.AdapterInfo().vendor != std::string_view{"apple"}; - if (!(accuracy_level == 4 && block_size % 32 == 0 && - components_k == 4 && K % 128 == 0 && N % 16 == 0 && - use_dp4a)) { + bool use_dp4a = + context.HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor != std::string_view{"apple"}; + if (!(accuracy_level == 4 && block_size % 32 == 0 && components_k == 4 && K % 128 == 0 && N % 16 == 0 && use_dp4a)) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 962295e1671eb..bb3a0debb5d51 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -30,37 +30,28 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false, bool lora_shared_b = false) : Program{has_lora_down ? (lora_quantized ? "DP4AMatMulNBitsQLoRADown" : "DP4AMatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBitsLoRAUp") - : "DP4AMatMulNBits"}, - tile_size_m_(tile_size_m), - tile_size_n_(tile_size_n), - tile_size_k_vec_(tile_size_k_vec), - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32), - has_lora_up_(has_lora_up), - has_lora_down_(has_lora_down), - lora_transposed_(lora_transposed), - lora_quantized_(lora_quantized), - lora_shared_b_(lora_shared_b) {} + DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, + uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32, bool has_qlora_up = false) + : Program{has_qlora_up ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBits"}, + tile_size_m_(tile_size_m), + tile_size_n_(tile_size_n), + tile_size_k_vec_(tile_size_k_vec), + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K8", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K8", ProgramUniformVariableDataType::Uint32}, {"K16", ProgramUniformVariableDataType::Uint32}, + {"num_M_tile", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"weight_idx", ProgramUniformVariableDataType::Uint32}); @@ -76,51 +67,38 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_weight_idx_indirect_; bool is_qualcomm_; bool acc_f32_; - bool has_lora_up_; - bool has_lora_down_; - bool lora_transposed_; - bool lora_quantized_; - bool lora_shared_b_; + bool has_qlora_up_; }; class DP4AMatMulNBitsSmallMProgram final : public Program { public: - DP4AMatMulNBitsSmallMProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, - bool broadcast_a_row = false, - bool acc_f32 = false, bool has_lora_up = false, bool lora_transposed = false, bool lora_quantized = false, - bool has_qlora_down = false, bool has_lora_down = false) : Program{has_lora_down ? "DP4AMatMulNBitsLoRADownSmallM" : has_qlora_down ? "DP4AMatMulNBitsQLoRADownSmallM" - : lora_quantized ? "DP4AMatMulNBitsQLoRAUpSmallM" - : "DP4AMatMulNBitsSmallMProgram"}, - tile_size_k_vec_(tile_size_k_vec), - tile_size_(tile_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - single_scale_weights_(single_scale_weights), - broadcast_a_row_(broadcast_a_row), - acc_f32_(acc_f32), - has_lora_up_(has_lora_up), - lora_transposed_(lora_transposed), - lora_quantized_(lora_quantized), - has_qlora_down_(has_qlora_down), - has_lora_down_(has_lora_down) {} + DP4AMatMulNBitsSmallMProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, + bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false, + bool has_qlora_up = false, bool has_qlora_down = false) + : Program{has_qlora_down ? "DP4AMatMulNBitsQLoRADownSmallM" + : has_qlora_up ? "DP4AMatMulNBitsQLoRAUpSmallM" + : "DP4AMatMulNBitsSmallMProgram"}, + tile_size_k_vec_(tile_size_k_vec), + tile_size_(tile_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + single_scale_weights_(single_scale_weights), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up), + has_qlora_down_(has_qlora_down) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"K32", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K16", ProgramUniformVariableDataType::Uint32}, {"K32", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, - {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_k_vec_; @@ -133,49 +111,31 @@ class DP4AMatMulNBitsSmallMProgram final : public Program::max(), - bool has_weight_idx_indirect = false, - const Tensor* y = nullptr); + bool has_weight_idx_indirect = false, const Tensor* y = nullptr); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template index be9143cae48ad..b5d71296e0b6f 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template @@ -9,9 +9,7 @@ #param n_bits #param broadcast_a_row #param has_zero_points -#param has_lora_up -#param lora_quantized -#param lora_transposed +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -19,27 +17,11 @@ #use .getByOffset .setByOffset -fn write_lora_output(offset: u32, value: output_element_t) { -#if has_lora_up +fn write_qlora_output(offset: u32, value: output_element_t) { +#if has_qlora_up let row = (offset * 1u) / uniforms.N; let col = (offset * 1u) % uniforms.N; - var delta = output_element_t(0); -#if lora_quantized - delta = qlora_delta(row, col); -#elif lora_transposed - for (var r4 = 0u; r4 < 8u; r4++) { - let bv = lora_b[col * 8u + r4]; - let tv = lora_t[row * 8u + r4]; - delta = bv.x * tv.x + delta; - delta = bv.y * tv.y + delta; - delta = bv.z * tv.z + delta; - delta = bv.w * tv.w + delta; - } -#else - for (var r = 0u; r < 32u; r++) { - delta = lora_b[r * uniforms.N + col] * lora_t[row * 32u + r] + delta; - } -#endif + let delta = qlora_delta(row, col); output.setByOffset(offset, value + delta); #else output.setByOffset(offset, value); @@ -292,9 +274,9 @@ $MAIN { #if has_bias let bias_value = acc_element_t(bias[b_global + b_bias_offset]); // Downcast to the output element type only at the final store. - write_lora_output(output_idx, output_element_t(output_value + bias_value)); + write_qlora_output(output_idx, output_element_t(output_value + bias_value)); #else - write_lora_output(output_idx, output_element_t(output_value)); + write_qlora_output(output_idx, output_element_t(output_value)); #endif } } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 0b5f7284b7b26..7131961c8a5ba 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -19,17 +19,13 @@ namespace onnxruntime { namespace contrib { namespace webgpu { -ONNX_OPERATOR_KERNEL_EX( - MatMulNBits, - kMSDomain, - 1, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", WebGpuSupportedFloatTypes()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) - .TypeConstraint("T3", DataTypeImpl::GetTensorType()) - .TypeConstraint("T4", DataTypeImpl::GetTensorType()), - MatMulNBits); +ONNX_OPERATOR_KERNEL_EX(MatMulNBits, kMSDomain, 1, kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", WebGpuSupportedFloatTypes()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("T3", DataTypeImpl::GetTensorType()) + .TypeConstraint("T4", DataTypeImpl::GetTensorType()), + MatMulNBits); Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -44,13 +40,11 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } - if (has_lora_up_) { + if (has_qlora_up_) { shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); - if (lora_quantized_) { - shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); - } + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); } const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -59,23 +53,15 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons ORT_ENFORCE(tile_m_ % (workgroup_size / 8) == 0, "tile_m must be a multiple of workgroup_size / 8."); ORT_ENFORCE(nbits_ == 4 || nbits_ == 8, "Only 4/8 bits are supported for webgpu matmulnbits."); - return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_wide_tile.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), - WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), - WGSL_TEMPLATE_PARAMETER(nbits, nbits_), - WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_), - WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_), - WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales, scales)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/matmul_nbits_wide_tile.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(nbits, nbits_), + WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_), WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_), + WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_), WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), + WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales)); } // Apply similar idea with DP4AMatMulNBitsSmallMProgram algorithm. @@ -92,26 +78,20 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } - if (has_lora_up_) { + if (has_qlora_up_) { shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); - if (lora_quantized_) { - shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); - } + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, acc_f32_); } - if (has_lora_down_) { + if (has_qlora_down_) { shader.AddInput("lora_at"); - if (lora_quantized_) shader.AddInput("lora_a_scales"); + shader.AddInput("lora_a_scales"); } const auto& output = shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); - if (has_lora_down_) { + if (has_qlora_down_) { shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - if (lora_quantized_) { - GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_, qlora_down_subgroups_); - } else { - GenerateLoRAGemv(shader, "uniforms.num_N_tile", a.GetByOffset("k4")); - } + GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_, qlora_down_subgroups_); } const uint32_t components_a = a.NumComponents(); @@ -122,31 +102,20 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const uint32_t a_length_per_tile = tile_size_k / components_a; uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec; - return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits.wgsl.template", - WGSL_TEMPLATE_PARAMETER(a_length_per_tile, a_length_per_tile), - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), - WGSL_TEMPLATE_PARAMETER(component_a, components_a), - WGSL_TEMPLATE_PARAMETER(component_b, components_b), - WGSL_TEMPLATE_PARAMETER(elements_in_value_b, elements_in_value_b), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_lora_up, has_lora_up_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(lora_quantized, lora_quantized_), - WGSL_TEMPLATE_PARAMETER(lora_transposed, lora_transposed_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, false), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k, tile_size_k), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/matmul_nbits.wgsl.template", WGSL_TEMPLATE_PARAMETER(a_length_per_tile, a_length_per_tile), + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), + WGSL_TEMPLATE_PARAMETER(component_a, components_a), WGSL_TEMPLATE_PARAMETER(component_b, components_b), + WGSL_TEMPLATE_PARAMETER(elements_in_value_b, elements_in_value_b), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, false), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k, tile_size_k), WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), + WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { @@ -161,7 +130,8 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context const bool has_zero_points = zero_points != nullptr; if (has_zero_points) { - ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), + "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); } MatMulComputeHelper helper; @@ -194,7 +164,8 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context * @param zero_points Pointer to the zero-point tensor for quantization; must be of type uint8 if provided. * weight_index > 0 is only supported when zero_points is nullptr. * @param bias Pointer to the bias tensor; optional. - * @param K_op The K dimension of the operation (number of columns in 'a' and rows in 'b' before quantization). + * @param K_op The K dimension of the operation (number of columns in 'a' and rows in 'b' before + * quantization). * @param N_op The N dimension of the operation (number of columns in 'b'). * @param block_size_op The block size used for quantization partitioning. * @param accuracy_level Accuracy level influencing the choice of optimized kernel. @@ -211,36 +182,29 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context * - A wide tile program is used when block size, component count, and other criteria are met. * - Otherwise, a default matmul program is used. */ -Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, const Tensor* bias, - int64_t K_op, - int64_t N_op, - int64_t block_size_op, - int64_t accuracy_level, - int64_t nbits, - onnxruntime::webgpu::ComputeContext& context, - Tensor* y, - const uint32_t weight_index, - const Tensor* weight_index_indirect, - uint32_t override_M, const Tensor* lora_t, const Tensor* lora_b, - const Tensor* lora_a_transposed, Tensor* lora_a_output, bool lora_transposed, - const Tensor* lora_b_scales, const Tensor* lora_a_scales, +Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, + int64_t nbits, onnxruntime::webgpu::ComputeContext& context, Tensor* y, + const uint32_t weight_index, const Tensor* weight_index_indirect, uint32_t override_M, + const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_a_transposed, + Tensor* lora_a_output, const Tensor* lora_b_scales, const Tensor* lora_a_scales, DP4AQuantizedInput* quantized_input) { TensorShape b_shape({N_op, K_op}); MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); - const bool has_lora_up = lora_t != nullptr; - const bool has_lora_down = lora_a_transposed != nullptr; - const bool lora_quantized = lora_b_scales != nullptr || lora_a_scales != nullptr; - ORT_RETURN_IF_NOT(!has_lora_down || (lora_a_output != nullptr && !has_lora_up && a->Shape().Size() == K_op && - nbits == 4 && block_size_op == 32), - "Unsupported joint LoRA projection."); + const bool has_qlora_up = lora_t != nullptr; + const bool has_qlora_down = lora_a_transposed != nullptr; + ORT_RETURN_IF_NOT(!has_qlora_down || (lora_a_output != nullptr && lora_a_scales != nullptr && !has_qlora_up && + a->Shape().Size() == K_op && nbits == 4 && block_size_op == 32), + "Unsupported joint QLoRA projection."); const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; const bool has_zero_points = zero_points != nullptr; if (has_zero_points) { - ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), + "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); } const uint32_t batch_count = onnxruntime::narrow(helper.OutputOffsets().size()); @@ -266,51 +230,43 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // n_blocks_per_col up to the next multiple of (8/nbits) — the number of zero point // values per byte — so that the linear stride correctly skips byte-boundary padding. const uint32_t zp_elements_per_byte = 8 / static_cast(nbits); - uint32_t zero_blocks_per_col = (n_blocks_per_col + zp_elements_per_byte - 1) / zp_elements_per_byte * zp_elements_per_byte; + uint32_t zero_blocks_per_col = + (n_blocks_per_col + zp_elements_per_byte - 1) / zp_elements_per_byte * zp_elements_per_byte; // apple|intel - Experimental dawn support for subgroup matrix matmul. std::optional subgroup_matrix_config; - if (!has_lora_down && CanApplySubgroupMatrixMatMulNBits(context, - accuracy_level, - block_size, - batch_count, - N, - K, - static_cast(nbits), - y->DataType() == DataTypeImpl::GetType(), - subgroup_matrix_config, - M, - has_weight_idx_indirect)) { - if (has_lora_up) { + if (!has_qlora_down && CanApplySubgroupMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, + static_cast(nbits), + y->DataType() == DataTypeImpl::GetType(), + subgroup_matrix_config, M, has_weight_idx_indirect)) { + if (has_qlora_up) { auto base_output = context.CreateGPUTensor(y->DataType(), y->Shape()); - ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, - static_cast(nbits), zero_blocks_per_col, - *subgroup_matrix_config, context, &base_output, weight_index, - weight_index_indirect)); - return ApplyLoRAOutput(context, &base_output, lora_t, lora_b, y, lora_transposed, lora_b_scales); + ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits( + a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), zero_blocks_per_col, + *subgroup_matrix_config, context, &base_output, weight_index, weight_index_indirect)); + return ApplyQLoRAOutput(context, &base_output, lora_t, lora_b, lora_b_scales, y); } - return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, - static_cast(nbits), zero_blocks_per_col, - *subgroup_matrix_config, context, y, weight_index, + return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), + zero_blocks_per_col, *subgroup_matrix_config, context, y, weight_index, weight_index_indirect); } - // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of length of M. - // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). - if (!has_lora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, - M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, lora_transposed, lora_b_scales, quantized_input); + // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of + // length of M. DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte + // values). + if (!has_qlora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, + has_weight_idx_indirect, y)) { + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, + zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), + context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, + lora_b_scales, quantized_input); } // WideTileProgram // This program is optimized for Block32 prefill. - const bool use_wide_tile_program = !has_lora_down && CanApplyWideTileMatMulNBits(M, - K, - block_size, - nbits, - has_weight_idx_indirect, - components_a, - components_b); + const bool use_wide_tile_program = + !has_qlora_down && + CanApplyWideTileMatMulNBits(M, K, block_size, nbits, has_weight_idx_indirect, components_a, components_b); const bool acc_f32 = context.EnableMatmulFp32Accumulation(); @@ -331,11 +287,19 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // width it supports; otherwise (including when Subgroups isn't supported) it falls // back to a direct reduction. // NOTE: NVIDIA GPUs prefer direct reduction. - const uint32_t subgroup_min_size = (context.HasFeature(wgpu::FeatureName::Subgroups) && !is_nvidia) - ? context.AdapterInfo().subgroupMinSize - : 0u; - - MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, static_cast(nbits), subgroup_min_size, acc_f32, has_lora_up, lora_transposed, lora_quantized}; + const uint32_t subgroup_min_size = + (context.HasFeature(wgpu::FeatureName::Subgroups) && !is_nvidia) ? context.AdapterInfo().subgroupMinSize : 0u; + + MatMulNBitsWideTileProgram program{has_zero_points, + has_bias, + has_weight_idx, + has_weight_idx_indirect, + tile_m, + tile_n, + static_cast(nbits), + subgroup_min_size, + acc_f32, + has_qlora_up}; program.SetWorkgroupSize(workgroup_size); program.SetDispatchGroupSize(num_N_tile, num_M_tile, batch_count); @@ -344,12 +308,9 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const uint32_t K_of_b = n_blocks_per_col * blob_size / components_b_with_u32; const uint32_t K_of_a = K / components_a; - program.AddInput({a, - ProgramTensorMetadataDependency::TypeAndRank, - onnxruntime::narrow(components_a)}); - program.AddInput({b, - ProgramTensorMetadataDependency::TypeAndRank, - onnxruntime::narrow(components_b_with_u32)}); + program.AddInput({a, ProgramTensorMetadataDependency::TypeAndRank, onnxruntime::narrow(components_a)}); + program.AddInput( + {b, ProgramTensorMetadataDependency::TypeAndRank, onnxruntime::narrow(components_b_with_u32)}); program.AddInput({scales, ProgramTensorMetadataDependency::TypeAndRank}); if (has_zero_points) { program.AddInput({zero_points, @@ -363,9 +324,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, if (has_weight_idx_indirect) { program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } - program.AddOutput({y, - ProgramTensorMetadataDependency::TypeAndRank, - onnxruntime::narrow(components)}); + program.AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, onnxruntime::narrow(components)}); program.AddUniformVariables({{batch_count}, {dispatch_M}, {N}, @@ -376,21 +335,20 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {num_N_tile}, {num_M_tile}, {weight_index}}); - program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32, has_lora_up, lora_transposed, lora_quantized); + program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, + acc_f32, has_qlora_up); - if (has_lora_up) { - program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 - : 1}); - if (lora_quantized) program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + if (has_qlora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(program); } // Use tile_size_k_vec=32 by default for better K-dimension parallelism. // Intel devices use 16 as they have different subgroup/cache characteristics. - const uint32_t tile_size_k_vec = - (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u; + const uint32_t tile_size_k_vec = (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u; constexpr uint32_t workgroup_size = 128; constexpr uint32_t tile_size = 8; @@ -399,14 +357,25 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; // The decode reduction assigns one rank to each 32-lane subgroup. const auto& adapter = context.AdapterInfo(); - const bool qlora_down_subgroups = has_lora_down && lora_quantized && - adapter.vendor == std::string_view{"nvidia"} && - context.HasFeature(wgpu::FeatureName::Subgroups) && - adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; - MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, qlora_down_subgroups}; + const bool qlora_down_subgroups = has_qlora_down && adapter.vendor == std::string_view{"nvidia"} && + context.HasFeature(wgpu::FeatureName::Subgroups) && adapter.subgroupMinSize == 32 && + adapter.subgroupMaxSize == 32; + MatMulNBitsProgram program{tile_size, + static_cast(nbits), + has_zero_points, + has_bias, + has_weight_idx, + has_weight_idx_indirect, + single_scale_weights, + tile_size_k_vec, + broadcast_a, + acc_f32, + has_qlora_up, + has_qlora_down, + qlora_down_subgroups}; program.SetWorkgroupSize(workgroup_size); uint32_t num_N_tile = (N + tile_size - 1) / tile_size; - program.SetDispatchGroupSize(num_N_tile + (has_lora_down ? (lora_quantized ? 4 : 8) : 0), dispatch_M, batch_count); + program.SetDispatchGroupSize(num_N_tile + (has_qlora_down ? 4 : 0), dispatch_M, batch_count); program .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_a)}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_b_with_u32)}, @@ -424,7 +393,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {batch_count}, {weight_index}, {dispatch_M}}) - .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32, has_lora_up, has_lora_down, lora_transposed, lora_quantized, qlora_down_subgroups); + .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, + tile_size_k_vec, broadcast_a, acc_f32, has_qlora_up, has_qlora_down, qlora_down_subgroups); if (has_zero_points) { program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -434,16 +404,15 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, if (has_weight_idx_indirect) { program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } - if (has_lora_up) { - program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, lora_transposed ? 4 : 1}); - program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : lora_transposed ? 4 - : 1}); - if (lora_quantized) program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + if (has_qlora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); } - if (has_lora_down) { - ORT_RETURN_IF_NOT(components_a == 4, "Joint LoRA projection requires vec4 X."); - program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, lora_quantized ? 16 : 4}); - if (lora_quantized) program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); + if (has_qlora_down) { + ORT_RETURN_IF_NOT(components_a == 4, "Joint QLoRA projection requires vec4 X."); + program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); program.AddOutput({lora_a_output, ProgramTensorMetadataDependency::Type}); } return context.RunProgram(program); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index 82be7cc1a6201..c5870a9edcce5 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -10,32 +10,33 @@ namespace onnxruntime { namespace contrib { namespace webgpu { struct DP4AQuantizedInput; -// Adds an independent FP16 rank-32 down-projection branch to a 256-thread shader. -void GenerateLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups); -void GenerateLoRADecode(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups); - -// Adds the same FP32 GEMV reduction to standalone and fused 128-thread shaders. -void GenerateLoRAGemv(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, - const std::string& x_value); - using namespace onnxruntime::webgpu; class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, bool has_lora_up = false, bool lora_transposed = false, bool lora_quantized = false) - : Program{lora_quantized ? "MatMulNBitsQLoRAUpWideTile" : "MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32), has_lora_up_(has_lora_up), lora_transposed_(lora_transposed), lora_quantized_(lora_quantized) {} + MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, + bool has_qlora_up = false) + : Program{has_qlora_up ? "MatMulNBitsQLoRAUpWideTile" : "MatMulNBitsWideTile"}, + has_zero_points_{has_zero_points}, + has_bias_{has_bias}, + has_weight_idx_{has_weight_idx}, + has_weight_idx_indirect_{has_weight_idx_indirect}, + tile_m_(tile_m), + tile_n_(tile_n), + nbits_(nbits), + subgroup_min_size_(subgroup_min_size), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up) {} Status GenerateShaderCode(ShaderHelper& sh) const override; - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"Batch", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K_of_a", ProgramUniformVariableDataType::Uint32}, - {"K_of_b", ProgramUniformVariableDataType::Uint32}, - {"n_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"Batch", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K_of_a", ProgramUniformVariableDataType::Uint32}, + {"K_of_b", ProgramUniformVariableDataType::Uint32}, {"n_blocks_per_col", ProgramUniformVariableDataType::Uint32}, + {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"num_M_tile", ProgramUniformVariableDataType::Uint32}, + {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: bool has_zero_points_; @@ -47,16 +48,18 @@ class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false, bool has_lora_up = false, bool has_lora_down = false, bool lora_transposed = false, bool lora_quantized = false, bool qlora_down_subgroups = false) - : Program{has_lora_down ? (lora_quantized ? (qlora_down_subgroups ? "MatMulNBitsQLoRADownSubgroup" : "MatMulNBitsQLoRADown") : "MatMulNBitsLoRADown") : has_lora_up ? (lora_quantized ? "MatMulNBitsQLoRAUp" : "MatMulNBitsLoRAUp") - : "MatMulNBits"}, + MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, + bool broadcast_a_row = false, bool acc_f32 = false, bool has_qlora_up = false, + bool has_qlora_down = false, bool qlora_down_subgroups = false) + : Program{has_qlora_down ? (qlora_down_subgroups ? "MatMulNBitsQLoRADownSubgroup" : "MatMulNBitsQLoRADown") + : has_qlora_up ? "MatMulNBitsQLoRAUp" + : "MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), @@ -67,25 +70,18 @@ class MatMulNBitsProgram final : public Program { tile_size_k_vec_(tile_size_k_vec), broadcast_a_row_(broadcast_a_row), acc_f32_(acc_f32), - has_lora_up_(has_lora_up), - has_lora_down_(has_lora_down), - lora_transposed_(lora_transposed), - lora_quantized_(lora_quantized), + has_qlora_up_(has_qlora_up), + has_qlora_down_(has_qlora_down), qlora_down_subgroups_(qlora_down_subgroups) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K_of_a", ProgramUniformVariableDataType::Uint32}, - {"K_of_b", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, {"K_of_a", ProgramUniformVariableDataType::Uint32}, + {"K_of_b", ProgramUniformVariableDataType::Uint32}, {"block_size", ProgramUniformVariableDataType::Uint32}, {"blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, - {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"batch_count", ProgramUniformVariableDataType::Uint32}, + {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_; @@ -98,10 +94,8 @@ class MatMulNBitsProgram final : public Program { uint32_t tile_size_k_vec_; bool broadcast_a_row_; bool acc_f32_; - bool has_lora_up_; - bool has_lora_down_; - bool lora_transposed_; - bool lora_quantized_; + bool has_qlora_up_; + bool has_qlora_down_; bool qlora_down_subgroups_; }; @@ -127,18 +121,15 @@ class MatMulNBits final : public WebGpuKernel { int64_t bits_; }; -Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, const Tensor* bias, - int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, int64_t bits_op, - onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index = 0, - const Tensor* weight_index_indirect = nullptr, - uint32_t override_M = 0, - const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, - const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, bool lora_transposed = false, +Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, + int64_t bits_op, onnxruntime::webgpu::ComputeContext& context, Tensor* y, + const uint32_t weight_index = 0, const Tensor* weight_index_indirect = nullptr, + uint32_t override_M = 0, const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, + const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, const Tensor* lora_b_scales = nullptr, const Tensor* lora_a_scales = nullptr, DP4AQuantizedInput* quantized_input = nullptr); -Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed = false, const Tensor* quantized_scales = nullptr); - } // namespace webgpu } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template index 4a5b40d3cd4de..69e320d1acde5 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template @@ -13,9 +13,7 @@ #param tile_size_k_vec #param tile_size_k #param tile_size -#param has_lora_up -#param lora_quantized -#param lora_transposed +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -267,24 +265,8 @@ $MAIN { output_value += acc_element_t(bias[b_global + b_bias_offset]); #endif // Downcast to the output type only at the final write. -#if has_lora_up - var delta = output_element_t(0); -#if lora_quantized - delta = qlora_delta(batch * uniforms.M + a_global, b_global); -#elif lora_transposed - for (var r4 = 0u; r4 < 8u; r4++) { - let bv = lora_b[b_global * 8u + r4]; - let tv = lora_t[(batch * uniforms.M + a_global) * 8u + r4]; - delta = bv.x * tv.x + delta; - delta = bv.y * tv.y + delta; - delta = bv.z * tv.z + delta; - delta = bv.w * tv.w + delta; - } -#else - for (var r = 0u; r < 32u; r++) { - delta = lora_b[r * uniforms.N + b_global] * lora_t[(batch * uniforms.M + a_global) * 32u + r] + delta; - } -#endif +#if has_qlora_up + let delta = qlora_delta(batch * uniforms.M + a_global, b_global); output.setByOffset(output_idx, output_element_t(output_value) + delta); #else output.setByOffset(output_idx, output_element_t(output_value)); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc deleted file mode 100644 index eb1c2b3966ea9..0000000000000 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_lora.cc +++ /dev/null @@ -1,410 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/webgpu/quantization/matmul_nbits.h" -#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" -#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" -#include "core/common/safeint.h" -#include "core/providers/webgpu/math/matmul.h" -#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" -#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" -#include "core/providers/webgpu/shader_helper.h" -#include "core/providers/webgpu/webgpu_supported_types.h" - -namespace onnxruntime::contrib::webgpu { - -void GenerateLoRADown(ShaderHelper& shader, const std::string& base_groups) { - shader.AddInput("lora_x"); - shader.AddInput("lora_at"); - shader.AddOutput("lora_down", ShaderUsage::UseValueTypeAlias); - shader.AdditionalImplementation() << R"( -var lora_x_tile: array; -var lora_a_tile: array; -)"; - shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" - << R"( -if (workgroup_idx >= lora_base_groups) { - let group = workgroup_idx - lora_base_groups; - let row_local = local_idx / 8u; - let k4_local = local_idx % 8u; - let row = group * 32u + row_local; - let rank_base = k4_local * 4u; - var acc = lora_down_value_t(0); - for (var start = 0u; start < uniforms.K / 4u; start += 8u) { - var xv = lora_down_value_t(0); - if (row < uniforms.M * uniforms.batch_count) { - xv = lora_x[row * (uniforms.K / 4u) + start + k4_local]; - } - lora_x_tile[local_idx] = xv; - lora_a_tile[local_idx] = lora_at[row_local * (uniforms.K / 4u) + start + k4_local]; - workgroupBarrier(); - for (var kk = 0u; kk < 8u; kk++) { - let x = lora_x_tile[row_local * 8u + kk]; - let a0 = lora_a_tile[rank_base * 8u + kk]; - let a1 = lora_a_tile[(rank_base + 1u) * 8u + kk]; - let a2 = lora_a_tile[(rank_base + 2u) * 8u + kk]; - let a3 = lora_a_tile[(rank_base + 3u) * 8u + kk]; - acc = lora_down_value_t(a0.x, a1.x, a2.x, a3.x) * x.x + acc; - acc = lora_down_value_t(a0.y, a1.y, a2.y, a3.y) * x.y + acc; - acc = lora_down_value_t(a0.z, a1.z, a2.z, a3.z) * x.z + acc; - acc = lora_down_value_t(a0.w, a1.w, a2.w, a3.w) * x.w + acc; - } - workgroupBarrier(); - } - if (row < uniforms.M * uniforms.batch_count) { - lora_down[row * 8u + k4_local] = acc; - } - return; -} -)"; -} - -void GenerateLoRAGemv(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value) { - shader.AdditionalImplementation() << "var lora_partials: array, 128>;\n"; - shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" - << R"( - // Extra workgroups compute each rank output once, alongside independent base - // tiles. The next dispatch supplies the global dependency before consuming T. - if (workgroup_idx >= lora_base_groups) { - let lane = local_idx % 32u; - let rank = (workgroup_idx - lora_base_groups) * 4u + local_idx / 32u; - var acc = vec4(0); - for (var k4 = lane; k4 < uniforms.K / 4u; k4 += 32u) { - acc = fma(vec4()" << x_value - << R"(), vec4(lora_at[rank * (uniforms.K / 4u) + k4]), acc); - } - lora_partials[local_idx] = acc; - workgroupBarrier(); - for (var stride = 16u; stride > 0u; stride /= 2u) { - if (lane < stride) { lora_partials[local_idx] += lora_partials[local_idx + stride]; } - workgroupBarrier(); - } - if (lane == 0u) { - let v = lora_partials[local_idx]; - lora_partials[local_idx] = vec4(v.x + v.z, v.y + v.w, 0, 0); - } - // Materialize the penultimate reduction to prevent reassociation of the - // final four terms, matching the original GEMV's workgroup barriers. - workgroupBarrier(); - if (lane == 0u) { - let v = lora_partials[local_idx]; - lora_down[rank] = lora_down_element_t(v.x + v.y); - } - return; - } -)"; -} - -class LoRAGemvProgram final : public Program { - public: - LoRAGemvProgram() : Program{"LoRADownGemv"} {} - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}); - Status GenerateShaderCode(ShaderHelper& shader) const override { - shader.AddInput("lora_x"); - shader.AddInput("lora_at"); - shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - GenerateLoRAGemv(shader, "0u", "lora_x[k4]"); - return Status::OK(); - } -}; - -class LoRADownProgram final : public Program { - public: - LoRADownProgram() : Program{"LoRADownVec4"} {} - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"batch_count", ProgramUniformVariableDataType::Uint32}); - Status GenerateShaderCode(ShaderHelper& shader) const override { - GenerateLoRADown(shader, "0u"); - return Status::OK(); - } -}; - -static void GenerateLoRAParallel(ShaderHelper& shader, const std::string& base_groups, uint32_t tile_rows) { - shader.AddInput("lora_x"); - shader.AddInput("lora_at"); - shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - auto& body = shader.MainFunctionBody(); - body << "let lora_base_groups = " << base_groups << ";\n" - << "if (workgroup_idx >= lora_base_groups) {\n" - << "let group = workgroup_idx - lora_base_groups;\n" - << "let row = (group / 8u) * " << tile_rows << "u;\n" - << R"( - let rank = (group % 8u) * 4u + local_idx / 32u; - let lane = local_idx % 32u; -)"; - // Reuse each contiguous weight vector across prompt rows. FP32 partials - // avoid accumulating FP16 rounding error over the long K dimension. - for (uint32_t row = 0; row < tile_rows; ++row) body << "var acc" << row << " = vec4(0);\n"; - body << "for (var k4 = lane; k4 < uniforms.K / 4u; k4 += 32u) {\n" - << "let av = vec4(lora_at[rank * (uniforms.K / 4u) + k4]);\n"; - for (uint32_t row = 0; row < tile_rows; ++row) { - body << "if (row + " << row << "u < uniforms.M * uniforms.batch_count) {\n" - << "acc" << row << " = fma(vec4(lora_x[(row + " << row - << "u) * (uniforms.K / 4u) + k4]), av, acc" << row << ");\n}\n"; - } - body << "}\n"; - for (uint32_t row = 0; row < tile_rows; ++row) { - body << "let sum" << row << " = subgroupAdd((acc" << row << ".x + acc" << row << ".z) + (acc" - << row << ".y + acc" << row << ".w));\n" - << "if (lane == 0u && row + " << row << "u < uniforms.M * uniforms.batch_count) {\n" - << "lora_down[(row + " << row << "u) * 32u + rank] = lora_down_element_t(sum" << row << ");\n}\n"; - } - body << "return;\n}\n"; -} - -void GenerateLoRADecode(ShaderHelper& shader, const std::string& base_groups) { - // One workgroup per rank exposes more parallelism than four ranks per - // workgroup; the four subgroup partials need only one workgroup barrier. - shader.AddInput("lora_x"); - shader.AddInput("lora_at"); - shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - shader.AdditionalImplementation() << "var lora_sums: array;\n"; - shader.MainFunctionBody() << "let lora_base_groups = " << base_groups << ";\n" - << R"( -if (workgroup_idx >= lora_base_groups) { - let group = workgroup_idx - lora_base_groups; - let row = group / 32u; - let rank = group % 32u; - var acc = vec4(0); - for (var k4 = local_idx; k4 < uniforms.K / 4u; k4 += 128u) { - acc = fma(vec4(lora_x[row * (uniforms.K / 4u) + k4]), - vec4(lora_at[rank * (uniforms.K / 4u) + k4]), acc); - } - let sum = subgroupAdd((acc.x + acc.z) + (acc.y + acc.w)); - if (local_idx % 32u == 0u) { lora_sums[local_idx / 32u] = sum; } - workgroupBarrier(); - if (local_idx == 0u) { - lora_down[row * 32u + rank] = lora_down_element_t((lora_sums[0] + lora_sums[2]) + (lora_sums[1] + lora_sums[3])); - } - return; -} -)"; -} - -class LoRADownParallelProgram final : public Program { - public: - explicit LoRADownParallelProgram(uint32_t tile_rows) : Program{"LoRADownParallel"}, tile_rows_(tile_rows) {} - ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(LoRADownParallelProgram); - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"batch_count", ProgramUniformVariableDataType::Uint32}); - Status GenerateShaderCode(ShaderHelper& shader) const override { - GenerateLoRAParallel(shader, "0u", tile_rows_); - return Status::OK(); - } - - private: - uint32_t tile_rows_; -}; - -// Fallback epilogue for base kernels without an integrated adapter output path. -class LoRAOutputProgram final : public Program { - public: - explicit LoRAOutputProgram(bool transposed = false) - : Program{transposed ? "LoRAOutputVec4" : "LoRAOutput"}, transposed_(transposed) {} - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}); - Status GenerateShaderCode(ShaderHelper& shader) const override { - shader.AddInput("base", ShaderUsage::UseElementTypeAlias); - shader.AddInput("t"); - shader.AddInput("b"); - shader.AddOutput("output"); - shader.MainFunctionBody() - << "if (global_idx >= uniforms.size) { return; }\n" - << "let row = global_idx / uniforms.N; let col = global_idx % uniforms.N;\n" - << "var delta = base_element_t(0);\n"; - if (transposed_) { - shader.MainFunctionBody() - << "for (var r4 = 0u; r4 < 8u; r4++) {\n" - << " let bv = b[col * 8u + r4]; let tv = t[row * 8u + r4];\n" - << " delta = bv.x * tv.x + delta; delta = bv.y * tv.y + delta;\n" - << " delta = bv.z * tv.z + delta; delta = bv.w * tv.w + delta;\n}\n"; - } else { - shader.MainFunctionBody() - << "for (var r = 0u; r < 32u; r++) { delta = b[r * uniforms.N + col] * t[row * 32u + r] + delta; }\n"; - } - shader.MainFunctionBody() << "output[global_idx] = base[global_idx] + delta;\n"; - return Status::OK(); - } - - private: - bool transposed_; -}; - -Status ApplyLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, Tensor* y, bool transposed, const Tensor* quantized_scales) { - if (quantized_scales) return ApplyQLoRAOutput(context, base, t, b, quantized_scales, y); - const auto size = narrow(y->Shape().Size()); - LoRAOutputProgram program(transposed); - program.AddInputs({{base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, - {t, ProgramTensorMetadataDependency::Type, transposed ? 4 : 1}, - {b, ProgramTensorMetadataDependency::Type, transposed ? 4 : 1}}) - .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) - .AddUniformVariables({{size}, {narrow(b->Shape()[transposed ? 0 : 1])}}) - .SetDispatchGroupSize(CeilDiv(size, 64u)); - return context.RunProgram(program); -} - -Status ApplyTransposedLoRA(onnxruntime::webgpu::ComputeContext& context, const Tensor* x, const Tensor* weight, - const Tensor* scales, const Tensor* at, const Tensor* bt, Tensor* y, - int64_t k, int64_t n, int64_t accuracy_level, bool fuse_up) { - const auto rows = narrow(x->Shape().Size() / k); - const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); - const bool use_gemv = rows == 1 && k >= 2048 && k <= 8192 && - context.AdapterInfo().vendor == std::string_view{"nvidia"}; - auto down_shape = x->Shape().AsShapeVector(); - down_shape.back() = 32; - auto t = context.CreateGPUTensor(x->DataType(), TensorShape(down_shape)); - const auto compute_down = [&]() -> Status { - const auto& adapter = context.AdapterInfo(); - if (rows > 1 && adapter.vendor == std::string_view{"nvidia"} && - context.HasFeature(wgpu::FeatureName::Subgroups) && - adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) { - constexpr uint32_t tile_rows = 4; - LoRADownParallelProgram down(tile_rows); - down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, - {at, ProgramTensorMetadataDependency::Type, 4}}) - .AddOutput({&t, ProgramTensorMetadataDependency::Type}) - .SetWorkgroupSize(128) - .SetDispatchGroupSize(CeilDiv(rows, tile_rows) * 8) - .AddUniformVariables({{narrow(k)}, {rows}, {1u}}) - .CacheHint(tile_rows); - ORT_RETURN_IF_ERROR(context.RunProgram(down)); - } else if (use_gemv) { - LoRAGemvProgram down; - down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, - {at, ProgramTensorMetadataDependency::Type, 4}}) - .AddOutput({&t, ProgramTensorMetadataDependency::Type}) - .SetWorkgroupSize(128) - .SetDispatchGroupSize(8) - .AddUniformVariables({{narrow(k)}}); - ORT_RETURN_IF_ERROR(context.RunProgram(down)); - } else { - LoRADownProgram down; - down.AddInputs({{x, ProgramTensorMetadataDependency::Type, 4}, - {at, ProgramTensorMetadataDependency::Type, 4}}) - .AddOutput({&t, ProgramTensorMetadataDependency::Type, 4}) - .SetWorkgroupSize(256) - .SetDispatchGroupSize(CeilDiv(rows, 32u)) - .AddUniformVariables({{rows}, {narrow(k)}, {1u}}); - ORT_RETURN_IF_ERROR(context.RunProgram(down)); - } - return Status::OK(); - }; - if (fuse_up) { - ORT_RETURN_IF_ERROR(compute_down()); - return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, - context, y, 0, nullptr, 0, &t, bt, nullptr, nullptr, true); - } - auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); - // FP32 can select DP4A even at M=1. Keep that base policy when the - // floating joint down kernel would otherwise replace it. - const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits( - context, accuracy_level, 32, narrow(n), narrow(k), 4, m, false, y); - const auto& adapter = context.AdapterInfo(); - if (m == 1 && k >= 2048 && k <= 8192 && base_dp4a && x->IsDataType() && - adapter.vendor == std::string_view{"nvidia"} && context.HasFeature(wgpu::FeatureName::Subgroups) && - adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) { - // Quantize only the base's X. The appended floating down projection - // reads original X and writes T for the subsequent up/add dispatch. - ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits(x, weight, scales, nullptr, nullptr, rows, m, m, - narrow(n), narrow(k), 32, 0, - kMinMForTileOptimization, 4, context, &base, 0, - nullptr, nullptr, nullptr, at, &t)); - } else if (use_gemv && !base_dp4a) { - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, - context, &base, 0, nullptr, 0, nullptr, nullptr, at, &t)); - } else if (m >= kMinMForTileOptimization && context.AdapterInfo().vendor == std::string_view{"nvidia"} && - base_dp4a) { - // Preserve DP4A's activation quantization and compute T from original X - // in extra workgroups of the base dispatch. - ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits(x, weight, scales, nullptr, nullptr, rows / m, m, m, - narrow(n), narrow(k), 32, 0, - kMinMForTileOptimization, 4, context, &base, 0, - nullptr, nullptr, nullptr, at, &t)); - } else { - ORT_RETURN_IF_ERROR(compute_down()); - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k, n, 32, accuracy_level, 4, - context, &base)); - } - return ApplyLoRAOutput(context, &base, &t, bt, y, true); -} - -class MatMulNBitsLoRA final : public WebGpuKernel { - public: - explicit MatMulNBitsLoRA(const OpKernelInfo& info) - : WebGpuKernel(info), k_(info.GetAttr("K")), n_(info.GetAttr("N")), accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)), transposed_(info.GetAttrOrDefault("adapter_transposed", 0)) { - const auto path = info.GetConfigOptions().GetConfigOrDefault("session.webgpu.lora_fusion_mode", "auto"); - ORT_ENFORCE(path == "auto" || path == "fuse_down" || path == "fuse_up", - "session.webgpu.lora_fusion_mode must be auto, fuse_down or fuse_up."); - fusion_mode_ = path == "auto" ? LoRAFusionMode::Auto : path == "fuse_up" ? LoRAFusionMode::FuseUp - : LoRAFusionMode::FuseDown; - ORT_ENFORCE(transposed_ == 0 || transposed_ == 1, "Invalid adapter_transposed value."); - ORT_ENFORCE(info.GetAttr("bits") == 4 && info.GetAttr("block_size") == 32, - "MatMulNBitsLoRA currently requires INT4 block size 32."); - ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0 && n_ % 4 == 0, "Unsupported projection dimensions."); - } - - Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override { - const auto* x = context.Input(0); - const auto* weight = context.Input(1); - const auto* scales = context.Input(2); - const auto* a = context.Input(3); - const auto* b = context.Input(4); - ORT_RETURN_IF_NOT(x->Shape().NumDimensions() >= 2 && a->Shape().NumDimensions() == 2 && - b->Shape().NumDimensions() == 2, - "Invalid LoRA input ranks."); - const auto adapter_axis = static_cast(transposed_); - ORT_RETURN_IF_NOT(a->Shape()[adapter_axis] == k_ && b->Shape()[1 - adapter_axis] == n_, "Invalid LoRA K/N dimensions."); - const int64_t rank_a = a->Shape()[1 - adapter_axis]; - const int64_t rank_b = b->Shape()[adapter_axis]; - ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (rank_b == 0 || rank_b == 32), - "MatMulNBitsLoRA currently supports rank 32 or an empty adapter."); - ORT_RETURN_IF_NOT(weight->Shape() == TensorShape({n_, k_ / 32, 16}) && - scales->Shape().Size() == SafeInt(n_) * (k_ / 32), - "Invalid packed weight/scales shape."); - MatMulComputeHelper helper; - ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n_, k_}), false, true)); - auto* y = context.Output(0, helper.OutputShape()); - if (y->Shape().Size() == 0) return Status::OK(); - if (rank_a == 0 || rank_b == 0) { - return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y); - } - if (transposed_) { - // Use the per-matrix row count so batched one-token decode still selects fuse_down. - const auto m = x->Shape()[x->Shape().NumDimensions() - 2]; - const bool fuse_up = fusion_mode_ == LoRAFusionMode::FuseUp || - (fusion_mode_ == LoRAFusionMode::Auto && m > 1); - return ApplyTransposedLoRA(context, x, weight, scales, a, b, y, k_, n_, accuracy_level_, fuse_up); - } - MatMulComputeHelper lora_helper; - ORT_RETURN_IF_ERROR(lora_helper.Compute(x->Shape(), a->Shape())); - auto t = context.CreateGPUTensor(x->DataType(), lora_helper.OutputShape()); - std::vector inputs{x, a}; - ORT_RETURN_IF_ERROR(compute_dispatcher_.Compute(context, Activation(), inputs, &t, - /*is_channels_last=*/true, /*b_is_constant=*/false)); - return ApplyMatMulNBits(x, weight, scales, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, - context, y, 0, nullptr, 0, &t, b); - } - - private: - enum class LoRAFusionMode { Auto, - FuseDown, - FuseUp }; - - int64_t k_; - int64_t n_; - int64_t accuracy_level_; - int64_t transposed_; - LoRAFusionMode fusion_mode_; - mutable MatMulComputeDispatcher compute_dispatcher_; -}; - -ONNX_OPERATOR_KERNEL_EX( - MatMulNBitsLoRA, kMSDomain, 1, kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T", WebGpuSupportedFloatTypes()) - .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), - MatMulNBitsLoRA); - -} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index d6e34ad20007a..8cad004d693c7 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -14,8 +14,8 @@ namespace onnxruntime::contrib::webgpu { -void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, - bool acc_f32, bool use_subgroups) { +void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32, + bool use_subgroups) { if (use_subgroups) { shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n"; auto& body = shader.MainFunctionBody(); @@ -41,8 +41,8 @@ if (workgroup_idx >= qlora_base_groups) { << "let scale" << r << " = lora_a_scales[rank" << r << " * (uniforms.K / 32u) + k / 32u];\n" << "var sum" << r << " = qlora_acc_t(0);\n"; for (uint32_t i = 0; i < 4; ++i) { - body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" - << r << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; + body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" << r + << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; } body << "partial" << r << " += sum" << r << ";\n"; } @@ -282,10 +282,9 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { context.AdapterInfo().vendor == std::string_view{"nvidia"}) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( - x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, - narrow(n_), narrow(k_), 32, narrow(k_ / 32), - kMinMForTileOptimization, 4, context, &base, 0, - nullptr, nullptr, nullptr, aq, &t, true, nullptr, nullptr, 0, as)); + x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, narrow(n_), + narrow(k_), 32, narrow(k_ / 32), kMinMForTileOptimization, 4, context, &base, 0, nullptr, + nullptr, nullptr, aq, &t, nullptr, nullptr, 0, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && @@ -293,25 +292,25 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { if (joint_down) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, - nullptr, 0, nullptr, nullptr, aq, &t, true, nullptr, as)); + nullptr, 0, nullptr, nullptr, aq, &t, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } // Allow activation quantization for prefill only when the model opts into it. const auto down_accuracy = fuse_up && m >= kMinMForTileOptimization ? accuracy_level_ : 0; DP4AQuantizedInput quantized_x; // Rank 32 exposes too few output tiles for the tiled DP4A kernel at these sizes. - const bool small_m_down = fuse_up && m <= 4096 && k_ >= 2048 && - context.AdapterInfo().vendor == std::string_view{"nvidia"} && - CanApplyDP4AMatrixMatMulNBits(context, down_accuracy, 32, 32, - narrow(k_), 4, m, false, &t); + const bool small_m_down = + fuse_up && m <= 4096 && k_ >= 2048 && context.AdapterInfo().vendor == std::string_view{"nvidia"} && + CanApplyDP4AMatrixMatMulNBits(context, down_accuracy, 32, 32, narrow(k_), 4, m, false, &t); if (small_m_down) { ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( - x, aq, as, nullptr, nullptr, narrow(x->Shape().Size() / k_ / m), m, m, 32, - narrow(k_), 32, narrow(k_ / 32), kMinMForTileOptimization, 8, context, &t, 0, - nullptr, nullptr, nullptr, nullptr, nullptr, false, nullptr, &quantized_x, 32)); + x, aq, as, nullptr, nullptr, narrow(x->Shape().Size() / k_ / m), m, m, 32, narrow(k_), 32, + narrow(k_ / 32), kMinMForTileOptimization, 8, context, &t, 0, nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, &quantized_x, 32)); } else { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, - nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); + nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + &quantized_x)); } std::optional up_config; const bool subgroup_up = @@ -319,11 +318,12 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { narrow(n_), 32, 8, x->IsDataType(), up_config, m); if (fuse_up && !subgroup_up) { return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, &t, - bq, nullptr, nullptr, true, bs, nullptr, &quantized_x); + bq, nullptr, nullptr, bs, nullptr, &quantized_x); } auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, - nullptr, 0, nullptr, nullptr, nullptr, nullptr, false, nullptr, nullptr, &quantized_x)); + nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + &quantized_x)); if (!subgroup_up) return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); // Keep the original up kernel on devices with a different matrix-accumulation path. auto delta = context.CreateGPUTensor(y->DataType(), y->Shape()); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template index 25379afea40e0..bc24ccccdac33 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template @@ -2,9 +2,7 @@ // Licensed under the MIT License. #param has_zero_points -#param has_lora_up -#param lora_quantized -#param lora_transposed +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -116,24 +114,8 @@ fn load_scale(row : u32, block_idx : u32) -> output_element_t { fn write_output(batch : u32, row : u32, col : u32, value : output_element_t) { if (batch < uniforms.Batch && row < uniforms.M && col < uniforms.N) { let offset = batch * uniforms.M * uniforms.N + row * uniforms.N + col; -#if has_lora_up - var delta = output_element_t(0); -#if lora_quantized - delta = qlora_delta(batch * uniforms.M + row, col); -#elif lora_transposed - for (var r4 = 0u; r4 < 8u; r4++) { - let bv = lora_b[col * 8u + r4]; - let tv = lora_t[(batch * uniforms.M + row) * 8u + r4]; - delta = bv.x * tv.x + delta; - delta = bv.y * tv.y + delta; - delta = bv.z * tv.z + delta; - delta = bv.w * tv.w + delta; - } -#else - for (var r = 0u; r < 32u; r++) { - delta = lora_b[r * uniforms.N + col] * lora_t[(batch * uniforms.M + row) * 32u + r] + delta; - } -#endif +#if has_qlora_up + let delta = qlora_delta(batch * uniforms.M + row, col); output.setByOffset(offset, value + delta); #else output.setByOffset(offset, value); diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index b098284c9bea3..f6b10778d6a09 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -62,7 +62,6 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 1f3fcc967144a..177774c22cb95 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3941,56 +3941,6 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by } }); - ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsLoRA) - .SetDomain(kMSDomain) - .SinceVersion(1) - .SetDoc( - "Computes MatMulNBits(X, W, scales) + MatMul(MatMul(X, lora_A), lora_B). " - "With adapter_transposed=1, transposes both stored adapter inputs in this formula. " - "The adapter remains replaceable. A zero rank in either adapter input disables the adapter. " - "Intermediate results are rounded to the input type before addition. " - "Supports FP16 or FP32, rank 32, INT4 block size 32, implicit zero point 8, and no bias.") - .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) - .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) - .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) - .Attr("block_size", "Base weight block size; must be 32.", AttributeProto::INT, int64_t{32}) - .Attr("accuracy_level", "MatMulNBits base accumulation policy.", AttributeProto::INT, int64_t{0}) - .Attr("adapter_transposed", "Whether adapters are stored as [R_A,K] and [N,R_B].", AttributeProto::INT, int64_t{0}) - .Input(0, "X", "Input with rank at least two and last dimension K.", "T") - .Input(1, "W", "INT4 packed weights of shape [N, K/32, 16].", "TQ") - .Input(2, "scales", "Block scales containing N * K/32 elements; may be flattened.", "T") - .Input(3, "lora_A", "Adapter of shape [K, R_A], or [R_A, K] if adapter_transposed=1; R_A is 0 or 32.", "T") - .Input(4, "lora_B", "Adapter of shape [R_B, N], or [N, R_B] if adapter_transposed=1; R_B is 0 or 32.", "T") - .Output(0, "Y", "Input shape with last dimension replaced by N.", "T") - .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, "FP16 or FP32 activations, block scales, and adapter weights.") - .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized base weights.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - const auto k = getAttribute(ctx, "K", -1); - const auto n = getAttribute(ctx, "N", -1); - const auto transposed = getAttribute(ctx, "adapter_transposed", 0); - if ((transposed != 0 && transposed != 1) || k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || - getAttribute(ctx, "bits", 4) != 4 || getAttribute(ctx, "block_size", 32) != 32) { - fail_shape_inference("Unsupported MatMulNBitsLoRA attributes."); - } - propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() < 2) { - fail_shape_inference("X must have rank at least two."); - } - MatmulWithQuantWeightShapeInference(ctx, k, n, true); - for (size_t i : {size_t{3}, size_t{4}}) { - if (!hasInputShape(ctx, i)) continue; - const auto& s = getInputShape(ctx, i); - if (s.dim_size() != 2) fail_shape_inference("Adapter inputs must have rank two."); - const int feature_axis = (i == 3 ? 0 : 1) ^ static_cast(transposed); - const auto& feature = s.dim(feature_axis); - const auto& rank = s.dim(1 - feature_axis); - if ((feature.has_dim_value() && feature.dim_value() != (i == 3 ? k : n)) || - (rank.has_dim_value() && rank.dim_value() != 0 && rank.dim_value() != 32)) { - fail_shape_inference("Invalid adapter dimensions."); - } - } - }); - ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsQLoRA) .SetDomain(kMSDomain) .SinceVersion(1) diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 800b0340e912c..60833655ef1f0 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -60,7 +60,6 @@ #include "core/optimizer/group_query_attention_pre_norm_fusion.h" #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/matmul_nbits_mlp_fusion.h" -#include "core/optimizer/matmul_nbits_lora_fusion.h" #include "core/optimizer/matmul_nbits_qlora_fusion.h" #include "core/optimizer/matmul_bn_fusion.h" #include "core/optimizer/matmul_integer_to_float.h" @@ -467,7 +466,6 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kCudaExecutionProvider, onnxruntime::kWebGpuExecutionProvider})); - bool has_matmul_nbits_lora_kernel = false; bool has_matmul_nbits_qlora_kernel = false; bool has_matmul_nbits_mlp_kernel = false; bool has_matmul_nbits_qkv_kernel = false; @@ -486,17 +484,12 @@ InlinedVector> GenerateTransformers( nullptr) .IsOK(); }; - has_matmul_nbits_lora_kernel = has_webgpu_kernel("MatMulNBitsLoRA"); has_matmul_nbits_qlora_kernel = has_webgpu_kernel("MatMulNBitsQLoRA"); has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp"); has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv"); } } } - if (has_matmul_nbits_lora_kernel && - session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { - transformers.emplace_back(std::make_unique()); - } if (has_matmul_nbits_qlora_kernel && session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { transformers.emplace_back(std::make_unique()); diff --git a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc deleted file mode 100644 index 8796ee372c6c0..0000000000000 --- a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.cc +++ /dev/null @@ -1,199 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/optimizer/matmul_nbits_lora_fusion.h" - -#include "core/graph/graph_utils.h" -#include "core/graph/node_attr_utils.h" -#include "core/optimizer/initializer.h" -#include "core/optimizer/matmul_nbits_fusion_utils.h" -#include "core/optimizer/utils.h" - -namespace onnxruntime { -namespace { - -const Node* Producer(const Graph& graph, const Node& node, int input) { - const auto* edge = graph_utils::GetInputEdge(node, input); - return edge ? graph.GetNode(edge->GetNode().Index()) : nullptr; -} - -bool SingleConsumer(const Graph& graph, const Node& node) { - return !graph.NodeProducesGraphOutput(node) && optimizer_utils::CheckOutputEdges(graph, node, 1); -} - -bool IsMatMul(const Node* node) { - return node && graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMul", {9, 13}) && - node->GetExecutionProviderType() == kWebGpuExecutionProvider; -} - -bool MatrixShape(const NodeArg& arg, int64_t rows, int64_t cols) { - const auto* s = arg.Shape(); - return s && s->dim_size() == 2 && s->dim(0).has_dim_value() && s->dim(1).has_dim_value() && - s->dim(0).dim_value() == rows && s->dim(1).dim_value() == cols; -} - -const Node* AdapterTranspose(const Graph& graph, const Node& matmul) { - const auto* node = Producer(graph, matmul, 1); - if (!node || !graph_utils::IsSupportedOptypeVersionAndDomain(*node, "Transpose", {1, 13, 21, 23, 24}) || - node->GetExecutionProviderType() != kWebGpuExecutionProvider || !SingleConsumer(graph, *node)) return nullptr; - const auto* shape = node->InputDefs()[0]->Shape(); - if (!shape || shape->dim_size() != 2) return nullptr; - auto perm = node->GetAttributes().find("perm"); - if (perm != node->GetAttributes().end() && - (perm->second.ints_size() != 2 || perm->second.ints(0) != 1 || perm->second.ints(1) != 0)) return nullptr; - return node; -} - -bool ZeroDefault(const Graph& graph, const NodeArg& arg) { - if (!graph.IsInputsIncludingInitializers(&arg) || graph.IsOutput(&arg)) { - return false; - } - const ONNX_NAMESPACE::TensorProto* proto = nullptr; - if (!graph.GetInitializedTensor(arg.Name(), proto)) { - return false; - } - // Only rewrite a private adapter parameter: another consumer may depend on - // its original dimensions or on IEEE zero-times-infinity semantics. - const auto consumers = graph.GetConsumerNodes(arg.Name()); - if (consumers.size() != 1) { - return false; - } - Initializer data(graph, *proto, graph.ModelPath()); - auto bytes = data.DataAsByteSpan(); - return std::all_of(bytes.begin(), bytes.end(), [](uint8_t v) { return v == 0; }); -} - -void EmptyDefault(Graph& graph, NodeArg& arg, int rank_axis) { - const ONNX_NAMESPACE::TensorProto* original = nullptr; - ORT_ENFORCE(graph.GetInitializedTensor(arg.Name(), original)); - ONNX_NAMESPACE::TensorProto empty; - empty.set_name(original->name()); - empty.set_data_type(original->data_type()); - for (int i = 0; i < original->dims_size(); ++i) { - empty.add_dims(i == rank_axis ? 0 : original->dims(i)); - } - auto shape = *arg.Shape(); - shape.mutable_dim(rank_axis)->Clear(); - arg.SetShape(shape); - graph.RemoveInitializedTensor(arg.Name()); - graph.AddInitializedTensor(empty); -} - -} // namespace - -Status MatMulNBitsLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, - const logging::Logger& logger) const { - GraphViewer viewer(graph); - for (auto index : viewer.GetNodesInTopologicalOrder()) { - auto* add = graph.GetNode(index); - if (!add) continue; - ORT_RETURN_IF_ERROR(Recurse(*add, modified, graph_level, logger)); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(*add, "Add", {7, 13, 14}) || - add->GetExecutionProviderType() != kWebGpuExecutionProvider) continue; - const Node* base = Producer(graph, *add, 0); - const Node* up = Producer(graph, *add, 1); - if (base && base->OpType() != "MatMulNBits") std::swap(base, up); - if (!base || !graph_utils::IsSupportedOptypeVersionAndDomain(*base, "MatMulNBits", {1}, kMSDomain) || - base->GetExecutionProviderType() != kWebGpuExecutionProvider || !IsMatMul(up)) continue; - const Node* down = Producer(graph, *up, 0); - if (!IsMatMul(down) || down->InputDefs()[0] != base->InputDefs()[0] || - !SingleConsumer(graph, *base) || !SingleConsumer(graph, *up) || !SingleConsumer(graph, *down)) continue; - - using namespace matmul_nbits_fusion_utils; - const int64_t k = GetIntAttr(*base, "K", -1, true); - const int64_t n = GetIntAttr(*base, "N", -1, true); - const auto* shape = down->InputDefs()[1]->Shape(); - if (!shape || shape->dim_size() != 2 || !shape->dim(1).has_dim_value()) continue; - const int64_t rank = shape->dim(1).dim_value(); - if (k <= 0 || n <= 0 || rank != 32 || k % 32 != 0 || n % 4 != 0 || - GetIntAttr(*base, "bits", 4) != 4 || GetIntAttr(*base, "block_size", -1, true) != 32 || - GetIntAttr(*base, "weight_prepacked", 0) != 0 || HasInput(*base, 3) || HasInput(*base, 4) || - HasInput(*base, 5) || !MatrixShape(*down->InputDefs()[1], k, rank) || - !MatrixShape(*up->InputDefs()[1], rank, n) || - !base->InputDefs()[0]->Shape() || base->InputDefs()[0]->Shape()->dim_size() < 2) continue; - - const auto* x_type = base->InputDefs()[0]->TypeAsProto(); - if (!x_type || !x_type->has_tensor_type()) continue; - const auto float_type = x_type->tensor_type().elem_type(); - if (float_type != ONNX_NAMESPACE::TensorProto::FLOAT16 && float_type != ONNX_NAMESPACE::TensorProto::FLOAT) - continue; - bool types_ok = true; - for (const auto* arg : {base->InputDefs()[2], down->InputDefs()[1], up->InputDefs()[1]}) { - const auto* type = arg->TypeAsProto(); - types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == float_type; - } - if (!types_ok) continue; - - // MatMulNBits also accepts legacy flattened weights. Leave those graphs - // alone until the fused kernel supports that representation. - const auto* weight_shape = base->InputDefs()[1]->Shape(); - if (!weight_shape || weight_shape->dim_size() != 3 || - !weight_shape->dim(0).has_dim_value() || weight_shape->dim(0).dim_value() != n || - !weight_shape->dim(1).has_dim_value() || weight_shape->dim(1).dim_value() != k / 32 || - !weight_shape->dim(2).has_dim_value() || weight_shape->dim(2).dim_value() != 16) continue; - const auto* scale_shape = base->InputDefs()[2]->Shape(); - const bool flat_scales = scale_shape && scale_shape->dim_size() == 1 && - scale_shape->dim(0).has_dim_value() && - scale_shape->dim(0).dim_value() / n == k / 32 && - scale_shape->dim(0).dim_value() % n == 0; - if (!flat_scales && !MatrixShape(*base->InputDefs()[2], n, k / 32)) continue; - - auto* a = const_cast(down->InputDefs()[1]); - auto* b = const_cast(up->InputDefs()[1]); - const auto* transpose_a = AdapterTranspose(graph, *down); - const auto* transpose_b = AdapterTranspose(graph, *up); - const bool transposed = transpose_a && transpose_b && - MatrixShape(*transpose_a->InputDefs()[0], rank, k) && - MatrixShape(*transpose_b->InputDefs()[0], n, rank); - if (transposed) { - a = const_cast(transpose_a->InputDefs()[0]); - b = const_cast(transpose_b->InputDefs()[0]); - } - const bool empty_defaults = ZeroDefault(graph, *a) && ZeroDefault(graph, *b); - InlinedVector inputs{const_cast(base->InputDefs()[0]), - const_cast(base->InputDefs()[1]), - const_cast(base->InputDefs()[2]), a, b}; - auto attrs = base->GetAttributes(); - attrs.erase("weight_prepacked"); - utils::SetNodeAttribute(utils::MakeAttribute("adapter_transposed", int64_t{transposed}), attrs); - Node& fused = graph.AddNode(graph.GenerateNodeName("MatMulNBitsLoRA"), "MatMulNBitsLoRA", - "fused unmerged rank-32 LoRA projection", inputs, - add->MutableOutputDefs(), &attrs, kMSDomain); - fused.SetExecutionProviderType(kWebGpuExecutionProvider); - // All input producer edges and final consumers must be preserved before - // removal, including weights supplied by upstream nodes. - const Node* sources[] = {base, base, base, transposed ? transpose_a : down, transposed ? transpose_b : up}; - const int slots[] = {0, 1, 2, transposed ? 0 : 1, transposed ? 0 : 1}; - for (int i = 0; i < 5; ++i) { - if (const auto* edge = graph_utils::GetInputEdge(*sources[i], slots[i])) { - graph.AddEdge(edge->GetNode().Index(), fused.Index(), edge->GetSrcArgIndex(), i); - } - } - const auto outputs = graph_utils::GraphEdge::GetNodeOutputEdges(*add); - for (const Node* node : {base, down, up, static_cast(add)}) { - graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(node->Index())); - } - const NodeIndex indices[] = {base->Index(), down->Index(), up->Index(), add->Index()}; - if (transposed) { - graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(transpose_a->Index())); - graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(transpose_b->Index())); - } - for (auto idx : indices) graph.RemoveNode(idx); - if (transposed) { - graph.RemoveNode(transpose_a->Index()); - graph.RemoveNode(transpose_b->Index()); - } - for (const auto& edge : outputs) { - graph.AddEdge(fused.Index(), edge.dst_node, 0, edge.dst_arg_index); - } - if (empty_defaults) { - EmptyDefault(graph, *a, transposed ? 0 : 1); - EmptyDefault(graph, *b, transposed ? 1 : 0); - } - LOGS(logger, VERBOSE) << "MatMulNBitsLoRAFusion: " << fused.Name() << " empty defaults=" << empty_defaults; - modified = true; - } - return Status::OK(); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h deleted file mode 100644 index 055f7cb72bcc4..0000000000000 --- a/onnxruntime/core/optimizer/matmul_nbits_lora_fusion.h +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "core/optimizer/graph_transformer.h" - -namespace onnxruntime { - -// Opt-in fusion for finite-valued inference. Zero overridable defaults can be -// represented with rank zero, so the kernel can skip inactive adapters. -class MatMulNBitsLoRAFusion final : public GraphTransformer { - public: - MatMulNBitsLoRAFusion() : GraphTransformer("MatMulNBitsLoRAFusion", {kWebGpuExecutionProvider}) {} - - private: - Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; -}; - -} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc deleted file mode 100644 index f09f780aee3c8..0000000000000 --- a/onnxruntime/test/contrib_ops/matmul_nbits_lora_test.cc +++ /dev/null @@ -1,217 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "test/providers/provider_test_utils.h" -#include "test/util/include/default_providers.h" - -namespace onnxruntime::test { -#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) - -template -void RunEmptyAdapterSkipsProjection() { - for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { - SCOPED_TRACE(mode[0] ? mode : "default"); - for (bool transposed : {false, true}) { - for (int64_t m : {1, 37}) { - for (auto ranks : {std::pair{0, 0}, {0, 32}, {32, 0}, {32, 32}}) { - auto ep = DefaultWebGpuExecutionProvider(); - ASSERT_NE(ep, nullptr); - OpTester test("MatMulNBitsLoRA", 1, kMSDomain); - test.AddAttribute("K", 64); - test.AddAttribute("N", 128); - test.AddAttribute("bits", 4); - test.AddAttribute("block_size", 32); - test.AddAttribute("adapter_transposed", transposed); - test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); - // Packed nibbles 9 => dequantized weight 1/32, so the base is exactly 1. - test.AddInput("W", {128, 2, 16}, std::vector(128 * 2 * 16, 0x99), true); - test.AddInput("scales", {128, 2}, std::vector(256, T(0.03125f)), true); - test.AddInput("lora_A", transposed ? std::vector{ranks.first, 64} : std::vector{64, ranks.first}, - std::vector(64 * ranks.first, T(0.0625f))); - test.AddInput("lora_B", transposed ? std::vector{128, ranks.second} : std::vector{ranks.second, 128}, - std::vector(128 * ranks.second, T(0.03125f))); - const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; - test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); - std::vector> eps; - eps.push_back(std::move(ep)); - SessionOptions options; - if (mode[0]) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - } - test.Config(options); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } - } - } - } -} - -TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjection) { RunEmptyAdapterSkipsProjection(); } -TEST(MatMulNBitsLoRA, EmptyAdapterSkipsProjectionFp32) { RunEmptyAdapterSkipsProjection(); } - -TEST(MatMulNBitsLoRA, Fp32Numerics) { - struct Shape { - int64_t m, k, n, batch; - }; - for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, - Shape{2, 2048, 36, 2}, Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { - const auto [m, k, n, batch] = shape; - SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); - std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); - std::vector aq(32 * k), bq(n * 32); - for (int64_t row = 0; row < batch * m; ++row) { - for (int64_t col = 0; col < k; ++col) { - // Exactly representable by DP4A quantization, including zero rows. - const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; - x[row * k + col] = row % 7 == 6 ? 0.0f : value / 128.0f; - } - } - for (int64_t r = 0; r < 32; ++r) { - for (int64_t col = 0; col < k; ++col) - aq[r * k + col] = static_cast(96 + (r * 17 + col * 13) % 65); - for (int64_t block = 0; block < k / 32; ++block) - as[r * (k / 32) + block] = 0.001234567f * (1 + (r + block) % 4); - } - for (int64_t col = 0; col < n; ++col) { - bs[col] = 0.002345678f * (1 + col % 3); - for (int64_t r = 0; r < 32; ++r) - bq[col * 32 + r] = static_cast(112 + (col * 7 + r * 3) % 33); - } - std::vector at(32 * k), bt(n * 32); - for (int64_t r = 0; r < 32; ++r) - for (int64_t col = 0; col < k; ++col) - at[r * k + col] = (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; - for (int64_t col = 0; col < n; ++col) - for (int64_t r = 0; r < 32; ++r) - bt[col * 32 + r] = (static_cast(bq[col * 32 + r]) - 128) * bs[col]; - // Independent dense FP64 reference, rounding intermediate T - // to FP32. Non-half-representable weights detect accidental FP16 arithmetic. - for (int64_t row = 0; row < batch * m; ++row) { - double base = 0; - for (int64_t col = 0; col < k; ++col) base += x[row * k + col] * 0.03125; - std::vector t(32); - for (int64_t r = 0; r < 32; ++r) { - double sum = 0; - for (int64_t col = 0; col < k; ++col) { - const float weight = at[r * k + col]; - sum += static_cast(x[row * k + col]) * weight; - } - t[r] = static_cast(sum); - } - for (int64_t col = 0; col < n; ++col) { - double delta = 0; - for (int64_t r = 0; r < 32; ++r) { - const float weight = bt[col * 32 + r]; - delta += static_cast(t[r]) * weight; - } - expected[row * n + col] = static_cast(base) + static_cast(delta); - } - } - for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { - for (int64_t accuracy : {0, 4}) { - SCOPED_TRACE(MakeString(mode, ", accuracy=", accuracy)); - OpTester test("MatMulNBitsLoRA", 1, kMSDomain); - test.AddAttribute("K", k); - test.AddAttribute("N", n); - test.AddAttribute("accuracy_level", accuracy); - test.AddAttribute("adapter_transposed", 1); - test.AddInput("X", {batch, m, k}, x); - test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); - test.AddInput("A", {32, k}, at); - test.AddInput("B", {n, 32}, bt); - test.AddOutput("Y", {batch, m, n}, expected); - test.SetOutputAbsErr("Y", 0.00001f); - test.SetOutputRelErr("Y", 0.00001f); - SessionOptions options; - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - test.Config(options); - std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } - } - } -} - -template -void RunParallelDownNumerics() { - struct Shape { - int64_t batch, m, k, n; - }; - for (const auto shape : {Shape{1, 1, 2048, 64}, Shape{2, 1, 8192, 64}, Shape{2, 33, 2176, 64}, - Shape{1, 128, 4096, 64}, Shape{1, 1024, 2048, 16}}) { - const auto [batch, m, k, n] = shape; - SCOPED_TRACE(MakeString("batch=", batch, ", M=", m, ", K=", k, ", N=", n)); - std::vector x(batch * m * k), a(32 * k), b(n * 32), expected(batch * m * n); - for (int64_t row = 0; row < batch * m; ++row) { - for (int64_t col = 0; col < k; ++col) { - // Off-grid floating X detects accidental reuse of the base's INT8 X. - x[row * k + col] = T(row % 7 == 6 ? 0.0f : (static_cast((row * 19 + col * 11) % 255) - 127) * 0.00712345f); - } - } - for (int64_t r = 0; r < 32; ++r) { - for (int64_t col = 0; col < k; ++col) - a[r * k + col] = T((static_cast((r * 17 + col * 13) % 65) - 32) * 0.001234567f); - } - for (int64_t col = 0; col < n; ++col) - for (int64_t r = 0; r < 32; ++r) - b[col * 32 + r] = T((static_cast((col * 7 + r * 3) % 33) - 16) * 0.002345678f); - // Dense FP64 reference with exactly one cast at each MatMul boundary. - // A zero base isolates adapter arithmetic from activation quantization. - for (int64_t row = 0; row < batch * m; ++row) { - std::vector t(32); - for (int64_t r = 0; r < 32; ++r) { - double sum = 0; - for (int64_t col = 0; col < k; ++col) - sum += static_cast(static_cast(x[row * k + col])) * static_cast(a[r * k + col]); - t[r] = T(static_cast(sum)); - } - for (int64_t col = 0; col < n; ++col) { - double sum = 0; - for (int64_t r = 0; r < 32; ++r) - sum += static_cast(static_cast(t[r])) * static_cast(b[col * 32 + r]); - expected[row * n + col] = T(static_cast(sum)); - } - } - OpTester test("MatMulNBitsLoRA", 1, kMSDomain); - test.AddAttribute("K", k); - test.AddAttribute("N", n); - test.AddAttribute("accuracy_level", 4); - test.AddAttribute("adapter_transposed", 1); - test.AddInput("X", {batch, m, k}, x); - test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x88), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, T(0.03125f)), true); - test.AddInput("A", {32, k}, a); - test.AddInput("B", {n, 32}, b); - test.AddOutput("Y", {batch, m, n}, expected); - test.SetOutputAbsErr("Y", std::is_same_v ? 2e-6f : 1e-4f); - test.SetOutputRelErr("Y", std::is_same_v ? 2e-5f : 0.003f); - std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } -} - -TEST(MatMulNBitsLoRA, ParallelDownNumericsFp16) { RunParallelDownNumerics(); } -TEST(MatMulNBitsLoRA, ParallelDownNumericsFp32) { RunParallelDownNumerics(); } - -TEST(MatMulNBitsLoRA, RejectInvalidAdapterShape) { - auto ep = DefaultWebGpuExecutionProvider(); - ASSERT_NE(ep, nullptr); - OpTester test("MatMulNBitsLoRA", 1, kMSDomain); - test.AddAttribute("K", 64); - test.AddAttribute("N", 128); - test.AddInput("X", {1, 64}, std::vector(64)); - test.AddInput("W", {128, 2, 16}, std::vector(4096), true); - test.AddInput("scales", {128, 2}, std::vector(256), true); - test.AddInput("lora_A", {63, 32}, std::vector(63 * 32)); - test.AddInput("lora_B", {32, 128}, std::vector(32 * 128)); - test.AddOutput("Y", {1, 128}, std::vector(128)); - std::vector> eps; - eps.push_back(std::move(ep)); - test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid adapter dimensions", {}, nullptr, &eps); -} - -#endif -} // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc deleted file mode 100644 index 4aa38acf3bc17..0000000000000 --- a/onnxruntime/test/optimizer/matmul_nbits_lora_fusion_test.cc +++ /dev/null @@ -1,341 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/graph/node_attr_utils.h" -#include "core/optimizer/matmul_nbits_lora_fusion.h" -#include "test/optimizer/graph_transform_test_fixture.h" -#include "test/optimizer/webgpu_fusion_test_util.h" -#include "test/util/include/default_providers.h" - -namespace onnxruntime::test { -#ifndef DISABLE_CONTRIB_OPS -namespace { - -// Fused FP16 reductions can differ by a few ULPs across GPU architectures. -// Use two FP16 epsilons plus a small floor for values near cancellation. -constexpr double kFp16FusionAbsTolerance = 1e-4; -constexpr double kFp16FusionRelTolerance = 2.0 / 1024.0; -// The K-parallel down projection uses FP32 partials instead of the unfused -// sequential FP16 sum over K=2048. Allow its accumulated rounding difference -// only for affected cases; ParallelDownNumerics checks a dense FP64 reference. -constexpr double kFp16ParallelDownAbsTolerance = 4.0 / 1024.0; - -template -void BuildLoRA(ModelTestBuilder& builder, int64_t m = 1, bool zeros = false, - bool expose_intermediate = false, bool different_x = false, int64_t rank = 32, - bool flattened_weights = false, bool transposed = false, int64_t n = 128, - int isolate = 0, int64_t k_override = 0, int64_t batch = 1) { - const int64_t k = k_override ? k_override : transposed ? 2048 - : 128; - auto* x = builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)); - auto weight_shape = flattened_weights ? std::vector{n * k / 2} : std::vector{n, k / 32, 16}; - auto* w = builder.MakeInitializer(weight_shape, isolate ? uint8_t{0x88} : uint8_t{0}, isolate ? uint8_t{0x89} : uint8_t{255}); - auto* scales = builder.MakeInitializer({n, k / 32}, T(0.03125f), T(0.0625f)); - auto* a = builder.MakeInput(transposed ? std::vector{rank, k} : std::vector{k, rank}, T(-0.0625f), T(0.0625f)); - auto* b = builder.MakeInput(transposed ? std::vector{n, rank} : std::vector{rank, n}, T(-0.0625f), T(0.0625f)); - if (isolate == 1 || isolate == 2 || isolate == 4) { - auto* arg = isolate == 2 ? a : b; - auto values = builder.feeds_.at(arg->Name()).GetMutable()->MutableDataAsSpan(); - std::fill(values.begin(), values.end(), T(0.0f)); - const int64_t cols = arg->Shape()->dim(1).dim_value(); - for (int64_t r = 0; r < rank; ++r) values[r * cols + r] = T(1.0f); - } - if (isolate == 4) { - auto xv = builder.feeds_.at(x->Name()).GetMutable()->MutableDataAsSpan(); - auto av = builder.feeds_.at(a->Name()).GetMutable()->MutableDataAsSpan(); - std::fill(xv.begin(), xv.end(), T(0.0f)); - std::fill(av.begin(), av.end(), T(0.0f)); - xv[0] = T(10000.0f); - xv[1] = T(0.001f); - xv[2] = T(-10000.0f); - xv[3] = T(0.001f); - for (int i = 0; i < 4; ++i) av[i] = T(1.0f); - } - if (zeros) { - for (auto* arg : {a, b}) { - ONNX_NAMESPACE::TensorProto init; - init.set_name(arg->Name()); - init.set_data_type(utils::ToTensorProtoElementType()); - int64_t count = 1; - for (const auto& dim : arg->Shape()->dim()) { - init.add_dims(dim.dim_value()); - count *= dim.dim_value(); - } - init.set_raw_data(std::string(static_cast(count) * sizeof(T), '\0')); - builder.graph_.AddInitializedTensor(init); - } - builder.graph_.SetInputs({x, a, b}); - } - auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* t = expose_intermediate ? builder.MakeOutput(std::vector{batch, m, rank}) - : builder.MakeIntermediate(std::vector{batch, m, rank}); - auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); - auto* y = builder.MakeOutput(std::vector{batch, m, n}); - NodeAttributes attrs; - for (auto [name, value] : {std::pair{"K", k}, {"N", n}, {"bits", int64_t{4}}, {"block_size", int64_t{32}}}) { - utils::SetNodeAttribute(utils::MakeAttribute(name, value), attrs); - } - auto& base_node = builder.AddNode("MatMulNBits", {x, w, scales}, {base}, kMSDomain); - for (const auto& [name, attr] : attrs) base_node.AddAttributeProto(attr); - base_node.AddAttribute("accuracy_level", int64_t{4}); - auto* lora_x = different_x ? builder.MakeInput({batch, m, k}, T(0.0f), T(1.0f)) : x; - auto* matmul_a = a; - auto* matmul_b = b; - if (transposed) { - matmul_a = builder.MakeIntermediate(std::vector{k, rank}); - matmul_b = builder.MakeIntermediate(std::vector{rank, n}); - auto& ta = builder.AddNode("Transpose", {a}, {matmul_a}); - auto& tb = builder.AddNode("Transpose", {b}, {matmul_b}); - for (auto* node : {&ta, &tb}) { - node->AddAttribute("perm", std::vector{1, 0}); - node->SetExecutionProviderType(kWebGpuExecutionProvider); - } - } - auto& down = builder.AddNode("MatMul", {lora_x, matmul_a}, {t}); - auto& up = builder.AddNode("MatMul", {t, matmul_b}, {delta}); - auto& add = builder.AddNode("Add", {delta, base}, {y}); - for (auto* node : {&base_node, &down, &up, &add}) node->SetExecutionProviderType(kWebGpuExecutionProvider); -} - -Status CheckCount(const Graph& graph, size_t count) { - size_t found = 0; - for (const auto& node : graph.Nodes()) found += node.OpType() == "MatMulNBitsLoRA"; - ORT_RETURN_IF_NOT(found == count, "Unexpected fused LoRA count: ", found); - return Status::OK(); -} - -} // namespace - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionPreservesOverridableInputs) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, true); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { - ORT_RETURN_IF_ERROR(CheckCount(g, 1)); - ORT_RETURN_IF_NOT(g.GetOverridableInitializers().size() == 2, "Adapter inputs were removed."); - for (const auto* arg : g.GetOverridableInitializers()) { - const ONNX_NAMESPACE::TensorProto* init = nullptr; - ORT_RETURN_IF_NOT(g.GetInitializedTensor(arg->Name(), init), "Missing adapter default."); - ORT_RETURN_IF_NOT(init->dims(0) == 0 || init->dims(1) == 0, "Default was not emptied."); - } - return Status::OK(); - })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedAdapters) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, true, false, false, 32, false, true); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { - ORT_RETURN_IF_ERROR(CheckCount(g, 1)); - for (const auto& node : g.Nodes()) { - ORT_RETURN_IF_NOT(node.OpType() != "Transpose", "Adapter transpose was not absorbed."); - if (node.OpType() == "MatMulNBitsLoRA") { - ORT_RETURN_IF_NOT(node.GetAttributes().at("adapter_transposed").i() == 1, "Missing layout flag."); - } - } - ORT_RETURN_IF_NOT(g.GetOverridableInitializers().size() == 2, "Adapter inputs were removed."); - return Status::OK(); - })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedNumerics) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { - SCOPED_TRACE(mode[0] ? mode : "default"); - for (int64_t m : {1, 2, 37}) { - for (int64_t n : {128, 132}) { - SCOPED_TRACE(m); - SCOPED_TRACE(n); - const bool parallel_down = m > 1 && (std::string_view(mode) != "fuse_down" || m < 4 || n % 16 != 0); - RunWebGpuFusionTransformerTest( - [m, n](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, n); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - parallel_down ? kFp16ParallelDownAbsTolerance : kFp16FusionAbsTolerance, kFp16FusionRelTolerance, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - if (mode[0]) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - } - }); - } - } - } -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedProjectionNumerics) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - // An identity up/down adapter isolates each projection. Zero base weights - // expose low-magnitude errors; the last case stresses FP32 cancellation. - for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { - SCOPED_TRACE(mode[0] ? mode : "default"); - for (int isolate : {1, 2, 3, 4}) { - for (int64_t k : {2048, 3072, 8192}) { - SCOPED_TRACE(isolate); - SCOPED_TRACE(k); - RunWebGpuFusionTransformerTest( - [isolate, k](ModelTestBuilder& b) { BuildLoRA(b, 1, true, false, false, 32, false, true, 2048, isolate, k); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - isolate == 4 ? 0.0 : kFp16FusionAbsTolerance, isolate == 4 ? 0.0 : kFp16FusionRelTolerance, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - if (mode[0]) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - } - }); - } - } - } -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionTransposedBatchedNumerics) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { - SCOPED_TRACE(mode[0] ? mode : "default"); - for (int64_t batch : {2, 4}) { - for (int64_t m : {1, 3, 37}) { - SCOPED_TRACE(batch); - SCOPED_TRACE(m); - const bool parallel_down = std::string_view(mode) != "fuse_down" || m < 4; - RunWebGpuFusionTransformerTest( - [batch, m](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, true, 128, 0, 0, batch); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - parallel_down ? kFp16ParallelDownAbsTolerance : kFp16FusionAbsTolerance, kFp16FusionRelTolerance, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - if (mode[0]) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - } - }); - } - } - } -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsObservableIntermediate) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, true); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { return CheckCount(g, 0); })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsDifferentActivation) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, true); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { return CheckCount(g, 0); })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionRejectsUnsupportedRank) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, false, 16); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { return CheckCount(g, 0); })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionLeavesLegacyPackedShapesAlone) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildLoRA(b, 1, false, false, false, 32, true); }, 21, *logger_, - std::make_unique(), TransformerLevel::Level2, 1, nullptr, - [](const Graph& g) { return CheckCount(g, 0); })); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionMatchesWebGpuDecodeAndPrefill) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (int64_t m : {1, 37}) { - SCOPED_TRACE(m); - RunWebGpuFusionTransformerTest( - [m](ModelTestBuilder& b) { BuildLoRA(b, m, true); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - kFp16FusionAbsTolerance, kFp16FusionRelTolerance, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); - } -} - -template -void RunLoRASwitchesFeedsAndDefaults(const logging::Logger& logger) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - Model model("LoRA switching", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), - {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); - auto& graph = model.MainGraph(); - ModelTestBuilder helper(graph); - BuildLoRA(helper, 1, true); - helper.SetGraphOutputs(); - ASSERT_STATUS_OK(graph.Resolve()); - const auto adapters = graph.GetOverridableInitializers(); - ASSERT_EQ(adapters.size(), 2u); - std::string data; - model.ToProto().SerializeToString(&data); - std::vector expected; - for (bool fused : {false, true}) { - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Level1; - InferenceSessionWrapper session(options, GetEnvironment()); - ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); - ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); - if (fused) { - ASSERT_STATUS_OK(session.RegisterGraphTransformer(std::make_unique(), - TransformerLevel::Level2)); - } - ASSERT_STATUS_OK(session.Initialize()); - // Off -> adapter A -> adapter B -> partially supplied -> off. - for (int state = 0; state < 5; ++state) { - auto feeds = helper.feeds_; - for (size_t i = 0; i < adapters.size(); ++i) { - const auto& name = adapters[i]->Name(); - if (state == 0 || state == 4 || (state == 3 && i == 1)) { - feeds.erase(name); - } else { - auto* tensor = feeds.at(name).GetMutable(); - auto values = tensor->MutableDataAsSpan(); - std::fill(values.begin(), values.end(), T(state == 2 ? -0.03125f : 0.0625f)); - } - } - std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, helper.output_names_, &outputs)); - ASSERT_EQ(outputs.size(), 1u); - if (fused) { - auto result = CompareOrtValue(outputs[0], expected[state], 0.0, 0.0, false); - EXPECT_EQ(result.first, COMPARE_RESULT::SUCCESS) << "state=" << state << " " << result.second; - } else { - expected.push_back(std::move(outputs[0])); - } - } - } -} -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaults) { - RunLoRASwitchesFeedsAndDefaults(*logger_); -} -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionSwitchesFeedsAndDefaultsFp32) { - RunLoRASwitchesFeedsAndDefaults(*logger_); -} - -TEST_F(GraphTransformationTests, MatMulNBitsLoRAFusionFp32Numerics) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (bool transposed : {false, true}) { - for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { - for (int64_t m : {1, 2, 37}) { - for (int64_t n : {128, 132}) { - for (int64_t batch : {1, 2}) { - SCOPED_TRACE(MakeString("transposed=", transposed, ", mode=", mode, ", M=", m, - ", N=", n, ", batch=", batch)); - RunWebGpuFusionTransformerTest( - [=](ModelTestBuilder& b) { BuildLoRA(b, m, true, false, false, 32, false, transposed, n, 0, 2048, batch); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckCount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, 1e-5, 1e-5, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - }); - } - } - } - } - } -} -#endif -} // namespace onnxruntime::test From 2ea9a834aaec6b157e8b0ec472e8836ff08fccf0 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Oct 2026 09:49:32 +0800 Subject: [PATCH 16/35] Fix QLoRA fusion test template parsing on GCC 11 --- onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 614cba24d6cd8..bca218223706e 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -292,7 +292,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = std::string a_name, b_name; // Dyadic inputs keep the switching checks exact on all reduction paths. for (auto& [name, value] : builder.feeds_) { - auto* tensor = value.GetMutable(); + Tensor* tensor = value.template GetMutable(); if (tensor->IsDataType()) { auto bytes = tensor->MutableDataAsSpan(); std::fill(bytes.begin(), bytes.end(), uint8_t{129}); @@ -390,7 +390,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = const bool adapter = std::any_of(adapters.begin(), adapters.end(), [&name](const NodeArg* arg) { return arg->Name() == name; }); if (state == 0 && adapter) continue; - const auto& source = value.Get(); + const Tensor& source = value.template Get(); OrtValue input; Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); From 9ea931eb45d0ba752917a131a30a066c47edf271 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Oct 2026 15:51:01 +0800 Subject: [PATCH 17/35] [WebGPU] Generalize K-parallel DP4A dispatch for narrow outputs --- docs/WebGPU_MatMulNBitsQLoRA_Performance.md | 8 +++-- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 22 +++++++++---- ...e => dp4a_matmul_k_parallel.wgsl.template} | 30 ++++++----------- .../webgpu/quantization/dp4a_matmul_nbits.cc | 32 ++++++++++--------- .../webgpu/quantization/dp4a_matmul_nbits.h | 19 ++++++----- .../webgpu/quantization/matmul_nbits.cc | 2 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 19 +++-------- .../test/contrib_ops/matmul_4bits_test.cc | 2 +- .../test/contrib_ops/matmul_8bits_test.cc | 2 +- .../contrib_ops/matmul_nbits_qlora_test.cc | 21 ++++++++---- 10 files changed, 78 insertions(+), 79 deletions(-) rename onnxruntime/contrib_ops/webgpu/quantization/{dp4a_matmul_small_m.wgsl.template => dp4a_matmul_k_parallel.wgsl.template} (87%) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md index 60bc4593a9cc9..5af7c4cc5367d 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md @@ -75,8 +75,12 @@ Skipping adapter arithmetic does not recover those composed base fusions. - Prefill defaults to `fuse_up`: compute T, then combine the base projection with the INT8 up projection and Add. Eligible DP4A projections share one - quantized X allocation. The narrow down projection reuses the small-M DP4A - kernel with 32 output columns. + quantized X allocation. The narrow down projection reuses the shared + K-parallel DP4A kernel with 32 output columns. Within eligible DP4A paths, + this kernel handles M < 4 or (N <= 32 and K >= 2048), as well as indirect + expert weights. The narrow-N threshold requires two full K-parallel + iterations to amortize reduction overhead; shorter K retains tiled reuse + across M. - Eligible FP16 decode defaults to `fuse_down`: one dispatch computes the base and down projections in independent workgroups, followed by up/Add. Packed adapter weights are dequantized as needed; X stays floating point. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 3f22bd2cb51ae..9ff459c1a475c 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -63,8 +63,9 @@ model's `accuracy_level` permits it. Eligible FP32 DP4A decode also shares the base projection's quantized X with the down projection. Other decode down projections and all up projections retain floating-point activations. The provider's `enableMatmulFp32Accumulation` option -controls floating-point lane accumulation for FP16 inputs. FP32 inputs always -use FP32 floating-point accumulation. Intermediate T uses the activation type; +controls floating-point lane accumulation for FP16 inputs, except that the +K-parallel DP4A kernel always accumulates in FP32 when N <= 32. FP32 inputs +always use FP32 floating-point accumulation. Intermediate T uses the activation type; the specialized decode kernel performs its final subgroup reduction in FP32. These paths are not claimed to be bitwise equivalent to the unfused graph. @@ -103,10 +104,19 @@ thread, loading packed weights as `vec4` and T as `vec4` along rank. DP4A prefill quantizes X once and shares the quantized tensor and scales between the down projection and base projection, making three dispatches per -adapted projection. For eligible NVIDIA inputs with K >= 2048 and M <= 4096, -the rank-32 down projection uses the small-M DP4A reduction with a 32-column -tile. This narrow output remains underutilized by the general tiled kernel -even at larger M. The general DP4A kernel also uses 32-column tiles when N < 64. +adapted projection. Within the shared DP4A implementation, +`DP4AMatMulNBitsKParallelProgram` handles M below +`min_M_for_tile_optimization` (currently 4), or N <= 32 with K >= 2048. +Indirect expert weights also retain this row-wise path. Small M retains four +output columns per tile; otherwise narrow projections use an N-column tile. +This kernel uses FP32 accumulation for N <= 32. The rank-32 down projection +uses this shared selection without a QLoRA-specific override. Each K-parallel iteration covers 1024 +elements; requiring two full iterations for narrow N conservatively amortizes +reduction overhead, while shorter K retains tiled reuse across M. This is an +initial policy, not a universal performance crossover. Existing DP4A +eligibility checks and subgroup-matrix preference still apply. Other DP4A +projections use the general tiled kernel, with 32-column tiles when N < 64 +and 64-column tiles otherwise. Explicit `fuse_down` and `fuse_up` remain available. Shapes outside the joint decode specialization calculate T separately. Devices that select subgroup matrix arithmetic for the original diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template similarity index 87% rename from onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template rename to onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template index b5d71296e0b6f..8c0313bf4d009 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template @@ -1,6 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +// K-parallel DP4A projection: each workgroup computes one input row and +// tile_size output columns. Its 128 threads process four columns concurrently, +// using 32 K lanes per column and 32 elements per lane in each K iteration. +// The group stages 1024 activations and their scales in workgroup memory, +// reusing them across output columns, including additional four-column groups. +// Each lane accumulates scaled integer dot products; the 32 lane partials are +// reduced through inter_results to produce each output value after processing K. +// This exposes parallelism along K when small M or N yields few output tiles. + #param acc_f32 #param tile_size #param tile_size_k_vec @@ -32,27 +41,6 @@ fn write_qlora_output(offset: u32, value: output_element_t) { #include "quantization/dp4a_matmul_common.wgsl.template" -// This algorithm works to compute dot product of k in parallel, by processing k at each step amongst tile_size_k_vec threads, -// and utilizing the remaining threads in the workgroup to process additional rows of b in parallel (such that the values in shared memory for A can be reused). -// For each load of k, the tile_size_k_vec threads also reload B tile_size/num_concurrent_b_rows times to compute partial dot products of other B rows -// in order to complete all tile_size b rows in this workgroup and also reusing the loaded in register values of a. - -// 1. Each workgroup handles tile_size_k_vec * k_vectorization_in_b (32) columns and num_concurrent_b_rows of matrix B at a time, -// iterating over the columns to compute a partial dot product. -// 2. Uses vec4 vectorization where each K represents 32 elements of matrix B - -// 1. Workgroup Responsibility: -// - Processes one row of matrix A -// - Handles tile_size rows of matrix B -// -// 2. Computation Process: -// - Reads [tile_size][tile_size_k_vec] block of B data at a time -// - Each thread within workgroup computes dot products of 32 A*B elements since each K represents 32 elements of matrix B -// - Stores intermediate results in shared memory (inter_results) -// - Iterates through columns accumulating results in inter_results -// - Performs final reduction sum in inter_results for output -// sub_tile_count is the number of concurrent b rows processed by the workgroup. - const double_tile_size_k_vec = 2 * tile_size_k_vec; // Precision of the cross-tile accumulator, selected by the diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 98a2bd2ca58dd..c688313a5870f 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -55,7 +55,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { } // scale_A components = 1, b components = 4, output components = 1 -Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) const { +Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); @@ -89,7 +89,7 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co } return WGSL_TEMPLATE_APPLY( - shader, "quantization/dp4a_matmul_small_m.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), WGSL_TEMPLATE_PARAMETER(has_qlora_down, has_qlora_down_), WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), @@ -110,15 +110,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_at, Tensor* lora_down, const Tensor* lora_b_scales, - DP4AQuantizedInput* quantized_input, uint32_t small_m_tile_n, - const Tensor* lora_a_scales) { + DP4AQuantizedInput* quantized_input, const Tensor* lora_a_scales) { const bool has_qlora_down = lora_at != nullptr; - ORT_RETURN_IF_NOT( - small_m_tile_n == 0 || (nbits == 8 && N == 32 && small_m_tile_n <= 32 && small_m_tile_n % 4 == 0 && - !zero_points && !bias && !weight_index && !weight_index_indirect && !lora_t && !lora_at), - "Small-M override requires a plain rank-32 INT8 projection."); ORT_RETURN_IF_NOT(!has_qlora_down || (lora_a_scales && lora_down && a->IsDataType() && nbits == 4 && - block_size == 32 && M == 1 && dispatch_M == 1 && !small_m_tile_n && + block_size == 32 && M == 1 && dispatch_M == 1 && !weight_index && !weight_index_indirect && !zero_points && !bias && !lora_t), "Unsupported joint DP4A QLoRA decode projection."); constexpr uint32_t kVec4Components = 4; @@ -154,16 +149,23 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index != 0 || has_weight_idx_indirect; const bool single_scale_weights = (block_size == K * N); - if (small_m_tile_n || has_weight_idx_indirect || M < min_M_for_tile_optimization) { - uint32_t tile_size_k_vec = 32; - uint32_t tile_size_n = small_m_tile_n ? small_m_tile_n : 4; + const bool small_n = N <= 32; + // Each K-parallel iteration covers 32 lanes * 32 elements = 1024 K values. + // For narrow N, conservatively require two iterations to amortize reduction + // overhead; shorter K benefits more from the tiled kernel's reuse across M. + const bool use_k_parallel = M < min_M_for_tile_optimization || (small_n && K >= 2048); + // Indirect expert weights can differ by row, so they also require the row-wise kernel. + if (use_k_parallel || has_weight_idx_indirect) { + const uint32_t tile_size_k_vec = 32; + // Keep more workgroups available when there are only a few input rows. + const uint32_t tile_size_n = small_n && M >= min_M_for_tile_optimization ? N : 4; const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components); const bool broadcast_a = dispatch_M > M; // The narrow projection changes the K reduction order. Keep partial sums - // and the final reduction in FP32, then cast T once when storing it. - const bool acc_f32 = small_m_tile_n != 0 || context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{ + // and the final reduction in FP32, then cast once when storing the output. + const bool acc_f32 = small_n || context.EnableMatmulFp32Accumulation(); + DP4AMatMulNBitsKParallelProgram mul_program{ tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_qlora_up, has_qlora_down}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index bb3a0debb5d51..3282ff96ce020 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -70,15 +70,15 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_qlora_up_; }; -class DP4AMatMulNBitsSmallMProgram final : public Program { +class DP4AMatMulNBitsKParallelProgram final : public Program { public: - DP4AMatMulNBitsSmallMProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, - bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, - bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false, - bool has_qlora_up = false, bool has_qlora_down = false) - : Program{has_qlora_down ? "DP4AMatMulNBitsQLoRADownSmallM" - : has_qlora_up ? "DP4AMatMulNBitsQLoRAUpSmallM" - : "DP4AMatMulNBitsSmallMProgram"}, + DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, + bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false, + bool has_qlora_up = false, bool has_qlora_down = false) + : Program{has_qlora_down ? "DP4AMatMulNBitsQLoRADownKParallel" + : has_qlora_up ? "DP4AMatMulNBitsQLoRAUpKParallel" + : "DP4AMatMulNBitsKParallelProgram"}, tile_size_k_vec_(tile_size_k_vec), tile_size_(tile_size), nbits_(nbits), @@ -124,8 +124,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const Tensor* lora_b = nullptr, const Tensor* lora_at = nullptr, Tensor* lora_down = nullptr, const Tensor* lora_b_scales = nullptr, DP4AQuantizedInput* quantized_input = nullptr, - // Nonzero selects the K-parallel kernel for a rank-32 INT8 projection. - uint32_t small_m_tile_n = 0, const Tensor* lora_a_scales = nullptr); + const Tensor* lora_a_scales = nullptr); // The optional M / has_weight_idx_indirect / y arguments fold the original // dispatch-precondition (DP4A is preferred when M is large enough, or diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 7131961c8a5ba..f8861648d77a8 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -64,7 +64,7 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales)); } -// Apply similar idea with DP4AMatMulNBitsSmallMProgram algorithm. +// Apply similar idea with DP4AMatMulNBitsKParallelProgram algorithm. Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias); const auto& b = shader.AddInput("input_b"); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 8cad004d693c7..b7d18fd691fce 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -284,7 +284,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, narrow(n_), narrow(k_), 32, narrow(k_ / 32), kMinMForTileOptimization, 4, context, &base, 0, nullptr, - nullptr, nullptr, aq, &t, nullptr, nullptr, 0, as)); + nullptr, nullptr, aq, &t, nullptr, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && @@ -298,20 +298,9 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // Allow activation quantization for prefill only when the model opts into it. const auto down_accuracy = fuse_up && m >= kMinMForTileOptimization ? accuracy_level_ : 0; DP4AQuantizedInput quantized_x; - // Rank 32 exposes too few output tiles for the tiled DP4A kernel at these sizes. - const bool small_m_down = - fuse_up && m <= 4096 && k_ >= 2048 && context.AdapterInfo().vendor == std::string_view{"nvidia"} && - CanApplyDP4AMatrixMatMulNBits(context, down_accuracy, 32, 32, narrow(k_), 4, m, false, &t); - if (small_m_down) { - ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( - x, aq, as, nullptr, nullptr, narrow(x->Shape().Size() / k_ / m), m, m, 32, narrow(k_), 32, - narrow(k_ / 32), kMinMForTileOptimization, 8, context, &t, 0, nullptr, nullptr, nullptr, nullptr, - nullptr, nullptr, &quantized_x, 32)); - } else { - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, - nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, - &quantized_x)); - } + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, + nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + &quantized_x)); std::optional up_config; const bool subgroup_up = CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 0fc9214b6c88f..f24e54c5a9705 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -1289,7 +1289,7 @@ TEST(MatMulNBits, Float16_LargeK_AccumulatorOverflow) { // accuracy_level 4, M = 8 -> dp4a_matmul.wgsl.template (where the adapter supports it) // CanApplyDP4AMatrixMatMulNBits also needs subgroups, a non-Apple vendor and M >= 4 // (kMinMForTileOptimization) for an fp16 output, so the M = 2 case below only reaches -// dp4a_matmul_small_m on an fp32-output or Qualcomm adapter; elsewhere it lands on the generic +// dp4a_matmul_k_parallel on an fp32-output or Qualcomm adapter; elsewhere it lands on the generic // kernel and is one more shape for it rather than dp4a coverage. // With the option on and an adapter that has subgroup matrices, the subgroup-matrix path declines // itself (its cooperative-matrix result type is f16 and cannot honour the request) and the dispatch diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index d03639e7a5e6f..a5b13c7f35e39 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1059,7 +1059,7 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { #if defined(USE_WEBGPU) TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { - for (int64_t m : {4, 65}) { + for (int64_t m : {3, 4, 65}) { for (int64_t n : {16, 32, 48, 64, 80}) { for (bool has_zero_point : {false, true}) { TestOptions8Bits opts{}; diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index b5680c4b5c981..215093050ec00 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -259,16 +259,14 @@ TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { } } -TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { +TEST(MatMulNBitsQLoRA, NarrowKParallelReduction) { #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; #else auto initial_ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(initial_ep, nullptr); - auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); - if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId()) || - gpu.AdapterInfo().vendor != std::string_view{"nvidia"}) { - GTEST_SKIP() << "The narrow small-M projection is selected on NVIDIA DP4A devices."; + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId())) { + GTEST_SKIP() << "Requires a device that selects DP4A for the rank-32 down projection."; } #endif struct Shape { @@ -276,7 +274,9 @@ TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { int64_t k; int64_t batch; }; - for (const auto shape : {Shape{4, 2048, 1}, Shape{37, 2176, 2}, Shape{128, 3072, 1}, Shape{65, 8192, 1}}) { + for (const auto shape : {Shape{4, 128, 1}, Shape{3, 1920, 1}, Shape{4, 1920, 1}, Shape{4, 2048, 1}, + Shape{37, 2176, 2}, Shape{128, 3072, 1}, + Shape{65, 8192, 1}, Shape{4097, 128, 1}}) { const auto [m, k, batch] = shape; SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", batch=", batch)); constexpr int64_t n = 64; @@ -332,8 +332,15 @@ TEST(MatMulNBitsQLoRA, NarrowSmallMReduction) { test.AddOutput("Y", {batch, m, n}, expected); test.SetOutputAbsErr("Y", 0.002f); test.SetOutputRelErr("Y", 0.001f); + ConfigOptions options; + // Short-K fallbacks do not force FP32 accumulation like the narrow DP4A + // kernel. Enable it explicitly for this tight dense-reference comparison. + if (k < 2048) { + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + } std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); + eps.push_back(WebGpuExecutionProviderWithOptions(options)); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); } } From 2ea6f9bb47cb02c0244d271b81ba06d21861622a Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Oct 2026 16:13:03 +0800 Subject: [PATCH 18/35] [WebGPU] Enable QLoRA fast paths by device capabilities --- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 16 +++-- .../webgpu/quantization/dp4a_matmul_nbits.cc | 4 +- .../webgpu/quantization/matmul_nbits.cc | 15 +++-- .../webgpu/quantization/matmul_nbits_qlora.cc | 11 ++-- .../contrib_ops/matmul_nbits_qlora_test.cc | 64 ++++++++++++++++--- 5 files changed, 85 insertions(+), 25 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 9ff459c1a475c..1ab70e347b9a1 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -77,14 +77,14 @@ insert a transpose or materialize a complete dequantized adapter. `session.webgpu.lora_fusion_mode=auto` is the default. M is the second-to-last dimension of X; a one-token prompt follows the decode selection. -| Phase | Selected mode | NVIDIA implementation | +| Phase | Selected mode | Implementation for eligible inputs | |---|---|---| | Prefill, M > 1 | `fuse_up` | INT8 down projection, then base with INT8 up projection and addition; eligible DP4A paths share quantized X | | Decode, M = 1 | `fuse_down` | Base and INT8 down projection in one dispatch, then INT8 up projection and addition | The joint decode kernel supports one total input row and K from 2048 through 8192, when the base projection does not select DP4A. FP32 base projections -can select DP4A even at M=1. On NVIDIA, eligible FP32 inputs with K from 2048 +can select DP4A even at M=1. Eligible FP32 inputs with K from 2048 through 8192 use three dispatches: quantize X once, jointly compute the INT4 base and INT8 down projections, then perform the INT8 up projection and addition. The joint DP4A shader retains the base's 128-thread workgroups and adds eight @@ -94,9 +94,15 @@ its own output. The next dispatch provides synchronization. The existing DP4A eligibility checks, including `accuracy_level=4`, still apply. Adapter-off, FP16 decode, and prefill selection are unchanged. +The joint decode paths use shape, type, and capability checks without a vendor +restriction. The subgroup reduction accepts a fixed width of 32, or explicitly +requests 32 with `SetSubgroupSize(32)` when `SubgroupSizeControl` is supported +and 32 is within the adapter's range. Other devices use the workgroup-memory +down reduction. Existing base MatMulNBits device policies still apply. + The floating-input joint kernel retains 128 threads per workgroup and adds four workgroups for the -32 down-projection outputs. On devices with fixed 32-lane subgroups, each lane +32 down-projection outputs. With fixed or explicitly selected 32-lane subgroups, each lane loads X into registers and reuses it for two rank outputs, then reduces lane partials with FP32 subgroup sums. The previous workgroup-memory reduction remains the fallback. `QLoRAOutputVec4` uses 64 threads and one output per @@ -116,7 +122,9 @@ reduction overhead, while shorter K retains tiled reuse across M. This is an initial policy, not a universal performance crossover. Existing DP4A eligibility checks and subgroup-matrix preference still apply. Other DP4A projections use the general tiled kernel, with 32-column tiles when N < 64 -and 64-column tiles otherwise. +and 64-column tiles otherwise. Narrow tiles use 32 rows when the device's +maximum subgroup size is below 64; otherwise they keep 64 rows to accommodate +the existing 64-lane subgroup variant. Explicit `fuse_down` and `fuse_up` remain available. Shapes outside the joint decode specialization calculate T separately. Devices that select subgroup matrix arithmetic for the original diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index c688313a5870f..5ac3fa164ad69 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -204,7 +204,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } - const uint32_t tile_size_m = N < 64 && context.AdapterInfo().vendor == std::string_view{"nvidia"} ? 32 : 64; + // The 64-lane subgroup variant spans 64 rows. Reserve 32-row tiles for + // devices whose subgroup range cannot select that variant. + const uint32_t tile_size_m = N < 64 && context.AdapterInfo().subgroupMaxSize < 64 ? 32 : 64; const uint32_t tile_size_n = N < 64 ? 32 : 64; const uint32_t tile_size_k_vec = N < 64 && nbits == 8 && block_size == 32 ? 8 : 2; TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components}; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index f8861648d77a8..c440c971122d3 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -355,11 +355,15 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, constexpr uint32_t kU32Components = 4; uint32_t components_b_with_u32 = components_b * kU32Components; uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; - // The decode reduction assigns one rank to each 32-lane subgroup. + // The decode reduction assigns two ranks to each 32-lane subgroup. Devices + // with a variable subgroup range must explicitly request that width. const auto& adapter = context.AdapterInfo(); - const bool qlora_down_subgroups = has_qlora_down && adapter.vendor == std::string_view{"nvidia"} && - context.HasFeature(wgpu::FeatureName::Subgroups) && adapter.subgroupMinSize == 32 && - adapter.subgroupMaxSize == 32; + const bool fixed_subgroup_32 = adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; + const bool request_subgroup_32 = has_qlora_down && !fixed_subgroup_32 && + context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) && + adapter.subgroupMinSize <= 32 && adapter.subgroupMaxSize >= 32; + const bool qlora_down_subgroups = has_qlora_down && context.HasFeature(wgpu::FeatureName::Subgroups) && + (fixed_subgroup_32 || request_subgroup_32); MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, @@ -374,6 +378,9 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, has_qlora_down, qlora_down_subgroups}; program.SetWorkgroupSize(workgroup_size); + if (qlora_down_subgroups && request_subgroup_32) { + program.SetSubgroupSize(32); + } uint32_t num_N_tile = (N + tile_size - 1) / tile_size; program.SetDispatchGroupSize(num_N_tile + (has_qlora_down ? 4 : 0), dispatch_M, batch_count); program diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index b7d18fd691fce..0fbc1b4092bae 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -18,13 +18,14 @@ void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, con bool use_subgroups) { if (use_subgroups) { shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n"; + // Subgroup membership need not match consecutive local invocation indices. auto& body = shader.MainFunctionBody(); body << "let qlora_base_groups = " << base_groups << ";\n" << R"( if (workgroup_idx >= qlora_base_groups) { let rank_base = (workgroup_idx - qlora_base_groups) * 8u; - let lane = local_idx % 32u; - let rank_local = local_idx / 32u; + let lane = sg_id; + let rank_local = subgroup_index; var partial0 = qlora_acc_t(0); var partial1 = qlora_acc_t(0); for (var k = lane * 16u; k < uniforms.K; k += 512u) { @@ -278,8 +279,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // shader must not replace that base projection's activation quantization. const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), narrow(k_), 4, m, false, y); - if (!fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192 && - context.AdapterInfo().vendor == std::string_view{"nvidia"}) { + if (!fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, narrow(n_), @@ -287,8 +287,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { nullptr, nullptr, aq, &t, nullptr, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } - const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && - context.AdapterInfo().vendor == std::string_view{"nvidia"} && !base_dp4a; + const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && !base_dp4a; if (joint_down) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 215093050ec00..5cdf10665c845 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -8,6 +8,9 @@ #if !defined(ORT_USE_EP_API_ADAPTERS) #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "core/providers/webgpu/webgpu_context.h" +#if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) +#include "dawn/native/DawnNative.h" +#endif #endif #include "core/providers/webgpu/webgpu_provider_options.h" #endif @@ -109,7 +112,7 @@ void RunQuantizedPrefillAndReusedInput() { TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { RunQuantizedPrefillAndReusedInput(); } TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInputFp32) { RunQuantizedPrefillAndReusedInput(); } -TEST(MatMulNBitsQLoRA, Fp32Numerics) { +void RunFp32Numerics(const ConfigOptions& provider_options = {}) { struct Shape { int64_t m, k, n, batch; }; @@ -182,23 +185,63 @@ TEST(MatMulNBitsQLoRA, Fp32Numerics) { ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); test.Config(options); std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); + eps.push_back(WebGpuExecutionProviderWithOptions(provider_options)); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); } } } } +TEST(MatMulNBitsQLoRA, Fp32Numerics) { RunFp32Numerics(); } + +TEST(MatMulNBitsQLoRA, Fp32NumericsWithoutSubgroups) { +#if defined(ORT_USE_EP_API_ADAPTERS) || defined(__wasm__) || defined(USE_EXTERNAL_DAWN) + GTEST_SKIP() << "Requires native Dawn to create a device without subgroups."; +#else + // Keep the default context alive to initialize Dawn's proc table. + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + const auto instance_feature = wgpu::InstanceFeatureName::TimedWaitAny; + wgpu::InstanceDescriptor instance_desc{}; + instance_desc.requiredFeatureCount = 1; + instance_desc.requiredFeatures = &instance_feature; + dawn::native::Instance instance{&instance_desc}; + wgpu::RequestAdapterOptions adapter_options{}; + adapter_options.backendType = static_cast(webgpu::WebGpuContextConfig{}.backend_type); + auto adapters = instance.EnumerateAdapters(&adapter_options); + ASSERT_FALSE(adapters.empty()); + + const auto device_feature = wgpu::FeatureName::ImplicitDeviceSynchronization; + wgpu::DeviceDescriptor device_desc{}; + device_desc.requiredFeatureCount = 1; + device_desc.requiredFeatures = &device_feature; + auto device = wgpu::Device::Acquire(adapters.front().CreateDevice(&device_desc)); + ASSERT_NE(device, nullptr); + ASSERT_FALSE(device.HasFeature(wgpu::FeatureName::Subgroups)); + + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kDeviceId, "1")); + ASSERT_STATUS_OK(options.AddConfigEntry( + webgpu::options::kWebGpuInstance, std::to_string(reinterpret_cast(instance.Get())).c_str())); + ASSERT_STATUS_OK(options.AddConfigEntry( + webgpu::options::kWebGpuDevice, std::to_string(reinterpret_cast(device.Get())).c_str())); + auto external_ep = WebGpuExecutionProviderWithOptions(options); + ASSERT_NE(external_ep, nullptr); + ASSERT_FALSE(webgpu::WebGpuContextFactory::GetContext(1).DeviceHasFeature(wgpu::FeatureName::Subgroups)); + // Exercise the joint decode workgroup-memory reduction, including K tails, + // and the prefill fallback against the same dense FP64 reference. + RunFp32Numerics(options); +#endif +} + TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; #else auto initial_ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(initial_ep, nullptr); - auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); - if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId()) || - gpu.AdapterInfo().vendor != std::string_view{"nvidia"}) { - GTEST_SKIP() << "The joint DP4A decode projection is selected on NVIDIA DP4A devices."; + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId())) { + GTEST_SKIP() << "The joint decode projection requires DP4A device support."; } #endif constexpr int64_t k = 2176, n = 64, batch = 3; @@ -353,10 +396,11 @@ TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { ASSERT_NE(initial_ep, nullptr); auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); const auto& adapter = gpu.AdapterInfo(); - if (!gpu.DeviceHasFeature(wgpu::FeatureName::Subgroups) || - adapter.vendor != std::string_view{"nvidia"} || - adapter.subgroupMinSize != 32 || adapter.subgroupMaxSize != 32) { - GTEST_SKIP() << "The decode specialization requires NVIDIA 32-lane subgroups."; + const bool subgroup_32 = (adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) || + (gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl) && + adapter.subgroupMinSize <= 32 && adapter.subgroupMaxSize >= 32); + if (!gpu.DeviceHasFeature(wgpu::FeatureName::Subgroups) || !subgroup_32) { + GTEST_SKIP() << "The decode specialization requires fixed or selectable 32-lane subgroups."; } #endif for (bool acc_f32 : {false, true}) { From 41037661bd44cdbbc262b382b3d55c0db6df3a6e Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Oct 2026 22:11:31 +0800 Subject: [PATCH 19/35] [WebGPU] Simplify QLoRA dispatch and reuse DP4A for short FP32 prefill --- docs/WebGPU_MatMulNBitsQLoRA_Performance.md | 8 +- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 21 +-- .../webgpu/quantization/matmul_nbits_qlora.cc | 15 +- .../contrib_ops/matmul_nbits_qlora_test.cc | 149 +++++++++--------- .../matmul_nbits_qlora_fusion_test.cc | 72 ++++----- 5 files changed, 121 insertions(+), 144 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md index 5af7c4cc5367d..78b40e980aca4 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md @@ -73,15 +73,17 @@ Skipping adapter arithmetic does not recover those composed base fusions. ## Retained execution paths -- Prefill defaults to `fuse_up`: compute T, then combine the base projection +- Prefill (M > 1) selects `fuse_up`: compute T, then combine the base projection with the INT8 up projection and Add. Eligible DP4A projections share one quantized X allocation. The narrow down projection reuses the shared - K-parallel DP4A kernel with 32 output columns. Within eligible DP4A paths, + K-parallel DP4A kernel with four output columns for M < 4 and 32 otherwise. + FP32 short prefill uses this path when the base also selects DP4A, preserving + activation reuse without changing fusion. Within eligible DP4A paths, this kernel handles M < 4 or (N <= 32 and K >= 2048), as well as indirect expert weights. The narrow-N threshold requires two full K-parallel iterations to amortize reduction overhead; shorter K retains tiled reuse across M. -- Eligible FP16 decode defaults to `fuse_down`: one dispatch computes the base +- Eligible FP16 decode selects `fuse_down`: one dispatch computes the base and down projections in independent workgroups, followed by up/Add. Packed adapter weights are dequantized as needed; X stays floating point. - Eligible FP32 decode uses quantize X, joint DP4A base/down, then up/Add. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 1ab70e347b9a1..14349b68af81d 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -11,8 +11,7 @@ the model's existing `model.decoder.session_options` object: ```json { - "session.enable_matmul_nbits_lora_fusion": "1", - "session.webgpu.lora_fusion_mode": "auto" + "session.enable_matmul_nbits_lora_fusion": "1" } ``` @@ -74,8 +73,8 @@ insert a transpose or materialize a complete dequantized adapter. ## Dispatch -`session.webgpu.lora_fusion_mode=auto` is the default. M is the second-to-last -dimension of X; a one-token prompt follows the decode selection. +Selection is automatic. M is the second-to-last dimension of X; a one-token +prompt follows the decode selection. | Phase | Selected mode | Implementation for eligible inputs | |---|---|---| @@ -84,15 +83,14 @@ dimension of X; a one-token prompt follows the decode selection. The joint decode kernel supports one total input row and K from 2048 through 8192, when the base projection does not select DP4A. FP32 base projections -can select DP4A even at M=1. Eligible FP32 inputs with K from 2048 +can select DP4A even at M=1. Eligible FP32 decode inputs with K from 2048 through 8192 use three dispatches: quantize X once, jointly compute the INT4 base and INT8 down projections, then perform the INT8 up projection and addition. The joint DP4A shader retains the base's 128-thread workgroups and adds eight workgroups per input row for rank 32, with four ranks per workgroup. Both projections read the same quantized activation buffer and scales; each writes its own output. The next dispatch provides synchronization. The existing DP4A -eligibility checks, including `accuracy_level=4`, still apply. Adapter-off, -FP16 decode, and prefill selection are unchanged. +eligibility checks, including `accuracy_level=4`, still apply. The joint decode paths use shape, type, and capability checks without a vendor restriction. The subgroup reduction accepts a fixed width of 32, or explicitly @@ -108,6 +106,10 @@ partials with FP32 subgroup sums. The previous workgroup-memory reduction remains the fallback. `QLoRAOutputVec4` uses 64 threads and one output per thread, loading packed weights as `vec4` and T as `vec4` along rank. +FP32 short prefill (1 < M < 4) also uses standalone DP4A for the down +projection when the base selects DP4A. This preserves the base/up fusion and +reuses its activation quantization. FP16 short prefill and inputs that do not +meet the base DP4A requirements retain the floating-input down projection. DP4A prefill quantizes X once and shares the quantized tensor and scales between the down projection and base projection, making three dispatches per adapted projection. Within the shared DP4A implementation, @@ -125,9 +127,8 @@ projections use the general tiled kernel, with 32-column tiles when N < 64 and 64-column tiles otherwise. Narrow tiles use 32 rows when the device's maximum subgroup size is below 64; otherwise they keep 64 rows to accommodate the existing 64-lane subgroup variant. -Explicit `fuse_down` and `fuse_up` remain available. Shapes outside the joint -decode specialization calculate T -separately. Devices that select subgroup matrix arithmetic for the original +Shapes outside the joint decode specialization calculate T separately. +Devices that select subgroup matrix arithmetic for the original INT8 up projection retain that kernel and a separate Add to preserve arithmetic. ## Adapter off and graph matching diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 0fbc1b4092bae..c88657a220136 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -231,11 +231,6 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { k_(info.GetAttr("K")), n_(info.GetAttr("N")), accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)) { - const auto mode = info.GetConfigOptions().GetConfigOrDefault("session.webgpu.lora_fusion_mode", "auto"); - ORT_ENFORCE(mode == "auto" || mode == "fuse_up" || mode == "fuse_down", - "session.webgpu.lora_fusion_mode must be auto, fuse_down or fuse_up."); - automatic_ = mode == "auto"; - fuse_up_ = mode == "fuse_up"; ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0 && n_ % 4 == 0 && info.GetAttr("bits") == 4 && info.GetAttr("block_size") == 32 && info.GetAttr("lora_bits") == 8 && info.GetAttr("lora_block_size") == 32, @@ -274,7 +269,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { t_shape.back() = 32; auto t = context.CreateGPUTensor(x->DataType(), TensorShape(t_shape)); const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); - const bool fuse_up = fuse_up_ || (automatic_ && m > 1); + const bool fuse_up = m > 1; // FP32 may select DP4A even for M=1. The joint floating-point down // shader must not replace that base projection's activation quantization. const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), @@ -294,8 +289,10 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { nullptr, 0, nullptr, nullptr, aq, &t, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } - // Allow activation quantization for prefill only when the model opts into it. - const auto down_accuracy = fuse_up && m >= kMinMForTileOptimization ? accuracy_level_ : 0; + // FP32 small-M prefill can use the existing DP4A down kernel and share + // quantized X with the base projection when accuracy_level permits it. + const auto down_accuracy = + fuse_up && (m >= kMinMForTileOptimization || (x->IsDataType() && base_dp4a)) ? accuracy_level_ : 0; DP4AQuantizedInput quantized_x; ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, @@ -331,8 +328,6 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { int64_t k_; int64_t n_; int64_t accuracy_level_; - bool automatic_; - bool fuse_up_; }; ONNX_OPERATOR_KERNEL_EX(MatMulNBitsQLoRA, kMSDomain, 1, kWebGpuExecutionProvider, diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 5cdf10665c845..44dd63f9883dd 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -20,30 +20,25 @@ namespace onnxruntime::test { template void RunEmptyAndActiveAdapters() { - for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { - for (int64_t m : {1, 37}) { - for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { - auto ep = DefaultWebGpuExecutionProvider(); - ASSERT_NE(ep, nullptr); - OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); - test.AddAttribute("K", 64); - test.AddAttribute("N", 128); - test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); - test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); - test.AddInput("S", {128, 2}, std::vector(256, T(0.03125f)), true); - test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); - test.AddInput("AS", {32, 2}, std::vector(64, T(0.03125f))); - test.AddInput("BQ", {128, ranks.second, 32}, std::vector(128 * ranks.second * 32, 129)); - test.AddInput("BS", {128, 1}, std::vector(128, T(0.0625f))); - const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; - test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); - SessionOptions options; - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - test.Config(options); - std::vector> eps; - eps.push_back(std::move(ep)); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } + for (int64_t m : {1, 37}) { + for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 128); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); + test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); + test.AddInput("S", {128, 2}, std::vector(256, T(0.03125f)), true); + test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); + test.AddInput("AS", {32, 2}, std::vector(64, T(0.03125f))); + test.AddInput("BQ", {128, ranks.second, 32}, std::vector(128 * ranks.second * 32, 129)); + test.AddInput("BS", {128, 1}, std::vector(128, T(0.0625f))); + const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; + test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); } } } @@ -78,7 +73,7 @@ TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { template void RunQuantizedPrefillAndReusedInput() { - for (int64_t m : {1, 3, 4, 65}) { + for (int64_t m : {1, 2, 3, 4, 65}) { for (int64_t accuracy : {0, 4}) { OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); constexpr int64_t k = 256; @@ -118,7 +113,8 @@ void RunFp32Numerics(const ConfigOptions& provider_options = {}) { }; for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, Shape{1, 2176, 64, 7}, Shape{1, 8192, 64, 1}, - Shape{2, 2048, 36, 2}, Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { + Shape{2, 2048, 36, 2}, Shape{2, 8192, 64, 1}, Shape{3, 3072, 64, 3}, + Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { const auto [m, k, n, batch] = shape; SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); @@ -164,30 +160,25 @@ void RunFp32Numerics(const ConfigOptions& provider_options = {}) { expected[row * n + col] = static_cast(base) + static_cast(delta); } } - for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { - for (int64_t accuracy : {0, 4}) { - SCOPED_TRACE(MakeString(mode, ", accuracy=", accuracy)); - OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); - test.AddAttribute("K", k); - test.AddAttribute("N", n); - test.AddAttribute("accuracy_level", accuracy); - test.AddInput("X", {batch, m, k}, x); - test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); - test.AddInput("AQ", {32, k / 32, 32}, aq); - test.AddInput("AS", {32, k / 32}, as); - test.AddInput("BQ", {n, 1, 32}, bq); - test.AddInput("BS", {n, 1}, bs); - test.AddOutput("Y", {batch, m, n}, expected); - test.SetOutputAbsErr("Y", 0.00001f); - test.SetOutputRelErr("Y", 0.00001f); - SessionOptions options; - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - test.Config(options); - std::vector> eps; - eps.push_back(WebGpuExecutionProviderWithOptions(provider_options)); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } + for (int64_t accuracy : {0, 4}) { + SCOPED_TRACE(accuracy); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", accuracy); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, bs); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + std::vector> eps; + eps.push_back(WebGpuExecutionProviderWithOptions(provider_options)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); } } } @@ -234,20 +225,22 @@ TEST(MatMulNBitsQLoRA, Fp32NumericsWithoutSubgroups) { #endif } -TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { +void RunFp32DP4AQuantization(int64_t m) { + SCOPED_TRACE(m); #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; #else auto initial_ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(initial_ep, nullptr); if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId())) { - GTEST_SKIP() << "The joint decode projection requires DP4A device support."; + GTEST_SKIP() << "Requires DP4A device support."; } #endif constexpr int64_t k = 2176, n = 64, batch = 3; - std::vector x(batch * k), quantized_x(batch * k), as(k), expected(batch * n); + const int64_t rows = batch * m; + std::vector x(rows * k), quantized_x(rows * k), as(k), expected(rows * n); std::vector aq(32 * k), bq(n * 32, 128); - for (int64_t row = 0; row < batch; ++row) { + for (int64_t row = 0; row < rows; ++row) { for (int64_t col = 0; col < k; ++col) { const float scale = 1.0f / static_cast(1 << (8 + (row + col / 128) % 3)); const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; @@ -266,7 +259,7 @@ TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { for (int64_t col = 0; col < n; ++col) bq[col * 32 + col % 32] = 129; // An identity up projection exposes every T value. The reference uses // dense FP64 products on the known grid, including a zero row and a K tail. - for (int64_t row = 0; row < batch; ++row) { + for (int64_t row = 0; row < rows; ++row) { double base = 0; for (int64_t col = 0; col < k; ++col) base += quantized_x[row * k + col] * 0.03125; for (int64_t r = 0; r < 32; ++r) { @@ -277,29 +270,29 @@ TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { expected[row * n + r] = expected[row * n + r + 32] = static_cast(base) + static_cast(sum); } } - for (const char* mode : {"auto", "fuse_down"}) { - SCOPED_TRACE(mode); - OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); - test.AddAttribute("K", k); - test.AddAttribute("N", n); - test.AddAttribute("accuracy_level", 4); - test.AddInput("X", {batch, 1, k}, x); - test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); - test.AddInput("AQ", {32, k / 32, 32}, aq); - test.AddInput("AS", {32, k / 32}, as); - test.AddInput("BQ", {n, 1, 32}, bq); - test.AddInput("BS", {n, 1}, std::vector(n, 1.0f)); - test.AddOutput("Y", {batch, 1, n}, expected); - test.SetOutputAbsErr("Y", 0.00001f); - test.SetOutputRelErr("Y", 0.00001f); - SessionOptions options; - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - test.Config(options); - std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, std::vector(n, 1.0f)); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + +TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { RunFp32DP4AQuantization(1); } + +TEST(MatMulNBitsQLoRA, SmallMPrefillDP4AQuantization) { + for (int64_t m : {2, 3}) RunFp32DP4AQuantization(m); } TEST(MatMulNBitsQLoRA, NarrowKParallelReduction) { diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index bca218223706e..8d59515e8d415 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -104,53 +104,39 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { bool zero_base; }; const ShapeCase shapes[] = {{1, 64, 128, 1, false}, {1, 2048, 128, 1, false}, {1, 3072, 132, 1, true}, {1, 8192, 128, 1, true}, {2, 2048, 132, 1, true}, {37, 2048, 128, 1, false}, {37, 3072, 132, 1, true}, {1, 2048, 128, 2, true}, {37, 2048, 132, 2, true}}; - for (const char* mode : {"", "auto", "fuse_down", "fuse_up"}) { - for (const auto& shape : shapes) { - SCOPED_TRACE(mode[0] ? mode : "default"); - SCOPED_TRACE(shape.m); - SCOPED_TRACE(shape.k); - SCOPED_TRACE(shape.n); - SCOPED_TRACE(shape.batch); - // DP4A prefill quantizes X; subgroup decode changes the K reduction - // order. Keep exact comparisons for the other dispatch paths. - const bool approximate = - (shape.m >= 4 && std::string_view{mode} != "fuse_down") || - (shape.m == 1 && shape.batch == 1 && shape.k >= 2048 && shape.k <= 8192 && - std::string_view{mode} != "fuse_up"); - RunWebGpuFusionTransformerTest( - [shape](ModelTestBuilder& b) { - BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); - }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), - [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - if (mode[0]) - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - }); - } + for (const auto& shape : shapes) { + SCOPED_TRACE(shape.m); + SCOPED_TRACE(shape.k); + SCOPED_TRACE(shape.n); + SCOPED_TRACE(shape.batch); + // DP4A prefill quantizes X; subgroup decode changes the K reduction + // order. Keep exact comparisons for the other dispatch paths. + const bool approximate = + shape.m >= 4 || + (shape.m == 1 && shape.batch == 1 && shape.k >= 2048 && shape.k <= 8192); + RunWebGpuFusionTransformerTest( + [shape](ModelTestBuilder& b) { + BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); + }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), + [] { return DefaultWebGpuExecutionProvider(); }); } } TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Numerics) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (const char* mode : {"auto", "fuse_down", "fuse_up"}) { - for (const auto& [m, n] : {std::pair{1, 128}, {1, 132}, {2, 128}, {37, 132}}) { - SCOPED_TRACE(MakeString(mode, ", M=", m, ", N=", n)); - // Eligible FP32 decode shares the base projection's quantized X. - const bool quantized_down = (m >= 4 && std::string_view{mode} != "fuse_down") || - (m == 1 && n % 16 == 0 && std::string_view{mode} != "fuse_up"); - RunWebGpuFusionTransformerTest( - [m, n](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, n); }, - [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, - TransformerLevel::Level1, TransformerLevel::Level2, 21, - quantized_down ? 0.002 : 0.00001, quantized_down ? 0.001 : 0.00001, - std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }, - [mode](SessionOptions& options) { - ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.webgpu.lora_fusion_mode", mode)); - }); - } + for (const auto& [m, n] : {std::pair{1, 128}, {1, 132}, {2, 128}, {2, 132}, {3, 128}, {37, 132}}) { + SCOPED_TRACE(MakeString("M=", m, ", N=", n)); + // Eligible FP32 decode and short prefill share quantized X with the base. + const bool quantized_down = m >= 4 || n % 16 == 0; + RunWebGpuFusionTransformerTest( + [m, n](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, n); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + quantized_down ? 0.002 : 0.00001, quantized_down ? 0.001 : 0.00001, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); } } @@ -278,7 +264,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { template void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (int64_t m : {1, 37}) { + for (int64_t m : {1, 2, 3, 37}) { SCOPED_TRACE(m); Model model("quantized LoRA empty defaults", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); From 51a970c240fb50576a5b4ec641831c40022226c6 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Oct 2026 22:21:11 +0800 Subject: [PATCH 20/35] [WebGPU] Fix QLoRA CI tolerances and execute FP16 INT2 tests --- .../test/contrib_ops/matmul_2bits_test.cc | 12 +++++++ .../matmul_nbits_qlora_fusion_test.cc | 35 ++++++++++++++----- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 26106177956d1..bf0fb13783713 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -255,6 +255,18 @@ void RunTest2Bits(const TestOptions2Bits& opts) { test.ConfigEps(std::move(execution_providers)); test.RunWithConfig(); } +#ifdef USE_WEBGPU + if constexpr (std::is_same::value) { + // FP16 WebGPU cases must execute too; an unused OpTester traps in Debug + // and silently omits coverage in Release. + execution_providers.push_back(DefaultWebGpuExecutionProvider()); + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.ConfigEps(std::move(execution_providers)); + test.RunWithConfig(); + } +#endif } template diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 8d59515e8d415..2ca6802fc3036 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -109,18 +109,15 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { SCOPED_TRACE(shape.k); SCOPED_TRACE(shape.n); SCOPED_TRACE(shape.batch); - // DP4A prefill quantizes X; subgroup decode changes the K reduction - // order. Keep exact comparisons for the other dispatch paths. - const bool approximate = - shape.m >= 4 || - (shape.m == 1 && shape.batch == 1 && shape.k >= 2048 && shape.k <= 8192); + // FP16 fusion can change rounding on Metal even without DP4A or subgroup + // reduction. Use the FP16 error budget for every dispatch path. RunWebGpuFusionTransformerTest( [shape](ModelTestBuilder& b) { BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, TransformerLevel::Level1, TransformerLevel::Level2, 21, - approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), + 0.002, 0.001, std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); } } @@ -150,12 +147,13 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Accumulation) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 2, 37}) { SCOPED_TRACE(m); - const bool approximate = m == 1 || m >= 4; + // FP32 accumulation still produces FP16 intermediates and output, so + // fused and unfused execution need not be bitwise identical. RunWebGpuFusionTransformerTest( [m](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, 132, 1, false, false, true); }, [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, TransformerLevel::Level1, TransformerLevel::Level2, 21, - approximate ? 0.002 : 0.0, approximate ? 0.001 : 0.0, std::make_unique(), + 0.002, 0.001, std::make_unique(), [] { ConfigOptions options; ORT_THROW_IF_ERROR(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, @@ -228,6 +226,20 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { ASSERT_STATUS_OK(graph.Resolve()); const auto adapters = graph.GetOverridableInitializers(); ASSERT_EQ(adapters.size(), 4u); + // Keep this state-transition test exact across backends. With X=1/4 and + // unit packed weights, the adapter delta is 512 with default scales and + // 1/2 with supplied scales; random-data roundoff is covered by Numerics. + for (auto& [name, value] : builder.feeds_) { + auto* tensor = value.GetMutable(); + if (tensor->IsDataType()) { + auto bytes = tensor->MutableDataAsSpan(); + std::fill(bytes.begin(), bytes.end(), uint8_t{129}); + } else { + auto values = tensor->MutableDataAsSpan(); + const float fill = tensor->Shape().NumDimensions() == 3 ? 0.25f : 0.03125f; + std::fill(values.begin(), values.end(), MLFloat16(fill)); + } + } std::string data; model.ToProto().SerializeToString(&data); std::vector expected; @@ -241,6 +253,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { ASSERT_STATUS_OK( session.RegisterGraphTransformer(std::make_unique(), TransformerLevel::Level2)); ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); // Off, packed weights only, scales only, all parameters, then off again. for (int state = 0; state < 5; ++state) { auto feeds = builder.feeds_; @@ -256,6 +269,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { } else { const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << "state=" << state << " " << comparison.second; + const auto base = expected[0].Get().DataAsSpan(); + const auto actual = output[0].Get().DataAsSpan(); + const float delta = state == 1 ? 512.0f : (state == 3 ? 0.5f : 0.0f); + for (size_t i = 0; i < actual.size(); ++i) { + EXPECT_EQ(actual[i], MLFloat16(base[i].ToFloat() + delta)) << "state=" << state << ", i=" << i; + } } } } From 2611096138e9f1aae7e7bfd17e54e6361bf18997 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 19 Oct 2026 17:25:14 +0800 Subject: [PATCH 21/35] [WebGPU] Note follow-up for arbitrary QLoRA ranks --- .../contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index c88657a220136..bf4f3fc9df32e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -251,6 +251,8 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { ORT_RETURN_IF_NOT( w->Shape() == TensorShape({n_, k_ / 32, 16}) && s->Shape().Size() == SafeInt(n_) * (k_ / 32), "Invalid base weight/scales shape."); + // TODO: Support arbitrary adapter ranks with a rank uniform, per-block up scales, + // and vector-tail handling while preserving rank-32 performance. const auto rank_a = aq->Shape()[0]; const auto blocks_b = bq->Shape()[1]; ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (blocks_b == 0 || blocks_b == 1) && aq->Shape()[1] == k_ / 32 && From 7b9880263fd9cdfb9deebd669e7b1dbfdeee9b80 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 19 Oct 2026 17:27:07 +0800 Subject: [PATCH 22/35] Remove duplicate telemetry test registration after merging main --- cmake/onnxruntime_unittests.cmake | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 8ae494fb513a1..4e630b5bb45d2 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -474,26 +474,6 @@ if(WIN32) if(onnxruntime_USE_1DS_TELEMETRY) list(APPEND onnxruntime_test_framework_src_patterns "${TEST_SRC_DIR}/platform/windows_telemetry/*.cc" ) - - if(NOT GDK_PLATFORM AND NOT CMAKE_SYSTEM_NAME STREQUAL "WindowsStore") - set(telemetry_lifecycle_test_dir "${TEST_SRC_DIR}/platform/windows_telemetry/lifecycle") - onnxruntime_add_shared_library(onnxruntime_telemetry_shutdown_test_library - "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.cc" - "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.def") - target_link_libraries(onnxruntime_telemetry_shutdown_test_library PRIVATE onnxruntime_common mat) - onnxruntime_add_executable(onnxruntime_telemetry_shutdown_test - "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test.cc") - add_dependencies(onnxruntime_telemetry_shutdown_test onnxruntime_telemetry_shutdown_test_library) - foreach(mode process_exit normal_destruction explicit_unload) - add_test(NAME onnxruntime_telemetry_${mode}_test - COMMAND onnxruntime_telemetry_shutdown_test - "$" ${mode} - WORKING_DIRECTORY "$") - set_tests_properties(onnxruntime_telemetry_${mode}_test PROPERTIES TIMEOUT 15) - endforeach() - set_target_properties(onnxruntime_telemetry_shutdown_test_library onnxruntime_telemetry_shutdown_test - PROPERTIES FOLDER "ONNXRuntimeTest") - endif() endif() endif() From a7900dc0525ac0238791a7c07346e5f87d673c5b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 19 Oct 2026 17:50:56 +0800 Subject: [PATCH 23/35] [WebGPU] Adapt QLoRA decode to subgroup ID opt-in --- onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index c440c971122d3..fd5b52830707a 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -363,6 +363,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) && adapter.subgroupMinSize <= 32 && adapter.subgroupMaxSize >= 32; const bool qlora_down_subgroups = has_qlora_down && context.HasFeature(wgpu::FeatureName::Subgroups) && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && (fixed_subgroup_32 || request_subgroup_32); MatMulNBitsProgram program{tile_size, static_cast(nbits), @@ -377,6 +378,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, has_qlora_up, has_qlora_down, qlora_down_subgroups}; + program.SetUseSubgroupId(qlora_down_subgroups); program.SetWorkgroupSize(workgroup_size); if (qlora_down_subgroups && request_subgroup_32) { program.SetSubgroupSize(32); From be2c49cf4827521fba40f0e34c19501fafc66cda Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 19 Oct 2026 18:11:33 +0800 Subject: [PATCH 24/35] [WebGPU] Guard QLoRA down projection against padded dispatches --- .../dp4a_matmul_k_parallel.wgsl.template | 4 + .../webgpu/quantization/matmul_nbits_qlora.cc | 13 +- .../webgpu/dispatch_normalization_test.cc | 115 ++++++++++++++++++ 3 files changed, 126 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template index 8c0313bf4d009..440ba6202913c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template @@ -109,6 +109,10 @@ fn qlora_down(group: u32, local_idx: u32) { $MAIN { #if has_qlora_down let base_groups = uniforms.batch_count * uniforms.dispatch_M * uniforms.num_N_tile; + // Normalization may pad the dispatch beyond the eight adapter groups per batch. + if (workgroup_idx >= base_groups + uniforms.batch_count * 8u) { + return; + } if (workgroup_idx >= base_groups) { qlora_down(workgroup_idx - base_groups, local_idx); return; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index bf4f3fc9df32e..3bdcb18340cb2 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -16,12 +16,14 @@ namespace onnxruntime::contrib::webgpu { void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32, bool use_subgroups) { + auto& body = shader.MainFunctionBody(); + // Dispatch normalization can append workgroups beyond the four adapter groups. + body << "let qlora_base_groups = " << base_groups << ";\n" + << "if (workgroup_idx >= qlora_base_groups + 4u) { return; }\n"; if (use_subgroups) { shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n"; // Subgroup membership need not match consecutive local invocation indices. - auto& body = shader.MainFunctionBody(); - body << "let qlora_base_groups = " << base_groups << ";\n" - << R"( + body << R"( if (workgroup_idx >= qlora_base_groups) { let rank_base = (workgroup_idx - qlora_base_groups) * 8u; let lane = sg_id; @@ -66,8 +68,7 @@ if (workgroup_idx >= qlora_base_groups) { var qlora_x_tile: array, 128>; var qlora_partials: array, 8>; )"; - shader.MainFunctionBody() << "let qlora_base_groups = " << base_groups << ";\n" - << R"( + body << R"( if (workgroup_idx >= qlora_base_groups) { let rank_base = (workgroup_idx - qlora_base_groups) * 8u; let lane = local_idx % 32u; @@ -77,7 +78,7 @@ if (workgroup_idx >= qlora_base_groups) { let k4 = start / 4u + local_idx; var xv = vec4(0); if (k4 < uniforms.K / 4u) { xv = )" - << x_value << R"(; } + << x_value << R"(; } qlora_x_tile[local_idx] = xv; workgroupBarrier(); for (var rank_offset = 0u; rank_offset < 8u; rank_offset += 4u) { diff --git a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc index a197551015f3e..a77a7ae791c1d 100644 --- a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc +++ b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc @@ -26,6 +26,11 @@ #include "test/test_environment.h" #include "test/util/include/asserts.h" +#if !defined(DISABLE_CONTRIB_OPS) +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#endif + namespace onnxruntime { namespace test { namespace { @@ -218,6 +223,116 @@ void CheckResultsAndPadding(gsl::span output_data, size_t output_si << "Normalization-added workgroups modified output padding."; } +#if !defined(DISABLE_CONTRIB_OPS) +void CheckQLoRADownPadding(WebGpuExecutionProvider& ep, bool dp4a, bool subgroups) { + auto& context = WebGpuContextFactory::GetContext(0); + const auto& adapter = context.AdapterInfo(); + const bool fixed_subgroup_32 = adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; + if (subgroups && + (!context.DeviceHasFeature(wgpu::FeatureName::Subgroups) || + !context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) || + (!fixed_subgroup_32 && + (!context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl) || + adapter.subgroupMinSize > 32 || adapter.subgroupMaxSize < 32)))) { + GTEST_SKIP() << "QLoRA subgroup decode requires 32-lane subgroups and subgroup_id."; + } + + constexpr uint32_t N = 24, K = 512, rank = 32, tile_n = 8; + const uint32_t batches = dp4a ? 2 : 1; + const uint32_t base_groups = batches * (N / tile_n); + const uint32_t down_groups = dp4a ? batches * 8 : 4; + const OrtMemoryInfo memory_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault); + // Exercise both cube and square normalization with small physical buffers. + for (uint32_t limit : {dp4a ? 3u : 2u, 5u}) { + SCOPED_TRACE(limit); + std::array physical{base_groups + down_groups, 1, 1}; + ASSERT_STATUS_OK(webgpu::detail::NormalizeDispatchGroupSize(physical[0], physical[1], physical[2], limit)); + const uint32_t physical_count = physical[0] * physical[1] * physical[2]; + ASSERT_GT(physical_count, base_groups + down_groups); + const uint32_t padded_batches = dp4a ? CeilDiv(physical_count - base_groups, 8u) : 1; + const uint32_t padded_ranks = dp4a ? rank : (physical_count - base_groups) * 8; + // Oversize the backing buffers so an unguarded shader writes the canary + // instead of accessing unallocated GPU memory, even with robustness enabled. + std::vector x_data(padded_batches * K, 0.5f); + std::vector xq_data(padded_batches * K / 4, 0x40404040u); + std::vector xs_data(padded_batches * K / 128, 1.0f / 128.0f); + std::vector w_data(N * K / 2, 0x99); + std::vector s_data(N * K / 32, 0.03125f); + std::vector aq_data(padded_ranks * K, 129); + std::vector as_data(padded_ranks * K / 32, 0.0625f); + std::vector output_data(batches * N + rank, kOutputPaddingCanary); + std::vector down_data(padded_batches * padded_ranks, kOutputPaddingCanary); + auto x_buffer = CreateStorageBuffer(ep, x_data); + auto xq_buffer = CreateStorageBuffer(ep, xq_data); + auto xs_buffer = CreateStorageBuffer(ep, xs_data); + auto w_buffer = CreateStorageBuffer(ep, w_data); + auto s_buffer = CreateStorageBuffer(ep, s_data); + auto aq_buffer = CreateStorageBuffer(ep, aq_data); + auto as_buffer = CreateStorageBuffer(ep, as_data); + auto output_buffer = CreateStorageBuffer(ep, output_data); + auto down_buffer = CreateStorageBuffer(ep, down_data); + Tensor x(DataTypeImpl::GetType(), {batches, 1, K}, x_buffer.Get(), memory_info); + Tensor xq(DataTypeImpl::GetType(), {batches, 1, K / 4}, xq_buffer.Get(), memory_info); + Tensor xs(DataTypeImpl::GetType(), {batches, 1, 1, K / 128}, xs_buffer.Get(), memory_info); + Tensor w(DataTypeImpl::GetType(), {N, K / 32, 16}, w_buffer.Get(), memory_info); + Tensor s(DataTypeImpl::GetType(), {N, K / 32}, s_buffer.Get(), memory_info); + Tensor aq(DataTypeImpl::GetType(), {rank, K / 32, 32}, aq_buffer.Get(), memory_info); + Tensor as(DataTypeImpl::GetType(), {rank, K / 32}, as_buffer.Get(), memory_info); + Tensor output(DataTypeImpl::GetType(), {batches, 1, N}, output_buffer.Get(), memory_info); + Tensor down(DataTypeImpl::GetType(), {batches, 1, rank}, down_buffer.Get(), memory_info); + if (dp4a) { + contrib::webgpu::DP4AMatMulNBitsKParallelProgram program{ + 32, tile_n, 4, false, false, false, false, false, false, true, false, true}; + program.AddInputs({{&xq, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {&xs, ProgramTensorMetadataDependency::TypeAndRank}, + {&w, ProgramTensorMetadataDependency::TypeAndRank, 16}, + {&s, ProgramTensorMetadataDependency::TypeAndRank}, + {&aq, ProgramTensorMetadataDependency::Type, 16}, + {&as, ProgramTensorMetadataDependency::Type}}) + .AddOutputs({{&output, ProgramTensorMetadataDependency::TypeAndRank}, + {&down, ProgramTensorMetadataDependency::Type}}) + .AddUniformVariables({{batches}, {1u}, {N}, {K}, {K / 16}, {K / 32}, {32u}, {N / tile_n}, {0u}, {0u}, {1u}}) + .SetWorkgroupSize(128) + .SetDispatchGroupSize(physical[0], physical[1], physical[2]); + ASSERT_NO_FATAL_FAILURE(RunAndReadProgram(ep, program, output_buffer.Get(), output_data)); + } else { + contrib::webgpu::MatMulNBitsProgram program{ + tile_n, 4, false, false, false, false, false, 32, false, false, false, true, subgroups}; + program.AddInputs({{&x, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {&w, ProgramTensorMetadataDependency::TypeAndRank, 16}, + {&s, ProgramTensorMetadataDependency::TypeAndRank}, + {&aq, ProgramTensorMetadataDependency::Type, 16}, + {&as, ProgramTensorMetadataDependency::Type}}) + .AddOutputs({{&output, ProgramTensorMetadataDependency::TypeAndRank}, + {&down, ProgramTensorMetadataDependency::Type}}) + .AddUniformVariables({{1u}, {N}, {K}, {K / 4}, {K / 32}, {32u}, {K / 32}, {0u}, {N / tile_n}, {1u}, {0u}, {1u}}) + .SetWorkgroupSize(128) + .SetUseSubgroupId(subgroups) + .SetDispatchGroupSize(physical[0], physical[1], physical[2]); + if (subgroups && !fixed_subgroup_32) { + program.SetSubgroupSize(32); + } + ASSERT_NO_FATAL_FAILURE(RunAndReadProgram(ep, program, output_buffer.Get(), output_data)); + } + ep.BufferManager().Download(ep.Recording(), down_buffer.Get(), down_data.data(), down_data.size() * sizeof(float)); + CheckResultsAndPadding(output_data, batches * N, 8.0f); + CheckResultsAndPadding(down_data, batches * rank, 16.0f); + } +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRASharedDownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), false, false); +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRASubgroupDownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), false, true); +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRADP4ADownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), true, false); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + class LogicalDispatchProbeProgram final : public Program { public: LogicalDispatchProbeProgram() : Program{"LogicalDispatchProbe"} {} From b684a410cafc74a62212a6b99c175cb9d0b4c3cf Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Oct 2026 08:23:41 +0800 Subject: [PATCH 25/35] webgpu: Fuse QLoRA subgroup prefill and share narrow down input Compute the adapter down projection with a narrow subgroup-matrix tile, reuse the packed input, and fuse the up projection and addition into the subgroup base kernel. Preserve the DP4A and tiled paths and cover subgroup dispatch selection in tests. Verified on Intel Vulkan: built onnxruntime and onnxruntime_provider_test; all 17 MatMulNBitsQLoRA and WebGpuDataTransfer tests passed. Changed-file lintrunner checks passed. --- .../webgpu/quantization/matmul_nbits.cc | 23 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 57 ++--- .../subgroup_matrix_matmul_nbits.cc | 215 ++++++++++++------ .../subgroup_matrix_matmul_nbits.h | 23 +- ...ix_matmul_nbits_16x16x16_128.wgsl.template | 84 +++++++ ..._matrix_matmul_nbits_8x16x16.wgsl.template | 86 ++++++- ...up_matrix_matmul_nbits_8x8x8.wgsl.template | 48 ++++ .../contrib_ops/matmul_nbits_qlora_test.cc | 75 ++++++ 8 files changed, 485 insertions(+), 126 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index c440c971122d3..aa259254f71eb 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -235,27 +235,22 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // apple|intel - Experimental dawn support for subgroup matrix matmul. std::optional subgroup_matrix_config; - if (!has_qlora_down && CanApplySubgroupMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, - static_cast(nbits), - y->DataType() == DataTypeImpl::GetType(), - subgroup_matrix_config, M, has_weight_idx_indirect)) { - if (has_qlora_up) { - auto base_output = context.CreateGPUTensor(y->DataType(), y->Shape()); - ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits( - a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), zero_blocks_per_col, - *subgroup_matrix_config, context, &base_output, weight_index, weight_index_indirect)); - return ApplyQLoRAOutput(context, &base_output, lora_t, lora_b, lora_b_scales, y); - } + if (!has_qlora_down && + CanApplySubgroupMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, + static_cast(nbits), + y->DataType() == DataTypeImpl::GetType(), + subgroup_matrix_config, M, has_weight_idx_indirect)) { return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), zero_blocks_per_col, *subgroup_matrix_config, context, y, weight_index, - weight_index_indirect); + weight_index_indirect, lora_t, lora_b, lora_b_scales); } // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of // length of M. DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte // values). - if (!has_qlora_down && CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, - has_weight_idx_indirect, y)) { + if (!has_qlora_down && + CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, + has_weight_idx_indirect, y)) { return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, lora_t, lora_b, nullptr, nullptr, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index c88657a220136..eaa797a5932d9 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -188,21 +188,6 @@ class QLoRAOutputProgram final : public Program { bool acc_f32_; }; -class QLoRAAddProgram final : public Program { - public: - QLoRAAddProgram() : Program{"QLoRAAdd"} {} - ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(QLoRAAddProgram); - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}); - Status GenerateShaderCode(ShaderHelper& shader) const override { - shader.AddInput("base"); - shader.AddInput("delta"); - shader.AddOutput("output"); - shader.MainFunctionBody() - << "if (global_idx < uniforms.size) { output[global_idx] = base[global_idx] + delta[global_idx]; }\n"; - return Status::OK(); - } -}; - } // namespace Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, @@ -274,6 +259,23 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // shader must not replace that base projection's activation quantization. const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), narrow(k_), 4, m, false, y); + std::optional base_subgroup_matrix_config; + const auto subgroup_m = narrow(helper.M()); + const bool base_subgroup_matrix = + fuse_up && CanApplySubgroupMatrixMatMulNBits( + context, accuracy_level_, 32, narrow(helper.OutputOffsets().size()), + narrow(n_), narrow(k_), 4, x->IsDataType(), + base_subgroup_matrix_config, subgroup_m, false); + if (base_subgroup_matrix && base_subgroup_matrix_config->Is(8, 16, 16)) { + const auto down_config = SelectSubgroupMatrixConfig( + context, {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, + 8, 16, 16, 32, true}}); + if (down_config.has_value()) { + return ApplySubgroupMatrixMatMulNBitsQLoRA( + x, w, s, aq, as, bq, bs, subgroup_m, narrow(n_), narrow(k_), + *base_subgroup_matrix_config, *down_config, context, &t, y); + } + } if (!fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( @@ -292,16 +294,15 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // FP32 small-M prefill can use the existing DP4A down kernel and share // quantized X with the base projection when accuracy_level permits it. const auto down_accuracy = - fuse_up && (m >= kMinMForTileOptimization || (x->IsDataType() && base_dp4a)) ? accuracy_level_ : 0; + fuse_up && + (m >= kMinMForTileOptimization || (x->IsDataType() && base_dp4a)) + ? accuracy_level_ + : 0; DP4AQuantizedInput quantized_x; ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, &quantized_x)); - std::optional up_config; - const bool subgroup_up = - CanApplySubgroupMatrixMatMulNBits(context, 0, 32, narrow(helper.OutputOffsets().size()), - narrow(n_), 32, 8, x->IsDataType(), up_config, m); - if (fuse_up && !subgroup_up) { + if (fuse_up) { return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, &t, bq, nullptr, nullptr, bs, nullptr, &quantized_x); } @@ -309,19 +310,7 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, nullptr, 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, &quantized_x)); - if (!subgroup_up) return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); - // Keep the original up kernel on devices with a different matrix-accumulation path. - auto delta = context.CreateGPUTensor(y->DataType(), y->Shape()); - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(&t, bq, bs, nullptr, nullptr, 32, n_, 32, 0, 8, context, &delta)); - const auto size = narrow(y->Shape().Size()); - QLoRAAddProgram add; - add.AddInputs({{&base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, - {&delta, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}}) - .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) - .AddUniformVariables({{size}}) - .SetWorkgroupSize(64) - .SetDispatchGroupSize(CeilDiv(size, 64u)); - return context.RunProgram(add); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } private: diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index 2fb980fa5737b..160de58d6a2af 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -2,7 +2,10 @@ // Licensed under the MIT License. #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" + +#include #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -61,10 +64,12 @@ Status GenerateShaderCode16x16x16(ShaderHelper& shader, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, const ShaderVariableHelper& output, - uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect) { + uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up) { // Use the 128x128 tile shader for the 16x16x16 config. return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_16x16x16_128.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points), @@ -82,10 +87,13 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, const ShaderVariableHelper& output, - uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up, uint32_t tile_m, uint32_t tile_n, bool has_tail_buffer) { return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_tail_buffer, has_tail_buffer), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), @@ -104,9 +112,11 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, Status GenerateShaderCode8x8x8(ShaderHelper& shader, const ShaderVariableHelper& a, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, - const ShaderVariableHelper& output, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect) { + const ShaderVariableHelper& output, uint32_t nbits, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up) { return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points), @@ -131,17 +141,27 @@ Status SubgroupMatrixMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b_scales"); + GenerateQLoRAUp(shader, false); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); if (has_tail_buffer_) { shader.AddOutput("tail_output", ShaderUsage::None); } if (config_.Is(8, 8, 8)) { - return GenerateShaderCode8x8x8(shader, a, b, scales_b, output, nbits_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_); + return GenerateShaderCode8x8x8(shader, a, b, scales_b, output, nbits_, has_zero_points_, has_bias_, + has_weight_idx_, has_weight_idx_indirect_, has_qlora_up_); } else if (config_.Is(8, 16, 16)) { - return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_, tile_m_, tile_n_, has_tail_buffer_); + return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, + has_weight_idx_, has_weight_idx_indirect_, has_qlora_up_, tile_m_, tile_n_, + has_tail_buffer_); } else if (config_.Is(16, 16, 16)) { - return GenerateShaderCode16x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_); + return GenerateShaderCode16x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, + has_weight_idx_, has_weight_idx_indirect_, has_qlora_up_); } else { return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, "Unsupported subgroup matrix config dimensions."); @@ -167,6 +187,53 @@ class SubgroupMatrixMatMulNBitsTailCopyProgram final : public Program= 2048 && N % 256 == 0) { + return {128, 256, 512}; + } + if (M >= 256 && N % 128 == 0) { + return {128, 128, 512}; + } + } + return {64, 64, 256}; + } + if (config.Is(16, 16, 16)) { + return {128, 128, 128}; + } + return {32, 64, 128}; +} + +Status PrepackMatMulNBitsInput(const Tensor* a, uint32_t M, uint32_t K, uint32_t padded_M, + const SubgroupMatrixConfig& config, onnxruntime::webgpu::ComputeContext& context, + Tensor& packed) { + PrepackProgram program{config.M, config.K}; + program.SetWorkgroupSize(config.subgroupSize); + if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + program.SetSubgroupSize(config.subgroupSize); + } + ORT_ENFORCE(K % config.K == 0 && padded_M % config.M == 0); + program.SetDispatchGroupSize(padded_M / config.M, K / config.K, 1); + packed = context.CreateGPUTensor(a->DataType(), TensorShape{padded_M, K}); + program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddOutputs({{&packed, ProgramTensorMetadataDependency::Rank, packed.Shape(), 1}}) + .AddUniformVariables({{M}, {K}}) + .CacheHint(config.M, config.K); + return context.RunProgram(program); +} + +} // namespace + Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, const Tensor* bias, uint32_t M, @@ -178,40 +245,30 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, - const Tensor* weight_index_indirect) { + const Tensor* weight_index_indirect, + const Tensor* lora_t, + const Tensor* lora_b, + const Tensor* lora_b_scales, + const Tensor* prepacked_a, + uint32_t narrow_tile_m) { const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; - - // Determine tile sizes first (needed for prepack padding). - uint32_t tile_m = 32; - uint32_t tile_n = 64; - uint32_t workgroup_size = 128; - if (config.Is(8, 16, 16)) { - // Default: 64x64 tile size, 256 thread. - tile_m = 64; - tile_n = 64; - workgroup_size = 256; - // Bias output uses a scratch tile, so keep it at 64x64 for workgroup memory limits. - if (!has_bias) { - if (M >= 2048 && N % 256 == 0) { - // Large M (M>=2048): 128x256 tile size, 512 threads. - tile_m = 128; - tile_n = 256; - workgroup_size = 512; - } else if (M >= 256 && N % 128 == 0) { - // Mid-sized M (256<=M<2048): 128x128 tile size, 512 threads. - tile_m = 128; - tile_n = 128; - workgroup_size = 512; - } - } - } else if (config.Is(16, 16, 16)) { - // 16x16x16 config: 4 subgroups, 128 threads, 128x128 tiles - tile_m = 128; - tile_n = 128; - workgroup_size = 128; + const bool has_qlora_up = lora_t != nullptr; + ORT_RETURN_IF_NOT((lora_b != nullptr) == has_qlora_up && (lora_b_scales != nullptr) == has_qlora_up, + "QLoRA up inputs must be provided together."); + + auto [tile_m, tile_n, workgroup_size] = SelectMatMulNBitsTiling(config, M, N, has_bias); + if (narrow_tile_m != 0) { + ORT_ENFORCE(config.Is(8, 16, 16) && nbits == 8 && N == 32 && + narrow_tile_m % config.M == 0 && !has_bias && !has_qlora_up); + tile_m = narrow_tile_m; + tile_n = 32; + workgroup_size = tile_m / config.M * config.subgroupSize; + } + if (config.Is(8, 16, 16) && config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32) { + ORT_ENFORCE(narrow_tile_m != 0, "FP32 subgroup accumulation is implemented for the narrow QLoRA down tile."); } ORT_ENFORCE(workgroup_size <= context.DeviceLimits().maxComputeWorkgroupSizeX && workgroup_size <= context.DeviceLimits().maxComputeInvocationsPerWorkgroup, @@ -219,32 +276,14 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // If applicable, layout optimization of input matrix A(MxK) can be used for SubgroupMatrixLoad. Tensor a_prepack; - if (config.needsPrepack) { - const auto m = config.M; - const auto k = config.K; - - // Optimize the layout of input matrix A(MxK) for SubgroupMatrixLoad. - PrepackProgram prepack_program{m, k}; - prepack_program.SetWorkgroupSize(config.subgroupSize); - if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { - prepack_program.SetSubgroupSize(config.subgroupSize); - } - - // Pad M to workgroup tile size so all subgroups read valid prepacked data. - const uint32_t padded_M = CeilDiv(M, tile_m) * tile_m; - const auto dispatch_group_size_x = padded_M / m; - ORT_ENFORCE(K % k == 0, "K must be a multiple of ", k); - const auto dispatch_group_size_y = K / k; - // Each workgroup will process one subgroup matrix of size m x k. - prepack_program.SetDispatchGroupSize(dispatch_group_size_x, dispatch_group_size_y, 1); - - TensorShape a_prepack_shape{padded_M, K}; - a_prepack = context.CreateGPUTensor(a->DataType(), a_prepack_shape); - prepack_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddOutputs({{&a_prepack, ProgramTensorMetadataDependency::Rank, a_prepack.Shape(), 1}}) - .AddUniformVariables({{M}, {K}}) - .CacheHint(m, k); - ORT_RETURN_IF_ERROR(context.RunProgram(prepack_program)); + if (prepacked_a != nullptr) { + ORT_ENFORCE(config.needsPrepack && prepacked_a->DataType() == a->DataType() && + prepacked_a->Shape().NumDimensions() == 2 && prepacked_a->Shape()[1] == K && + prepacked_a->Shape()[0] >= CeilDiv(M, tile_m) * tile_m); + a = prepacked_a; + } else if (config.needsPrepack) { + ORT_RETURN_IF_ERROR(PrepackMatMulNBitsInput(a, M, K, CeilDiv(M, tile_m) * tile_m, + config, context, a_prepack)); a = &a_prepack; } @@ -255,8 +294,12 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // padded buffer instead of `y` (which is too short by this point); the crop-copy // dispatch below moves the valid rows into `y`. This keeps the no-bias write-out // free of any bounds-checked workgroup-scratch store. - const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && (M % tile_m != 0); - SubgroupMatrixMatMulNBitsProgram matmul_program{nbits, config, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, has_tail_buffer}; + const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && + config.resultComponentType != wgpu::SubgroupMatrixComponentType::F32 && + (M % tile_m != 0); + SubgroupMatrixMatMulNBitsProgram matmul_program{nbits, config, has_zero_points, has_bias, has_weight_idx, + has_weight_idx_indirect, has_qlora_up, tile_m, tile_n, + has_tail_buffer}; matmul_program.SetWorkgroupSize(workgroup_size); // Pin kernels running on variable-size adapters to the subgroup size they were written for. @@ -278,7 +321,7 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te static_cast(config.componentType), static_cast(config.resultComponentType), config.M, config.N, config.K, config.subgroupSize, - has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, + has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, has_qlora_up, tile_m, tile_n, has_tail_buffer); if (has_zero_points) { matmul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); @@ -289,6 +332,11 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te if (has_weight_idx_indirect) { matmul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_qlora_up) { + matmul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + matmul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + matmul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } Tensor tail_buffer; if (has_tail_buffer) { @@ -322,6 +370,34 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te return Status::OK(); } +Status ApplySubgroupMatrixMatMulNBitsQLoRA( + const Tensor* x, const Tensor* w, const Tensor* scales, + const Tensor* lora_a, const Tensor* lora_a_scales, + const Tensor* lora_b, const Tensor* lora_b_scales, + uint32_t M, uint32_t N, uint32_t K, const SubgroupMatrixConfig& config, + const SubgroupMatrixConfig& down_config, + onnxruntime::webgpu::ComputeContext& context, Tensor* t, Tensor* y) { + ORT_ENFORCE(config.Is(8, 16, 16) && config.needsPrepack); + ORT_ENFORCE(down_config.Is(8, 16, 16) && down_config.needsPrepack && + down_config.componentType == config.componentType && + down_config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32 && + down_config.subgroupSize == config.subgroupSize); + constexpr uint32_t kDownTileM = 32; + const auto base_tiling = SelectMatMulNBitsTiling(config, M, N, false); + const auto padded_M = std::max(CeilDiv(M, base_tiling.m) * base_tiling.m, + CeilDiv(M, kDownTileM) * kDownTileM); + // The layout depends on the subgroup matrix dimensions, not the workgroup tile. + // Both projections can read one packed X, padded for the larger workgroup tile. + Tensor packed_x; + ORT_RETURN_IF_ERROR(PrepackMatMulNBitsInput(x, M, K, padded_M, config, context, packed_x)); + ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits( + x, lora_a, lora_a_scales, nullptr, nullptr, M, 32, K, 8, K / 32, down_config, context, t, + 0, nullptr, nullptr, nullptr, nullptr, &packed_x, kDownTileM)); + return ApplySubgroupMatrixMatMulNBits( + x, w, scales, nullptr, nullptr, M, N, K, 4, CeilDiv(K / 32, 2u) * 2u, config, context, y, + 0, nullptr, t, lora_b, lora_b_scales, &packed_x); +} + bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, uint32_t block_size, @@ -358,12 +434,9 @@ bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& cont return false; } - // Every fp16 preference below has resultComponentType == F16, and the kernels declare - // subgroup_matrix_result to match, so the accumulation inside - // subgroupMatrixMultiplyAccumulate is f16 and there is no variant of this kernel that can - // honour an f32 accumulator request. Decline the path instead of ignoring the option: the - // caller then falls through to a kernel that does honour it. Only fp16 outputs are affected; - // the fp32 config accumulates in f32 already. + // The FP16 configurations below accumulate the base projection in FP16. The separately + // selected FP32-accumulating QLoRA down tile does not support a general base projection. + // Fall back when FP32 accumulation is requested instead of ignoring the option. if (is_fp16 && context.EnableMatmulFp32Accumulation()) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h index 2f3c49171b80c..94a4d1e29ae74 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h @@ -20,14 +20,17 @@ using namespace onnxruntime::webgpu; class SubgroupMatrixMatMulNBitsProgram final : public Program { public: - SubgroupMatrixMatMulNBitsProgram(uint32_t nbits, SubgroupMatrixConfig config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, bool has_tail_buffer) - : Program{"SubgroupMatrixMatMulNBits"}, + SubgroupMatrixMatMulNBitsProgram(uint32_t nbits, SubgroupMatrixConfig config, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up, + uint32_t tile_m, uint32_t tile_n, bool has_tail_buffer) + : Program{has_qlora_up ? "SubgroupMatrixMatMulNBitsQLoRAUp" : "SubgroupMatrixMatMulNBits"}, nbits_(nbits), config_(config), has_zero_points_(has_zero_points), has_bias_(has_bias), has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, + has_qlora_up_{has_qlora_up}, tile_m_(tile_m), tile_n_(tile_n), has_tail_buffer_{has_tail_buffer} {}; @@ -48,6 +51,7 @@ class SubgroupMatrixMatMulNBitsProgram final : public Program buf_b: array; // 128 * 40 = 5120 element var coopmat_stage: array; // 16*16*4 = 1024 elements = 2048 bytes // Total shared memory: 10240 + 2048 = 12288 bytes +#if has_qlora_up +const kQLoRARank: u32 = 32; +var lora_a: array; +#endif + #if n_bits == 4 // Dequantize 4-bit quantized B weights into buf_b with padded stride. // First half: rows 0..63, Second half: rows 64..127 @@ -312,6 +318,84 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + for (var load_unit = local_idx; load_unit < BM * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let tile_m = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_m = global_base_m + tile_m; + let offset = tile_m * kQLoRARank + r4 * 4u; + if (global_m < uniforms.M) { + let values = lora_t[global_m * (kQLoRARank / 4u) + r4]; + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(values[i]); + } + } else { + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(0.0); + } + } + } + for (var load_unit = local_idx; load_unit < BN * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let tile_n = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_n = global_base_n + tile_n; + let offset = tile_n * SHMEM_STRIDE + r4 * 4u; + if (global_n < uniforms.N) { + let values = qlora_weight(global_n, r4); + for (var i = 0u; i < 4u; i++) { + buf_b[offset + i] = f16(values[i]); + } + } else { + for (var i = 0u; i < 4u; i++) { + buf_b[offset + i] = f16(0.0); + } + } + } + workgroupBarrier(); + + for (var k = 0u; k < kQLoRARank; k += kSgMatK) { + let b_base = (warp_c * WN) * SHMEM_STRIDE + k; + var b0: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base, SHMEM_STRIDE); + var b1: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 16u * SHMEM_STRIDE, SHMEM_STRIDE); + var b2: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 32u * SHMEM_STRIDE, SHMEM_STRIDE); + var b3: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 48u * SHMEM_STRIDE, SHMEM_STRIDE); + + let a_base = (warp_r * WM) * kQLoRARank + k; + var a0: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base, kQLoRARank); + var a1: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 16u * kQLoRARank, kQLoRARank); + var a2: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 32u * kQLoRARank, kQLoRARank); + var a3: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 48u * kQLoRARank, kQLoRARank); + + c00 = subgroupMatrixMultiplyAccumulate(a0, b0, c00); + c01 = subgroupMatrixMultiplyAccumulate(a0, b1, c01); + c02 = subgroupMatrixMultiplyAccumulate(a0, b2, c02); + c03 = subgroupMatrixMultiplyAccumulate(a0, b3, c03); + c10 = subgroupMatrixMultiplyAccumulate(a1, b0, c10); + c11 = subgroupMatrixMultiplyAccumulate(a1, b1, c11); + c12 = subgroupMatrixMultiplyAccumulate(a1, b2, c12); + c13 = subgroupMatrixMultiplyAccumulate(a1, b3, c13); + c20 = subgroupMatrixMultiplyAccumulate(a2, b0, c20); + c21 = subgroupMatrixMultiplyAccumulate(a2, b1, c21); + c22 = subgroupMatrixMultiplyAccumulate(a2, b2, c22); + c23 = subgroupMatrixMultiplyAccumulate(a2, b3, c23); + c30 = subgroupMatrixMultiplyAccumulate(a3, b0, c30); + c31 = subgroupMatrixMultiplyAccumulate(a3, b1, c31); + c32 = subgroupMatrixMultiplyAccumulate(a3, b2, c32); + c33 = subgroupMatrixMultiplyAccumulate(a3, b3, c33); + } + workgroupBarrier(); +#endif + // Write out results let dr = global_base_m + warp_r * WM; let dc = global_base_n + warp_c * WN; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template index 312196af6986b..f4fdab6ec96b3 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template @@ -13,8 +13,10 @@ // K dimension iterated in kChunkSizeK(32) blocks #param n_bits +#param acc_f32 #param has_zero_points #param has_bias +#param has_qlora_up #param has_tail_buffer #param has_weight_idx #param has_weight_idx_indirect @@ -55,17 +57,27 @@ const kSgMatCountK : u32 = kChunkSizeK / kSgMatDimK; // Number of subgroup_mat alias SgMatLeft = subgroup_matrix_left; alias SgMatRight = subgroup_matrix_right; +#if acc_f32 +alias SgMatResult = subgroup_matrix_result; +// The narrow down projection accumulates over long K in FP32, then rounds T to FP16. +var result_f32 : array; +#else alias SgMatResult = subgroup_matrix_result; +#endif // Shared local memory for dequantized B weights [kTileN, kChunkSizeK], column-major var tile_b : array; #if has_bias -// Scratch space for bounds-checked write-out with bias addition, [kTileM, kTileN], -// row-major. The no-bias path never needs this - see `has_tail_buffer` below. +// Scratch space for bounds-checked write-out with bias addition, [kTileM, kTileN], row-major. var scratch : array; #endif +#if has_qlora_up +const kQLoRARank : u32 = 32; +var lora_a : array; +#endif + // --- Data Loading Functions --- // scale_offset is resolved before the K loop, not recomputed for each load. @@ -175,8 +187,74 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + // Reuse the subgroup-matrix accumulators for the rank-32 adapter-up product. + // Loading the row-major adapter activation into workgroup memory avoids the + // standalone subgroup layout and adapter-up dispatches. + for (var load_unit = local_idx; load_unit < kTileM * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let row_in_tile = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_m = global_base_m + row_in_tile; + if (global_m < uniforms.M) { + let values = lora_t[global_m * (kQLoRARank / 4u) + r4]; + let offset = row_in_tile * kQLoRARank + r4 * 4u; + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(values[i]); + } + } + } + for (var load_unit = local_idx; load_unit < kTileN * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let col_in_tile = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_n = global_base_n + col_in_tile; + if (global_n < uniforms.N) { + let values = qlora_weight(global_n, r4); + let offset = col_in_tile * kQLoRARank + r4 * 4u; + for (var i = 0u; i < 4u; i++) { + tile_b[offset + i] = f16(values[i]); + } + } + } + workgroupBarrier(); + + let lora_a_base = sg_idx * kSgMatDimM * kQLoRARank; + for (var k_idx = 0u; k_idx < kQLoRARank; k_idx += kSgMatDimK) { + var sg_mat_a: SgMatLeft = subgroupMatrixLoad(&lora_a, + lora_a_base + k_idx, + kQLoRARank); + var sg_mat_offset_b = k_idx; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + var sg_mat_b: SgMatRight = subgroupMatrixLoad(&tile_b, + sg_mat_offset_b, + kQLoRARank); + sg_mat_offset_b += kSgMatStrideN; + sg_mat_acc[sg_mat_n_idx] = + subgroupMatrixMultiplyAccumulate(sg_mat_a, sg_mat_b, sg_mat_acc[sg_mat_n_idx]); + } + } + workgroupBarrier(); +#endif + // Write out -#if has_bias +#if acc_f32 + let scratch_store_base = sg_idx * kSgMatDimM * kTileN; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&result_f32, + scratch_store_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + kTileN); + } + workgroupBarrier(); + for (var i = local_idx; i < kTileM * kTileN; i += workgroup_size_x) { + let row = global_base_m + i / kTileN; + let col = global_base_n + i % kTileN; + if (row < uniforms.M && col < uniforms.N) { + output.setByOffset(row * uniforms.N + col, output_element_t(result_f32[i])); + } + } +#elif has_bias let scratch_store_base = sg_idx * kSgMatDimM * kTileN; for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { subgroupMatrixStore(&scratch, @@ -204,7 +282,7 @@ $MAIN { #endif for (var i : u32 = 0; i < ElemsPerThreadRow; i++) { if (out_global_n_base + i < uniforms.N) { - let val = f16(scratch[scratch_base + i]) + bias[bias_offset + out_global_n_base + i]; + let val = output_element_t(scratch[scratch_base + i]) + bias[bias_offset + out_global_n_base + i]; output.setByOffset(out_base + i, val); } } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template index f28360af79c45..6f8927903124e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template @@ -7,6 +7,7 @@ #param n_bits #param has_zero_points #param has_bias +#param has_qlora_up #param has_weight_idx #param has_weight_idx_indirect @@ -211,6 +212,53 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + for (var idx = local_idx; idx < tile_rows * tile_k; idx += workgroup_size_x) { + let row = idx / tile_k; + let col = idx % tile_k; + let global_row = a_global_base + row; + if (global_row < uniforms.M) { + tile_A[idx] = compute_precision(lora_t[global_row * 8u + col / 4u][col % 4u]); + } + } + for (var idx = local_idx; idx < tile_cols * tile_k; idx += workgroup_size_x) { + let row = idx / tile_k; + let col = idx % tile_k; + let global_col = b_global_base + row; + if (global_col < uniforms.N) { + tile_B[idx] = compute_precision(qlora_weight(global_col, col / 4u)[col % 4u]); + } + } + workgroupBarrier(); + + for (var step: u32 = 0; step < tile_k; step += 8u) { + let matrix_a_offset = subtile_idy * subtile_rows * tile_k + step; + var matA0: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &tile_A, matrix_a_offset, tile_k); + var matA1: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &tile_A, matrix_a_offset + 8u * tile_k, tile_k); + let matrix_b_offset = subtile_idx * subtile_cols * tile_k + step; + var matB0: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset, tile_k); + var matB1: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 8u * tile_k, tile_k); + var matB2: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 16u * tile_k, tile_k); + var matB3: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 24u * tile_k, tile_k); + + matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00); + matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01); + matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02); + matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03); + matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10); + matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11); + matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12); + matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13); + } + workgroupBarrier(); +#endif + // Write out // Write out top block subgroupMatrixStore(&scratch[subtile_id][0], 0, matC00, 8); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 44dd63f9883dd..4d93d77f9ce98 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -7,6 +7,7 @@ #if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) #if !defined(ORT_USE_EP_API_ADAPTERS) #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/webgpu_context.h" #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) #include "dawn/native/DawnNative.h" @@ -107,6 +108,80 @@ void RunQuantizedPrefillAndReusedInput() { TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { RunQuantizedPrefillAndReusedInput(); } TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInputFp32) { RunQuantizedPrefillAndReusedInput(); } +TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixUsesFusedUp) { +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific dispatch checks require direct adapter capability inspection."; +#else + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()); + constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; + const auto& adapter = gpu.AdapterInfo(); + const auto& device_configs = gpu.SubgroupMatrixConfigs(); + const auto config = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {device_configs.configs, device_configs.configCount}, + adapter.subgroupMinSize, adapter.subgroupMaxSize, + gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{kF16, kF16, 16, 16, 16, 32, true}, + {kF16, kF16, 8, 16, 16, 32, true}, + {kF16, kF16, 8, 8, 8, 32, false}}); + if (!gpu.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix) || !config.has_value()) { + GTEST_SKIP() << "Requires a compatible FP16 subgroup-matrix configuration."; + } + const auto down_config = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {device_configs.configs, device_configs.configCount}, + adapter.subgroupMinSize, adapter.subgroupMaxSize, + gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{kF16, wgpu::SubgroupMatrixComponentType::F32, 8, 16, 16, 32, true}}); + + constexpr int64_t m = 256, k = 256, n = 128; + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {1, m, k}, std::vector(m * k, MLFloat16(0.5f))); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, + std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {32, k / 32, 32}, std::vector(32 * k, 129)); + test.AddInput("AS", {32, k / 32}, std::vector(k, MLFloat16(0.03125f))); + test.AddInput("BQ", {n, 1, 32}, std::vector(n * 32, 129)); + test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(0.0625f))); + test.AddOutput("Y", {1, m, n}, std::vector(m * n, MLFloat16(12.0f))); + + SessionOptions session_options; + session_options.session_log_severity_level = static_cast(logging::Severity::kINFO); + std::vector> eps; + eps.push_back(std::move(ep)); + testing::internal::CaptureStderr(); + test.Run(session_options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + const auto logs = testing::internal::GetCapturedStderr(); + EXPECT_NE(logs.find("SubgroupMatrixMatMulNBitsQLoRAUp"), std::string::npos) << logs; + if (config->Is(8, 16, 16) && down_config.has_value()) { + auto count_program = [&logs](const std::string& name) { + const auto marker = "Starting program \"" + name + "["; + size_t count = 0; + for (auto pos = logs.find(marker); pos != std::string::npos; pos = logs.find(marker, pos + marker.size())) { + ++count; + } + return count; + }; + const auto down_count = count_program("SubgroupMatrixMatMulNBits"); + EXPECT_GT(down_count, 0u) << logs; + EXPECT_NE(logs.find("SubgroupMatrixMatMulNBits[8|2|1|8|16|16|32|"), std::string::npos) << logs; + EXPECT_EQ(count_program("SubgroupMatrixMatMulLayout"), down_count) << logs; + EXPECT_EQ(count_program("SubgroupMatrixMatMulNBitsQLoRAUp"), down_count) << logs; + EXPECT_NE(logs.find("|32|32|0]:128"), std::string::npos) << logs; + EXPECT_NE(logs.find("|128|128|0]:512"), std::string::npos) << logs; + EXPECT_EQ(logs.find("Starting program \"MatMulNBits["), std::string::npos) << logs; + EXPECT_EQ(logs.find("DP4AMatMulQuantize"), std::string::npos) << logs; + EXPECT_EQ(logs.find("Starting program \"DP4AMatMulNBits"), std::string::npos) << logs; + } + EXPECT_EQ(logs.find("QLoRAOutputVec4"), std::string::npos) << logs; + EXPECT_EQ(logs.find("QLoRAAdd"), std::string::npos) << logs; +#endif +} + void RunFp32Numerics(const ConfigOptions& provider_options = {}) { struct Shape { int64_t m, k, n, batch; From 034acaaf84755b01a57e57f64ddfcda72aa02879 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Oct 2026 08:23:41 +0800 Subject: [PATCH 26/35] webgpu: Batch unmapped input uploads on UMA devices Use the existing bounded staging upload path for multi-tensor copies on UMA too. UMA capability alone does not mean destination buffers are mapped; serial uploads unnecessarily submit each tensor separately. Existing mapped-buffer and single-copy fallbacks remain unchanged. Verified on Intel Vulkan: all 5 WebGpuDataTransfer tests passed, including chunked uploads and buffer reuse; all 12 MatMulNBitsQLoRA tests also passed. Build and changed-file lint passed. --- onnxruntime/core/providers/webgpu/data_transfer.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 0bf7b0319a7ea..7e0998b605621 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -72,7 +72,7 @@ common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { } common::Status DataTransfer::CopyTensors(const std::vector& src_dst_pairs) const { - if (src_dst_pairs.size() < 2 || buffer_manager_.SupportsUMA()) { + if (src_dst_pairs.size() < 2) { return IDataTransfer::CopyTensors(src_dst_pairs); } InlinedVector uploads; From a49f0196ec6e3fd172023d1b6e42d1a2ad48b328 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Oct 2026 08:43:39 +0800 Subject: [PATCH 27/35] webgpu: Batch plugin input uploads on the owning session stream Batch same-device CPU-to-GPU inputs on a shared non-null WebGPU session stream using the existing bounded staging uploader, including on UMA. Preserve mixed-transfer, mapped-buffer, single-copy, and streamless fallbacks, and retain allocation clears and deferred command ordering. Add real submission-count and data-integrity tests for repeated uploads, staging chunk boundaries, empty inputs, oversized tensors, and single-input fallback. The batching tests fail on the control plugin with four submissions instead of one or two. Verified Intel Vulkan: production and provider-test builds, lint, 67 plugin tests passed with 6 intentional capability-specific skips; all 17 built-in QLoRA/data-transfer tests passed. Fixed 1024-token adapter-on prefill improved from 848.1 to 716.7 ms across paired run orders (15.5%); decode remained about 35 tokens/s. All work remains local. --- .../core/providers/webgpu/ep/sync_stream.cc | 7 + .../core/providers/webgpu/ep/sync_stream.h | 3 + .../webgpu/webgpu_provider_factory.cc | 33 ++++ .../webgpu/plugin/test_webgpu_allocators.cc | 165 ++++++++++++++++++ 4 files changed, 208 insertions(+) diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc index f7831672483b3..396084276095d 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc @@ -87,6 +87,13 @@ common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void* return transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, bytes); } +void UploadTensorsOnWebGpuStream(const OrtSyncStream* stream, int device_id, + gsl::span uploads) { + auto& ep = WebGpuSyncStream::From(stream).ep_; + ORT_ENFORCE(ep.GetDeviceId() == device_id, "WebGPU upload destination does not match the stream device."); + ep.BufferManager().UploadBatch(ep.Recording(), uploads); +} + } // namespace ep } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.h b/onnxruntime/core/providers/webgpu/ep/sync_stream.h index 043f0655ec074..aec71587aebae 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.h +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.h @@ -6,6 +6,7 @@ #include #include "core/common/status.h" +#include "core/providers/webgpu/buffer_manager.h" #include "core/session/onnxruntime_c_api.h" namespace onnxruntime { @@ -18,6 +19,8 @@ OrtSyncStreamImpl* CreateWebGpuSyncStream(WebGpuExecutionProvider& ep); CommandRecordingState& GetWebGpuStreamCommandState(const OrtSyncStream* stream); common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void* src_data, bool src_is_gpu, void* dst_data, bool dst_is_gpu, size_t bytes); +void UploadTensorsOnWebGpuStream(const OrtSyncStream* stream, int device_id, + gsl::span uploads); } // namespace ep } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 520103e687241..3af2a76ff361d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -580,6 +580,39 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { return nullptr; } +#if defined(ORT_USE_EP_API_ADAPTERS) + // Batch unmapped inputs on the owning Session timeline, including on UMA devices. + // Check the entire list before copying so mixed transfers retain their original ordering. + if (num_tensors > 1 && streams != nullptr && streams[0] != nullptr) { + InlinedVector uploads; + uploads.reserve(num_tensors); + int device_id = -1; + for (size_t idx = 0; idx < num_tensors; ++idx) { + if (streams[idx] != streams[0]) break; + Ort::ConstValue src_value{src_tensors[idx]}; + Ort::UnownedValue dst_value{dst_tensors[idx]}; + auto src_info = src_value.GetTensorMemoryInfo(); + auto dst_info = dst_value.GetTensorMemoryInfo(); + if (src_info.GetDeviceType() != OrtMemoryInfoDeviceType_CPU || + dst_info.GetDeviceType() != OrtMemoryInfoDeviceType_GPU) break; + const auto* src_device = impl.ep_api.MemoryInfo_GetMemoryDevice(src_info); + const auto* dst_device = impl.ep_api.MemoryInfo_GetMemoryDevice(dst_info); + if (!CanCopyImpl(this_ptr, src_device, dst_device)) break; + const int dst_device_id = impl.ep_api.MemoryDevice_GetDeviceId(dst_device); + if (idx == 0) device_id = dst_device_id; + if (dst_device_id != device_id) break; + const size_t size = src_value.GetTensorSizeInBytes(); + auto target = static_cast(dst_value.GetTensorMutableRawData()); + if (size != 0 && wgpuBufferGetMapState(target) != WGPUBufferMapState_Unmapped) break; + uploads.push_back({src_value.GetTensorRawData(), target, size}); + } + if (uploads.size() == num_tensors) { + webgpu::ep::UploadTensorsOnWebGpuStream(streams[0], device_id, uploads); + return nullptr; + } + } +#endif + for (size_t idx = 0; idx < num_tensors; ++idx) { const OrtMemoryInfo* src_memory_info = nullptr; const OrtMemoryInfo* dst_memory_info = nullptr; diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc index 2bf2e6b3ab7d1..107446a80b07f 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc @@ -506,6 +506,119 @@ int VerifyInstanceNormScratchAllocationsStayBatched(Ort::Env& env, Ort::ConstEpD return 0; } +int VerifySessionInputUploadBatching(Ort::Env& env, Ort::ConstEpDevice ep_device, + gsl::span input_sizes, size_t expected_submissions) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(18); + auto* graph = model.mutable_graph(); + graph->set_name("webgpu_batched_session_inputs"); + auto* node = graph->add_node(); + node->set_op_type("Concat"); + node->add_output("Y"); + auto* axis = node->add_attribute(); + axis->set_name("axis"); + axis->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis->set_i(0); + InlinedVector names; + names.reserve(input_sizes.size()); + int64_t output_size = 0; + for (size_t i = 0; i < input_sizes.size(); ++i) { + names.push_back("X" + std::to_string(i)); + node->add_input(names.back()); + auto* input = graph->add_input(); + input->set_name(names.back()); + auto* type = input->mutable_type()->mutable_tensor_type(); + type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + type->mutable_shape()->add_dim()->set_dim_value(input_sizes[i]); + output_size += input_sizes[i]; + } + auto* output = graph->add_output(); + output->set_name("Y"); + auto* output_type = output->mutable_type()->mutable_tensor_type(); + output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output_type->mutable_shape()->add_dim()->set_dim_value(output_size); + + struct SubmissionProbe { + SubmissionProbe() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SubmissionProbe); + + bool armed{false}; + bool saw_program{false}; + size_t baseline{0}; + size_t input_submissions{0}; + + static void ORT_API_CALL Log(void* param, OrtLoggingLevel, const char*, const char*, const char*, + const char* message) { + auto& probe = *static_cast(param); + if (probe.armed && !probe.saw_program && + std::string_view{message}.find("Starting program") != std::string_view::npos) { + probe.saw_program = true; + probe.input_submissions = nonempty_queue_submissions.load(std::memory_order_relaxed) - probe.baseline; + } + } + } probe; + auto procs = dawn::native::GetProcs(); + procs.queueSubmit = CountNonemptyQueueSubmissions; + dawnProcSetProcs(&procs); + Ort::SessionOptions options; + options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); + options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); + options.SetLogSeverityLevel(ORT_LOGGING_LEVEL_INFO); + Ort::ThrowOnError(Ort::GetApi().SetUserLoggingFunction(options, SubmissionProbe::Log, &probe)); + options.AppendExecutionProvider_V2( + env, {ep_device}, {{"dawnProcTable", std::to_string(reinterpret_cast(&procs))}}); + const auto model_bytes = model.SerializeAsString(); + Ort::Session session(env, model_bytes.data(), model_bytes.size(), options); + const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + InlinedVector> input_data; + InlinedVector inputs; + InlinedVector input_names; + input_data.reserve(input_sizes.size()); + inputs.reserve(input_sizes.size()); + input_names.reserve(input_sizes.size()); + for (size_t i = 0; i < input_sizes.size(); ++i) { + input_data.emplace_back(static_cast(input_sizes[i])); + inputs.push_back(Ort::Value::CreateTensor(cpu_memory, input_data.back().data(), input_data.back().size(), + &input_sizes[i], 1)); + input_names.push_back(names[i].c_str()); + } + const char* output_names[] = {"Y"}; + // Warm pipelines/caches, then measure actual input submissions before the first program. + // Changing every input also catches stale uploads and clears submitted after the copy. + for (int iteration = 0; iteration < 6; ++iteration) { + for (size_t i = 0; i < input_data.size(); ++i) { + for (size_t j = 0; j < input_data[i].size(); ++j) { + input_data[i][j] = static_cast(iteration * 100 + i * 10 + j % 7); + } + } + probe.baseline = nonempty_queue_submissions.load(std::memory_order_relaxed); + probe.saw_program = false; + probe.armed = true; + auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names.data(), inputs.data(), inputs.size(), + output_names, 1); + probe.armed = false; + if (iteration >= 3 && (!probe.saw_program || probe.input_submissions != expected_submissions)) { + std::fprintf(stderr, "Input batching failed: saw_program=%d, submissions=%zu, expected=%zu\n", + static_cast(probe.saw_program), probe.input_submissions, expected_submissions); + return 11; + } + if (outputs[0].GetTensorTypeAndShapeInfo().GetElementCount() != static_cast(output_size)) return 12; + const float* actual = outputs[0].GetTensorData(); + size_t offset = 0; + for (const auto& values : input_data) { + for (float expected : values) { + if (actual[offset++] != expected) { + std::fputs("Batched Session input data was corrupted or stale\n", stderr); + return 13; + } + } + } + } + std::fputs("Session input uploads preserved data and expected submission count\n", stderr); + return 0; +} + WGPUBufferMapState InjectMappedBufferOnFree(WGPUBuffer buffer) { if (buffer == free_failure_buffer) { ++free_failure_probes; @@ -627,6 +740,58 @@ TEST_F(WebGpuSessionAllocatorDeathTest, WarmedInstanceNormalizationScratchAlloca #endif } +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsAreBatched) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16, 16, 16, 16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsSpanMultipleChunks) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + // Four 5 MiB + 4 byte inputs require two staging chunks at the 16 MiB limit. + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching( + Env(), EpDevice(), std::array{1310721, 1310721, 1310721, 1310721}, 2)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsIncludeEmptyTensors) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{0, 16, 0, 16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsHandleOversizedTensor) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{4194305, 16}, 2)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionSingleInputUploadRetainsFallback) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + #endif } // namespace From 4625103acbc3977a7f8be61fd84f056c0fd48b33 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 22 Oct 2026 14:48:56 +0800 Subject: [PATCH 28/35] webgpu: Generalize QLoRA schema and simplify correctness tests --- docs/ContribOperators.md | 32 +-- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 22 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 16 +- .../core/graph/contrib_ops/contrib_defs.cc | 102 ++++--- .../optimizer/matmul_nbits_qlora_fusion.cc | 2 +- .../test/contrib_ops/matmul_2bits_test.cc | 3 +- .../test/contrib_ops/matmul_4bits_test.cc | 3 +- .../test/contrib_ops/matmul_8bits_test.cc | 29 +- ...matmul_nbits_qlora_shape_inference_test.cc | 139 ++++++++++ .../contrib_ops/matmul_nbits_qlora_test.cc | 165 +++++------- .../matmul_nbits_qlora_fusion_test.cc | 95 ++----- .../webgpu/plugin/test_webgpu_allocators.cc | 252 ++++++------------ 12 files changed, 448 insertions(+), 412 deletions(-) create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 8b69cd2b0814f..b0a56088c71aa 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4280,7 +4280,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulNBitsQLoRA** - Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights stored as unsigned bytes with implicit zero point 128 and block scales matching the activation type. Eligible adapter down projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension in either packed adapter disables the adapter. Intermediate projections use the activation type (FP16 or FP32) before the final addition. Adapters remain replaceable. + Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). Each projection follows MatMulNBits block quantization and bit packing, with implicit zero points 2^(bits-1) for the base and 2^(lora_bits-1) for the adapters. No bias or explicit zero points are applied. The adapter rank R may be any positive integer and need not be a multiple of lora_block_size. Partial final blocks are padded; only the logical K or R elements participate in each projection. An empty rank dimension in either packed adapter disables the adapter, while scale defaults are retained. Intermediate projections and block scales use the activation type before the final addition. Adapters remain replaceable. Execution providers may support a subset of the defined types, quantization settings, and ranks. #### Version @@ -4290,38 +4290,38 @@ This version of the operator has been available since version 1 of the 'com.micr
K : int (required)
-
Input feature dimension; a positive multiple of 32.
+
Positive input feature dimension of the base and down-projection weight matrices.
N : int (required)
-
Output feature dimension; a positive multiple of 4.
+
Positive output feature dimension of the base and up-projection weight matrices.
accuracy_level : int
-
Accuracy policy for the base and eligible adapter down projections.
+
Accuracy policy following MatMulNBits: 0 (unset), 1 (fp32), 2 (fp16), 3 (bf16), or 4 (int8). Applies to the base and eligible adapter down projections; the up projection uses level 0.
bits : int
-
Base weight bit width; must be 4.
+
Base weight bit width, following MatMulNBits (2, 4, or 8).
block_size : int
-
Base weight block size; must be 32.
+
Base quantization block size along K, following MatMulNBits: a power of two, at least 16.
lora_bits : int
-
Adapter weight bit width; must be 8.
+
Adapter weight bit width, following MatMulNBits (2, 4, or 8).
lora_block_size : int
-
Adapter weight block size; must be 32.
+
Adapter quantization block size along K for lora_A and R for lora_B: a power of two, at least 16.
#### Inputs
X : T
-
Input of rank at least two, with last dimension K.
+
Input of rank at least one, with last dimension K, following MatMulNBits.
W : TQ
-
Packed INT4 base weights [N,K/32,16].
+
Packed base weights [N,ceil(K/block_size),block_size*bits/8].
scales : T
-
Base block scales with N*K/32 elements; may be flattened.
+
Base block scales [N,ceil(K/block_size)]; may be flattened.
lora_A : TQ
-
Packed INT8 down weights [R_A,K/32,32], R_A=0 or 32.
+
Packed down weights [R_A,ceil(K/lora_block_size),lora_block_size*lora_bits/8], R_A=R or 0.
lora_A_scales : T
-
Down block scales [32,K/32]; retained even when the adapter is empty.
+
Down block scales [R,ceil(K/lora_block_size)]; retained when the adapter is empty.
lora_B : TQ
-
Packed INT8 up weights [N,R_B/32,32], R_B=0 or 32.
+
Packed up weights [N,R_blocks,lora_block_size*lora_bits/8], R_blocks=ceil(R/lora_block_size) or 0.
lora_B_scales : T
-
Up block scales [N,1]; retained even when the adapter is empty.
+
Up block scales [N,ceil(R/lora_block_size)]; retained when the adapter is empty.
#### Outputs @@ -4335,7 +4335,7 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(float16), tensor(float)
-
FP16 or FP32 activations and block scales.
+
Activations and block scales share the same floating-point type (FP16 or FP32).
TQ : tensor(uint8)
Packed quantized weights.
diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 14349b68af81d..60f253ebd217f 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -1,7 +1,7 @@ # WebGPU MatMulNBitsQLoRA prototype -`com.microsoft.MatMulNBitsQLoRA` fuses an INT4 base projection, two INT8 adapter -projections, and their residual Add. Adapter weights remain replaceable. This +`com.microsoft.MatMulNBitsQLoRA` fuses a quantized base projection, two quantized +adapter projections, and their residual Add. Adapter weights remain replaceable. This is an inference operator for quantized adapters; it does not merge adapter weights into the base or perform adapter training. @@ -29,8 +29,22 @@ the transformer does not undo those explicit nodes. ## Weight format and arithmetic +The [operator definition](ContribOperators.md#com.microsoft.matmulnbitsqlora) +follows MatMulNBits quantization settings and permits any positive adapter rank, +including ranks that are not divisible by 8 or by the quantization block size. +Packed storage pads the final block; the logical rank determines how many +elements participate in the up projection. + +The current WebGPU kernel and graph matcher support a subset: FP16/FP32 +activations and scales, INT4 base weights, INT8 adapter weights, block size 32 +for all projections, active adapter rank 32, positive K divisible by 32, +positive N, and X of rank at least two. Other configurations +are not supported by this kernel. Explicit nodes reject unsupported configurations, +and the matcher leaves them unfused. These are +implementation limits rather than restrictions on the operator definition. + The operator takes `X, W, scales, lora_A, lora_A_scales, lora_B, lora_B_scales`. -For rank 32, positive K divisible by 32, and positive N divisible by 4: +For the current WebGPU implementation: | Tensor | Type | Shape | |---|---|---| @@ -54,7 +68,7 @@ delta = F(T * Bt^T) Y = F(MatMulNBits(X,W,scales) + delta) ``` -Base quantization stays INT4 with block size 32 and implicit zero point 8. +This implementation uses INT4 base quantization with block size 32 and implicit zero point 8. The equations describe the weight interpretation and mathematical projections; the optimized kernels can change reduction order and activation precision. Eligible prefill down projections use DP4A activation quantization when the diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 678ea9467f883..377b866118bc9 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -217,10 +217,12 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { k_(info.GetAttr("K")), n_(info.GetAttr("N")), accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)) { - ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0 && n_ % 4 == 0 && info.GetAttr("bits") == 4 && + ORT_ENFORCE(info.GetAttr("bits") == 4 && info.GetAttr("block_size") == 32 && info.GetAttr("lora_bits") == 8 && info.GetAttr("lora_block_size") == 32, - "MatMulNBitsQLoRA requires INT4 base and INT8 adapters with block size 32."); + "WebGPU MatMulNBitsQLoRA supports only bits=4, block_size=32, lora_bits=8, and lora_block_size=32."); + ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0, + "WebGPU MatMulNBitsQLoRA requires positive K divisible by 32 and positive N."); } Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override { @@ -231,9 +233,9 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { const auto* as = context.Input(4); const auto* bq = context.Input(5); const auto* bs = context.Input(6); - ORT_RETURN_IF_NOT( - x->Shape().NumDimensions() >= 2 && aq->Shape().NumDimensions() == 3 && bq->Shape().NumDimensions() == 3, - "Invalid quantized LoRA input ranks."); + ORT_RETURN_IF_NOT(x->Shape().NumDimensions() >= 2, "WebGPU MatMulNBitsQLoRA requires X rank at least two."); + ORT_RETURN_IF_NOT(aq->Shape().NumDimensions() == 3 && bq->Shape().NumDimensions() == 3, + "Invalid quantized LoRA input ranks."); ORT_RETURN_IF_NOT( w->Shape() == TensorShape({n_, k_ / 32, 16}) && s->Shape().Size() == SafeInt(n_) * (k_ / 32), "Invalid base weight/scales shape."); @@ -241,7 +243,9 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { // and vector-tail handling while preserving rank-32 performance. const auto rank_a = aq->Shape()[0]; const auto blocks_b = bq->Shape()[1]; - ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (blocks_b == 0 || blocks_b == 1) && aq->Shape()[1] == k_ / 32 && + ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (blocks_b == 0 || blocks_b == 1), + "WebGPU MatMulNBitsQLoRA supports only adapter rank 32 (or empty adapters)."); + ORT_RETURN_IF_NOT(aq->Shape()[1] == k_ / 32 && aq->Shape()[2] == 32 && bq->Shape()[0] == n_ && bq->Shape()[2] == 32, "Invalid packed adapter dimensions."); MatMulComputeHelper helper; diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f3fd12b16bdc5..4c398ff505007 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3953,52 +3953,94 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .SinceVersion(1) .SetDoc( "Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). " - "The base uses INT4 block-32 weights with implicit zero point 8. Adapters use INT8 block-32 weights " - "stored as unsigned bytes with implicit zero point 128 and block scales matching the activation type. Eligible adapter down " - "projections may quantize activations according to accuracy_level. Rank is 32; an empty rank dimension " - "in either packed adapter disables the adapter. " - "Intermediate projections use the activation type (FP16 or FP32) before the final addition. Adapters remain replaceable.") - .Attr("K", "Input feature dimension; a positive multiple of 32.", AttributeProto::INT) - .Attr("N", "Output feature dimension; a positive multiple of 4.", AttributeProto::INT) - .Attr("bits", "Base weight bit width; must be 4.", AttributeProto::INT, int64_t{4}) - .Attr("block_size", "Base weight block size; must be 32.", AttributeProto::INT, int64_t{32}) - .Attr("accuracy_level", "Accuracy policy for the base and eligible adapter down projections.", AttributeProto::INT, - int64_t{0}) - .Attr("lora_bits", "Adapter weight bit width; must be 8.", AttributeProto::INT, int64_t{8}) - .Attr("lora_block_size", "Adapter weight block size; must be 32.", AttributeProto::INT, int64_t{32}) - .Input(0, "X", "Input of rank at least two, with last dimension K.", "T") - .Input(1, "W", "Packed INT4 base weights [N,K/32,16].", "TQ") - .Input(2, "scales", "Base block scales with N*K/32 elements; may be flattened.", "T") - .Input(3, "lora_A", "Packed INT8 down weights [R_A,K/32,32], R_A=0 or 32.", "TQ") - .Input(4, "lora_A_scales", "Down block scales [32,K/32]; retained even when the adapter is empty.", "T") - .Input(5, "lora_B", "Packed INT8 up weights [N,R_B/32,32], R_B=0 or 32.", "TQ") - .Input(6, "lora_B_scales", "Up block scales [N,1]; retained even when the adapter is empty.", "T") + "Each projection follows MatMulNBits block quantization and bit packing, with implicit zero points " + "2^(bits-1) for the base and 2^(lora_bits-1) for the adapters. No bias or explicit zero points are applied. " + "The adapter rank R may be any positive integer and need not be a multiple of lora_block_size. " + "Partial final blocks are padded; only the logical K or R elements participate in each projection. " + "An empty rank dimension in either packed adapter disables the adapter, while scale defaults are retained. " + "Intermediate projections and block scales use the activation type before the final addition. " + "Adapters remain replaceable. Execution providers may support a subset of the defined types, " + "quantization settings, and ranks.") + .Attr("K", "Positive input feature dimension of the base and down-projection weight matrices.", + AttributeProto::INT) + .Attr("N", "Positive output feature dimension of the base and up-projection weight matrices.", + AttributeProto::INT) + .Attr("bits", "Base weight bit width, following MatMulNBits (2, 4, or 8).", AttributeProto::INT, int64_t{4}) + .Attr("block_size", "Base quantization block size along K, following MatMulNBits: a power of two, at least 16.", + AttributeProto::INT, int64_t{32}) + .Attr("accuracy_level", + "Accuracy policy following MatMulNBits: 0 (unset), 1 (fp32), 2 (fp16), 3 (bf16), or 4 (int8). " + "Applies to the base and eligible adapter down projections; the up projection uses level 0.", + AttributeProto::INT, int64_t{0}) + .Attr("lora_bits", "Adapter weight bit width, following MatMulNBits (2, 4, or 8).", + AttributeProto::INT, int64_t{8}) + .Attr("lora_block_size", + "Adapter quantization block size along K for lora_A and R for lora_B: a power of two, at least 16.", + AttributeProto::INT, int64_t{32}) + .Input(0, "X", "Input of rank at least one, with last dimension K, following MatMulNBits.", "T") + .Input(1, "W", "Packed base weights [N,ceil(K/block_size),block_size*bits/8].", "TQ") + .Input(2, "scales", "Base block scales [N,ceil(K/block_size)]; may be flattened.", "T") + .Input(3, "lora_A", + "Packed down weights [R_A,ceil(K/lora_block_size),lora_block_size*lora_bits/8], R_A=R or 0.", "TQ") + .Input(4, "lora_A_scales", + "Down block scales [R,ceil(K/lora_block_size)]; retained when the adapter is empty.", "T") + .Input(5, "lora_B", + "Packed up weights [N,R_blocks,lora_block_size*lora_bits/8], R_blocks=ceil(R/lora_block_size) or 0.", "TQ") + .Input(6, "lora_B_scales", + "Up block scales [N,ceil(R/lora_block_size)]; retained when the adapter is empty.", "T") .Output(0, "Y", "Input shape with the last dimension replaced by N.", "T") - .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, "FP16 or FP32 activations and block scales.") + .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, + "Activations and block scales share the same floating-point type (FP16 or FP32).") .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized weights.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { const auto k = getAttribute(ctx, "K", -1); const auto n = getAttribute(ctx, "N", -1); - if (k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || - getAttribute(ctx, "bits", 4) != 4 || getAttribute(ctx, "block_size", 32) != 32 || - getAttribute(ctx, "lora_bits", 8) != 8 || getAttribute(ctx, "lora_block_size", 32) != 32) { - fail_shape_inference("Unsupported MatMulNBitsQLoRA attributes."); + const auto lora_bits = getAttribute(ctx, "lora_bits", 8); + const auto lora_block_size = getAttribute(ctx, "lora_block_size", 32); + const auto valid_quantization = [](int64_t bits, int64_t block_size) { + return (bits == 2 || bits == 4 || bits == 8) && block_size >= 16 && + (block_size & (block_size - 1)) == 0; + }; + if (k <= 0 || n <= 0 || + !valid_quantization(getAttribute(ctx, "bits", 4), getAttribute(ctx, "block_size", 32)) || + !valid_quantization(lora_bits, lora_block_size)) { + fail_shape_inference( + "MatMulNBitsQLoRA requires positive K and N, bit widths 2, 4, or 8, " + "and power-of-two block sizes of at least 16."); } propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() < 2) { - fail_shape_inference("X must have rank at least two."); - } MatmulWithQuantWeightShapeInference(ctx, k, n, true); + int64_t rank = -1; for (size_t index : {size_t{3}, size_t{4}, size_t{5}, size_t{6}}) { if (!hasInputShape(ctx, index)) continue; const auto& shape = getInputShape(ctx, index); const bool packed = index == 3 || index == 5; if (shape.dim_size() != (packed ? 3 : 2)) fail_shape_inference("Invalid quantized adapter rank."); - const int64_t dims[4][3] = {{32, k / 32, 32}, {32, k / 32, 0}, {n, 1, 32}, {n, 1, 0}}; + if ((index == 3 || index == 4) && shape.dim(0).has_dim_value()) { + const auto value = shape.dim(0).dim_value(); + if (index == 3 && value == 0) continue; // Empty packed defaults retain the scale tensor's rank. + if (value <= 0 || (rank > 0 && value != rank)) { + fail_shape_inference("Inconsistent adapter rank between lora_A and lora_A_scales."); + } + rank = value; + } + } + const auto k_blocks = k / lora_block_size + (k % lora_block_size != 0); + const auto r_blocks = rank < 0 ? -1 : rank / lora_block_size + (rank % lora_block_size != 0); + const auto blob_size = (lora_block_size / 8) * lora_bits; + const int64_t dims[4][3] = { + {rank, k_blocks, blob_size}, + {rank, k_blocks, 0}, + {n, r_blocks, blob_size}, + {n, r_blocks, 0}, + }; + for (size_t index : {size_t{3}, size_t{4}, size_t{5}, size_t{6}}) { + if (!hasInputShape(ctx, index)) continue; + const auto& shape = getInputShape(ctx, index); for (int axis = 0; axis < shape.dim_size(); ++axis) { const auto& dim = shape.dim(axis); const bool empty_axis = (index == 3 && axis == 0) || (index == 5 && axis == 1); - if (dim.has_dim_value() && dim.dim_value() != dims[index - 3][axis] && + if (dim.has_dim_value() && dims[index - 3][axis] >= 0 && dim.dim_value() != dims[index - 3][axis] && !(empty_axis && dim.dim_value() == 0)) { fail_shape_inference("Invalid quantized adapter dimensions."); } diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc index aa7e6ec9ecb5d..9a568c4efe795 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -98,7 +98,7 @@ Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph continue; const int64_t k = GetIntAttr(*base, "K", -1, true); const int64_t n = GetIntAttr(*base, "N", -1, true); - if (k <= 0 || k % 32 != 0 || n <= 0 || n % 4 != 0 || GetIntAttr(*down, "K", -1, true) != k || + if (k <= 0 || k % 32 != 0 || n <= 0 || GetIntAttr(*down, "K", -1, true) != k || GetIntAttr(*down, "N", -1, true) != 32 || GetIntAttr(*up, "K", -1, true) != 32 || GetIntAttr(*up, "N", -1, true) != n) continue; diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index bf0fb13783713..26a955b118e68 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1534,7 +1534,8 @@ TEST(MatMul2Bits, Float32_2b_BlkLen32_Accuracy0) { #ifdef USE_WEBGPU TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int2) { - for (int64_t n : {16, 32, 48, 64, 80}) { + // Partial narrow tile, multiple narrow tiles, and a partial wide tile. + for (int64_t n : {16, 48, 80}) { for (bool has_zero_point : {false, true}) { TestOptions2Bits opts{}; opts.M = 65; diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index a9f28a1311cfc..a1987ea1ce7c6 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -454,7 +454,8 @@ TEST(MatMulNBits, Float32_4b_Accuracy0) { #ifdef USE_WEBGPU TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int4) { - for (int64_t n : {16, 32, 48, 64, 80}) { + // Partial narrow tile, multiple narrow tiles, and a partial wide tile. + for (int64_t n : {16, 48, 80}) { for (bool has_zero_point : {false, true}) { TestOptions opts{}; opts.batch_count = 2; diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index a5b13c7f35e39..120dd87e15f28 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1059,21 +1059,20 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { #if defined(USE_WEBGPU) TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { - for (int64_t m : {3, 4, 65}) { - for (int64_t n : {16, 32, 48, 64, 80}) { - for (bool has_zero_point : {false, true}) { - TestOptions8Bits opts{}; - opts.M = m; - opts.N = n; - opts.K = 256; - opts.accuracy_level = 4; - opts.has_zero_point = has_zero_point; - opts.has_bias = true; - opts.output_abs_error = 0.05f; - opts.output_rel_error = 0.02f; - RunTest8Bits(opts); - RunTest8Bits(opts); - } + // Small-M reduction, the tiled threshold, and row/column tails in both tile widths. + for (const auto& [m, n] : {std::pair{3, 16}, {4, 32}, {65, 48}, {65, 80}}) { + for (bool has_zero_point : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m; + opts.N = n; + opts.K = 256; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.05f; + opts.output_rel_error = 0.02f; + RunTest8Bits(opts); + RunTest8Bits(opts); } } } diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc new file mode 100644 index 0000000000000..d6c39aed2d4e7 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc @@ -0,0 +1,139 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#include "test/util/include/test_utils.h" + +namespace onnxruntime::test { +namespace { + +struct QLoRAShapeConfig { + int64_t k = 64; + int64_t n = 128; + int64_t rank = 32; + int64_t bits = 4; + int64_t block_size = 32; + int64_t lora_bits = 8; + int64_t lora_block_size = 32; +}; + +template +Status ResolveQLoRAShape(const QLoRAShapeConfig& c, const std::function& edit = nullptr) { + Model model("qlora_shape", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + const auto k_blocks = (c.k + c.block_size - 1) / c.block_size; + const auto lora_k_blocks = (c.k + c.lora_block_size - 1) / c.lora_block_size; + const auto r_blocks = (c.rank + c.lora_block_size - 1) / c.lora_block_size; + const auto lora_blob = c.lora_block_size * c.lora_bits / 8; + auto* x = builder.MakeInput(std::vector{2, 3, c.k}); + auto* w = builder.MakeInput(std::vector{c.n, k_blocks, c.block_size * c.bits / 8}); + auto* s = builder.MakeInput(std::vector{c.n, k_blocks}); + auto* a = builder.MakeInput(std::vector{c.rank, lora_k_blocks, lora_blob}); + auto* as = builder.MakeInput(std::vector{c.rank, lora_k_blocks}); + auto* b = builder.MakeInput(std::vector{c.n, r_blocks, lora_blob}); + auto* bs = builder.MakeInput(std::vector{c.n, r_blocks}); + auto* y = builder.MakeOutput(); + auto& node = builder.AddNode("MatMulNBitsQLoRA", {x, w, s, a, as, b, bs}, {y}, kMSDomain); + node.AddAttribute("K", c.k); + node.AddAttribute("N", c.n); + node.AddAttribute("bits", c.bits); + node.AddAttribute("block_size", c.block_size); + node.AddAttribute("lora_bits", c.lora_bits); + node.AddAttribute("lora_block_size", c.lora_block_size); + if (edit) edit(node); + builder.SetGraphOutputs(); + auto status = model.MainGraph().Resolve(); + if (status.IsOK()) { + EXPECT_EQ(y->TypeAsProto()->tensor_type().elem_type(), utils::ToTensorProtoElementType()); + auto expected = *x->Shape(); + expected.mutable_dim(expected.dim_size() - 1)->set_dim_value(c.n); + CheckShapeEquality(y->Shape(), &expected); + } + return status; +} + +void SetInputDimension(Node& node, size_t input, int axis, int64_t value) { + auto* arg = node.MutableInputDefs()[input]; + auto shape = *arg->Shape(); + auto* dim = shape.mutable_dim(axis); + dim->Clear(); + if (value >= 0) dim->set_dim_value(value); + arg->SetShape(shape); +} + +} // namespace + +TEST(MatMulNBitsQLoRASchema, GeneralQuantizationAndArbitraryRank) { + for (const QLoRAShapeConfig& c : { + QLoRAShapeConfig{63, 65, 7, 8, 64, 4, 16}, + {65, 33, 64, 2, 16, 8, 128}, + {256, 128, 129, 4, 128, 2, 64}, + {64, 128, 1}, + {64, 128, 16}, + {64, 128, 32}}) { + SCOPED_TRACE(testing::Message() << "K=" << c.k << " N=" << c.n << " R=" << c.rank); + ASSERT_STATUS_OK(ResolveQLoRAShape(c)); + ASSERT_STATUS_OK(ResolveQLoRAShape(c)); + } +} + +TEST(MatMulNBitsQLoRASchema, EmptyAdaptersRetainScaleDefaults) { + for (int mask : {1, 2, 3}) { + ASSERT_STATUS_OK(ResolveQLoRAShape({63, 65, 7, 8, 64, 4, 16}, [mask](Node& node) { + if (mask & 1) SetInputDimension(node, 3, 0, 0); + if (mask & 2) SetInputDimension(node, 5, 1, 0); + })); + } +} + +TEST(MatMulNBitsQLoRASchema, DynamicAdapterRankAndVectorInput) { + ASSERT_STATUS_OK(ResolveQLoRAShape({}, [](Node& node) { + SetInputDimension(node, 3, 0, -1); + SetInputDimension(node, 4, 0, -1); + SetInputDimension(node, 5, 1, -1); + SetInputDimension(node, 6, 1, -1); + ONNX_NAMESPACE::TensorShapeProto x_shape; + x_shape.add_dim()->set_dim_value(64); + node.MutableInputDefs()[0]->SetShape(x_shape); + })); +} + +#ifndef ORT_NO_EXCEPTIONS +TEST(MatMulNBitsQLoRASchema, RejectsInconsistentAdapterShapes) { + for (const auto& [input, axis, value] : {std::tuple{4, 0, 16}, + {3, 1, 3}, + {3, 2, 16}, + {5, 1, 2}, + {6, 1, 2}}) { + const auto status = ResolveQLoRAShape({}, [=](Node& node) { SetInputDimension(node, input, axis, value); }); + EXPECT_FALSE(status.IsOK()); + } +} + +TEST(MatMulNBitsQLoRASchema, RejectsInvalidQuantizationAttributes) { + const std::pair invalid_attributes[] = { + {"bits", 3}, + {"lora_bits", 16}, + {"block_size", 8}, + {"block_size", 48}, + {"lora_block_size", 0}, + {"lora_block_size", 24}, + {"K", 0}, + {"N", -1}, + }; + for (const auto& [name, value] : invalid_attributes) { + SCOPED_TRACE(name); + const auto status = ResolveQLoRAShape({}, [=](Node& node) { node.AddAttribute(name, value); }); + EXPECT_FALSE(status.IsOK()); + } +} +#endif // ORT_NO_EXCEPTIONS + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 6f2482e64a7be..daa9a064aab3f 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -21,22 +21,24 @@ namespace onnxruntime::test { template void RunEmptyAndActiveAdapters() { - for (int64_t m : {1, 37}) { + // Scalar output, a two-column tail, and a partial wide tile. + for (const auto& [m, n] : {std::pair{1, 1}, {2, 2}, {37, 129}}) { + SCOPED_TRACE(MakeString("M=", m, ", N=", n)); for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { auto ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(ep, nullptr); OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); test.AddAttribute("K", 64); - test.AddAttribute("N", 128); + test.AddAttribute("N", n); test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); - test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); - test.AddInput("S", {128, 2}, std::vector(256, T(0.03125f)), true); + test.AddInput("W", {n, 2, 16}, std::vector(n * 32, 0x99), true); + test.AddInput("S", {n, 2}, std::vector(n * 2, T(0.03125f)), true); test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); test.AddInput("AS", {32, 2}, std::vector(64, T(0.03125f))); - test.AddInput("BQ", {128, ranks.second, 32}, std::vector(128 * ranks.second * 32, 129)); - test.AddInput("BS", {128, 1}, std::vector(128, T(0.0625f))); + test.AddInput("BQ", {n, ranks.second, 32}, std::vector(n * ranks.second * 32, 129)); + test.AddInput("BS", {n, 1}, std::vector(n, T(0.0625f))); const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; - test.AddOutput("Y", {1, m, 128}, std::vector(m * 128, T(expected))); + test.AddOutput("Y", {1, m, n}, std::vector(m * n, T(expected))); std::vector> eps; eps.push_back(std::move(ep)); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); @@ -47,68 +49,53 @@ void RunEmptyAndActiveAdapters() { TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { RunEmptyAndActiveAdapters(); } TEST(MatMulNBitsQLoRA, EmptyAndActiveAdaptersFp32) { RunEmptyAndActiveAdapters(); } -TEST(MatMulNBitsQLoRA, RejectsUnsupportedAdapterShapeAndBits) { - for (bool wrong_bits : {false, true}) { +#ifndef ORT_NO_EXCEPTIONS +TEST(MatMulNBitsQLoRA, RejectsUnsupportedWebGpuConfigurations) { + struct Config { + int64_t k = 64, n = 128, rank = 32, bits = 4, block_size = 32, lora_bits = 8, lora_block_size = 32; + const char* error = "WebGPU MatMulNBitsQLoRA supports only bits=4"; + }; + const Config configs[] = { + {64, 128, 32, 8}, + {64, 128, 32, 4, 64}, + {64, 128, 32, 4, 32, 4}, + {64, 128, 32, 4, 32, 8, 64}, + {64, 128, 7, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA supports only adapter rank 32"}, + {64, 128, 64, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA supports only adapter rank 32"}, + {63, 128, 32, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA requires positive K"}, + }; + for (const Config& c : configs) { auto ep = DefaultWebGpuExecutionProvider(); ASSERT_NE(ep, nullptr); OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); - test.AddAttribute("K", 64); - test.AddAttribute("N", 128); - if (wrong_bits) test.AddAttribute("lora_bits", 4); - test.AddInput("X", {1, 1, 64}, std::vector(64, MLFloat16(0.5f))); - test.AddInput("W", {128, 2, 16}, std::vector(4096, 0x99), true); - test.AddInput("S", {128, 2}, std::vector(256, MLFloat16(0.03125f)), true); - const int64_t rank = wrong_bits ? 32 : 16; - test.AddInput("AQ", {rank, 2, 32}, std::vector(rank * 64, 129)); - test.AddInput("AS", {32, 2}, std::vector(64, MLFloat16(0.03125f))); - test.AddInput("BQ", {128, 1, 32}, std::vector(4096, 129)); - test.AddInput("BS", {128, 1}, std::vector(128, MLFloat16(0.0625f))); - test.AddOutput("Y", {1, 1, 128}, std::vector(128, MLFloat16(3.0f))); + test.AddAttribute("K", c.k); + test.AddAttribute("N", c.n); + test.AddAttribute("bits", c.bits); + test.AddAttribute("block_size", c.block_size); + test.AddAttribute("lora_bits", c.lora_bits); + test.AddAttribute("lora_block_size", c.lora_block_size); + const auto k_blocks = (c.k + c.block_size - 1) / c.block_size; + const auto lora_k_blocks = (c.k + c.lora_block_size - 1) / c.lora_block_size; + const auto r_blocks = (c.rank + c.lora_block_size - 1) / c.lora_block_size; + const auto blob = c.block_size * c.bits / 8; + const auto lora_blob = c.lora_block_size * c.lora_bits / 8; + test.AddInput("X", {1, 1, c.k}, std::vector(c.k, 0.5f)); + test.AddInput("W", {c.n, k_blocks, blob}, std::vector(c.n * k_blocks * blob), true); + test.AddInput("S", {c.n, k_blocks}, std::vector(c.n * k_blocks, 1.0f), true); + test.AddInput("AQ", {c.rank, lora_k_blocks, lora_blob}, + std::vector(c.rank * lora_k_blocks * lora_blob)); + test.AddInput("AS", {c.rank, lora_k_blocks}, std::vector(c.rank * lora_k_blocks, 1.0f)); + test.AddInput("BQ", {c.n, r_blocks, lora_blob}, std::vector(c.n * r_blocks * lora_blob)); + test.AddInput("BS", {c.n, r_blocks}, std::vector(c.n * r_blocks, 1.0f)); + test.AddOutput("Y", {1, 1, c.n}, std::vector(c.n)); std::vector> eps; eps.push_back(std::move(ep)); - test.Run(OpTester::ExpectResult::kExpectFailure, - wrong_bits ? "Unsupported MatMulNBitsQLoRA attributes" : "Invalid quantized adapter dimensions", {}, - nullptr, &eps); + test.Run(OpTester::ExpectResult::kExpectFailure, c.error, {}, nullptr, &eps); } } +#endif // ORT_NO_EXCEPTIONS -template -void RunQuantizedPrefillAndReusedInput() { - for (int64_t m : {1, 2, 3, 4, 65}) { - for (int64_t accuracy : {0, 4}) { - OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); - constexpr int64_t k = 256; - constexpr int64_t n = 128; - constexpr int64_t batches = 2; - test.AddAttribute("K", k); - test.AddAttribute("N", n); - test.AddAttribute("accuracy_level", accuracy); - std::vector x(batches * m * k); - std::vector expected(batches * m * n); - for (int64_t row = 0; row < batches * m; ++row) { - const float value = row % 2 == 0 ? 0.5f : -0.25f; - std::fill_n(x.begin() + row * k, k, T(value)); - std::fill_n(expected.begin() + row * n, n, T(value * 24.0f)); - } - test.AddInput("X", {batches, m, k}, x); - test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); - test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, T(0.03125f)), true); - test.AddInput("AQ", {32, k / 32, 32}, std::vector(32 * k, 129)); - test.AddInput("AS", {32, k / 32}, std::vector(k, T(0.03125f))); - test.AddInput("BQ", {n, 1, 32}, std::vector(n * 32, 129)); - test.AddInput("BS", {n, 1}, std::vector(n, T(0.0625f))); - test.AddOutput("Y", {batches, m, n}, expected); - std::vector> eps; - eps.push_back(DefaultWebGpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); - } - } -} - -TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInput) { RunQuantizedPrefillAndReusedInput(); } -TEST(MatMulNBitsQLoRA, QuantizedPrefillAndReusedInputFp32) { RunQuantizedPrefillAndReusedInput(); } - -TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixUsesFusedUp) { +TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixNumerics) { #if defined(ORT_USE_EP_API_ADAPTERS) GTEST_SKIP() << "Hardware-specific dispatch checks require direct adapter capability inspection."; #else @@ -128,13 +115,9 @@ TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixUsesFusedUp) { if (!gpu.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix) || !config.has_value()) { GTEST_SKIP() << "Requires a compatible FP16 subgroup-matrix configuration."; } - const auto down_config = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( - {device_configs.configs, device_configs.configCount}, - adapter.subgroupMinSize, adapter.subgroupMaxSize, - gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), - {{kF16, wgpu::SubgroupMatrixComponentType::F32, 8, 16, 16, 32, true}}); - - constexpr int64_t m = 256, k = 256, n = 128; + // Include a partial row tile; the dispatch assertion below prevents a fallback + // from hiding numerical failures in the subgroup-matrix kernel. + constexpr int64_t m = 65, k = 256, n = 128; OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); test.AddAttribute("K", k); test.AddAttribute("N", n); @@ -157,28 +140,6 @@ TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixUsesFusedUp) { test.Run(session_options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); const auto logs = testing::internal::GetCapturedStderr(); EXPECT_NE(logs.find("SubgroupMatrixMatMulNBitsQLoRAUp"), std::string::npos) << logs; - if (config->Is(8, 16, 16) && down_config.has_value()) { - auto count_program = [&logs](const std::string& name) { - const auto marker = "Starting program \"" + name + "["; - size_t count = 0; - for (auto pos = logs.find(marker); pos != std::string::npos; pos = logs.find(marker, pos + marker.size())) { - ++count; - } - return count; - }; - const auto down_count = count_program("SubgroupMatrixMatMulNBits"); - EXPECT_GT(down_count, 0u) << logs; - EXPECT_NE(logs.find("SubgroupMatrixMatMulNBits[8|2|1|8|16|16|32|"), std::string::npos) << logs; - EXPECT_EQ(count_program("SubgroupMatrixMatMulLayout"), down_count) << logs; - EXPECT_EQ(count_program("SubgroupMatrixMatMulNBitsQLoRAUp"), down_count) << logs; - EXPECT_NE(logs.find("|32|32|32|0|0|0|0|0|0]:128"), std::string::npos) << logs; - EXPECT_NE(logs.find("|128|128|32|0|0|0|0|1|0]:512"), std::string::npos) << logs; - EXPECT_EQ(logs.find("Starting program \"MatMulNBits["), std::string::npos) << logs; - EXPECT_EQ(logs.find("DP4AMatMulQuantize"), std::string::npos) << logs; - EXPECT_EQ(logs.find("Starting program \"DP4AMatMulNBits"), std::string::npos) << logs; - } - EXPECT_EQ(logs.find("QLoRAOutputVec4"), std::string::npos) << logs; - EXPECT_EQ(logs.find("QLoRAAdd"), std::string::npos) << logs; #endif } @@ -186,10 +147,17 @@ void RunFp32Numerics(const ConfigOptions& provider_options = {}) { struct Shape { int64_t m, k, n, batch; }; - for (const auto shape : {Shape{1, 64, 128, 1}, Shape{1, 2048, 36, 1}, Shape{1, 2048, 64, 1}, Shape{1, 8192, 132, 1}, - Shape{1, 2176, 64, 7}, Shape{1, 8192, 64, 1}, - Shape{2, 2048, 36, 2}, Shape{2, 8192, 64, 1}, Shape{3, 3072, 64, 3}, - Shape{37, 2048, 132, 2}, Shape{128, 3072, 64, 1}}) { + // Small-K decode, joint decode with a K tail, long reductions, and batched + // generic/wide/DP4A prefill. Odd N covers scalar output tails. + const Shape shapes[] = { + {1, 64, 1, 1}, + {1, 2176, 127, 1}, + {1, 8192, 64, 1}, + {2, 2048, 129, 2}, + {37, 2176, 127, 2}, + {37, 2176, 64, 2}, + }; + for (const auto& shape : shapes) { const auto [m, k, n, batch] = shape; SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); @@ -367,7 +335,7 @@ void RunFp32DP4AQuantization(int64_t m) { TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { RunFp32DP4AQuantization(1); } TEST(MatMulNBitsQLoRA, SmallMPrefillDP4AQuantization) { - for (int64_t m : {2, 3}) RunFp32DP4AQuantization(m); + RunFp32DP4AQuantization(2); } TEST(MatMulNBitsQLoRA, NarrowKParallelReduction) { @@ -385,9 +353,8 @@ TEST(MatMulNBitsQLoRA, NarrowKParallelReduction) { int64_t k; int64_t batch; }; - for (const auto shape : {Shape{4, 128, 1}, Shape{3, 1920, 1}, Shape{4, 1920, 1}, Shape{4, 2048, 1}, - Shape{37, 2176, 2}, Shape{128, 3072, 1}, - Shape{65, 8192, 1}, Shape{4097, 128, 1}}) { + // Cover both sides of the K-parallel threshold and a batched K/M tail. + for (const auto shape : {Shape{4, 1920, 1}, Shape{4, 2048, 1}, Shape{37, 2176, 2}}) { const auto [m, k, batch] = shape; SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", batch=", batch)); constexpr int64_t n = 64; @@ -472,7 +439,7 @@ TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { } #endif for (bool acc_f32 : {false, true}) { - for (const auto& [k, n] : {std::pair{2048, 36}, {4096, 132}, {8192, 64}}) { + for (const auto& [k, n] : {std::pair{2176, 127}, {8192, 64}}) { SCOPED_TRACE(MakeString("K=", k, ", N=", n, ", acc_f32=", acc_f32)); std::vector x(k); std::vector aq(32 * k); diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 7040b0f59e2cc..7812fd5c65078 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -105,7 +105,13 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { int64_t m, k, n, batch; bool zero_base; }; - const ShapeCase shapes[] = {{1, 64, 128, 1, false}, {1, 2048, 128, 1, false}, {1, 3072, 132, 1, true}, {1, 8192, 128, 1, true}, {2, 2048, 132, 1, true}, {37, 2048, 128, 1, false}, {37, 3072, 132, 1, true}, {1, 2048, 128, 2, true}, {37, 2048, 132, 2, true}}; + const ShapeCase shapes[] = { + {1, 64, 1, 1, false}, + {1, 2176, 127, 1, true}, + {2, 2048, 129, 2, true}, + {37, 2048, 128, 1, false}, + {37, 2176, 127, 2, true}, + }; for (const auto& shape : shapes) { SCOPED_TRACE(shape.m); SCOPED_TRACE(shape.k); @@ -126,7 +132,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Numerics) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (const auto& [m, n] : {std::pair{1, 128}, {1, 132}, {2, 128}, {2, 132}, {3, 128}, {37, 132}}) { + for (const auto& [m, n] : {std::pair{1, 128}, {2, 128}, {37, 128}, {2, 129}}) { SCOPED_TRACE(MakeString("M=", m, ", N=", n)); // Eligible FP32 decode and short prefill share quantized X with the base. const bool quantized_down = m >= 4 || n % 16 == 0; @@ -217,75 +223,10 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentAdapterAc [](const Graph& g) { return CheckQLoRACount(g, 0); })); } -TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPartialOverrides) { - if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - Model model("quantized LoRA switching", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), - {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); - auto& graph = model.MainGraph(); - ModelTestBuilder builder(graph); - BuildQLoRA(builder, 1, 64); - builder.SetGraphOutputs(); - ASSERT_STATUS_OK(graph.Resolve()); - const auto adapters = graph.GetOverridableInitializers(); - ASSERT_EQ(adapters.size(), 4u); - // Keep this state-transition test exact across backends. With X=1/4 and - // unit packed weights, the adapter delta is 512 with default scales and - // 1/2 with supplied scales; random-data roundoff is covered by Numerics. - for (auto& [name, value] : builder.feeds_) { - auto* tensor = value.GetMutable(); - if (tensor->IsDataType()) { - auto bytes = tensor->MutableDataAsSpan(); - std::fill(bytes.begin(), bytes.end(), uint8_t{129}); - } else { - auto values = tensor->MutableDataAsSpan(); - const float fill = tensor->Shape().NumDimensions() == 3 ? 0.25f : 0.03125f; - std::fill(values.begin(), values.end(), MLFloat16(fill)); - } - } - std::string data; - model.ToProto().SerializeToString(&data); - std::vector expected; - for (bool fused : {false, true}) { - SessionOptions options; - options.graph_optimization_level = TransformerLevel::Level1; - InferenceSessionWrapper session(options, GetEnvironment()); - ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); - ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); - if (fused) - ASSERT_STATUS_OK( - session.RegisterGraphTransformer(std::make_unique(), TransformerLevel::Level2)); - ASSERT_STATUS_OK(session.Initialize()); - ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); - // Off, packed weights only, scales only, all parameters, then off again. - for (int state = 0; state < 5; ++state) { - auto feeds = builder.feeds_; - for (const auto* input : adapters) { - const bool packed = input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto::UINT8; - if (state == 0 || state == 4 || (state == 1 && !packed) || (state == 2 && packed)) feeds.erase(input->Name()); - } - std::vector output; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, builder.output_names_, &output)); - ASSERT_EQ(output.size(), 1u); - if (!fused) { - expected.push_back(std::move(output[0])); - } else { - const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); - EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << "state=" << state << " " << comparison.second; - const auto base = expected[0].Get().DataAsSpan(); - const auto actual = output[0].Get().DataAsSpan(); - const float delta = state == 1 ? 512.0f : (state == 3 ? 0.5f : 0.0f); - for (size_t i = 0; i < actual.size(); ++i) { - EXPECT_EQ(actual[i], MLFloat16(base[i].ToFloat() + delta)) << "state=" << state << ", i=" << i; - } - } - } - } -} - template void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; - for (int64_t m : {1, 2, 3, 37}) { + for (int64_t m : {1, 2, 37}) { SCOPED_TRACE(m); Model model("quantized LoRA empty defaults", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); @@ -315,15 +256,18 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = ASSERT_FALSE(a_name.empty()); ASSERT_FALSE(b_name.empty()); std::vector expected; - for (bool fused : {false, true}) { - if (fused) { + // Compare unfused, load-time fusion, and serialized fusion using the same + // partial-override states and exact reference values. + for (int mode : {0, 1, 2}) { + const bool fused = mode != 0; + if (mode == 2) { bool modified = false; ASSERT_STATUS_OK(MatMulNBitsQLoRAFusion().Apply(graph, modified, logger)); ASSERT_TRUE(modified); ASSERT_STATUS_OK(graph.Resolve()); } const auto proto = model.ToProto(); - if (fused) { + if (mode == 2) { ASSERT_EQ(proto.graph().node_size(), 1); ASSERT_EQ(proto.graph().node(0).op_type(), "MatMulNBitsQLoRA"); int empty_count = 0; @@ -344,9 +288,12 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = options.graph_optimization_level = TransformerLevel::Level1; InferenceSessionWrapper session(options, GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); - // No LoRA transformer is registered at load time: the serialized model - // itself must retain empty defaults and accept nonempty rank-32 overrides. ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + if (mode == 1) { + ASSERT_STATUS_OK( + session.RegisterGraphTransformer(std::make_unique(), TransformerLevel::Level2)); + } + // Mode 2 loads the serialized operator without a LoRA transformer. ASSERT_STATUS_OK(session.Initialize()); ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); // Off, A only, B only, scales only, packed weights only, all, off again. @@ -424,7 +371,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = Tensor::InitOrtValue(reference.DataType(), reference.Shape(), CPUAllocator::DefaultInstance(), actual); // Each adapter state owns a capture ID and stable GPU bindings. Replay // skips branch selection, so one ID cannot be reused across on/off states. - for (int state : {0, 0, 0, 1, 1, 1, 0, 1, 0, 1}) { + for (int state : {0, 0, 0, 1, 1, 1, 0, 1}) { SCOPED_TRACE(MakeString("capture state=", state)); auto* cpu_output = actual.GetMutable(); auto values = cpu_output->MutableDataAsSpan(); diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc index 107446a80b07f..cfd7509b5ae3c 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc @@ -196,6 +196,93 @@ class WebGpuPluginSharedAllocatorTest : public ::testing::Test { std::vector ep_devices_; }; +void CheckSessionInputUploads(Ort::Env& env, Ort::ConstEpDevice ep_device, + gsl::span input_sizes) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(18); + auto* graph = model.mutable_graph(); + graph->set_name("webgpu_batched_session_inputs"); + auto* node = graph->add_node(); + node->set_op_type("Concat"); + node->add_output("Y"); + auto* axis = node->add_attribute(); + axis->set_name("axis"); + axis->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis->set_i(0); + InlinedVector names; + names.reserve(input_sizes.size()); + int64_t output_size = 0; + for (size_t i = 0; i < input_sizes.size(); ++i) { + names.push_back("X" + std::to_string(i)); + node->add_input(names.back()); + auto* input = graph->add_input(); + input->set_name(names.back()); + auto* type = input->mutable_type()->mutable_tensor_type(); + type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + type->mutable_shape()->add_dim()->set_dim_value(input_sizes[i]); + output_size += input_sizes[i]; + } + auto* output = graph->add_output(); + output->set_name("Y"); + auto* output_type = output->mutable_type()->mutable_tensor_type(); + output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output_type->mutable_shape()->add_dim()->set_dim_value(output_size); + + Ort::SessionOptions options; + options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); + options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); + options.AppendExecutionProvider_V2(env, {ep_device}, {}); + const auto model_bytes = model.SerializeAsString(); + Ort::Session session(env, model_bytes.data(), model_bytes.size(), options); + const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + InlinedVector> input_data; + InlinedVector inputs; + InlinedVector input_names; + input_data.reserve(input_sizes.size()); + inputs.reserve(input_sizes.size()); + input_names.reserve(input_sizes.size()); + for (size_t i = 0; i < input_sizes.size(); ++i) { + input_data.emplace_back(static_cast(input_sizes[i])); + inputs.push_back(Ort::Value::CreateTensor(cpu_memory, input_data.back().data(), input_data.back().size(), + &input_sizes[i], 1)); + input_names.push_back(names[i].c_str()); + } + const char* output_names[] = {"Y"}; + // Change every input on the second run to catch stale copies or a reused + // buffer being cleared after its upload. + for (int iteration = 0; iteration < 2; ++iteration) { + SCOPED_TRACE(iteration); + for (size_t i = 0; i < input_data.size(); ++i) { + for (size_t j = 0; j < input_data[i].size(); ++j) { + input_data[i][j] = static_cast(iteration * 100 + i * 10 + j % 7); + } + } + auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names.data(), inputs.data(), inputs.size(), + output_names, 1); + ASSERT_EQ(outputs[0].GetTensorTypeAndShapeInfo().GetElementCount(), static_cast(output_size)); + const float* actual = outputs[0].GetTensorData(); + size_t offset = 0; + for (size_t i = 0; i < input_data.size(); ++i) { + const auto& values = input_data[i]; + EXPECT_TRUE(std::equal(values.begin(), values.end(), actual + offset)) << "input " << i; + offset += values.size(); + } + } +} + +TEST_F(WebGpuPluginSharedAllocatorTest, SessionInputUploadsPreserveData) { + for (const auto& sizes : { + InlinedVector{0, 16, 0, 16}, // Empty and nonempty inputs in one batch. + {1310721, 1310721, 1310721, 1310721}, // Cross the 16 MiB staging-chunk boundary. + {4194305, 16}, // One input larger than a staging chunk. + {16}, // Single-input fallback. + }) { + SCOPED_TRACE(testing::PrintToString(sizes)); + ASSERT_NO_FATAL_FAILURE(CheckSessionInputUploads(Env(), EpDevice(), sizes)); + } +} + #if !defined(USE_EXTERNAL_DAWN) #if !defined(BUILD_DAWN_SHARED_LIBRARY) namespace { @@ -506,119 +593,6 @@ int VerifyInstanceNormScratchAllocationsStayBatched(Ort::Env& env, Ort::ConstEpD return 0; } -int VerifySessionInputUploadBatching(Ort::Env& env, Ort::ConstEpDevice ep_device, - gsl::span input_sizes, size_t expected_submissions) { - ONNX_NAMESPACE::ModelProto model; - model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - model.add_opset_import()->set_version(18); - auto* graph = model.mutable_graph(); - graph->set_name("webgpu_batched_session_inputs"); - auto* node = graph->add_node(); - node->set_op_type("Concat"); - node->add_output("Y"); - auto* axis = node->add_attribute(); - axis->set_name("axis"); - axis->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); - axis->set_i(0); - InlinedVector names; - names.reserve(input_sizes.size()); - int64_t output_size = 0; - for (size_t i = 0; i < input_sizes.size(); ++i) { - names.push_back("X" + std::to_string(i)); - node->add_input(names.back()); - auto* input = graph->add_input(); - input->set_name(names.back()); - auto* type = input->mutable_type()->mutable_tensor_type(); - type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - type->mutable_shape()->add_dim()->set_dim_value(input_sizes[i]); - output_size += input_sizes[i]; - } - auto* output = graph->add_output(); - output->set_name("Y"); - auto* output_type = output->mutable_type()->mutable_tensor_type(); - output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - output_type->mutable_shape()->add_dim()->set_dim_value(output_size); - - struct SubmissionProbe { - SubmissionProbe() = default; - ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SubmissionProbe); - - bool armed{false}; - bool saw_program{false}; - size_t baseline{0}; - size_t input_submissions{0}; - - static void ORT_API_CALL Log(void* param, OrtLoggingLevel, const char*, const char*, const char*, - const char* message) { - auto& probe = *static_cast(param); - if (probe.armed && !probe.saw_program && - std::string_view{message}.find("Starting program") != std::string_view::npos) { - probe.saw_program = true; - probe.input_submissions = nonempty_queue_submissions.load(std::memory_order_relaxed) - probe.baseline; - } - } - } probe; - auto procs = dawn::native::GetProcs(); - procs.queueSubmit = CountNonemptyQueueSubmissions; - dawnProcSetProcs(&procs); - Ort::SessionOptions options; - options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); - options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); - options.SetLogSeverityLevel(ORT_LOGGING_LEVEL_INFO); - Ort::ThrowOnError(Ort::GetApi().SetUserLoggingFunction(options, SubmissionProbe::Log, &probe)); - options.AppendExecutionProvider_V2( - env, {ep_device}, {{"dawnProcTable", std::to_string(reinterpret_cast(&procs))}}); - const auto model_bytes = model.SerializeAsString(); - Ort::Session session(env, model_bytes.data(), model_bytes.size(), options); - const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); - InlinedVector> input_data; - InlinedVector inputs; - InlinedVector input_names; - input_data.reserve(input_sizes.size()); - inputs.reserve(input_sizes.size()); - input_names.reserve(input_sizes.size()); - for (size_t i = 0; i < input_sizes.size(); ++i) { - input_data.emplace_back(static_cast(input_sizes[i])); - inputs.push_back(Ort::Value::CreateTensor(cpu_memory, input_data.back().data(), input_data.back().size(), - &input_sizes[i], 1)); - input_names.push_back(names[i].c_str()); - } - const char* output_names[] = {"Y"}; - // Warm pipelines/caches, then measure actual input submissions before the first program. - // Changing every input also catches stale uploads and clears submitted after the copy. - for (int iteration = 0; iteration < 6; ++iteration) { - for (size_t i = 0; i < input_data.size(); ++i) { - for (size_t j = 0; j < input_data[i].size(); ++j) { - input_data[i][j] = static_cast(iteration * 100 + i * 10 + j % 7); - } - } - probe.baseline = nonempty_queue_submissions.load(std::memory_order_relaxed); - probe.saw_program = false; - probe.armed = true; - auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names.data(), inputs.data(), inputs.size(), - output_names, 1); - probe.armed = false; - if (iteration >= 3 && (!probe.saw_program || probe.input_submissions != expected_submissions)) { - std::fprintf(stderr, "Input batching failed: saw_program=%d, submissions=%zu, expected=%zu\n", - static_cast(probe.saw_program), probe.input_submissions, expected_submissions); - return 11; - } - if (outputs[0].GetTensorTypeAndShapeInfo().GetElementCount() != static_cast(output_size)) return 12; - const float* actual = outputs[0].GetTensorData(); - size_t offset = 0; - for (const auto& values : input_data) { - for (float expected : values) { - if (actual[offset++] != expected) { - std::fputs("Batched Session input data was corrupted or stale\n", stderr); - return 13; - } - } - } - } - std::fputs("Session input uploads preserved data and expected submission count\n", stderr); - return 0; -} - WGPUBufferMapState InjectMappedBufferOnFree(WGPUBuffer buffer) { if (buffer == free_failure_buffer) { ++free_failure_probes; @@ -740,58 +714,6 @@ TEST_F(WebGpuSessionAllocatorDeathTest, WarmedInstanceNormalizationScratchAlloca #endif } -TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsAreBatched) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; -#else - ASSERT_EXIT( - std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16, 16, 16, 16}, 1)), - ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); -#endif -} - -TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsSpanMultipleChunks) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; -#else - // Four 5 MiB + 4 byte inputs require two staging chunks at the 16 MiB limit. - ASSERT_EXIT( - std::_Exit(VerifySessionInputUploadBatching( - Env(), EpDevice(), std::array{1310721, 1310721, 1310721, 1310721}, 2)), - ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); -#endif -} - -TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsIncludeEmptyTensors) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; -#else - ASSERT_EXIT( - std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{0, 16, 0, 16}, 1)), - ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); -#endif -} - -TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsHandleOversizedTensor) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; -#else - ASSERT_EXIT( - std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{4194305, 16}, 2)), - ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); -#endif -} - -TEST_F(WebGpuSessionAllocatorDeathTest, SessionSingleInputUploadRetainsFallback) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; -#else - ASSERT_EXIT( - std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16}, 1)), - ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); -#endif -} - #endif } // namespace From f248e4120e3a22cb590eaf0cdbf964bd0656ccfb Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 22 Oct 2026 16:18:56 +0800 Subject: [PATCH 29/35] webgpu: Extract QLoRA shader helpers into WGSL templates --- .../webgpu/quantization/dp4a_matmul_nbits.cc | 4 +- .../webgpu/quantization/matmul_nbits.cc | 6 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 165 ++---------------- .../webgpu/quantization/matmul_nbits_qlora.h | 24 ++- .../matmul_nbits_qlora_down.wgsl.template | 119 +++++++++++++ .../matmul_nbits_qlora_up.wgsl.template | 72 ++++++++ .../subgroup_matrix_matmul_nbits.cc | 2 +- 7 files changed, 237 insertions(+), 155 deletions(-) create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template create mode 100644 onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 5ac3fa164ad69..9485908c49a46 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -38,7 +38,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("lora_t"); shader.AddInput("lora_b"); shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_, true); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_, true)); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); return WGSL_TEMPLATE_APPLY( @@ -73,7 +73,7 @@ Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) shader.AddInput("lora_t"); shader.AddInput("lora_b"); shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index de90756a0bfab..8bfa23fb1baf6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -44,7 +44,7 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); } const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -82,7 +82,7 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, acc_f32_); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); } if (has_qlora_down_) { shader.AddInput("lora_at"); @@ -91,7 +91,7 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& output = shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); if (has_qlora_down_) { shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); - GenerateQLoRADown(shader, "uniforms.num_N_tile", a.GetByOffset("k4"), acc_f32_, qlora_down_subgroups_); + ORT_RETURN_IF_ERROR(GenerateQLoRADown(shader, a, acc_f32_, qlora_down_subgroups_)); } const uint32_t components_a = a.NumComponents(); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 377b866118bc9..b30a11c379dc0 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -14,157 +14,24 @@ namespace onnxruntime::contrib::webgpu { -void GenerateQLoRADown(ShaderHelper& shader, const std::string& base_groups, const std::string& x_value, bool acc_f32, - bool use_subgroups) { - auto& body = shader.MainFunctionBody(); - // Dispatch normalization can append workgroups beyond the four adapter groups. - body << "let qlora_base_groups = " << base_groups << ";\n" - << "if (workgroup_idx >= qlora_base_groups + 4u) { return; }\n"; - if (use_subgroups) { - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n"; - // Subgroup membership need not match consecutive local invocation indices. - body << R"( -if (workgroup_idx >= qlora_base_groups) { - let rank_base = (workgroup_idx - qlora_base_groups) * 8u; - let lane = sg_id; - let rank_local = subgroup_index; - var partial0 = qlora_acc_t(0); - var partial1 = qlora_acc_t(0); - for (var k = lane * 16u; k < uniforms.K; k += 512u) { - var k4 = k / 4u; -)"; - // Keep each X vector in a register for both ranks assigned to this lane. - for (uint32_t i = 0; i < 4; ++i) { - body << "let x" << i << " = vec4(" << x_value << ");\n"; - if (i != 3) body << "k4 += 1u;\n"; - } - for (uint32_t r = 0; r < 2; ++r) { - body << "let rank" << r << " = rank_base + rank_local + " << r * 4 << "u;\n" - << "let packed" << r << " = lora_at[rank" << r << " * (uniforms.K / 16u) + k / 16u];\n" - << "let scale" << r << " = lora_a_scales[rank" << r << " * (uniforms.K / 32u) + k / 32u];\n" - << "var sum" << r << " = qlora_acc_t(0);\n"; - for (uint32_t i = 0; i < 4; ++i) { - body << "sum" << r << " += dot(x" << i << ", vec4((vec4(unpack4xU8(packed" << r - << "[" << i << "])) - vec4(128)) * scale" << r << "));\n"; - } - body << "partial" << r << " += sum" << r << ";\n"; - } - // FP32 avoids half-precision rounding at each step of the subgroup tree. - body << R"( - } - let result0 = subgroupAdd(f32(partial0)); - let result1 = subgroupAdd(f32(partial1)); - if (lane == 0u) { - lora_down[rank_base + rank_local] = output_element_t(result0); - lora_down[rank_base + rank_local + 4u] = output_element_t(result1); - } - return; -} -)"; - return; - } - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n" - << R"( -var qlora_x_tile: array, 128>; -var qlora_partials: array, 8>; -)"; - body << R"( -if (workgroup_idx >= qlora_base_groups) { - let rank_base = (workgroup_idx - qlora_base_groups) * 8u; - let lane = local_idx % 32u; - let rank_local = local_idx / 32u; - // Match the generic INT8 MatMulNBits tile, dot grouping and reduction order. - for (var start = 0u; start < uniforms.K; start += 512u) { - let k4 = start / 4u + local_idx; - var xv = vec4(0); - if (k4 < uniforms.K / 4u) { xv = )" - << x_value << R"(; } - qlora_x_tile[local_idx] = xv; - workgroupBarrier(); - for (var rank_offset = 0u; rank_offset < 8u; rank_offset += 4u) { - let rank = rank_base + rank_offset + rank_local; - let k = start + lane * 16u; - if (k < uniforms.K) { - let packed = lora_at[rank * (uniforms.K / 16u) + k / 16u]; - let scale = lora_a_scales[rank * (uniforms.K / 32u) + k / 32u]; - var sum = qlora_acc_t(0); - for (var i = 0u; i < 4u; i++) { - let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; - sum += dot(vec4(qlora_x_tile[lane * 4u + i]), vec4(bv)); - } - qlora_partials[rank_offset + rank_local][lane] += sum; - } - } - workgroupBarrier(); - } - if (local_idx < 8u) { - var value = qlora_acc_t(0); - for (var i = 0u; i < 32u; i++) { value += qlora_partials[local_idx][i]; } - lora_down[rank_base + local_idx] = output_element_t(value); - } - return; -} -)"; +Status GenerateQLoRADown(ShaderHelper& shader, const ShaderVariableHelper& x, bool acc_f32, + bool use_subgroups) { + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_qlora_down.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32), + WGSL_TEMPLATE_PARAMETER(use_subgroups, use_subgroups), + WGSL_TEMPLATE_VARIABLE(x, x)); } -void GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32, bool vector_output) { - shader.AdditionalImplementation() << "alias qlora_acc_t = " << (acc_f32 ? "f32" : "output_element_t") << ";\n" - << R"( -fn qlora_weight(col: u32, r4: u32) -> vec4 { - let packed = lora_b[col * 2u + r4 / 4u][r4 % 4u]; - return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; -} -fn qlora_delta(row: u32, col: u32) -> output_element_t { - var delta = qlora_acc_t(0); - // The original K=32 INT8 up projection uses wide tiles at M>=4. - if (uniforms.M >= 4u) { - for (var r4 = 0u; r4 < 8u; r4 += 2u) { - let a0 = lora_t[row * 8u + r4]; - let a1 = lora_t[row * 8u + r4 + 1u]; - delta += dot(vec4(a0), vec4(qlora_weight(col, r4))) + - dot(vec4(a1), vec4(qlora_weight(col, r4 + 1u))); - } - } else { - for (var block = 0u; block < 2u; block++) { - var sum = qlora_acc_t(0); - for (var i = 0u; i < 4u; i++) { - let r4 = block * 4u + i; - sum += dot(vec4(lora_t[row * 8u + r4]), vec4(qlora_weight(col, r4))); - } - delta += sum; - } - // The remaining generic-kernel partials are zero, including their sign. - delta += qlora_acc_t(0); - } - return output_element_t(delta); -} -)"; - if (vector_output) { - shader.AdditionalImplementation() << R"( -fn qlora_delta4(row: u32, col: u32) -> vec4 { - var delta = vec4(0); - for (var r4 = 0u; r4 < 8u; r4 += 2u) { - let a0 = vec4(lora_t[row * 8u + r4]); - let a1 = vec4(lora_t[row * 8u + r4 + 1u]); - delta += vec4( - dot(a0, vec4(qlora_weight(col, r4))), - dot(a0, vec4(qlora_weight(col + 1u, r4))), - dot(a0, vec4(qlora_weight(col + 2u, r4))), - dot(a0, vec4(qlora_weight(col + 3u, r4)))) + - vec4( - dot(a1, vec4(qlora_weight(col, r4 + 1u))), - dot(a1, vec4(qlora_weight(col + 1u, r4 + 1u))), - dot(a1, vec4(qlora_weight(col + 2u, r4 + 1u))), - dot(a1, vec4(qlora_weight(col + 3u, r4 + 1u)))); - } - return vec4(delta); -} -)"; - } +Status GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32, bool vector_output) { + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_qlora_up.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32), + WGSL_TEMPLATE_PARAMETER(vector_output, vector_output)); } namespace { +// Single-row execution stores the base projection separately. This shader computes +// the rank-32 up projection and adds it to that base output in one dispatch. class QLoRAOutputProgram final : public Program { public: explicit QLoRAOutputProgram(bool acc_f32) : Program{"QLoRAOutputVec4"}, acc_f32_(acc_f32) {} @@ -178,7 +45,7 @@ class QLoRAOutputProgram final : public Program { shader.AddInput("lora_b"); shader.AddInput("lora_b_scales"); shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); - GenerateQLoRAUp(shader, acc_f32_); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); shader.MainFunctionBody() << "if (global_idx >= uniforms.size) { return; }\n" << "output[global_idx] = base[global_idx] + qlora_delta(global_idx / uniforms.N, global_idx % uniforms.N);\n"; @@ -257,10 +124,14 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { } ORT_RETURN_IF_NOT(as->Shape() == TensorShape({32, k_ / 32}) && bs->Shape() == TensorShape({n_, 1}), "Invalid adapter scales shape."); + // With logical dequantized weights, Y = X * W + T * B_up, where T = X * A_down. + // T keeps X's leading dimensions and replaces K with the adapter rank (32). auto t_shape = x->Shape().AsShapeVector(); t_shape.back() = 32; auto t = context.CreateGPUTensor(x->DataType(), TensorShape(t_shape)); const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); + // Prefill fuses T * B_up into the base shader's output stage. Single-row + // decode uses ApplyQLoRAOutput after computing the base projection and T. const bool fuse_up = m > 1; // FP32 may select DP4A even for M=1. The joint floating-point down // shader must not replace that base projection's activation quantization. @@ -291,6 +162,8 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { nullptr, nullptr, aq, &t, nullptr, nullptr, as)); return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); } + // For one input vector, append four rank-reduction workgroups to the generic + // base dispatch. GenerateQLoRADown supplies their shader body; DP4A has its own. const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && !base_dp4a; if (joint_down) { auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h index 66148cdebf8ee..667b37c3b52b1 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -5,11 +5,29 @@ #include "core/providers/webgpu/program.h" +namespace onnxruntime::webgpu { +class ShaderVariableHelper; +} + namespace onnxruntime::contrib::webgpu { -void GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const std::string& base_groups, - const std::string& x_value, bool acc_f32, bool use_subgroups = false); -void GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32, bool vector_output = false); +// These C++ generation helpers append reusable WGSL fragments to the caller's +// shader. They do not create standalone GPU programs or dispatch GPU work. +// +// Append the rank-32 down-projection body to the generic single-vector decode +// shader. The base projection and T = X * A_down run in the same dispatch using +// separate workgroups. X uses vec4 loads; uniforms.num_N_tile marks the four +// appended adapter workgroups. +Status GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const onnxruntime::webgpu::ShaderVariableHelper& x, + bool acc_f32, bool use_subgroups = false); + +// Append WGSL helper functions for up-weight dequantization and delta = T * B_up. +// Generic and DP4A fused prefill shaders and the separate decode output shader +// reuse these functions. Subgroup-matrix shaders reuse only the dequantization +// helper and implement the multiplication with their own matrix instructions. +Status GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32, bool vector_output = false); + +// Dispatch the decode output shader to compute Y = base + T * B_up. Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, const Tensor* scales, Tensor* y); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template new file mode 100644 index 0000000000000..f1854a3e2495f --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#param acc_f32 +#param use_subgroups + +#use .getByOffset + +// Joint single-vector decode: the base projection and T = X * A_down share +// one dispatch. lora_at stores rank-major packed INT8 weights; lora_down is +// the rank-32 intermediate T consumed by the later up projection. +#if acc_f32 +alias qlora_acc_t = f32; +#else +alias qlora_acc_t = output_element_t; +#endif + +#if !use_subgroups +var qlora_x_tile: array, 128>; +var qlora_partials: array, 8>; +#endif + +$MAIN { + // Four appended workgroups compute eight ranks each. Normalizing the dispatch + // to multiple dimensions may add padding workgroups, which must not write T. + let qlora_base_groups = uniforms.num_N_tile; + if (workgroup_idx >= qlora_base_groups + 4u) { return; } + if (workgroup_idx >= qlora_base_groups) { + let rank_base = (workgroup_idx - qlora_base_groups) * 8u; +#if use_subgroups + // Each of the four 32-lane subgroups computes two ranks. Subgroup membership + // need not match consecutive local invocation indices, so use subgroup IDs. + let lane = sg_id; + let rank_local = subgroup_index; + var partial0 = qlora_acc_t(0); + var partial1 = qlora_acc_t(0); + for (var k = lane * 16u; k < uniforms.K; k += 512u) { + var k4 = k / 4u; + // Unrolling keeps the four X vectors in registers for reuse by both ranks + // and fixes the dot-product grouping that determines FP16 rounding. + let x0 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x1 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x2 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x3 = vec4(x.getByOffset(k4)); + let rank0 = rank_base + rank_local + 0u; + let packed0 = lora_at[rank0 * (uniforms.K / 16u) + k / 16u]; + let scale0 = lora_a_scales[rank0 * (uniforms.K / 32u) + k / 32u]; + var sum0 = qlora_acc_t(0); + sum0 += dot(x0, vec4( + (vec4(unpack4xU8(packed0[0])) - vec4(128)) * scale0)); + sum0 += dot(x1, vec4( + (vec4(unpack4xU8(packed0[1])) - vec4(128)) * scale0)); + sum0 += dot(x2, vec4( + (vec4(unpack4xU8(packed0[2])) - vec4(128)) * scale0)); + sum0 += dot(x3, vec4( + (vec4(unpack4xU8(packed0[3])) - vec4(128)) * scale0)); + partial0 += sum0; + let rank1 = rank_base + rank_local + 4u; + let packed1 = lora_at[rank1 * (uniforms.K / 16u) + k / 16u]; + let scale1 = lora_a_scales[rank1 * (uniforms.K / 32u) + k / 32u]; + var sum1 = qlora_acc_t(0); + sum1 += dot(x0, vec4( + (vec4(unpack4xU8(packed1[0])) - vec4(128)) * scale1)); + sum1 += dot(x1, vec4( + (vec4(unpack4xU8(packed1[1])) - vec4(128)) * scale1)); + sum1 += dot(x2, vec4( + (vec4(unpack4xU8(packed1[2])) - vec4(128)) * scale1)); + sum1 += dot(x3, vec4( + (vec4(unpack4xU8(packed1[3])) - vec4(128)) * scale1)); + partial1 += sum1; + } + // FP32 avoids half-precision rounding at each step of the subgroup tree. + let result0 = subgroupAdd(f32(partial0)); + let result1 = subgroupAdd(f32(partial1)); + if (lane == 0u) { + lora_down[rank_base + rank_local] = output_element_t(result0); + lora_down[rank_base + rank_local + 4u] = output_element_t(result1); + } +#else + let lane = local_idx % 32u; + let rank_local = local_idx / 32u; + // Match the generic INT8 MatMulNBits tile, dot grouping and reduction order + // so the FP16 rounding behavior is preserved in this shared-memory fallback. + for (var start = 0u; start < uniforms.K; start += 512u) { + let k4 = start / 4u + local_idx; + var xv = vec4(0); + if (k4 < uniforms.K / 4u) { xv = x.getByOffset(k4); } + qlora_x_tile[local_idx] = xv; + workgroupBarrier(); + for (var rank_offset = 0u; rank_offset < 8u; rank_offset += 4u) { + let rank = rank_base + rank_offset + rank_local; + let k = start + lane * 16u; + if (k < uniforms.K) { + let packed = lora_at[rank * (uniforms.K / 16u) + k / 16u]; + let scale = lora_a_scales[rank * (uniforms.K / 32u) + k / 32u]; + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; + sum += dot(vec4(qlora_x_tile[lane * 4u + i]), vec4(bv)); + } + qlora_partials[rank_offset + rank_local][lane] += sum; + } + } + workgroupBarrier(); + } + if (local_idx < 8u) { + var value = qlora_acc_t(0); + for (var i = 0u; i < 32u; i++) { value += qlora_partials[local_idx][i]; } + lora_down[rank_base + local_idx] = output_element_t(value); + } +#endif + // Adapter workgroups stop here; the remaining workgroups execute the base + // MatMulNBits body appended after this template. + return; + } +} diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template new file mode 100644 index 0000000000000..e44cd7957490f --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#param acc_f32 +#param vector_output + +// T = X * A_down is stored in lora_t as eight vec4 values per row (rank 32). +// These helpers compute the adapter contribution delta = T * B_up. The caller +// adds delta to the base projection, either in its output stage or a separate +// decode dispatch. Subgroup-matrix templates reuse only qlora_weight and perform +// the up multiplication with their own matrix instructions. +#if acc_f32 +alias qlora_acc_t = f32; +#else +alias qlora_acc_t = output_element_t; +#endif + +fn qlora_weight(col: u32, r4: u32) -> vec4 { + // Each output column has 32 packed INT8 weights in two vec4 values. + // r4 selects four consecutive ranks; 128 is the implicit zero point. + let packed = lora_b[col * 2u + r4 / 4u][r4 % 4u]; + return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; +} + +fn qlora_delta(row: u32, col: u32) -> output_element_t { + var delta = qlora_acc_t(0); + // Match the unfused INT8 up projection's reduction order: wide tiles at M>=4, + // generic tiles otherwise. Regrouping these sums can change FP16 rounding. + if (uniforms.M >= 4u) { + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = lora_t[row * 8u + r4]; + let a1 = lora_t[row * 8u + r4 + 1u]; + delta += dot(vec4(a0), vec4(qlora_weight(col, r4))) + + dot(vec4(a1), vec4(qlora_weight(col, r4 + 1u))); + } + } else { + for (var block = 0u; block < 2u; block++) { + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let r4 = block * 4u + i; + sum += dot(vec4(lora_t[row * 8u + r4]), vec4(qlora_weight(col, r4))); + } + delta += sum; + } + // The remaining generic-kernel partials are zero, including their sign. + delta += qlora_acc_t(0); + } + return output_element_t(delta); +} + +#if vector_output +// Tiled DP4A writes four adjacent output columns at once. Reuse each T vector +// across those columns while keeping the scalar wide-tile dot grouping. +fn qlora_delta4(row: u32, col: u32) -> vec4 { + var delta = vec4(0); + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = vec4(lora_t[row * 8u + r4]); + let a1 = vec4(lora_t[row * 8u + r4 + 1u]); + delta += vec4( + dot(a0, vec4(qlora_weight(col, r4))), + dot(a0, vec4(qlora_weight(col + 1u, r4))), + dot(a0, vec4(qlora_weight(col + 2u, r4))), + dot(a0, vec4(qlora_weight(col + 3u, r4)))) + + vec4( + dot(a1, vec4(qlora_weight(col, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 1u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 2u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 3u, r4 + 1u)))); + } + return vec4(delta); +} +#endif diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index b0f77e61f5f42..4fe60588b47ad 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -149,7 +149,7 @@ Status SubgroupMatrixMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); shader.AddInput("lora_b_scales"); - GenerateQLoRAUp(shader, false); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, false)); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); if (has_tail_buffer_) { From 40fc0c4f11c8b4dfaedb0c5d562dea0d9cb95a72 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 22 Oct 2026 16:22:44 +0800 Subject: [PATCH 30/35] webgpu: Disambiguate plugin upload test provider options --- .../test/providers/webgpu/plugin/test_webgpu_allocators.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc index cfd7509b5ae3c..6018c46a20c64 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc @@ -232,7 +232,7 @@ void CheckSessionInputUploads(Ort::Env& env, Ort::ConstEpDevice ep_device, Ort::SessionOptions options; options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); - options.AppendExecutionProvider_V2(env, {ep_device}, {}); + options.AppendExecutionProvider_V2(env, {ep_device}, Ort::KeyValuePairs{}); const auto model_bytes = model.SerializeAsString(); Ort::Session session(env, model_bytes.data(), model_bytes.size(), options); const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); From 171cbe3ecf3ba59d0e23e177e498fcd222cb05c0 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 23 Oct 2026 14:07:00 +0800 Subject: [PATCH 31/35] Support empty MatMulNBits weights on CPU and WebGPU Accept zero-sized packed defaults while retaining logical output shapes, and skip unused parameter packing and computation. Allow QLoRA fusion from empty adapter defaults with dynamic axes, and cover correctness, overrides, and graph replay. --- docs/ContribOperators.md | 11 +- .../cpu/quantization/matmul_nbits.cc | 35 ++- .../cpu/quantization/matmul_nbits_helper.h | 23 ++ .../webgpu/quantization/matmul_nbits.cc | 26 +- .../core/graph/contrib_ops/contrib_defs.cc | 10 +- .../optimizer/matmul_nbits_qlora_fusion.cc | 28 ++- .../matmul_nbits_empty_weights_test.cc | 227 ++++++++++++++++++ .../matmul_nbits_qlora_fusion_test.cc | 44 +++- 8 files changed, 384 insertions(+), 20 deletions(-) create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index b0a56088c71aa..418003822af2f 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4112,7 +4112,7 @@ This version of the operator has been available since version 1 of the 'com.micr dequantized_weight = (quantized_weight - zero_point) * scale 2. Matrix multiplication between the input matrix A and the dequantized weight matrix. - The weight matrix is a 2D constant matrix with the input feature count and output feature count specified by attributes 'K' and 'N'. + The logical weight matrix is 2D, with the input feature count and output feature count specified by attributes 'K' and 'N'. It is quantized block-wise along the K dimension with a block size specified by the 'block_size' attribute. The block size must be a power of 2 and not smaller than 16 (e.g., 16, 32, 64, 128). Each block has its own scale and zero-point. The quantization is performed using a bit-width specified by the 'bits' attribute, which can take values from 2 to 8. @@ -4120,6 +4120,13 @@ This version of the operator has been available since version 1 of the 'com.micr The quantized weights are stored in a bit-packed format along the K dimension, with each block being represented by a blob of uint8. For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a byte, and the second 4 bits are stored in the higher 4 bits of a byte. + When weight_prepacked is 0, B may represent an omitted projection by replacing one or more dimensions of its + packed shape (N, k_blocks, blob_size) with 0. All remaining dimensions must retain their normal values. + In this case, Y is filled with zeros (or the broadcast bias, if provided). Its shape is still determined by A, K and N. + The contents of A, scales, zero_points and g_idx are not used; the quantization parameters may also be empty. + This permits zero-sized defaults for optional weights, with dynamic input dimensions allowing nonempty overrides. + Empty packed weights are currently supported by the CPU and WebGPU execution providers. + #### Version This version of the operator has been available since version 1 of the 'com.microsoft' operator set. @@ -4147,7 +4154,7 @@ This version of the operator has been available since version 1 of the 'com.micr
A : T1
The input tensor, not quantized.
B : T2
-
Packed uint8 tensor of shape (N, k_blocks, blob_size), where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.
+
Packed uint8 tensor of shape (N, k_blocks, blob_size), where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). With weight_prepacked=0, one or more dimensions may be 0 to omit the projection as described above. The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.
scales : T1
Per-block scaling factors for dequantization with shape (N, k_blocks) and same data type as input A.
zero_points (optional) : T3
diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index d42dbae635047..46f5c81ea698c 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -221,6 +221,12 @@ class MatMulNBits final : public OpKernel { const Tensor* tensor_zero_point = nullptr; has_zp_input_ = info.TryGetConstantInput(InputIndex::zero_points, &tensor_zero_point); + + const Tensor* weight = nullptr; + // Runtime B may be empty. Defer parameter validation/conversion until Compute + // in that case, and never pack unused parameters for a constant empty B. + skip_prepack_ = !info.TryGetConstantInput(InputIndex::B, &weight) || + weight == nullptr || weight->Shape().Size() == 0; } Status Compute(OpKernelContext* context) const override; @@ -242,6 +248,7 @@ class MatMulNBits final : public OpKernel { const bool has_g_idx_; const bool has_bias_; bool scales_are_packed_{false}; + bool skip_prepack_{false}; const bool force_fp32_; const bool prefer_lut_gemm_{false}; const MLAS_QNBIT_GEMM_COMPUTE_TYPE compute_type_; @@ -344,6 +351,8 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; + if (skip_prepack_) return Status::OK(); + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, N_, K_, block_size_, nbits_, has_zp_arg_, has_zp_input_)); @@ -728,6 +737,9 @@ template <> Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc, /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) { + is_packed = false; + if (skip_prepack_) return Status::OK(); + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, N_, K_, block_size_, nbits_, has_zp_arg_, has_zp_input_)); @@ -1526,8 +1538,14 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { const Tensor* reorder_idx = ctx->Input(InputIndex::g_idx); const Tensor* bias = ctx->Input(InputIndex::bias); - ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( - a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); + const bool empty_weights = b != nullptr && b->Shape().Size() == 0; + if (empty_weights) { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_, K_, block_size_, nbits_, + Info().GetAttrOrDefault("weight_prepacked", 0))); + } else { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( + a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); + } TensorShape b_shape({static_cast(N_), static_cast(K_)}); MatMulComputeHelper helper; @@ -1540,6 +1558,19 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { return Status::OK(); } + if (empty_weights) { + auto output = y->MutableDataAsSpan(); + if (bias == nullptr) { + std::fill(output.begin(), output.end(), T1(0.0f)); + } else { + const auto bias_values = bias->DataAsSpan(); + for (size_t offset = 0; offset < output.size(); offset += N_) { + std::copy(bias_values.begin(), bias_values.end(), output.begin() + offset); + } + } + return Status::OK(); + } + AllocatorPtr allocator; ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator)); diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h index 0114e20182308..eeb3117d21d45 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h @@ -34,6 +34,29 @@ constexpr int64_t GetUint8ZeroPointBlobSize(int64_t k_blocks, int64_t bits) { return (k_blocks * bits + 7) / 8; } +// Empty packed weights represent an omitted projection, not a change to logical K/N. +// Keep this separate from CheckInputs: EPs must implement the zero/bias fast path +// before accepting empty weights, and must not dequantize unused parameters. +inline Status CheckEmptyInputs(const Tensor& quantized_weight, const Tensor* bias, + int64_t n, int64_t k, int64_t block_size, int64_t bits, + int64_t weight_prepacked) { + ORT_RETURN_IF_NOT(weight_prepacked == 0, "MatMulNBits: empty B requires weight_prepacked=0"); + ORT_RETURN_IF_NOT(bits == 2 || bits == 4 || bits == 8, "MatMulNBits: unsupported bits: ", bits); + ORT_RETURN_IF_NOT(block_size >= 16 && (block_size & (block_size - 1)) == 0 && + (bits != 2 || block_size <= 256), + "MatMulNBits: unsupported block_size: ", block_size); + const auto& shape = quantized_weight.Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() == 3 && shape.Size() == 0, + "MatMulNBits: empty B must have rank 3"); + const int64_t expected[] = {n, GetKBlocks(k, block_size), GetBlobSize(block_size, bits)}; + for (size_t i = 0; i < 3; ++i) { + ORT_RETURN_IF_NOT(shape[i] == 0 || shape[i] == expected[i], + "MatMulNBits: empty B dimension ", i, " must be 0 or ", expected[i], ", got ", shape[i]); + } + ASSERT_TENSOR_SHAPE(bias, make_shape(n)); + return Status::OK(); +} + template Status CheckInputs(const T* /*activation*/, const T* quantized_weight, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 8bfa23fb1baf6..122aee0a2ef45 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -4,6 +4,7 @@ #include #include +#include "contrib_ops/cpu/quantization/matmul_nbits_helper.h" #include "contrib_ops/webgpu/quantization/matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" @@ -11,6 +12,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/math/matmul.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -242,12 +244,16 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context const Tensor* g_idx = context.Input(4); const Tensor* bias = context.Input(5); - ORT_ENFORCE(g_idx == nullptr, "group_idx as input is not supported yet."); - - const bool has_zero_points = zero_points != nullptr; - if (has_zero_points) { - ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), - "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + const bool empty_weights = b->Shape().Size() == 0; + if (empty_weights) { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_, K_, block_size_, bits_, + Info().GetAttrOrDefault("weight_prepacked", 0))); + } else { + ORT_ENFORCE(g_idx == nullptr, "group_idx as input is not supported yet."); + if (zero_points != nullptr) { + ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), + "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + } } MatMulComputeHelper helper; @@ -260,6 +266,14 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context return Status::OK(); } + if (empty_weights) { + // Record a fill dispatch so graph replay also overwrites the output. The + // MatMul shader uses [batch, row, column] coordinates, even for a vector A. + auto program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( + bias, y, false, 1.0f, 1, TensorShape({1, y->Shape().Size() / N_, N_})); + return context.RunProgram(program); + } + return ApplyMatMulNBits(a, b, scales, zero_points, bias, K_, N_, block_size_, accuracy_level_, bits_, context, y, 0); } diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 4c398ff505007..62edef6d7ed23 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3870,13 +3870,20 @@ It is a fusion of two operations: dequantized_weight = (quantized_weight - zero_point) * scale 2. Matrix multiplication between the input matrix A and the dequantized weight matrix. -The weight matrix is a 2D constant matrix with the input feature count and output feature count specified by attributes 'K' and 'N'. +The logical weight matrix is 2D, with the input feature count and output feature count specified by attributes 'K' and 'N'. It is quantized block-wise along the K dimension with a block size specified by the 'block_size' attribute. The block size must be a power of 2 and not smaller than 16 (e.g., 16, 32, 64, 128). Each block has its own scale and zero-point. The quantization is performed using a bit-width specified by the 'bits' attribute, which can take values from 2 to 8. The quantized weights are stored in a bit-packed format along the K dimension, with each block being represented by a blob of uint8. For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a byte, and the second 4 bits are stored in the higher 4 bits of a byte. + +When weight_prepacked is 0, B may represent an omitted projection by replacing one or more dimensions of its +packed shape (N, k_blocks, blob_size) with 0. All remaining dimensions must retain their normal values. +In this case, Y is filled with zeros (or the broadcast bias, if provided). Its shape is still determined by A, K and N. +The contents of A, scales, zero_points and g_idx are not used; the quantization parameters may also be empty. +This permits zero-sized defaults for optional weights, with dynamic input dimensions allowing nonempty overrides. +Empty packed weights are currently supported by the CPU and WebGPU execution providers. )DOC"; ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBits) @@ -3907,6 +3914,7 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .Input(1, "B", "Packed uint8 tensor of shape (N, k_blocks, blob_size), " "where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). " + "With weight_prepacked=0, one or more dimensions may be 0 to omit the projection as described above. " "The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.", "T2") .Input(2, "scales", "Per-block scaling factors for dequantization with shape (N, k_blocks) and same data type as input A.", "T1") diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc index 9a568c4efe795..60ffd0b73eea4 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -23,17 +23,35 @@ bool Private(const Graph& graph, const Node& node) { return !graph.NodeProducesGraphOutput(node) && optimizer_utils::CheckOutputEdges(graph, node, 1); } -bool Shape(const NodeArg& arg, std::initializer_list expected) { +bool Shape(const NodeArg& arg, std::initializer_list expected, int empty_axis = -1) { const auto* shape = arg.Shape(); if (!shape || shape->dim_size() != static_cast(expected.size())) return false; int index = 0; for (auto value : expected) { - const auto& dim = shape->dim(index++); + const auto& dim = shape->dim(index); + if (index++ == empty_axis && (!dim.has_dim_value() || dim.dim_value() == 0)) continue; if (!dim.has_dim_value() || dim.dim_value() != value) return false; } return true; } +bool AdapterShape(const Graph& graph, const NodeArg& arg, std::initializer_list expected, int empty_axis) { + // K/N attributes establish the logical rank even when an overridable default + // is empty. Only the axis used by the QLoRA empty-adapter convention may vary. + if (!Shape(arg, expected, empty_axis)) return false; + const ONNX_NAMESPACE::TensorProto* init = nullptr; + if (graph.GetInitializedTensor(arg.Name(), init)) { + if (init->dims_size() != static_cast(expected.size())) return false; + int index = 0; + for (auto value : expected) { + const auto dim = init->dims(index); + if (dim != value && !(index == empty_axis && dim == 0)) return false; + ++index; + } + } + return true; +} + bool Projection(const Node* node, int64_t bits) { using namespace matmul_nbits_fusion_utils; return node && graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMulNBits", {1}, kMSDomain) && @@ -52,6 +70,7 @@ bool ZeroDefault(const Graph& graph, const NodeArg& q, const NodeArg& scales) { Initializer qdata(graph, *qp, graph.ModelPath()); Initializer sdata(graph, *sp, graph.ModelPath()); const auto bytes = qdata.DataAsByteSpan(); + if (!bytes.empty() && sdata.DataAsByteSpan().empty()) return false; if (!std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 128; })) return false; if (sp->data_type() == ONNX_NAMESPACE::TensorProto::FLOAT) { const auto values = sdata.DataAsSpan(); @@ -120,8 +139,9 @@ Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == expected; } if (!types_ok || !inputs[0]->Shape() || inputs[0]->Shape()->dim_size() < 2 || !Shape(*inputs[1], {n, k / 32, 16}) || - !Shape(*inputs[3], {32, k / 32, 32}) || !Shape(*inputs[4], {32, k / 32}) || !Shape(*inputs[5], {n, 1, 32}) || - !Shape(*inputs[6], {n, 1})) + !AdapterShape(graph, *inputs[3], {32, k / 32, 32}, 0) || + !AdapterShape(graph, *inputs[4], {32, k / 32}, 0) || + !AdapterShape(graph, *inputs[5], {n, 1, 32}, 1) || !AdapterShape(graph, *inputs[6], {n, 1}, 1)) continue; const auto* s = inputs[2]->Shape(); const bool flat_scales = s && s->dim_size() == 1 && s->dim(0).has_dim_value() && diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc new file mode 100644 index 0000000000000..4bc982fa64248 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc @@ -0,0 +1,227 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) + +#include + +#include "core/session/IOBinding.h" +#include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/providers/provider_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +class MatMulNBitsEmptyWeights : public ::testing::TestWithParam { + void SetUp() override { + if (GetParam() && !DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + } +}; + +template +void RunEmptyCases(bool webgpu) { + struct Case { + int64_t bits, block_size; + std::vector a_shape; + int empty_axis; + bool constant, bias; + }; + const Case cases[] = { + {8, 32, {2, 3, 64}, 0, false, false}, // down projection, batched + {8, 32, {2, 64}, 1, true, false}, // up projection, constant empty B + {4, 16, {64}, 1, true, true}, // vector A and bias + {2, 64, {2, 64}, 2, false, true}, // empty packed-blob axis + {4, 32, {0, 64}, 0, false, false}, // empty output + }; + for (const auto& c : cases) { + SCOPED_TRACE(MakeString("bits=", c.bits, ", block_size=", c.block_size, ", axis=", c.empty_axis)); + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 3); + test.AddAttribute("bits", c.bits); + test.AddAttribute("block_size", c.block_size); + test.AddInput("A", c.a_shape, + std::vector(TensorShape(c.a_shape).Size(), T(std::numeric_limits::quiet_NaN()))); + std::vector b_shape{3, 64 / c.block_size, c.block_size * c.bits / 8}; + b_shape[c.empty_axis] = 0; + test.AddInput("B", b_shape, {}, c.constant); + test.AddInput("scales", {0}, {}, true); + // Unused quantization inputs need neither their normal shapes nor data. + if (!c.constant) { + test.AddInput("zero_points", {0}, {}); + test.AddInput("g_idx", {0}, {}); + } else { + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + } + if (c.bias) test.AddInput("bias", {3}, {T(1.0f), T(-2.0f), T(3.0f)}, c.constant); + auto y_shape = c.a_shape; + y_shape.back() = 3; + std::vector expected(TensorShape(y_shape).Size(), T(0.0f)); + const T bias[] = {T(1.0f), T(-2.0f), T(3.0f)}; + if (c.bias) { + for (size_t i = 0; i < expected.size(); ++i) expected[i] = bias[i % 3]; + } + test.AddOutput("Y", y_shape, expected); + test.SetOutputTolerance(0.0f, 0.0f); + SessionOptions options; + options.use_per_session_threads = false; + if (webgpu) ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + auto ep = webgpu ? DefaultWebGpuExecutionProvider() : DefaultCpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + size_t packed_count = 0; + test.Config(options).ConfigEp(std::move(ep)).RunWithConfig(&packed_count); + EXPECT_EQ(packed_count, 0u); + } +} + +TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp32) { RunEmptyCases(GetParam()); } +TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp16) { RunEmptyCases(GetParam()); } + +TEST_P(MatMulNBitsEmptyWeights, RejectMalformedPackedShape) { + for (bool constant : {false, true}) { + for (const auto& shape : {std::vector{0, 2, 33}, std::vector{0, 64}}) { + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 3); + test.AddAttribute("bits", 8); + test.AddAttribute("block_size", 32); + test.AddInput("A", {1, 64}, std::vector(64, 1.0f)); + test.AddInput("B", shape, {}, constant); + test.AddInput("scales", {0}, {}, constant); + test.AddOutput("Y", {1, 3}, std::vector(3)); + SessionOptions options; + options.use_per_session_threads = false; + if (GetParam()) + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options) + .Config(OpTester::ExpectResult::kExpectFailure, "MatMulNBits: empty B") + .ConfigEp(GetParam() ? DefaultWebGpuExecutionProvider() : DefaultCpuExecutionProvider()) + .RunWithConfig(); + } + } +} + +template +void RunSwitching(bool webgpu, bool capture, bool with_bias = true) { + Model model("empty MatMulNBits defaults", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + auto* a = builder.MakeInput({2, 64}, T(0.5f), T(0.5f)); + auto* b = builder.MakeInput({3, 2, 32}, uint8_t{129}, uint8_t{130}); + auto* s = builder.MakeInput({3, 2}, T(0.25f), T(0.25f)); + for (auto* arg : {b, s}) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(arg->Name()); + init.set_data_type(arg == b ? ONNX_NAMESPACE::TensorProto::UINT8 : utils::ToTensorProtoElementType()); + init.add_dims(3); + init.add_dims(0); + if (arg == b) init.add_dims(32); + graph.AddInitializedTensor(init); + auto shape = *arg->Shape(); + shape.mutable_dim(1)->Clear(); + arg->SetShape(shape); + } + graph.SetInputs({a, b, s}); + auto* y = builder.MakeOutput(std::vector{2, 3}); + std::vector inputs{a, b, s}; + if (with_bias) { + auto* bias = builder.MakeInitializer({3}, T(1.0f), T(1.0f)); + inputs.insert(inputs.end(), {builder.MakeEmptyInput(), builder.MakeEmptyInput(), bias}); + } + auto& node = builder.AddNode("MatMulNBits", inputs, {y}, kMSDomain); + node.AddAttribute("K", int64_t{64}); + node.AddAttribute("N", int64_t{3}); + node.AddAttribute("bits", int64_t{8}); + node.AddAttribute("block_size", int64_t{32}); + node.AddAttribute("accuracy_level", int64_t{0}); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + + ConfigOptions config; + if (capture) ASSERT_STATUS_OK(config.AddConfigEntry("ep.webgpuexecutionprovider.enableGraphCapture", "1")); + auto ep = webgpu ? WebGpuExecutionProviderWithOptions(config) : DefaultCpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + auto* provider = ep.get(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + if (webgpu) ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + const auto data = model.ToProto().SerializeAsString(); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + auto off_feeds = builder.feeds_; + off_feeds.erase(b->Name()); + off_feeds.erase(s->Name()); + auto check = [with_bias](const OrtValue& output, int state) { + const T expected((with_bias ? 1.0f : 0.0f) + (state ? 8.0f : 0.0f)); + for (T value : output.Get().DataAsSpan()) EXPECT_EQ(value, expected); + }; + if (!capture) { + for (int state : {0, 1, 0, 1}) { + std::vector output; + ASSERT_STATUS_OK(session.Run(RunOptions{}, state ? builder.feeds_ : off_feeds, builder.output_names_, &output)); + check(output[0], state); + } + return; + } + + const auto allocator = provider->CreatePreferredAllocators().back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr bindings[2]; + OrtValue outputs[2]; + for (int state : {0, 1}) { + ASSERT_STATUS_OK(session.NewIOBinding(&bindings[state])); + for (const auto& [name, value] : state ? builder.feeds_ : off_feeds) { + const auto& source = value.template Get(); + OrtValue input; + Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); + ASSERT_STATUS_OK(bindings[state]->BindInput(name, input)); + } + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), allocator, outputs[state]); + ASSERT_STATUS_OK(bindings[state]->BindOutput(y->Name(), outputs[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeInputs()); + } + OrtValue actual; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), CPUAllocator::DefaultInstance(), actual); + for (int state : {0, 0, 0, 1, 1, 1, 0, 1}) { + auto* cpu = actual.GetMutable(); + auto values = cpu->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(-99.0f)); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(*cpu, *outputs[state].GetMutable())); + RunOptions run; + ASSERT_STATUS_OK(run.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, state ? "2" : "1")); + ASSERT_STATUS_OK(session.Run(run, *bindings[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeOutputs()); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(outputs[state].Get(), *cpu)); + check(actual, state); + } + EXPECT_TRUE(provider->IsGraphCaptured(1)); + EXPECT_TRUE(provider->IsGraphCaptured(2)); +} + +TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp32) { RunSwitching(GetParam(), false); } +TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp16) { RunSwitching(GetParam(), false); } + +INSTANTIATE_TEST_SUITE_P(Cpu, MatMulNBitsEmptyWeights, ::testing::Values(false)); +#if defined(USE_WEBGPU) +INSTANTIATE_TEST_SUITE_P(WebGpu, MatMulNBitsEmptyWeights, ::testing::Values(true)); +TEST(MatMulNBitsEmptyWeightsCapture, ZeroAndBiasReplay) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (bool bias : {false, true}) { + RunSwitching(true, true, bias); + RunSwitching(true, true, bias); + } +} +#endif + +} // namespace +} // namespace onnxruntime::test +#endif diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 7812fd5c65078..7eee55b472c50 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -15,7 +15,8 @@ namespace { template void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int64_t n = 128, int64_t batch = 1, - bool expose = false, bool different_x = false, bool zero_base = false, int64_t adapter_accuracy = 0) { + bool expose = false, bool different_x = false, bool zero_base = false, int64_t adapter_accuracy = 0, + bool empty_defaults = false) { auto* x = builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)); // Integer random ranges exclude the upper bound, even for constant weights. const uint8_t base_weight_min = zero_base ? uint8_t{0x88} : (ExactBase ? uint8_t{0x99} : uint8_t{0}); @@ -42,6 +43,14 @@ void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int6 const std::vector ones(static_cast(count), T(1.0f)); init.set_raw_data(ones.data(), ones.size() * sizeof(T)); } + if (empty_defaults) { + const int axis = arg == aq || arg == as ? 0 : 1; + init.set_dims(axis, 0); + init.clear_raw_data(); + auto shape = *arg->Shape(); + shape.mutable_dim(axis)->Clear(); + arg->SetShape(shape); + } builder.graph_.AddInitializedTensor(init); } auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); @@ -224,7 +233,7 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentAdapterAc } template -void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64) { +void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64, bool raw_empty_defaults = false) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 2, 37}) { SCOPED_TRACE(m); @@ -234,7 +243,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = ModelTestBuilder builder(graph); // A random base can round differently when the subgroup kernel accumulates // adapter-up into it. Keep this state/capture test exact; Numerics covers rounding. - BuildQLoRA(builder, m, k); + BuildQLoRA(builder, m, k, 128, 1, false, false, false, 0, raw_empty_defaults); builder.SetGraphOutputs(); ASSERT_STATUS_OK(graph.Resolve()); const auto adapters = graph.GetOverridableInitializers(); @@ -255,7 +264,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = } ASSERT_FALSE(a_name.empty()); ASSERT_FALSE(b_name.empty()); - std::vector expected; + std::vector expected(7); // Compare unfused, load-time fusion, and serialized fusion using the same // partial-override states and exact reference values. for (int mode : {0, 1, 2}) { @@ -286,6 +295,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = ASSERT_TRUE(proto.SerializeToString(&data)); SessionOptions options; options.graph_optimization_level = TransformerLevel::Level1; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); InferenceSessionWrapper session(options, GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); @@ -298,6 +308,8 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); // Off, A only, B only, scales only, packed weights only, all, off again. for (int state = 0; state < 7; ++state) { + // Empty scale defaults cannot be used with nonempty packed overrides. + if (raw_empty_defaults && (state == 1 || state == 2 || state == 4)) continue; SCOPED_TRACE(state); auto feeds = builder.feeds_; for (const auto* input : adapters) { @@ -320,7 +332,7 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = [base_value, delta](T value) { return value == T(base_value + delta); })) << "Expected every output to equal " << base_value + delta; if (!fused) { - expected.push_back(std::move(output[0])); + expected[state] = std::move(output[0]); } else { const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; @@ -403,5 +415,27 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionJointDP4ACapture) { RunQLoRASerializedEmptyDefaults(*logger_, 2048); } +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFromEmptyRawDefaults) { + RunQLoRASerializedEmptyDefaults(*logger_, 64, true); + RunQLoRASerializedEmptyDefaults(*logger_, 64, true); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsInvalidDefaultRank) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { + BuildQLoRA(b, 1, 64, 128, 1, false, false, false, 0, true); + auto* arg = b.graph_.GetInputsIncludingInitializers()[1]; + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(b.graph_.GetInitializedTensor(arg->Name(), original)); + auto invalid = *original; + invalid.set_dims(0, 16); + invalid.set_raw_data(std::string(16 * 2 * 32, static_cast(128))); + b.graph_.RemoveInitializedTensor(arg->Name()); + b.graph_.AddInitializedTensor(invalid); + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + #endif } // namespace onnxruntime::test From 12d74efad8aa2ff336a3419f5a842ea77a3b2afc Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 23 Oct 2026 15:55:03 +0800 Subject: [PATCH 32/35] Fix dependent template call in MatMulNBits empty-weight path --- onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index 46f5c81ea698c..e781655f287e0 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -1541,7 +1541,7 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { const bool empty_weights = b != nullptr && b->Shape().Size() == 0; if (empty_weights) { ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_, K_, block_size_, nbits_, - Info().GetAttrOrDefault("weight_prepacked", 0))); + Info().template GetAttrOrDefault("weight_prepacked", 0))); } else { ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); From 5f3272634297ec791ea8cbef4e05005a37a6da37 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 23 Oct 2026 16:55:25 +0800 Subject: [PATCH 33/35] Enable WebGPU QLoRA fusion automatically and preserve defaults --- docs/WebGPU_MatMulNBitsQLoRA_Performance.md | 2 +- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 70 ++++++------ .../core/optimizer/graph_transformer_utils.cc | 3 +- .../optimizer/matmul_nbits_qlora_fusion.cc | 47 +------- .../matmul_nbits_qlora_fusion_test.cc | 100 +++++++++++------- 5 files changed, 101 insertions(+), 121 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md index 78b40e980aca4..58aa03021f091 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md @@ -4,7 +4,7 @@ base projection with packed INT8 adapter projections. Both FP16 and FP32 models are supported, with scales and intermediate tensors matching the model's floating type. See the [operator guide](WebGPU_MatMulNBitsQLoRA_Prototype.md) -for layouts, opt-in graph matching, explicit export, and capture requirements. +for layouts, automatic graph matching, explicit export, and capture requirements. ## QLoRA-only split regression (2026-10-15) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 60f253ebd217f..26c4e5e7df9d7 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -5,21 +5,22 @@ adapter projections, and their residual Add. Adapter weights remain replaceable. is an inference operator for quantized adapters; it does not merge adapter weights into the base or perform adapter training. -Enable graph matching with `session.enable_matmul_nbits_lora_fusion=1` in -the decoder session options. For native GenAI, merge these entries into -the model's existing `model.decoder.session_options` object: +Graph matching runs automatically at the extended graph-optimization level +(`ORT_ENABLE_EXTENDED`) or higher when WebGPU provides the fused kernel. +CPU execution and unsupported patterns retain the original graph. For a comparison +with the original graph, disable this transformer using the standard session option. +For native GenAI, merge this entry into `model.decoder.session_options`: ```json { - "session.enable_matmul_nbits_lora_fusion": "1" + "optimization.disable_specified_optimizers": "MatMulNBitsQLoRAFusion" } ``` -This transformer handles packed INT8 adapters. An explicitly exported -`MatMulNBitsQLoRA` node does not require the graph-matching option. +Append the transformer name with a semicolon if other optimizers are already disabled. +This transformer handles packed INT8 adapters. Both built-in and plugin WebGPU are supported; plugin use requires the host -and provider library from this branch. Disable the option for a comparison -with the original graph. +and provider library from this branch. The transformer matches packed INT8 MatMulNBits adapter projections, not arbitrary DequantizeLinear/Cast/MatMul graphs. Those graphs require an offline @@ -153,16 +154,14 @@ must run on WebGPU; the down projection must consume the base's X. Intermediate outputs must be private. Bias, explicit zero points, group indices, prepacked weights, and incompatible shapes prevent this fusion. -Private overridable defaults containing only byte value 128, with finite -default scales, are replaced by empty packed adapter tensors in the session -graph. Scale defaults remain unchanged, preserving partial override semantics. -Empty A rank or empty B rank selects the ordinary base kernel before adapter -allocation or dispatch. Nonzero and shared packed defaults retain their data. +The matcher preserves initializer shapes and data, including full-size zero +adapter defaults and scale defaults. It does not replace zero weights with empty +tensors: doing so would skip arithmetic and can change NaN/Inf propagation. +Already-empty A rank or B rank selects the ordinary base kernel before adapter +allocation or dispatch. Full-size zero weights still execute the adapter projections. Input names and the adapter file stay unchanged. -Eliding a zero branch assumes finite activations, -adapter scales, and adapter intermediates. NaN/Inf propagation through zero multiplication can differ; -the transformer remains opt-in. GenAI builds without adapter-aware capture +GenAI builds without adapter-aware capture identity require separate generators when changing the adapter selection. [GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740) assigns separate capture IDs after additive adapter activation. Activation does not @@ -182,10 +181,11 @@ layout, and bound buffers stay compatible. ### Compact model packaging -An explicitly exported `MatMulNBitsQLoRA` node accepts empty packed defaults -directly in the serialized model. Full-size buffers filled with 128 are only -needed by the current three-MatMulNBits fusion pattern, not by the fused operator. -For an adapter-off default, export overridable initializers as follows: +Both an explicitly exported `MatMulNBitsQLoRA` node and the raw three-MatMulNBits +pattern accept empty packed defaults. The raw pattern requires a runtime with +empty-weight MatMulNBits support on CPU and WebGPU; its K/N attributes retain the +logical projection sizes. For an adapter-off default, export overridable +initializers as follows: | Initializer | Default shape | Default payload | |---|---|---| @@ -199,17 +199,14 @@ packed graph input as dynamic, so a rank-32 adapter can override its empty initializer. Either empty packed weight disables the adapter. Scale-only and single-weight overrides therefore continue to produce the base result; supplying both packed weights uses the supplied scales or their retained defaults. -Saving a graph after offline QLoRA fusion also produces this representation. - -This compact export requires a runtime with the custom operator. Disabling -fusion does not convert it to a portable fallback graph. Do not place empty -packed weights into the existing fixed-rank MatMulNBits chain: its N/K -attributes still require nonempty rank-32 projections. For an unfused fallback, -retain the original empty-default DequantizeLinear/MatMul graph. WebGPU's -empty-tensor handling supports that graph, but this transformer does not yet -match it. Empty DequantizeLinear outputs avoid dequantization work; a zero-K -MatMul with a nonempty output still initializes that output, and the residual -Add remains. +Scale defaults can also be empty, but activating the adapter then requires +supplying both packed weights and their scales. Fusion preserves whichever +representation the model exports; it does not compact full-size zero defaults. + +Use the raw MatMulNBits pattern for a model shared by CPU and WebGPU. WebGPU +fuses eligible patterns automatically; CPU executes the original operators. +Saving a fused graph serializes explicit MatMulNBitsQLoRA nodes and requires +the WebGPU kernel to execute them. Disabling fusion does not expand explicit nodes. ### FP32 packages and composed base fusions @@ -226,7 +223,7 @@ file across different activation types still requires scale conversion at the graph boundary or a separate mixed-scale type contract. Casting only initializer defaults would not handle later adapter overrides. -Adapter-off execution skips the adapter projections but calls a standalone +Empty adapter weights skip the adapter projections and call a standalone base MatMulNBits. The current QKV/MLP transformers do not compose with the QLoRA operator, so this can lose the throughput benefit of those base fusions. Empty defaults alone do not resolve that difference. A follow-up composed @@ -264,10 +261,11 @@ separately checks the specialized decode reduction. That reference test allows 0.003 absolute error for K=8192 with FP16 lane accumulation; shorter K and FP32 accumulation retain 0.002. Its relative tolerance stays 0.001. Deliberately exact switching checks remain exact. -A serialization regression -checks that empty packed defaults have no payload and that a reloaded explicit -operator accepts full and partial overrides without a load-time QLoRA transform. -It also captures adapter-off and adapter-on under separate IDs and alternates +Serialization regressions check that full-size zero defaults retain their payload, +empty packed defaults remain empty, and a reloaded explicit operator accepts +full and partial overrides without a load-time QLoRA transform. Session tests +exercise automatic fusion and the standard optimizer-disable option. +They also capture adapter-off and adapter-on under separate IDs and alternate replay with sentinel-filled outputs, checking that both captures execute and produce the correct results with CPU fallback disabled. See [the performance evaluation](WebGPU_MatMulNBitsQLoRA_Performance.md) for model-level diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 60833655ef1f0..64b456a5c9f77 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -490,8 +490,7 @@ InlinedVector> GenerateTransformers( } } } - if (has_matmul_nbits_qlora_kernel && - session_options.config_options.GetConfigOrDefault("session.enable_matmul_nbits_lora_fusion", "0") == "1") { + if (has_matmul_nbits_qlora_kernel) { transformers.emplace_back(std::make_unique()); } if (has_matmul_nbits_mlp_kernel) { diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc index 60ffd0b73eea4..971e13221d6fe 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -3,11 +3,8 @@ #include "core/optimizer/matmul_nbits_qlora_fusion.h" -#include - #include "core/graph/graph_utils.h" #include "core/graph/node_attr_utils.h" -#include "core/optimizer/initializer.h" #include "core/optimizer/matmul_nbits_fusion_utils.h" #include "core/optimizer/utils.h" @@ -61,39 +58,6 @@ bool Projection(const Node* node, int64_t bits) { (bits != 8 || GetIntAttr(*node, "accuracy_level", 0) == 0); } -bool ZeroDefault(const Graph& graph, const NodeArg& q, const NodeArg& scales) { - if (!graph.IsInputsIncludingInitializers(&q) || graph.IsOutput(&q) || graph.GetConsumerNodes(q.Name()).size() != 1) - return false; - const ONNX_NAMESPACE::TensorProto* qp = nullptr; - const ONNX_NAMESPACE::TensorProto* sp = nullptr; - if (!graph.GetInitializedTensor(q.Name(), qp) || !graph.GetInitializedTensor(scales.Name(), sp)) return false; - Initializer qdata(graph, *qp, graph.ModelPath()); - Initializer sdata(graph, *sp, graph.ModelPath()); - const auto bytes = qdata.DataAsByteSpan(); - if (!bytes.empty() && sdata.DataAsByteSpan().empty()) return false; - if (!std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 128; })) return false; - if (sp->data_type() == ONNX_NAMESPACE::TensorProto::FLOAT) { - const auto values = sdata.DataAsSpan(); - return std::all_of(values.begin(), values.end(), [](float value) { return std::isfinite(value); }); - } - const auto values = sdata.DataAsSpan(); - return std::all_of(values.begin(), values.end(), [](MLFloat16 value) { return std::isfinite(value.ToFloat()); }); -} - -void EmptyDefault(Graph& graph, NodeArg& arg, int axis) { - const ONNX_NAMESPACE::TensorProto* original = nullptr; - ORT_ENFORCE(graph.GetInitializedTensor(arg.Name(), original)); - ONNX_NAMESPACE::TensorProto empty; - empty.set_name(arg.Name()); - empty.set_data_type(original->data_type()); - for (int i = 0; i < original->dims_size(); ++i) empty.add_dims(i == axis ? 0 : original->dims(i)); - auto shape = *arg.Shape(); - shape.mutable_dim(axis)->Clear(); - arg.SetShape(shape); - graph.RemoveInitializedTensor(arg.Name()); - graph.AddInitializedTensor(empty); -} - } // namespace Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, @@ -147,8 +111,6 @@ Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph const bool flat_scales = s && s->dim_size() == 1 && s->dim(0).has_dim_value() && s->dim(0).dim_value() / n == k / 32 && s->dim(0).dim_value() % n == 0; if (!flat_scales && !Shape(*inputs[2], {n, k / 32})) continue; - const bool empty_defaults = - ZeroDefault(graph, *inputs[3], *inputs[4]) && ZeroDefault(graph, *inputs[5], *inputs[6]); auto attrs = base->GetAttributes(); attrs.erase("weight_prepacked"); utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", int64_t{8}), attrs); @@ -171,13 +133,10 @@ Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph const NodeIndex removed[] = {base->Index(), down->Index(), up->Index(), add->Index()}; for (auto node : removed) graph.RemoveNode(node); for (const auto& edge : outputs) graph.AddEdge(fused.Index(), edge.dst_node, 0, edge.dst_arg_index); - if (empty_defaults) { - // Keep scale defaults: overriding only a packed weight must still use its original scale. - EmptyDefault(graph, *inputs[3], 0); - EmptyDefault(graph, *inputs[5], 1); - } + // Preserve initializer shapes and data. Replacing full-size zero weights with + // empty adapters would bypass arithmetic, including NaN/Inf propagation. modified = true; - LOGS(logger, VERBOSE) << "MatMulNBitsQLoRAFusion: " << fused.Name() << " empty defaults=" << empty_defaults; + LOGS(logger, VERBOSE) << "MatMulNBitsQLoRAFusion: " << fused.Name(); } return Status::OK(); } diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc index 7eee55b472c50..de62dbee98119 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -5,6 +5,7 @@ #include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/IOBinding.h" #include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/optimizer/graph_transform_test_fixture.h" #include "test/optimizer/webgpu_fusion_test_util.h" #include "test/util/include/default_providers.h" @@ -88,24 +89,45 @@ Status CheckQLoRACount(const Graph& graph, size_t expected) { } // namespace TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPreservesDefaults) { - ASSERT_STATUS_OK(TestGraphTransformer( - [](ModelTestBuilder& b) { BuildQLoRA(b); }, 21, *logger_, std::make_unique(), - TransformerLevel::Level2, 1, nullptr, - [](const Graph& graph) { - ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); - const auto inputs = graph.GetOverridableInitializers(); - ORT_RETURN_IF_NOT(inputs.size() == 4, "Adapter parameters were removed."); - for (const auto* input : inputs) { - const ONNX_NAMESPACE::TensorProto* init = nullptr; - ORT_RETURN_IF_NOT(graph.GetInitializedTensor(input->Name(), init), "Missing default."); - if (init->data_type() == ONNX_NAMESPACE::TensorProto::UINT8) { - ORT_RETURN_IF_NOT(init->dims(0) == 0 || init->dims(1) == 0, "Packed default was not emptied."); - } else { - ORT_RETURN_IF_NOT(init->dims(0) > 0 && init->dims(1) > 0, "Scale default was emptied."); + for (bool empty_defaults : {false, true}) { + SCOPED_TRACE(empty_defaults); + InlinedHashMap initializers; + ASSERT_STATUS_OK(TestGraphTransformer( + [&](ModelTestBuilder& b) { + BuildQLoRA(b, 1, 64, 128, 1, false, false, false, 0, empty_defaults); + for (const auto& [name, init] : b.graph_.GetAllInitializedTensors()) { + initializers.emplace(name, init->SerializeAsString()); } - } - return Status::OK(); - })); + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [&](const Graph& graph) { + ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); + ORT_RETURN_IF_NOT(graph.GetOverridableInitializers().size() == 4, "Adapter parameters were removed."); + for (const auto& [name, original] : initializers) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(graph.GetInitializedTensor(name, init), "Missing initializer."); + ORT_RETURN_IF_NOT(init->SerializeAsString() == original, "Initializer shape or data changed."); + } + return Status::OK(); + })); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRequiresWebGpu) { + Model model("CPU quantized LoRA", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + ModelTestBuilder builder(model.MainGraph()); + BuildQLoRA(builder, 1, 64); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + const auto data = model.ToProto().SerializeAsString(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level2; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), 0)); } TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { @@ -233,11 +255,11 @@ TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentAdapterAc } template -void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = 64, bool raw_empty_defaults = false) { +void RunQLoRASerializedDefaults(const logging::Logger& logger, int64_t k = 64, bool raw_empty_defaults = false) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (int64_t m : {1, 2, 37}) { SCOPED_TRACE(m); - Model model("quantized LoRA empty defaults", false, ModelMetaData(), PathString(), + Model model("quantized LoRA defaults", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); auto& graph = model.MainGraph(); ModelTestBuilder builder(graph); @@ -265,8 +287,8 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = ASSERT_FALSE(a_name.empty()); ASSERT_FALSE(b_name.empty()); std::vector expected(7); - // Compare unfused, load-time fusion, and serialized fusion using the same - // partial-override states and exact reference values. + // Compare the standard optimizer opt-out, automatic fusion, and serialized + // fusion using the same partial-override states and exact reference values. for (int mode : {0, 1, 2}) { const bool fused = mode != 0; if (mode == 2) { @@ -279,30 +301,32 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = if (mode == 2) { ASSERT_EQ(proto.graph().node_size(), 1); ASSERT_EQ(proto.graph().node(0).op_type(), "MatMulNBitsQLoRA"); - int empty_count = 0; + int packed_count = 0; for (const auto& init : proto.graph().initializer()) { if (init.name() != a_name && init.name() != b_name) continue; - ++empty_count; + ++packed_count; ASSERT_EQ(init.dims_size(), 3); - EXPECT_EQ(init.dims(init.name() == a_name ? 0 : 1), 0); - EXPECT_TRUE(init.raw_data().empty()); + const int axis = init.name() == a_name ? 0 : 1; + EXPECT_EQ(init.dims(axis), raw_empty_defaults ? 0 : (axis == 0 ? 32 : 1)); + const int64_t size = init.dims(0) * init.dims(1) * init.dims(2); + EXPECT_EQ(init.raw_data(), std::string(static_cast(size), static_cast(128))); EXPECT_EQ(init.int32_data_size(), 0); EXPECT_EQ(init.external_data_size(), 0); } - ASSERT_EQ(empty_count, 2); + ASSERT_EQ(packed_count, 2); } std::string data; ASSERT_TRUE(proto.SerializeToString(&data)); SessionOptions options; - options.graph_optimization_level = TransformerLevel::Level1; + options.graph_optimization_level = mode == 2 ? TransformerLevel::Level1 : TransformerLevel::Level2; + if (mode == 0) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableSpecifiedOptimizers, + "MatMulNBitsQLoRAFusion")); + } ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); InferenceSessionWrapper session(options, GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); - if (mode == 1) { - ASSERT_STATUS_OK( - session.RegisterGraphTransformer(std::make_unique(), TransformerLevel::Level2)); - } // Mode 2 loads the serialized operator without a LoRA transformer. ASSERT_STATUS_OK(session.Initialize()); ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); @@ -403,21 +427,21 @@ void RunQLoRASerializedEmptyDefaults(const logging::Logger& logger, int64_t k = } } -TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaults) { - RunQLoRASerializedEmptyDefaults(*logger_); +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedDefaults) { + RunQLoRASerializedDefaults(*logger_); } -TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedEmptyDefaultsFp32) { - RunQLoRASerializedEmptyDefaults(*logger_); +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedDefaultsFp32) { + RunQLoRASerializedDefaults(*logger_); } TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionJointDP4ACapture) { - RunQLoRASerializedEmptyDefaults(*logger_, 2048); + RunQLoRASerializedDefaults(*logger_, 2048); } TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFromEmptyRawDefaults) { - RunQLoRASerializedEmptyDefaults(*logger_, 64, true); - RunQLoRASerializedEmptyDefaults(*logger_, 64, true); + RunQLoRASerializedDefaults(*logger_, 64, true); + RunQLoRASerializedDefaults(*logger_, 64, true); } TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsInvalidDefaultRank) { From 1360fa84c2fad07126fdc40f94c77d9ba7920495 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 28 Oct 2026 13:19:55 +0800 Subject: [PATCH 34/35] [WebGPU] Simplify DP4A tiling for MatMulNBits QLoRA Restore upstream tiled DP4A configuration while preserving the fused LoRA-up output stage and K-parallel narrow-N path. Update documentation and test comments to match. Validated with D3D12 provider tests, Muffin performance comparisons, and full-model token/logit checks with adapters off/on and graph capture off/on. --- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 8 +- .../quantization/dp4a_matmul.wgsl.template | 438 ++++++++++-------- .../webgpu/quantization/dp4a_matmul_nbits.cc | 22 +- .../webgpu/quantization/dp4a_matmul_nbits.h | 9 +- .../test/contrib_ops/matmul_2bits_test.cc | 2 +- .../test/contrib_ops/matmul_4bits_test.cc | 2 +- .../test/contrib_ops/matmul_8bits_test.cc | 2 +- 7 files changed, 251 insertions(+), 232 deletions(-) diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md index 26c4e5e7df9d7..d0f2d17867ff0 100644 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md @@ -138,10 +138,10 @@ elements; requiring two full iterations for narrow N conservatively amortizes reduction overhead, while shorter K retains tiled reuse across M. This is an initial policy, not a universal performance crossover. Existing DP4A eligibility checks and subgroup-matrix preference still apply. Other DP4A -projections use the general tiled kernel, with 32-column tiles when N < 64 -and 64-column tiles otherwise. Narrow tiles use 32 rows when the device's -maximum subgroup size is below 64; otherwise they keep 64 rows to accommodate -the existing 64-lane subgroup variant. +projections retain the upstream tiled kernel: 64x64 tiles, or 128x32 tiles on +eligible Qualcomm devices, with 256 threads and 32-element K steps. This also +handles rank-32 down projections with K < 2048. The base prefill projection +keeps its fused LoRA-up output stage without changing the matrix tiling. Shapes outside the joint decode specialization calculate T separately. Devices that select subgroup matrix arithmetic for the original INT8 up projection retain that kernel and a separate Add to preserve arithmetic. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index e773f36814e1c..5459a918e44b0 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -4,20 +4,16 @@ #param acc_f32 #param block_size #param n_bits -#param has_qlora_up #param has_bias +#param has_qlora_up #param has_zero_points #param is_qualcomm +#param qc_wave128 #param has_weight_idx #param has_weight_idx_indirect -#param tile_size_k_vec -#param tile_size_m -#param tile_size_n -#param workgroup_size #use .getByOffset .setByOffset - fn write_qlora_output(offset: u32, value: vec4) { #if has_qlora_up let row = (offset * 4u) / uniforms.N; @@ -29,28 +25,54 @@ fn write_qlora_output(offset: u32, value: vec4) { #endif } - #include "quantization/dp4a_matmul_common.wgsl.template" -// Each 16-thread group computes a 16x16 output subtile. Rectangular workgroup -// tiles avoid wasting half the columns on narrow projections. Wider K loads -// amortize barriers while retaining the original 32-element dot/scale order. -// A is dynamically INT8-quantized per 128 elements; B uses its model block size. - -const subtiles_n = tile_size_n / 16; +// This shader implements co-operative matrix multiply. The key idea here is to +// assume there is a primitive for medium size matrix multiply a subgroup can perform, +// using all its lanes and pooling all its registers to keep the values in registry. +// +// The entire workgroup which has N subgroups first loads a tile into shared memory, +// Then each subgroup loads a subtile from shared memory into registers and uses +// the medium size matrix multiply primitive to perform the math. +// The values for tile/subtile size are chosen to conform to the resource limits +// of an alderlake/tiger lake gpu. A tile is 64x64, workgroup is 256 threads - +// therefore there are 16 subgroups and 16 lanes in each subgroup. +// K the hidden dimension is paged in from RAM at k tile size which is 64. +// All this puts the shared memory requirement slightly above 16KB. +// WebGPU limit is 16KB, output is moved to registers instead of SHM to make +// everything fit in shared memory. +// +// Each subgroup performs a 16 x 64 x 16 multiply which is implemented with +// subgroup shuffle as a placeholder for the day the medium matrix mul primitive +// becomes available in WGSL. The registry requirements is ~2KB per subgroup, on +// Alderlake/Tigerlake subgroup has 8KB of registry space pooling the +// 512B of registry from each lane. +// +// The medium size matmul is implemented using dot4I8Packed, so the inputs for +// this shader require A to be int8 quantized with block size 128. B is regular +// matmulnbits input with block size 32. + +#if qc_wave128 +const tile_m = 128; +const tile_n = 32; +#else +const tile_m = 64; +const tile_n = 64; +#endif const subtile_size = 16; +const tile_size_k = 32; const vec_factor = 4; const u32_factor = 4; - +const tile_size_k_vec = 2; // Shared memory -var tile_A : array, tile_size_m>, tile_size_k_vec>; -var scale_A : array; -var tile_B : array, tile_size_n>, tile_size_k_vec>; -var scale_B : array; +var tile_A : array, tile_m>, tile_size_k_vec>; // tile_m x 32 +var scale_A : array; // tile_m x 1 +var tile_B : array, tile_n>, tile_size_k_vec>; // tile_n x 32 +var scale_B : array; // tile_n x 1 #if has_zero_points - var zeroes : array; + var zeroes : array; #endif fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) @@ -88,14 +110,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const actual_weight_idx : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16+kidx_v + col); #endif - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom4BitsTo8Bits(b_value, zero); - if (col % 2u == 0u) + if (col == 0) { // kidx_v - each kidx_v covers 16 values of k let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } } #endif @@ -123,14 +145,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif tile_B[col][row] = AlignWithZeroPoint(b_value); - if (col % 2u == 0u) + if (col == 0) { // kidx_v - each kidx_v covers 16 values of k - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); #if has_zero_points - zeroes[(col / 2u) * tile_size_n + row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); + zeroes[row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); #endif } } @@ -157,17 +179,15 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const b_weight_offset : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); #if has_zero_points let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value, zero); #else tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value); #endif - if (col % 2u == 0u) { - let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); - } + let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } #endif @@ -182,27 +202,31 @@ alias acc_element_t = output_element_t; $MAIN { #if n_bits == 2 + LoadDequantizationTable(local_idx); #if has_zero_points - const table_size = 1024u; -#else - const table_size = 256u; + LoadDequantizationTable(local_idx + 256); + LoadDequantizationTable(local_idx + 512); + LoadDequantizationTable(local_idx + 768); #endif - for (var i = local_idx; i < table_size; i += workgroup_size) { - LoadDequantizationTable(i); - } workgroupBarrier(); #endif + // During the load phase we use all 256 threads to load 64 rows of A/B. + // For each row we load tile_size_k_vec (2) vectorized elements, which are 32 elements of K. let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; } - let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size_m; - let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size_n; - - // Each group of 16 threads computes one 16x16 output subtile. + let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_m; + let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_n; + let load_AorB = u32(local_idx/128); + let load_row = u32((local_idx%128)/2); + let load_col = u32(local_idx%2); + + // During the compute phase, we have the 64x64 tile split into + // subtiles of 16x16. We have a grid of 4x4 subtiles. var subtile_id = u32(local_idx / subtile_size); - var subtile_idx = u32(subtile_id / subtiles_n); - var subtile_idy = u32(subtile_id % subtiles_n); + var subtile_idx = u32(subtile_id / 4); + var subtile_idy = u32(subtile_id % 4); var base_A = subtile_idx * 16; var base_B = subtile_idy * 16; // For each subtile we have 16 threads assigned. @@ -211,10 +235,10 @@ $MAIN { #if is_qualcomm // Qualcomm uses four 64x16 subtiles or two 128x16 subtiles. Each lane // owns one activation row and accumulates 16 output columns. - if (sg_size == tile_size_m) { + if (sg_size == tile_m) { subtile_id = u32(local_idx / sg_size); - subtile_idx = u32(subtile_id / subtiles_n); - subtile_idy = u32(subtile_id % subtiles_n); + subtile_idx = u32(subtile_id / 4); + subtile_idy = u32(subtile_id % 4); base_A = subtile_idx * sg_size; base_B = subtile_idy * 16; a_idx = sg_id; @@ -227,190 +251,196 @@ $MAIN { var lane_output4: vec4; #endif // K's vectorization is 16 items per index. See input_a/input_b. - // Load 32 or 128 K elements, represented by 2 or 8 vec4 values. + // tile_size_k_vec - is the k tile size in vectorized space (1/16). That is + // k tile size is 32. In vectorized space that is 32/16 = 2. for (var kidx_v:u32 = 0; kidx_v < uniforms.K16; kidx_v+=tile_size_k_vec) { // Load Phase: Populate shared memory for the workgroup. - for (var i = local_idx; i < (tile_size_m + tile_size_n) * tile_size_k_vec; i += workgroup_size) { - if (i < tile_size_m * tile_size_k_vec) { - loadSHMA(batch, a_global_base, kidx_v, i / tile_size_k_vec, i % tile_size_k_vec); - } else { - let j = i - tile_size_m * tile_size_k_vec; - loadSHMB(b_global_base, kidx_v, j / tile_size_k_vec, j % tile_size_k_vec); - } +#if qc_wave128 + loadSHMA(batch, a_global_base, kidx_v, local_idx / 2u, local_idx % 2u); + if (local_idx < 64u) { + loadSHMB(b_global_base, kidx_v, local_idx / 2u, local_idx % 2u); + } +#else + if (load_AorB == 0) + { + loadSHMA(batch, a_global_base, kidx_v, load_row, load_col); } + else + { + loadSHMB(b_global_base, kidx_v, load_row, load_col); + } +#endif workgroupBarrier(); - // Keep the original per-block scaling and accumulation order. - for (var k_sub = 0u; k_sub < tile_size_k_vec / 2u; k_sub++) { - // Step 1: Load from shared memory into registers across entire subgroup. - var own_a0: vec4 = tile_A[k_sub * 2u][base_A + a_idx]; - var own_a1: vec4 = tile_A[k_sub * 2u + 1u][base_A + a_idx]; - var own_scale_a: output_element_t = scale_A[base_A + a_idx]; + // Compute phase: Perform matmul for this subtile 16 x 32 x 16. + // Step 1: Load from shared memory into registers across entire subgroup. + var own_a0: vec4 = tile_A[0][base_A + a_idx]; + var own_a1: vec4 = tile_A[1][base_A + a_idx]; + var own_scale_a: output_element_t = scale_A[base_A + a_idx]; #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == tile_size_m) + if (sg_size == tile_m) + { + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + var zero: i32; + if (sg_id < 16) { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - var zero: i32; - if (sg_id < 16) - { - own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; - own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; - zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; - } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); - } + own_b0 = tile_B[0][base_B + sg_id]; + own_b1 = tile_B[1][base_B + sg_id]; + own_scale_b = scale_B[base_B + sg_id]; + zero = zeroes[base_B + sg_id]; } - else + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) { - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i], zeroes[k_sub * tile_size_n + base_B + i]); - } + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); } - #else - if (sg_size == 16) - { - var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; - var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; - var zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); - - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); - - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); - - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); - } - else + } + else + { + for (var i = 0u; i < 16u; i++) { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0], zeroes[k_sub * tile_size_n + base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1], zeroes[k_sub * tile_size_n + base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2], zeroes[k_sub * tile_size_n + base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3], zeroes[k_sub * tile_size_n + base_B + 3]); - - lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4], zeroes[k_sub * tile_size_n + base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5], zeroes[k_sub * tile_size_n + base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6], zeroes[k_sub * tile_size_n + base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7], zeroes[k_sub * tile_size_n + base_B + 7]); - - lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8], zeroes[k_sub * tile_size_n + base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9], zeroes[k_sub * tile_size_n + base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10], zeroes[k_sub * tile_size_n + base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11], zeroes[k_sub * tile_size_n + base_B + 11]); - - lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12], zeroes[k_sub * tile_size_n + base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13], zeroes[k_sub * tile_size_n + base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14], zeroes[k_sub * tile_size_n + base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15], zeroes[k_sub * tile_size_n + base_B + 15]); + lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i], zeroes[base_B + i]); } + } + #else + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[0][base_B + sg_id]; + var own_b1: vec4 = tile_B[1][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[base_B + sg_id]; + var zero = zeroes[base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); + + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); + + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); + + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0], zeroes[base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1], zeroes[base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2], zeroes[base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3], zeroes[base_B + 3]); + + lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4], zeroes[base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5], zeroes[base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6], zeroes[base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7], zeroes[base_B + 7]); + + lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8], zeroes[base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9], zeroes[base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10], zeroes[base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11], zeroes[base_B + 11]); + + lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12], zeroes[base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13], zeroes[base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14], zeroes[base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15], zeroes[base_B + 15]); + } #endif #else #if is_qualcomm - if (sg_size == tile_size_m) - { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - if (sg_id < 16) - { - own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; - own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; - } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); - } - } - else + if (sg_size == tile_m) + { + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + if (sg_id < 16) { - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i]); - } + own_b0 = tile_B[0][base_B + sg_id]; + own_b1 = tile_B[1][base_B + sg_id]; + own_scale_b = scale_B[base_B + sg_id]; } - #else - if (sg_size == 16) + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) { - var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; - var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); - - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); - - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); - - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); } - else + } + else + { + for (var i = 0u; i < 16u; i++) { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3]); - - lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7]); - - lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11]); - - lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15]); + lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i]); } + } + #else + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[0][base_B + sg_id]; + var own_b1: vec4 = tile_B[1][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); + + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); + + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); + + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3]); + + lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7]); + + lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11]); + + lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15]); + } #endif #endif - } workgroupBarrier(); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 6c770d3c9e66b..d52baf102b76c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -48,8 +48,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), - WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), + WGSL_TEMPLATE_PARAMETER(qc_wave128, SubgroupSize() == 128), WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } @@ -210,20 +209,17 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor context.AdapterInfo().subgroupMinSize <= 128 && context.AdapterInfo().subgroupMaxSize >= 128 && nbits == 4 && block_size == 32 && M >= 128; - // The 64-lane subgroup variant spans 64 rows. Reserve 32-row tiles for - // devices whose subgroup range cannot select that variant. - const uint32_t tile_size_m = qc_wave128 ? 128 : (N < 64 && context.AdapterInfo().subgroupMaxSize < 64 ? 32 : 64); - const uint32_t tile_size_n = qc_wave128 || N < 64 ? 32 : 64; - const uint32_t tile_size_k_vec = N < 64 && nbits == 8 && block_size == 32 ? 8 : 2; + const uint32_t tile_m = qc_wave128 ? 128 : 64; + const uint32_t tile_n = qc_wave128 ? 32 : 64; TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components}; - uint32_t num_M_tile = CeilDiv(M, tile_size_m); - uint32_t num_N_tile = CeilDiv(N, tile_size_n); + uint32_t num_M_tile = (M + tile_m - 1) / tile_m; + uint32_t num_N_tile = (N + tile_n - 1) / tile_n; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); DP4AMatMulNBitsProgram mul_program{ - tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, + block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_qlora_up}; - mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); + mul_program.SetWorkgroupSize(256); if (qc_wave128) { mul_program.SetSubgroupSize(128); } @@ -244,8 +240,8 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, - has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, + .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, + has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, has_qlora_up); if (has_zero_points) { mul_program.AddInput( diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 3282ff96ce020..967b71a3555e8 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -30,13 +30,9 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, - uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32, bool has_qlora_up = false) : Program{has_qlora_up ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBits"}, - tile_size_m_(tile_size_m), - tile_size_n_(tile_size_n), - tile_size_k_vec_(tile_size_k_vec), block_size_(block_size), nbits_(nbits), has_bias_(has_bias), @@ -56,9 +52,6 @@ class DP4AMatMulNBitsProgram final : public Program { {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: - uint32_t tile_size_m_; - uint32_t tile_size_n_; - uint32_t tile_size_k_vec_; uint32_t block_size_; uint32_t nbits_; bool has_bias_; diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 26a955b118e68..bb20adab1babb 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1534,7 +1534,7 @@ TEST(MatMul2Bits, Float32_2b_BlkLen32_Accuracy0) { #ifdef USE_WEBGPU TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int2) { - // Partial narrow tile, multiple narrow tiles, and a partial wide tile. + // Tiled row/column tails with N below and above the 64-column tile width. for (int64_t n : {16, 48, 80}) { for (bool has_zero_point : {false, true}) { TestOptions2Bits opts{}; diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index a1987ea1ce7c6..cc80dffdeb3e8 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -454,7 +454,7 @@ TEST(MatMulNBits, Float32_4b_Accuracy0) { #ifdef USE_WEBGPU TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int4) { - // Partial narrow tile, multiple narrow tiles, and a partial wide tile. + // Tiled row/column tails with N below and above the 64-column tile width. for (int64_t n : {16, 48, 80}) { for (bool has_zero_point : {false, true}) { TestOptions opts{}; diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 120dd87e15f28..fc1568600a87e 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1059,7 +1059,7 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { #if defined(USE_WEBGPU) TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { - // Small-M reduction, the tiled threshold, and row/column tails in both tile widths. + // Small-M reduction and tiled row/column tails with N below and above 64. for (const auto& [m, n] : {std::pair{3, 16}, {4, 32}, {65, 48}, {65, 80}}) { for (bool has_zero_point : {false, true}) { TestOptions8Bits opts{}; From 38638f614a6195f0ca889e66aca85015cd9c625f Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 29 Oct 2026 17:18:45 +0800 Subject: [PATCH 35/35] [WebGPU] Refresh MatMulNBitsQLoRA documentation --- docs/WebGPU_MatMulNBitsQLoRA.md | 236 +++++++++++++++++ docs/WebGPU_MatMulNBitsQLoRA_Performance.md | 122 --------- docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | 274 -------------------- 3 files changed, 236 insertions(+), 396 deletions(-) create mode 100644 docs/WebGPU_MatMulNBitsQLoRA.md delete mode 100644 docs/WebGPU_MatMulNBitsQLoRA_Performance.md delete mode 100644 docs/WebGPU_MatMulNBitsQLoRA_Prototype.md diff --git a/docs/WebGPU_MatMulNBitsQLoRA.md b/docs/WebGPU_MatMulNBitsQLoRA.md new file mode 100644 index 0000000000000..895c588986e61 --- /dev/null +++ b/docs/WebGPU_MatMulNBitsQLoRA.md @@ -0,0 +1,236 @@ +# WebGPU MatMulNBitsQLoRA + +`com.microsoft.MatMulNBitsQLoRA` combines a quantized base projection with a +quantized LoRA adapter during inference. Adapter weights remain separate and +replaceable; the operator does not merge them into the base weights or train them. +It reduces dispatches and intermediate memory traffic by sharing input preparation +and computing the adapter contribution inside existing MatMulNBits shaders. + +## Operator contract and current support + +The [operator definition](ContribOperators.md#com.microsoft.matmulnbitsqlora) +follows MatMulNBits quantization and packing rules. It allows arbitrary positive +adapter ranks, including ranks not divisible by eight or the block size. Partial +blocks are padded, and only logical elements participate in the projections. + +The current WebGPU kernel and graph matcher implement this subset: + +- FP16 or FP32 activations, scales, intermediate projections, and output, all + using the same floating type. +- INT4 base weights and INT8 adapter weights, all with block size 32. +- Active adapter rank 32, or empty packed adapter weights as described below. +- Input `X` of rank at least two, positive `K` divisible by 32, and positive `N`. + `N` need not be divisible by four; scalar output tails are supported. +- Implicit zero points: 8 for the base and 128 for the adapters. No bias, + explicit zero points, or group indices. + +These are implementation limits, not additional schema restrictions. Explicit +nodes reject unsupported configurations; the transformer leaves unsupported +patterns unfused. The fused kernel is available only on WebGPU. Both built-in and +plugin WebGPU can use it; plugin deployments need an ORT host with the schema and +transformer as well as a provider library with the kernel. + +For the current implementation, the active inputs are: + +| Input | Storage type | Shape | +| --- | --- | --- | +| `X` | FP16 or FP32 | `[..., M, K]` | +| `W` | Packed INT4 in UINT8 | `[N, K/32, 16]` | +| `scales` | Same as `X` | `[N, K/32]` or flattened | +| `lora_A` | Packed INT8 in UINT8 | `[32, K/32, 32]` | +| `lora_A_scales` | Same as `X` | `[32, K/32]` | +| `lora_B` | Packed INT8 in UINT8 | `[N, 1, 32]` | +| `lora_B_scales` | Same as `X` | `[N, 1]` | + +Let `A_down` and `B_up` be the logical dequantized matrices of shapes `[K, 32]` +and `[32, N]`. The adapter bytes represent signed values offset by 128, multiplied +by their block scales. The computation is: + +```text +base = MatMulNBits(X, W, scales) +T = X * A_down +delta = T * B_up +Y = base + delta +``` + +"Down" projects the input features to the adapter rank; "up" projects that rank +to the output features. `delta` is the adapter's contribution to the output. +`T` has shape `[..., M, 32]`, and `Y` has shape `[..., M, N]`. Dequantized weights +and intermediate projection results use the activation type before addition. +Any LoRA scaling factor must already be represented in the weights/scales; the +operator has no separate adapter multiplier. + +## Automatic graph fusion + +[MatMulNBitsQLoRAFusion](../onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc) +runs at `ORT_ENABLE_EXTENDED` or higher when WebGPU provides the fused kernel. +It matches this graph, accepting either order of the final Add inputs: + +```mermaid +flowchart LR + X --> Base["MatMulNBits: INT4 base"] + X --> Down["MatMulNBits: INT8 down"] + Down --> Up["MatMulNBits: INT8 up"] + Base --> Add + Up --> Add + Add --> Y +``` + +All four nodes must be assigned to WebGPU. The down projection must consume the +same `X` as the base projection. The base, down, and up outputs must each have +one consumer and must not be graph outputs. Packed shapes must match the layout +above, with logical dimensions established by the MatMulNBits `K`/`N` attributes. +Both adapter nodes must use `accuracy_level=0`; the fused node inherits the base +node's accuracy setting. Bias, explicit zero points, group indices, or prepacked +weights (`weight_prepacked != 0`) prevent matching. + +Fusion reuses the input NodeArgs and preserves initializer shapes, contents, and +overridability. It does not quantize adapters, compact zero defaults, or match +arbitrary DequantizeLinear/Cast/MatMul graphs. Such graphs need an offline +conversion to the packed MatMulNBits representation first. + +No QLoRA-specific enable option is needed. To compare with the raw graph, set +the session configuration entry `optimization.disable_specified_optimizers` to +`MatMulNBitsQLoRAFusion` (append it with a semicolon if other names are present). +Disabling the transformer does not expand explicitly exported fused nodes. + +## Execution and optimization + +[MatMulNBitsQLoRA::ComputeInternal](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc) +selects a path from the input shape, accuracy policy, and device capabilities. +Here, `M` is the second-to-last dimension of `X`; a one-token prompt uses the +same selection as decode. Graph fusion does not imply a single GPU dispatch. + +| Case | Execution | +| --- | --- | +| Either packed adapter is empty | Run ordinary base MatMulNBits, before allocating or dispatching adapter work. | +| Prefill (`M > 1`) | Compute `T`, then fuse the up projection and addition into the base shader's output stage. | +| Eligible decode (`M = 1`) | Compute base and down projections in separate workgroups of one dispatch, then run an up-plus-add output shader. | +| Other decode shapes | Compute down and base projections separately, then run the up-plus-add output shader. | + +Packed weights are dequantized where consumed, avoiding full floating-point +adapter-weight tensors. Prefill retains the rank-sized `T` intermediate while +avoiding separate base and adapter-delta output tensors. Decode retains separate +base and `T` outputs; the following dispatch synchronizes them before addition. + +### DP4A + +The [shared DP4A implementation](../onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc) +can quantize `X` once and reuse its packed values and scales for the base and +down projections within an invocation. Its normal eligibility checks, including +`accuracy_level=4`, still apply. Eligible prefill therefore uses quantization, +down projection, and base-plus-up dispatches. FP32 short prefill can also share +this preparation when the base selects DP4A. + +The rank-32 down projection uses the existing K-parallel kernel selection: +`M < 4`, or narrow output (`N <= 32`) with `K >= 2048`. Narrow reductions use +FP32 accumulation. Other shapes use the existing tiled DP4A kernel; the LoRA +output stage extends its computation without changing its matrix tiling. + +For FP32 decode with `2048 <= K <= 8192` and a DP4A base, the joint kernel adds +down-projection workgroups to the base dispatch. Both projections read the same +quantized `X`; up-plus-add is a separate dispatch. Prepared inputs are scoped to +the invocation and are not reused as cached values across different inputs. + +### Subgroup matrix + +The [subgroup-matrix implementation](../onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc) +adds the up projection using matrix instructions in the base shader's output +stage. On a compatible FP16 `8x16x16` configuration, a specialized path packs `X` +once for both projections, with row padding sufficient for both tile sizes. +The narrow down projection uses FP32 matrix accumulation and stores `T` in the +activation type; the base projection then consumes `T` for the fused up stage. + +Selection requires supported component types, matrix dimensions, subgroup size, +and device limits. Choosing a Vulkan backend alone does not guarantee support. +If the shared-packing specialization is unavailable, normal MatMulNBits selection +computes the down projection and selects an eligible base-plus-up shader. + +### Floating-point decode and shader helpers + +For one total input row, `2048 <= K <= 8192`, and a base that does not select +DP4A, the generic base shader appends four down-projection workgroups. A compatible +32-lane subgroup path reuses input values in registers and reduces partial sums; +other devices use workgroup memory. A variable subgroup range requires explicit +subgroup-size control before requesting 32 lanes. + +[GenerateQLoRADown and GenerateQLoRAUp](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h) +append reusable WGSL fragments to a caller's shader; they do not dispatch GPU work. +The [down template](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template) +supplies the generic joint-decode body. The +[up template](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template) +supplies dequantization and adapter-delta helpers used by fused prefill and the +separate decode output shader. Subgroup-matrix shaders reuse the dequantization +helper and implement the multiplication with their own matrix instructions. + +## Optional adapters and portable models + +Use the raw three-MatMulNBits-plus-Add pattern for a model shared by CPU and +WebGPU. CPU executes the original operators; WebGPU fuses eligible patterns. +Saving the optimized graph serializes explicit MatMulNBitsQLoRA nodes, which +require its WebGPU kernel. + +To export a compact adapter-off default, use overridable initializers with these +shapes and also list them as graph inputs: + +| Initializer | Default shape | Contents | +| --- | --- | --- | +| `lora_A` | `[0, K/32, 32]` | Empty packed weights | +| `lora_B` | `[N, 0, 32]` | Empty packed weights | +| `lora_A_scales` | `[32, K/32]` | Retained scale defaults | +| `lora_B_scales` | `[N, 1]` | Retained scale defaults | + +Declare the packed graph inputs' adapter axes as dynamic: axis 0 for `lora_A` +and axis 1 for `lora_B`. This allows active rank-32 weights to override the empty +defaults. Either empty packed weight disables the adapter. Scale-only or +single-weight overrides therefore retain the base result; both packed weights +activate the adapter using supplied scales or their defaults. Scale defaults may +also be empty, but activation then requires supplying their active shapes and data. + +This convention relies on empty-weight MatMulNBits support in the raw graph. +Full-size zero adapter weights still execute arithmetic, including its NaN/Inf +behavior. The transformer deliberately preserves them instead of converting them +to empty tensors. + +## Numerical behavior, capture, and limitations + +`accuracy_level` follows MatMulNBits policy for the base and eligible down +projections. Level 4 permits DP4A activation quantization, including in an +optimized down projection whose raw node uses level 0. Quantization, reduction +order, and accumulation precision can change results relative to the raw graph; +fusion does not promise bitwise equivalence. `enableMatmulFp32Accumulation` +controls accumulation in the shaders that honor that option; it does not disable +activation quantization. Evaluate fused-versus-raw differences separately from +the change caused by quantizing an originally floating-point adapter. + +Graph replay does not re-evaluate the empty-adapter branch. Native callers must +use distinct `gpu_graph_id` values for adapter-off and adapter-on captures, and +for other changes to shapes or execution paths. Each capture needs compatible, +stable GPU bindings and buffer lifetimes. Updating adapter values in the same +buffers can reuse a capture when shapes, layout, and dispatches remain compatible; +replacing a buffer or changing the path requires recapture. WebGPU does not derive +the capture ID from adapter state. + +The current QKV/MLP fusions do not compose with MatMulNBitsQLoRA. An empty adapter +uses standalone base MatMulNBits, so it may have different performance from a +separate base-only graph that benefits from those fusions. Supporting arbitrary +active ranks, additional quantization settings, and composed base fusions requires +further kernel/transformer work; the broader operator schema already allows rank +and quantization extensions. + +## Validation + +From a native WebGPU build's test output directory, run: + +```text +onnxruntime_provider_test.exe --gtest_filter="*QLoRA*:*MatMulNBitsInputPreparation*" +onnxruntime_test_all.exe --gtest_filter="*MatMulNBitsQLoRA*" +``` + +[Provider tests](../onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc) and +[fusion tests](../onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc) +cover FP16/FP32 numerics, prefill/decode, batching and tails, empty and partial +adapters, serialization, capture replay, and patterns that must remain unfused. +Hardware-specific tests skip when the required features are absent; those skips +do not validate the corresponding shader path. Full-model evaluation should also +check logits and generated tokens with each intended adapter and capture mode. diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md b/docs/WebGPU_MatMulNBitsQLoRA_Performance.md deleted file mode 100644 index 58aa03021f091..0000000000000 --- a/docs/WebGPU_MatMulNBitsQLoRA_Performance.md +++ /dev/null @@ -1,122 +0,0 @@ -# WebGPU quantized LoRA evaluation - -`MatMulNBitsQLoRA` keeps rank-32 adapters replaceable while combining an INT4 -base projection with packed INT8 adapter projections. Both FP16 and FP32 -models are supported, with scales and intermediate tensors matching the -model's floating type. See the [operator guide](WebGPU_MatMulNBitsQLoRA_Prototype.md) -for layouts, automatic graph matching, explicit export, and capture requirements. - -## QLoRA-only split regression (2026-10-15) - -A fresh paired comparison against the saved pre-split runtime (`0243e40267`) -uses the benchmark protocol below. Only the ORT DLL changes; GenAI, DXC, models -and adapters are identical. Builds, tests and profiling finished before timing. - -| Type | Prompt | Adapter | TTFT before / after (ms) | Change | Decode before / after (tok/s) | Change | -|---|---:|---|---:|---:|---:|---:| -| FP16 | 128 | off | 31.30 / 31.32 | +0.04% | 224.34 / 227.71 | +1.50% | -| FP16 | 128 | on | 40.48 / 40.20 | -0.70% | 205.08 / 206.54 | +0.71% | -| FP16 | 1024 | off | 163.33 / 163.04 | -0.18% | 222.47 / 225.62 | +1.42% | -| FP16 | 1024 | on | 173.43 / 173.15 | -0.16% | 194.33 / 193.15 | -0.61% | -| FP32 | 128 | off | 29.83 / 29.73 | -0.33% | 195.47 / 192.94 | -1.29% | -| FP32 | 128 | on | 43.41 / 43.51 | +0.25% | 180.45 / 179.47 | -0.54% | -| FP32 | 1024 | off | 162.45 / 162.55 | +0.06% | 192.53 / 192.49 | -0.02% | -| FP32 | 1024 | on | 195.10 / 192.97 | -1.09% | 166.45 / 169.01 | +1.54% | - -No repeatable material regression was observed: TTFT changed by -1.09% to -+0.25%, and decode throughput by -1.29% to +1.54% in these cells. Adapter-on -checks at 512/2048 tokens remained within 1.42%. A noisy 32-token first pass -showed a 2.48% FP16 decode loss and an 8.33% FP32 gain; a reverse-order repeat -measured -0.62%/-1.35% decode changes and -0.19%/+1.55% TTFT changes for -FP16/FP32. The experiment supports preserving performance, not a claimed speedup. -Run-to-run absolute rates varied, so compare paired runtimes rather than -historical absolute timings. - -All 38 native QLoRA/fusion/base regression tests passed with existing tolerances. -All 32 full-model before/after comparisons had exact logits and greedy tokens, -covering both types, two adapters, adapter on/off, and capture on/off. Native -build, Emscripten syntax checks and lint passed. These exact comparisons validate -the split; they do not remove the original-versus-optimized accuracy caveats -below. Original model/adapter hashes are unchanged. - -## Performance reference - -Native GenAI C++, built-in WebGPU, Windows D3D12, RTX 5080. Batch 1, fixed -prompt IDs, 128 generated tokens, capacity 4096, accuracy level 4, graph -capture, generator reuse, and token readback. Three warmups and five measured -requests per cell; medians. Timings were isolated from builds, tests, profiling, -and competing GPU evaluation. These are saved measurements of the QLoRA -implementation before floating-adapter support was separated into a follow-up. - -| Model type | Prompt tokens | TTFT off / on (ms) | Decode off / on (tokens/s) | -|---|---:|---:|---:| -| FP16 | 128 | 31.30 / 40.28 | 235.74 / 204.78 | -| FP16 | 1024 | 163.10 / 172.80 | 232.17 / 204.25 | -| FP32 | 128 | 29.62 / 43.47 | 202.64 / 180.06 | -| FP32 | 1024 | 162.50 / 192.40 | 198.13 / 176.73 | - -TTFT includes prompt processing and first-token sampling/readback. Adapter-off -uses the same optional graph's empty-adapter bypass. The measured TTFT overhead -is 28.7%/5.9% for FP16 and 46.7%/18.4% for FP32 at 128/1024 tokens. Decode -throughput loss is 10.8-13.1%; per-token latency overhead is 12.1-15.1%. -These are local device measurements, not a guarantee for other GPUs. - -An earlier paired comparison against upstream main measured FP16 adapter-on -TTFT reductions of 56-59% at 128 tokens and 29-30% at 1024 tokens, with decode -throughput increasing 75-78%. Those runs used two adapters and ten samples per -configuration on earlier runtime revisions. The original adapter-off graph -failed on that main revision because empty DequantizeLinear was unsupported; -its published base reference was a separately pruned graph. Later main fixes -addressed that empty-tensor issue. The optimized optional graph's adapter-off -decode remained about 10% below the separate base graph with QKV/MLP fusion. -Skipping adapter arithmetic does not recover those composed base fusions. - -## Retained execution paths - -- Prefill (M > 1) selects `fuse_up`: compute T, then combine the base projection - with the INT8 up projection and Add. Eligible DP4A projections share one - quantized X allocation. The narrow down projection reuses the shared - K-parallel DP4A kernel with four output columns for M < 4 and 32 otherwise. - FP32 short prefill uses this path when the base also selects DP4A, preserving - activation reuse without changing fusion. Within eligible DP4A paths, - this kernel handles M < 4 or (N <= 32 and K >= 2048), as well as indirect - expert weights. The narrow-N threshold requires two full K-parallel - iterations to amortize reduction overhead; shorter K retains tiled reuse - across M. -- Eligible FP16 decode selects `fuse_down`: one dispatch computes the base - and down projections in independent workgroups, followed by up/Add. Packed - adapter weights are dequantized as needed; X stays floating point. -- Eligible FP32 decode uses quantize X, joint DP4A base/down, then up/Add. - Quantized X and its scales are reused by the two projections. This increased - decode throughput by approximately 34% against the preceding four-dispatch - FP32 QLoRA implementation in the paired local experiment. -- Empty adapters return through the ordinary base implementation before - adapter allocation or dispatch. Unsupported specializations retain existing - kernels, including subgroup-matrix up projections where selected. - -Workgroup sizes, reduction order, quantization, and dispatch selection are -preserved when removing floating-adapter-only branches. The deferred row-tile -experiment is not included. - -## Numerical interpretation and coverage - -Accuracy level 4 permits DP4A activation quantization. Quantization and changed -reduction order can alter logits and generated tokens relative to the original -unfused package; these performance results do not assert bitwise equivalence. -In the FP32 support evaluation, adapter-off logits were exact. Adapter-on -comparison with the original package reached a maximum logit difference of -1.3800392, 71/72 forced-prefix next-token matches, and 3/6 identical generated -sequences. With activation quantization disabled in both configurations, the -maximum difference was 0.0000234097 and all tested choices and sequences matched. - -The subsequent joint DP4A decode change matched all six generated sequences -and all 72 forced-prefix choices against its preceding QLoRA implementation; -the maximum decode-logit difference was 0.822086. This is a small numerical -regression set, not a broad model-quality evaluation. Adapter-off and prefill -logits were exact across that decode change. - -Provider and fusion coverage includes FP16/FP32, dense FP64 references, K tails, -batching, prefill/decode, empty and partial adapters, compact serialization, -unsupported patterns, and separate adapter-on/off capture IDs. Full-model -capture OFF/ON and generator rewind checks use the GenAI adapter capture-identity -fix from [GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740). diff --git a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md b/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md deleted file mode 100644 index d0f2d17867ff0..0000000000000 --- a/docs/WebGPU_MatMulNBitsQLoRA_Prototype.md +++ /dev/null @@ -1,274 +0,0 @@ -# WebGPU MatMulNBitsQLoRA prototype - -`com.microsoft.MatMulNBitsQLoRA` fuses a quantized base projection, two quantized -adapter projections, and their residual Add. Adapter weights remain replaceable. This -is an inference operator for quantized adapters; it does not merge adapter -weights into the base or perform adapter training. - -Graph matching runs automatically at the extended graph-optimization level -(`ORT_ENABLE_EXTENDED`) or higher when WebGPU provides the fused kernel. -CPU execution and unsupported patterns retain the original graph. For a comparison -with the original graph, disable this transformer using the standard session option. -For native GenAI, merge this entry into `model.decoder.session_options`: - -```json -{ - "optimization.disable_specified_optimizers": "MatMulNBitsQLoRAFusion" -} -``` - -Append the transformer name with a semicolon if other optimizers are already disabled. -This transformer handles packed INT8 adapters. -Both built-in and plugin WebGPU are supported; plugin use requires the host -and provider library from this branch. - -The transformer matches packed INT8 MatMulNBits adapter projections, not -arbitrary DequantizeLinear/Cast/MatMul graphs. Those graphs require an offline -conversion to the weight format below. A model exported with explicit -MatMulNBitsQLoRA nodes does not need the transformer enabled, and disabling -the transformer does not undo those explicit nodes. - -## Weight format and arithmetic - -The [operator definition](ContribOperators.md#com.microsoft.matmulnbitsqlora) -follows MatMulNBits quantization settings and permits any positive adapter rank, -including ranks that are not divisible by 8 or by the quantization block size. -Packed storage pads the final block; the logical rank determines how many -elements participate in the up projection. - -The current WebGPU kernel and graph matcher support a subset: FP16/FP32 -activations and scales, INT4 base weights, INT8 adapter weights, block size 32 -for all projections, active adapter rank 32, positive K divisible by 32, -positive N, and X of rank at least two. Other configurations -are not supported by this kernel. Explicit nodes reject unsupported configurations, -and the matcher leaves them unfused. These are -implementation limits rather than restrictions on the operator definition. - -The operator takes `X, W, scales, lora_A, lora_A_scales, lora_B, lora_B_scales`. -For the current WebGPU implementation: - -| Tensor | Type | Shape | -|---|---|---| -| X | FP16 or FP32 | [..., M, K] | -| W | Packed INT4 in UINT8 | [N, K/32, 16] | -| scales | Same as X | [N, K/32] or flattened | -| lora_A | Packed INT8 in UINT8 | [32, K/32, 32] | -| lora_A_scales | Same as X | [32, K/32] | -| lora_B | Packed INT8 in UINT8 | [N, 1, 32] | -| lora_B_scales | Same as X | [N, 1] | - -Packed adapter bytes encode signed integers with an implicit zero point of -128. For down weight index `(r,k)` and up weight index `(n,r)`, with `F` -denoting the activation type (FP16 or FP32): - -```text -At[r,k] = F((F(Aq[r,k]) - 128) * As[r,floor(k/32)]) -Bt[n,r] = F((F(Bq[n,r]) - 128) * Bs[n,0]) -T = F(X * At^T) -delta = F(T * Bt^T) -Y = F(MatMulNBits(X,W,scales) + delta) -``` - -This implementation uses INT4 base quantization with block size 32 and implicit zero point 8. -The equations describe the weight interpretation and mathematical projections; -the optimized kernels can change reduction order and activation precision. -Eligible prefill down projections use DP4A activation quantization when the -model's `accuracy_level` permits it. Eligible FP32 DP4A decode also shares the -base projection's quantized X with the down projection. Other decode down -projections and all up projections retain floating-point activations. -The provider's `enableMatmulFp32Accumulation` option -controls floating-point lane accumulation for FP16 inputs, except that the -K-parallel DP4A kernel always accumulates in FP32 when N <= 32. FP32 inputs -always use FP32 floating-point accumulation. Intermediate T uses the activation type; -the specialized decode kernel performs its final subgroup reduction in FP32. -These paths are not claimed to be bitwise equivalent to the unfused graph. - -The input layout is already transposed and block packed. The fusion does not -insert a transpose or materialize a complete dequantized adapter. - -## Dispatch - -Selection is automatic. M is the second-to-last dimension of X; a one-token -prompt follows the decode selection. - -| Phase | Selected mode | Implementation for eligible inputs | -|---|---|---| -| Prefill, M > 1 | `fuse_up` | INT8 down projection, then base with INT8 up projection and addition; eligible DP4A paths share quantized X | -| Decode, M = 1 | `fuse_down` | Base and INT8 down projection in one dispatch, then INT8 up projection and addition | - -The joint decode kernel supports one total input row and K from 2048 through -8192, when the base projection does not select DP4A. FP32 base projections -can select DP4A even at M=1. Eligible FP32 decode inputs with K from 2048 -through 8192 use three dispatches: quantize X once, jointly compute the INT4 -base and INT8 down projections, then perform the INT8 up projection and addition. -The joint DP4A shader retains the base's 128-thread workgroups and adds eight -workgroups per input row for rank 32, with four ranks per workgroup. Both -projections read the same quantized activation buffer and scales; each writes -its own output. The next dispatch provides synchronization. The existing DP4A -eligibility checks, including `accuracy_level=4`, still apply. - -The joint decode paths use shape, type, and capability checks without a vendor -restriction. The subgroup reduction accepts a fixed width of 32, or explicitly -requests 32 with `SetSubgroupSize(32)` when `SubgroupSizeControl` is supported -and 32 is within the adapter's range. Other devices use the workgroup-memory -down reduction. Existing base MatMulNBits device policies still apply. - -The floating-input joint kernel retains 128 threads -per workgroup and adds four workgroups for the -32 down-projection outputs. With fixed or explicitly selected 32-lane subgroups, each lane -loads X into registers and reuses it for two rank outputs, then reduces lane -partials with FP32 subgroup sums. The previous workgroup-memory reduction -remains the fallback. `QLoRAOutputVec4` uses 64 threads and one output per -thread, loading packed weights as `vec4` and T as `vec4` along rank. - -FP32 short prefill (1 < M < 4) also uses standalone DP4A for the down -projection when the base selects DP4A. This preserves the base/up fusion and -reuses its activation quantization. FP16 short prefill and inputs that do not -meet the base DP4A requirements retain the floating-input down projection. -DP4A prefill quantizes X once and shares the quantized tensor and scales -between the down projection and base projection, making three dispatches per -adapted projection. Within the shared DP4A implementation, -`DP4AMatMulNBitsKParallelProgram` handles M below -`min_M_for_tile_optimization` (currently 4), or N <= 32 with K >= 2048. -Indirect expert weights also retain this row-wise path. Small M retains four -output columns per tile; otherwise narrow projections use an N-column tile. -This kernel uses FP32 accumulation for N <= 32. The rank-32 down projection -uses this shared selection without a QLoRA-specific override. Each K-parallel iteration covers 1024 -elements; requiring two full iterations for narrow N conservatively amortizes -reduction overhead, while shorter K retains tiled reuse across M. This is an -initial policy, not a universal performance crossover. Existing DP4A -eligibility checks and subgroup-matrix preference still apply. Other DP4A -projections retain the upstream tiled kernel: 64x64 tiles, or 128x32 tiles on -eligible Qualcomm devices, with 256 threads and 32-element K steps. This also -handles rank-32 down projections with K < 2048. The base prefill projection -keeps its fused LoRA-up output stage without changing the matrix tiling. -Shapes outside the joint decode specialization calculate T separately. -Devices that select subgroup matrix arithmetic for the original -INT8 up projection retain that kernel and a separate Add to preserve arithmetic. - -## Adapter off and graph matching - -The matcher requires two INT8 MatMulNBits nodes with block size 32 and accuracy -level 0, an INT4 base MatMulNBits node, and a final Add. All three projections -must run on WebGPU; the down projection must consume the base's X. Intermediate -outputs must be private. Bias, explicit zero points, group indices, prepacked -weights, and incompatible shapes prevent this fusion. - -The matcher preserves initializer shapes and data, including full-size zero -adapter defaults and scale defaults. It does not replace zero weights with empty -tensors: doing so would skip arithmetic and can change NaN/Inf propagation. -Already-empty A rank or B rank selects the ordinary base kernel before adapter -allocation or dispatch. Full-size zero weights still execute the adapter projections. -Input names and the adapter file stay unchanged. - -GenAI builds without adapter-aware capture -identity require separate generators when changing the adapter selection. -[GenAI #2740](https://github.com/microsoft/onnxruntime-genai/pull/2740) assigns -separate capture IDs after additive adapter activation. Activation does not -recompute an existing KV cache: rewind to zero and append the prompt again -if the adapter should apply to the earlier tokens. This does not add an -adapter-deactivation API; use a new base generator to return to adapter-off. - -Graph replay does not re-evaluate the empty-rank branch. Native callers must -use distinct `gpu_graph_id` run options for adapter-off and adapter-on captures, -with stable GPU bindings and buffer lifetimes for each ID. Captures also depend -on tensor shapes and execution mode; prefill and decode cannot share an ID -when their dispatches differ. The provider does not automatically derive an -ID from adapter state. Release and recapture an ID before changing the execution -path or bindings associated with it. Replacing active adapter values in place -is a separate case: the captured path can remain valid only when its shapes, -layout, and bound buffers stay compatible. - -### Compact model packaging - -Both an explicitly exported `MatMulNBitsQLoRA` node and the raw three-MatMulNBits -pattern accept empty packed defaults. The raw pattern requires a runtime with -empty-weight MatMulNBits support on CPU and WebGPU; its K/N attributes retain the -logical projection sizes. For an adapter-off default, export overridable -initializers as follows: - -| Initializer | Default shape | Default payload | -|---|---|---| -| lora_A | [0, K/32, 32] | Empty | -| lora_B | [N, 0, 32] | Empty | -| lora_A_scales | [32, K/32] | Retain the scale defaults | -| lora_B_scales | [N, 1] | Retain the scale defaults | - -Keep these tensors in the graph inputs. Declare the rank dimension of each -packed graph input as dynamic, so a rank-32 adapter can override its empty -initializer. Either empty packed weight disables the adapter. Scale-only and -single-weight overrides therefore continue to produce the base result; supplying -both packed weights uses the supplied scales or their retained defaults. -Scale defaults can also be empty, but activating the adapter then requires -supplying both packed weights and their scales. Fusion preserves whichever -representation the model exports; it does not compact full-size zero defaults. - -Use the raw MatMulNBits pattern for a model shared by CPU and WebGPU. WebGPU -fuses eligible patterns automatically; CPU executes the original operators. -Saving a fused graph serializes explicit MatMulNBitsQLoRA nodes and requires -the WebGPU kernel to execute them. Disabling fusion does not expand explicit nodes. - -### FP32 packages and composed base fusions - -The schema, kernel, and transformer support both FP16 and FP32 activations, -base scales, adapter scales, and outputs. All floating tensors in a node must -have the same type. FP32 packages retain FP32 intermediate projections and -scales; they do not convert the model to FP16. Shader cache keys include the -input/output types, so FP16 and FP32 programs have distinct cache entries. - -The packed INT8 weights use the same layout for both types, but an adapter -file with FP16 scales cannot directly override FP32 scale inputs (or vice -versa). Export scales in the package's activation type. Sharing one adapter -file across different activation types still requires scale conversion at -the graph boundary or a separate mixed-scale type contract. Casting only -initializer defaults would not handle later adapter overrides. - -Empty adapter weights skip the adapter projections and call a standalone -base MatMulNBits. The current QKV/MLP transformers do not compose with the -QLoRA operator, so this can lose the throughput benefit of those base fusions. -Empty defaults alone do not resolve that difference. A follow-up composed -QKV/MLP+QLoRA implementation should dispatch the existing fused base kernels -when every relevant adapter is empty, and preserve per-projection adapter -semantics when any adapter is active. Merely reordering the transformers does -not provide both behaviors. - -## Validation - -Run from the native build output directory: - -```text -onnxruntime_provider_test.exe --gtest_filter=MatMulNBitsQLoRA.* -onnxruntime_test_all.exe --gtest_filter=GraphTransformationTests.MatMulNBitsQLoRAFusion* -``` - -The quantized tests check original-versus-fused FP16 and FP32 results, FP32 accumulation, -batched and partial-tile shapes, empty adapters, partial overrides, -nonzero/shared defaults, unsupported shapes and attributes, and -patterns that must remain unfused. Full-model evaluation must separately -compare fusion equivalence against the original INT8 model and quantization -differences against FP16. Performance runs should use no profiling and record -the actual runtime libraries, provider, capture setting, and adapter state. - -FP32 tests include dense FP64 references with non-half-representable scales, -both fusion modes, DP4A and floating-point down projections, serialized empty -defaults, partial overrides, and separate adapter-on/off graph captures. - -Quantized-prefill and reordered-decode FP16 fusion comparisons use 0.002 absolute -and 0.001 relative tolerances; unaffected FP16 paths retain zero tolerance. -FP32 quantized-prefill comparisons use the same approximation bounds; other -FP32 comparisons use 1e-5 absolute/relative bounds. A dense reference test -separately checks the specialized decode reduction. -That reference test allows 0.003 absolute error for K=8192 with FP16 lane -accumulation; shorter K and FP32 accumulation retain 0.002. Its relative -tolerance stays 0.001. Deliberately exact switching checks remain exact. -Serialization regressions check that full-size zero defaults retain their payload, -empty packed defaults remain empty, and a reloaded explicit operator accepts -full and partial overrides without a load-time QLoRA transform. Session tests -exercise automatic fusion and the standard optimizer-disable option. -They also capture adapter-off and adapter-on under separate IDs and alternate -replay with sentinel-filled outputs, checking that both captures execute and -produce the correct results with CPU fallback disabled. -See [the performance evaluation](WebGPU_MatMulNBitsQLoRA_Performance.md) for model-level -accuracy limits and the main-versus-optimized performance comparison. WebGPU -also batches CPU-to-GPU input uploads to amortize per-tensor adapter transfer -cost; this is independent of the shader fusion.