From bd90fcf759cd2af877332c8fef2102fdb82745e2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 12:21:18 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 18 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 2 +- .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 60 ++- .../runtime/ops/sdpa_fd_decode/SdpaFdDecode.h | 33 +- .../sdpa_fd_reduce_gqa2_f16.wgsl | 83 ++++ .../sdpa_fd_reduce_gqa2_f16_wgsl.h | 107 +++++ .../sdpa_fd_split_gqa2_f16.wgsl | 164 +++++++ .../sdpa_fd_split_gqa2_f16_wgsl.h | 188 ++++++++ backends/webgpu/test/BUCK | 19 + .../webgpu/test/native/test_dynamic_shape.cpp | 41 +- .../webgpu/test/native/test_webgpu_utils.cpp | 56 +++ backends/webgpu/test/ops/test_sdpa.py | 118 ++--- backends/webgpu/test/test_webgpu_native.cpp | 424 ++++++++++++------ backends/webgpu/test/test_wgsl_codegen.py | 6 +- examples/models/qwen3/README.md | 28 ++ .../qwen3/config/qwen3_webgpu_q4gsw.yaml | 32 ++ .../qwen3/manifests/qwen3_0_6b_webgpu.json | 35 ++ examples/models/qwen3/targets.bzl | 16 + examples/models/qwen3/tests/BUCK | 15 + examples/models/qwen3/tests/targets.bzl | 22 + .../tests/test_webgpu_export_manifest.py | 352 +++++++++++++++ .../models/qwen3/webgpu_artifact_manifest.py | 309 +++++++++++++ 22 files changed, 1909 insertions(+), 219 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h create mode 100644 examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml create mode 100644 examples/models/qwen3/manifests/qwen3_0_6b_webgpu.json create mode 100644 examples/models/qwen3/tests/BUCK create mode 100644 examples/models/qwen3/tests/targets.bzl create mode 100644 examples/models/qwen3/tests/test_webgpu_export_manifest.py create mode 100644 examples/models/qwen3/webgpu_artifact_manifest.py diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index ebc3dcc6cd7..a1e6300b4ad 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -119,7 +119,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -155,7 +157,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -961,6 +963,13 @@ constexpr std::array kShaderRegistry = {{ kSdpaFdReduceWorkgroupSizeY, kSdpaFdReduceWorkgroupSizeZ, }, + { + "sdpa_fd_reduce_gqa2_f16", + kSdpaFdReduceGqa2F16WGSL, + kSdpaFdReduceGqa2F16WorkgroupSizeX, + kSdpaFdReduceGqa2F16WorkgroupSizeY, + kSdpaFdReduceGqa2F16WorkgroupSizeZ, + }, { "sdpa_fd_split", kSdpaFdSplitWGSL, @@ -968,6 +977,13 @@ constexpr std::array kShaderRegistry = {{ kSdpaFdSplitWorkgroupSizeY, kSdpaFdSplitWorkgroupSizeZ, }, + { + "sdpa_fd_split_gqa2_f16", + kSdpaFdSplitGqa2F16WGSL, + kSdpaFdSplitGqa2F16WorkgroupSizeX, + kSdpaFdSplitGqa2F16WorkgroupSizeY, + kSdpaFdSplitGqa2F16WorkgroupSizeZ, + }, { "sdpa_fd_split_half", kSdpaFdSplitHalfWGSL, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index bcb5deaa770..47b6467f990 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -694,7 +694,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { // lambda runs, so eager construction here can never throw on it. if (fd_eligible) { state.fd = make_sdpa_fd_decode_state( - gr.device(), Hq, Hkv, D, state.context_len, g, scale); + gr.device(), Hq, Hkv, D, state.context_len, g, scale, gr.kv_f16()); } return state; }; diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index b39c2ae9014..ed0e40f0942 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -11,7 +11,9 @@ #include #include #include +#include #include +#include #include #include @@ -27,14 +29,23 @@ namespace executorch::backends::webgpu { namespace { -// MUST match the .wgsl: MAX_SPLITS and WG_SIZE*MAX_D_PER_LANE. -constexpr uint32_t kSdpaFdSplitTile = 64; // KV positions per split -constexpr uint32_t kSdpaFdMaxSplits = 128; // == MAX_SPLITS in both .wgsl files // Public head-dim limit (kSdpaFdMaxHeadDim) must equal the kernel's lane-owns-D // reach; tie them so a WG_SIZE change can't silently desync the Sdpa.cpp gate. static_assert( kSdpaFdMaxHeadDim == kSdpaFdSplitWorkgroupSizeX * 2u, "kSdpaFdMaxHeadDim must match WG_SIZE * MAX_D_PER_LANE"); +static_assert( + kSdpaFdSplitGqa2F16WorkgroupSizeX == kSdpaFdSplitWorkgroupSizeX, + "Qwen GQA2 and generic split workgroup widths must match"); +static_assert( + kSdpaFdReduceGqa2F16WorkgroupSizeX == kSdpaFdReduceWorkgroupSizeX, + "Qwen GQA2 and generic reduce workgroup widths must match"); +static_assert( + kSdpaFdMaxHeadDim == kSdpaFdSplitGqa2F16WorkgroupSizeX * 2u, + "Qwen GQA2 split must reach D=128"); +static_assert( + kSdpaFdMaxSplits == 128u, + "MAX_SPLITS must match generic and GQA2 split/reduce shaders"); struct FdSplitParams { uint32_t _pad0; // 16B-alignment pad (head index derived from workgroup_id) @@ -138,7 +149,8 @@ SdpaFdDecodeState make_sdpa_fd_decode_state( int64_t D, int64_t context_len, int64_t g, - float scale) { + float scale, + bool kv_f16) { if (Hq <= 0 || Hkv <= 0 || D <= 0 || context_len <= 0 || g <= 0) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: dimensions must be positive"); @@ -162,17 +174,24 @@ SdpaFdDecodeState make_sdpa_fd_decode_state( "WebGPU sdpa FlashDecoding: head dim must be a multiple of 4"); } - uint32_t num_splits = static_cast( - (context_len + kSdpaFdSplitTile - 1) / kSdpaFdSplitTile); - num_splits = std::min(num_splits, kSdpaFdMaxSplits); + const bool qwen_gqa2_f16 = is_qwen_gqa2_f16_fd_route(kv_f16, Hq, Hkv, D, g); + const uint32_t num_splits = + sdpa_fd_num_splits(static_cast(context_len), qwen_gqa2_f16); const uint32_t split_len = static_cast((context_len + num_splits - 1) / num_splits); - const uint64_t split_threads = static_cast(Hq) * + const uint32_t split_heads = sdpa_fd_split_head_count(Hq, Hkv, qwen_gqa2_f16); + const uint32_t split_workgroup_size = qwen_gqa2_f16 + ? kSdpaFdSplitGqa2F16WorkgroupSizeX + : kSdpaFdSplitWorkgroupSizeX; + const uint32_t reduce_workgroup_size = qwen_gqa2_f16 + ? kSdpaFdReduceGqa2F16WorkgroupSizeX + : kSdpaFdReduceWorkgroupSizeX; + const uint64_t split_threads = static_cast(split_heads) * static_cast(num_splits) * - static_cast(kSdpaFdSplitWorkgroupSizeX); + static_cast(split_workgroup_size); const uint64_t reduce_threads = - static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX; + static_cast(Hq) * reduce_workgroup_size; if (split_threads > UINT32_MAX || reduce_threads > UINT32_MAX) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: thread count exceeds uint32 max"); @@ -181,12 +200,12 @@ SdpaFdDecodeState make_sdpa_fd_decode_state( const uint32_t split_wgc = utils::compute_1d_workgroup_count( device, static_cast(split_threads), - kSdpaFdSplitWorkgroupSizeX, + split_workgroup_size, "fd_split"); const uint32_t reduce_wgc = utils::compute_1d_workgroup_count( device, static_cast(reduce_threads), - kSdpaFdReduceWorkgroupSizeX, + reduce_workgroup_size, "fd_reduce"); return { static_cast(Hq), @@ -197,6 +216,7 @@ SdpaFdDecodeState make_sdpa_fd_decode_state( num_splits, split_len, scale, + qwen_gqa2_f16, {split_wgc, 1u}, {reduce_wgc, 1u}}; } @@ -230,7 +250,11 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( {k_cache.buffer, k_cache.nbytes}, {v_cache.buffer, v_cache.nbytes}}; const char* split_shader = kSdpaFdSplitWGSL; - if (graph.kv_f16()) { + const char* split_label = "fd_split"; + if (state.qwen_gqa2_f16) { + split_shader = kSdpaFdSplitGqa2F16WGSL; + split_label = "fd_split_gqa2_f16"; + } else if (graph.kv_f16()) { split_shader = kSdpaFdSplitHalfWGSL; } build_dispatch( @@ -243,7 +267,7 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( sizeof(sp), state.split_grid.x, true, - "fd_split"); + split_label); // Pass 2: reduce (Hq WGs) -> reads part_o, part_ml; writes out. FdReduceParams rp = make_reduce_params(state); @@ -252,9 +276,13 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( {out.buffer, out.nbytes}, {part_o, po_floats * sizeof(float)}, {part_ml, pml_floats * sizeof(float)}}; + const char* reduce_shader = + state.qwen_gqa2_f16 ? kSdpaFdReduceGqa2F16WGSL : kSdpaFdReduceWGSL; + const char* reduce_label = + state.qwen_gqa2_f16 ? "fd_reduce_gqa2_f16" : "fd_reduce"; build_dispatch( graph, - kSdpaFdReduceWGSL, + reduce_shader, reduce_bindings, 3, 1, @@ -262,7 +290,7 @@ SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( sizeof(rp), state.reduce_grid.x, true, - "fd_reduce"); + reduce_label); return {ub_split, ub_reduce, {dispatch_begin, graph.num_dispatches()}}; } diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h index d86b4c624b8..6735c07c678 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h @@ -19,6 +19,35 @@ namespace executorch::backends::webgpu { // MAX_D_PER_LANE(2). Decode shapes above this fall through to the materialized // SDPA path (the FD selection predicate in Sdpa.cpp checks this). constexpr int64_t kSdpaFdMaxHeadDim = 128; +constexpr uint32_t kSdpaFdGenericSplitTile = 64u; +constexpr uint32_t kSdpaFdQwenGqa2SplitTile = 128u; +constexpr uint32_t kSdpaFdMaxSplits = 128u; + +constexpr bool is_qwen_gqa2_f16_fd_route( + bool kv_f16, + int64_t hq, + int64_t hkv, + int64_t d, + int64_t g) { + return kv_f16 && hq == 16 && hkv == 8 && d == 128 && g == 2; +} + +constexpr uint32_t sdpa_fd_split_tile(bool qwen_gqa2_f16) { + return qwen_gqa2_f16 ? kSdpaFdQwenGqa2SplitTile : kSdpaFdGenericSplitTile; +} + +constexpr uint32_t sdpa_fd_num_splits( + uint32_t context_len, + bool qwen_gqa2_f16) { + const uint32_t tile = sdpa_fd_split_tile(qwen_gqa2_f16); + const uint32_t splits = context_len / tile + (context_len % tile != 0u); + return splits > kSdpaFdMaxSplits ? kSdpaFdMaxSplits : splits; +} + +constexpr uint32_t +sdpa_fd_split_head_count(int64_t hq, int64_t hkv, bool qwen_gqa2_f16) { + return static_cast(qwen_gqa2_f16 ? hkv : hq); +} struct SdpaFdDecodeState { uint32_t Hq; @@ -29,6 +58,7 @@ struct SdpaFdDecodeState { uint32_t num_splits; uint32_t split_len; float scale; + bool qwen_gqa2_f16; utils::WgCount split_grid; utils::WgCount reduce_grid; }; @@ -46,7 +76,8 @@ SdpaFdDecodeState make_sdpa_fd_decode_state( int64_t D, int64_t context_len, int64_t g, - float scale); + float scale, + bool kv_f16); // Records split + reduce with retained UBOs. Route selection is owned by the // caller so this helper never mutates recorded dispatch counts. diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16.wgsl new file mode 100644 index 00000000000..8d8f993b259 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16.wgsl @@ -0,0 +1,83 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_part_o: array; +@group(0) @binding(2) var t_part_ml: array; + +struct Params { + D: u32, + num_splits: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +// w_i = exp(m_i - M) per split, computed once and reused for the L-sum and every output dim. +var sh_w: array; +var sh_reduce: array; + +// FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h = wid.x; + let t = lid.x; + let D = params.D; + let ns = params.num_splits; + let head_base = h * MAX_SPLITS; + + var local_m: f32 = NEG_INF; + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + local_m = max(local_m, t_part_ml[(head_base + i) * 2u + 0u]); + } + sh_reduce[t] = local_m; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_reduce[t] = max(sh_reduce[t], sh_reduce[t + stride]); + } + workgroupBarrier(); + } + let M = sh_reduce[0]; + // Compute w_i = exp(m_i - M) once per split into shared memory (was recomputed per output dim). + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + sh_w[i] = exp(t_part_ml[(head_base + i) * 2u + 0u] - M); + } + workgroupBarrier(); + + var local_l: f32 = 0.0; + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + local_l = local_l + sh_w[i] * t_part_ml[(head_base + i) * 2u + 1u]; + } + sh_reduce[t] = local_l; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_reduce[t] = sh_reduce[t] + sh_reduce[t + stride]; + } + workgroupBarrier(); + } + let L = sh_reduce[0]; + let inv = select(0.0, 1.0 / L, L > 0.0); + + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + var acc: f32 = 0.0; + for (var i: u32 = 0u; i < ns; i = i + 1u) { + acc = acc + sh_w[i] * t_part_o[(head_base + i) * D + d]; + } + t_out[h * D + d] = acc * inv; + } + } +} diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h new file mode 100644 index 00000000000..43b52eeb81f --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_gqa2_f16_wgsl.h @@ -0,0 +1,107 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from sdpa_fd_reduce_gqa2_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: 1de12c2e2285186de29a81cc300ce7dba89c6652b56543f5ded4352284729646 +inline constexpr const char* kSdpaFdReduceGqa2F16WGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_part_o: array; +@group(0) @binding(2) var t_part_ml: array; + +struct Params { + D: u32, + num_splits: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +// w_i = exp(m_i - M) per split, computed once and reused for the L-sum and every output dim. +var sh_w: array; +var sh_reduce: array; + +// FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h = wid.x; + let t = lid.x; + let D = params.D; + let ns = params.num_splits; + let head_base = h * MAX_SPLITS; + + var local_m: f32 = NEG_INF; + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + local_m = max(local_m, t_part_ml[(head_base + i) * 2u + 0u]); + } + sh_reduce[t] = local_m; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_reduce[t] = max(sh_reduce[t], sh_reduce[t + stride]); + } + workgroupBarrier(); + } + let M = sh_reduce[0]; + // Compute w_i = exp(m_i - M) once per split into shared memory (was recomputed per output dim). + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + sh_w[i] = exp(t_part_ml[(head_base + i) * 2u + 0u] - M); + } + workgroupBarrier(); + + var local_l: f32 = 0.0; + for (var i: u32 = t; i < ns; i = i + WG_SIZE) { + local_l = local_l + sh_w[i] * t_part_ml[(head_base + i) * 2u + 1u]; + } + sh_reduce[t] = local_l; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_reduce[t] = sh_reduce[t] + sh_reduce[t + stride]; + } + workgroupBarrier(); + } + let L = sh_reduce[0]; + let inv = select(0.0, 1.0 / L, L > 0.0); + + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + var acc: f32 = 0.0; + for (var i: u32 = 0u; i < ns; i = i + 1u) { + acc = acc + sh_w[i] * t_part_o[(head_base + i) * D + d]; + } + t_out[h * D + d] = acc * inv; + } + } +} +)"; + +inline constexpr uint32_t kSdpaFdReduceGqa2F16WorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaFdReduceGqa2F16WorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaFdReduceGqa2F16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16.wgsl new file mode 100644 index 00000000000..6aef6689479 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16.wgsl @@ -0,0 +1,164 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +enable f16; + +@group(0) @binding(0) var t_part_o: array; +@group(0) @binding(1) var t_part_ml: array; +@group(0) @binding(2) var t_q: array; +@group(0) @binding(3) var t_k_cache: array; +@group(0) @binding(4) var t_v_cache: array; + +struct Params { + _pad0: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + num_splits: u32, + split_len: u32, + scale: f32, +} +@group(0) @binding(5) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const G: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +var sh_p: array; +var sh_red: array; + +// Qwen3 GQA=2 f16 FlashDecoding split. A workgroup covers both query heads +// associated with one KV head. QK remains independent for each query head; +// the V row is loaded once and applied to both heads' softmax weights. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h_kv = wid.x / params.num_splits; + let split_i = wid.x % params.num_splits; + let t = lid.x; + let D = params.D; + let D4 = D / 4u; + let ctx = params.context_len; + let kv_row_stride = params.Hkv * D; + + let c0 = split_i * params.split_len; + var c1 = c0 + params.split_len; + if (c1 > ctx) { c1 = ctx; } + + var m: array; + var l: array; + var o_acc: array, G>; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + m[group_head] = NEG_INF; + l[group_head] = 0.0; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + o_acc[group_head][nd] = 0.0; + } + } + + var block: u32 = c0; + loop { + if (block >= c1) { break; } + var n: u32 = c1 - block; + if (n > WG_SIZE) { n = WG_SIZE; } + + if (t < n) { + let kv_base = (block + t) * kv_row_stride + h_kv * D; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + let q_base = (h_kv * G + group_head) * D; + var acc4 = vec4(0.0, 0.0, 0.0, 0.0); + for (var i4: u32 = 0u; i4 < D4; i4 = i4 + 1u) { + let qi = q_base + i4 * 4u; + let ki = kv_base + i4 * 4u; + let qv = vec4( + t_q[qi], t_q[qi + 1u], t_q[qi + 2u], t_q[qi + 3u]); + let kv = vec4( + f32(t_k_cache[ki]), f32(t_k_cache[ki + 1u]), + f32(t_k_cache[ki + 2u]), f32(t_k_cache[ki + 3u])); + acc4 = acc4 + qv * kv; + } + sh_p[group_head * WG_SIZE + t] = + (acc4.x + acc4.y + acc4.z + acc4.w) * params.scale; + } + } else { + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + sh_p[group_head * WG_SIZE + t] = NEG_INF; + } + } + workgroupBarrier(); + + var rescale: array; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + sh_red[t] = sh_p[group_head * WG_SIZE + t]; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_red[t] = max(sh_red[t], sh_red[t + stride]); + } + workgroupBarrier(); + } + let m_new = max(m[group_head], sh_red[0]); + rescale[group_head] = exp(m[group_head] - m_new); + + var p_t: f32 = 0.0; + if (t < n) { + p_t = exp(sh_p[group_head * WG_SIZE + t] - m_new); + } + workgroupBarrier(); + sh_p[group_head * WG_SIZE + t] = p_t; + sh_red[t] = p_t; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_red[t] = sh_red[t] + sh_red[t + stride]; + } + workgroupBarrier(); + } + l[group_head] = rescale[group_head] * l[group_head] + sh_red[0]; + m[group_head] = m_new; + workgroupBarrier(); + } + + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + o_acc[group_head][nd] = + rescale[group_head] * o_acc[group_head][nd]; + } + for (var j: u32 = 0u; j < n; j = j + 1u) { + let v_base = (block + j) * kv_row_stride + h_kv * D; + let v_value = f32(t_v_cache[v_base + d]); + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + o_acc[group_head][nd] = o_acc[group_head][nd] + + sh_p[group_head * WG_SIZE + j] * v_value; + } + } + } + } + workgroupBarrier(); + block = block + WG_SIZE; + } + + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + let h = h_kv * G + group_head; + let part = h * MAX_SPLITS + split_i; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + t_part_o[part * D + d] = o_acc[group_head][nd]; + } + } + if (t == 0u) { + t_part_ml[part * 2u] = m[group_head]; + t_part_ml[part * 2u + 1u] = l[group_head]; + } + } +} diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h new file mode 100644 index 00000000000..7276feb47e8 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_gqa2_f16_wgsl.h @@ -0,0 +1,188 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from sdpa_fd_split_gqa2_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: 9571dc5a5b05c4f9d1f39c37e18100c3de3fe65573ae501f9b3442c8c21d9632 +inline constexpr const char* kSdpaFdSplitGqa2F16WGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +enable f16; + +@group(0) @binding(0) var t_part_o: array; +@group(0) @binding(1) var t_part_ml: array; +@group(0) @binding(2) var t_q: array; +@group(0) @binding(3) var t_k_cache: array; +@group(0) @binding(4) var t_v_cache: array; + +struct Params { + _pad0: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + num_splits: u32, + split_len: u32, + scale: f32, +} +@group(0) @binding(5) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const G: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +var sh_p: array; +var sh_red: array; + +// Qwen3 GQA=2 f16 FlashDecoding split. A workgroup covers both query heads +// associated with one KV head. QK remains independent for each query head; +// the V row is loaded once and applied to both heads' softmax weights. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h_kv = wid.x / params.num_splits; + let split_i = wid.x % params.num_splits; + let t = lid.x; + let D = params.D; + let D4 = D / 4u; + let ctx = params.context_len; + let kv_row_stride = params.Hkv * D; + + let c0 = split_i * params.split_len; + var c1 = c0 + params.split_len; + if (c1 > ctx) { c1 = ctx; } + + var m: array; + var l: array; + var o_acc: array, G>; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + m[group_head] = NEG_INF; + l[group_head] = 0.0; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + o_acc[group_head][nd] = 0.0; + } + } + + var block: u32 = c0; + loop { + if (block >= c1) { break; } + var n: u32 = c1 - block; + if (n > WG_SIZE) { n = WG_SIZE; } + + if (t < n) { + let kv_base = (block + t) * kv_row_stride + h_kv * D; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + let q_base = (h_kv * G + group_head) * D; + var acc4 = vec4(0.0, 0.0, 0.0, 0.0); + for (var i4: u32 = 0u; i4 < D4; i4 = i4 + 1u) { + let qi = q_base + i4 * 4u; + let ki = kv_base + i4 * 4u; + let qv = vec4( + t_q[qi], t_q[qi + 1u], t_q[qi + 2u], t_q[qi + 3u]); + let kv = vec4( + f32(t_k_cache[ki]), f32(t_k_cache[ki + 1u]), + f32(t_k_cache[ki + 2u]), f32(t_k_cache[ki + 3u])); + acc4 = acc4 + qv * kv; + } + sh_p[group_head * WG_SIZE + t] = + (acc4.x + acc4.y + acc4.z + acc4.w) * params.scale; + } + } else { + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + sh_p[group_head * WG_SIZE + t] = NEG_INF; + } + } + workgroupBarrier(); + + var rescale: array; + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + sh_red[t] = sh_p[group_head * WG_SIZE + t]; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_red[t] = max(sh_red[t], sh_red[t + stride]); + } + workgroupBarrier(); + } + let m_new = max(m[group_head], sh_red[0]); + rescale[group_head] = exp(m[group_head] - m_new); + + var p_t: f32 = 0.0; + if (t < n) { + p_t = exp(sh_p[group_head * WG_SIZE + t] - m_new); + } + workgroupBarrier(); + sh_p[group_head * WG_SIZE + t] = p_t; + sh_red[t] = p_t; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { + sh_red[t] = sh_red[t] + sh_red[t + stride]; + } + workgroupBarrier(); + } + l[group_head] = rescale[group_head] * l[group_head] + sh_red[0]; + m[group_head] = m_new; + workgroupBarrier(); + } + + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + o_acc[group_head][nd] = + rescale[group_head] * o_acc[group_head][nd]; + } + for (var j: u32 = 0u; j < n; j = j + 1u) { + let v_base = (block + j) * kv_row_stride + h_kv * D; + let v_value = f32(t_v_cache[v_base + d]); + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + o_acc[group_head][nd] = o_acc[group_head][nd] + + sh_p[group_head * WG_SIZE + j] * v_value; + } + } + } + } + workgroupBarrier(); + block = block + WG_SIZE; + } + + for (var group_head: u32 = 0u; group_head < G; group_head = group_head + 1u) { + let h = h_kv * G + group_head; + let part = h * MAX_SPLITS + split_i; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + t_part_o[part * D + d] = o_acc[group_head][nd]; + } + } + if (t == 0u) { + t_part_ml[part * 2u] = m[group_head]; + t_part_ml[part * 2u + 1u] = l[group_head]; + } + } +} +)"; + +inline constexpr uint32_t kSdpaFdSplitGqa2F16WorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaFdSplitGqa2F16WorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaFdSplitGqa2F16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index bcddaa0f66b..e65810bf33f 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -60,3 +60,22 @@ fbcode_target( "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", ], ) + +fbcode_target( + _kind = runtime.python_test, + name = "test_sdpa", + srcs = ["ops/test_sdpa.py"], + typing = True, + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib_mkl_noomp", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/exir:lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/extension/pybindings:portable_lib", + ], +) diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index cf16cd7e87c..a772b64c000 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -746,7 +746,9 @@ void expect_sdpa_route( const std::vector& names, int s, bool expect_k16, - const char* k16_kernel_name = "sdpa_streaming_attention_k16_causal_bound") { + const char* k16_kernel_name = "sdpa_streaming_attention_k16_causal_bound", + const char* fd_split_kernel_name = "fd_split", + const char* fd_reduce_kernel_name = "fd_reduce") { const bool expect_fd = s == 1; const bool expect_materialized = !expect_fd && !expect_k16; EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); @@ -754,9 +756,11 @@ void expect_sdpa_route( std::count(names.begin(), names.end(), k16_kernel_name), expect_k16 ? 1 : 0); EXPECT_EQ( - std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); + std::count(names.begin(), names.end(), fd_split_kernel_name), + expect_fd ? 1 : 0); EXPECT_EQ( - std::count(names.begin(), names.end(), "fd_reduce"), expect_fd ? 1 : 0); + std::count(names.begin(), names.end(), fd_reduce_kernel_name), + expect_fd ? 1 : 0); EXPECT_EQ( std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), expect_materialized ? 1 : 0); @@ -770,6 +774,19 @@ void expect_sdpa_route( names.size(), static_cast(2 + (expect_k16 ? 1 : (expect_fd ? 2 : 3)))); } + +TEST(DynamicShape, Qwen3FdProfileLabelsContract) { + expect_sdpa_route( + {"update_cache", + "update_cache", + "fd_split_gqa2_f16", + "fd_reduce_gqa2_f16"}, + 1, + false, + "sdpa_streaming_attention_qwen3_k16_causal_bound", + "fd_split_gqa2_f16", + "fd_reduce_gqa2_f16"); +} #endif void run_combined_routes(Module& m, int s) { @@ -1690,11 +1707,19 @@ TEST(DynamicShape, Qwen3K16CausalLiveRoutesProfile) { constexpr float kQwen3MaxError = 1e-2f; constexpr const char* kQwen3Kernel = "sdpa_streaming_attention_qwen3_k16_causal_bound"; + constexpr const char* kQwen3FdSplitKernel = "fd_split_gqa2_f16"; + constexpr const char* kQwen3FdReduceKernel = "fd_reduce_gqa2_f16"; Module module(g_dir + "/sdpa_k16_qwen3.pte"); load_sdpa_module(module, true); prime_k16_sdpa( module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); - expect_sdpa_route(current_profile_names(), 12, true, kQwen3Kernel); + expect_sdpa_route( + current_profile_names(), + 12, + true, + kQwen3Kernel, + kQwen3FdSplitKernel, + kQwen3FdReduceKernel); for (int s : {128, 1, 17, 1, 128}) { run_k16_sdpa( module, @@ -1705,7 +1730,13 @@ TEST(DynamicShape, Qwen3K16CausalLiveRoutesProfile) { kQwen3D, false, kQwen3MaxError); - expect_sdpa_route(current_profile_names(), s, s > 1, kQwen3Kernel); + expect_sdpa_route( + current_profile_names(), + s, + s > 1, + kQwen3Kernel, + kQwen3FdSplitKernel, + kQwen3FdReduceKernel); } } diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index f60dc7be2fd..87ca2542b4a 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -11,9 +11,13 @@ #include #include +#include +#include #include +#include + #include using namespace executorch::backends::webgpu; @@ -147,3 +151,55 @@ TEST(WebGPUUtils, ArgReduceResizeRejectsScratchGrowth) { EXPECT_FALSE(arg_reduce_partial_slots_fit(256u, 257u, 1u)); EXPECT_FALSE(arg_reduce_partial_slots_fit(256u, 129u, 2u)); } + +TEST(WebGPUUtils, QwenGqa2FdRouteRequiresEveryExactPredicate) { + EXPECT_TRUE(is_qwen_gqa2_f16_fd_route(true, 16, 8, 128, 2)); + EXPECT_FALSE(is_qwen_gqa2_f16_fd_route(false, 16, 8, 128, 2)); + EXPECT_FALSE(is_qwen_gqa2_f16_fd_route(true, 15, 8, 128, 2)); + EXPECT_FALSE(is_qwen_gqa2_f16_fd_route(true, 16, 7, 128, 2)); + EXPECT_FALSE(is_qwen_gqa2_f16_fd_route(true, 16, 8, 64, 2)); + EXPECT_FALSE(is_qwen_gqa2_f16_fd_route(true, 16, 8, 128, 1)); +} + +TEST(WebGPUUtils, QwenGqa2FdRouteUsesTile128Schedule) { + EXPECT_EQ(sdpa_fd_split_tile(true), 128u); + EXPECT_EQ(sdpa_fd_num_splits(1u, true), 1u); + EXPECT_EQ(sdpa_fd_num_splits(129u, true), 2u); + EXPECT_EQ(sdpa_fd_num_splits(7937u, true), 63u); + EXPECT_EQ(sdpa_fd_num_splits(8192u, true), 64u); + EXPECT_EQ(sdpa_fd_num_splits(8193u, true), 65u); + EXPECT_EQ(sdpa_fd_num_splits(8960u, true), 70u); + EXPECT_EQ(sdpa_fd_num_splits(16385u, true), 128u); +} + +TEST(WebGPUUtils, GenericFdFallbackKeepsTile64Schedule) { + EXPECT_EQ(sdpa_fd_split_tile(false), 64u); + EXPECT_EQ(sdpa_fd_num_splits(128u, false), 2u); + EXPECT_EQ(sdpa_fd_num_splits(8192u, false), 128u); + EXPECT_EQ(sdpa_fd_num_splits(8960u, false), 128u); +} + +TEST(WebGPUUtils, QwenGqa2FdUsesOneSplitWorkgroupPerKvHead) { + EXPECT_EQ(sdpa_fd_split_head_count(16u, 8u, true), 8u); + EXPECT_EQ(sdpa_fd_split_head_count(16u, 8u, false), 16u); +} + +TEST(WebGPUUtils, QwenGqa2FdRouteReentersAfterLongContext) { + const uint32_t splits[] = { + sdpa_fd_num_splits(128u, true), + sdpa_fd_num_splits(8192u, true), + sdpa_fd_num_splits(128u, true), + }; + EXPECT_EQ(splits[0], 1u); + EXPECT_EQ(splits[1], 64u); + EXPECT_EQ(splits[2], 1u); +} + +TEST(WebGPUUtils, QwenGqa2FdShaderLoadsEachVElementOnceForTheHeadPair) { + const std::string_view source(kSdpaFdSplitGqa2F16WGSL); + const std::string_view load = "t_v_cache["; + const size_t first = source.find(load); + ASSERT_NE(first, std::string_view::npos); + EXPECT_EQ(source.find(load, first + 1), std::string_view::npos); + EXPECT_EQ(source.find("atomic"), std::string_view::npos); +} diff --git a/backends/webgpu/test/ops/test_sdpa.py b/backends/webgpu/test/ops/test_sdpa.py index b2606256e28..3a7cee6e0e7 100644 --- a/backends/webgpu/test/ops/test_sdpa.py +++ b/backends/webgpu/test/ops/test_sdpa.py @@ -69,6 +69,36 @@ class SdpaConfig: # execution exercises the real fp32->fp16->fp32 cache conversion path. SdpaConfig("qwen3_prefill", 16, 8, 128, 128, 256, 0, 10.0, kv_f16=True), SdpaConfig("qwen3_odd_boundary", 16, 8, 128, 17, 64, 31, 10.0, kv_f16=True), + # Exact f16 GQA2 FlashDecoding coverage. These contexts exercise every + # reducer boundary required by the tile-128 route plus the production + # P8192/P8960 shapes. The final case keeps the same Qwen geometry but + # disables fp16 KV storage, proving the generic FlashDecoding fallback. + SdpaConfig("qwen3_fd_splits_1", 16, 8, 128, 1, 1, 0, 10.0, kv_f16=True), + SdpaConfig("qwen3_fd_splits_2", 16, 8, 128, 1, 129, 128, 10.0, kv_f16=True), + SdpaConfig( + "qwen3_fd_splits_63", 16, 8, 128, 1, 7937, 7936, 10.0, kv_f16=True + ), + SdpaConfig( + "qwen3_fd_splits_64", 16, 8, 128, 1, 8192, 8191, 10.0, kv_f16=True + ), + SdpaConfig( + "qwen3_fd_splits_65", 16, 8, 128, 1, 8193, 8192, 10.0, kv_f16=True + ), + SdpaConfig( + "qwen3_fd_splits_70", 16, 8, 128, 1, 8960, 8959, 10.0, kv_f16=True + ), + SdpaConfig( + "qwen3_fd_splits_128", + 16, + 8, + 128, + 1, + 16385, + 16384, + 10.0, + kv_f16=True, + ), + SdpaConfig("qwen3_fd_fallback_fp32", 16, 8, 128, 1, 8192, 8191, 10.0), ] @@ -264,6 +294,39 @@ def test_qwen3_fixture_contract(self) -> None: ) self.assertTrue(cfg.kv_f16) + expected_fd_contexts = { + "qwen3_fd_splits_1": 1, + "qwen3_fd_splits_2": 129, + "qwen3_fd_splits_63": 7937, + "qwen3_fd_splits_64": 8192, + "qwen3_fd_splits_65": 8193, + "qwen3_fd_splits_70": 8960, + "qwen3_fd_splits_128": 16385, + } + for name, context in expected_fd_contexts.items(): + with self.subTest(config=name): + self.assertIn(name, configs) + cfg = configs[name] + self.assertEqual( + (cfg.hq, cfg.hkv, cfg.d, cfg.s, cfg.cmax, cfg.input_pos), + (16, 8, 128, 1, context, context - 1), + ) + self.assertTrue(cfg.kv_f16) + + fallback = configs["qwen3_fd_fallback_fp32"] + self.assertEqual( + ( + fallback.hq, + fallback.hkv, + fallback.d, + fallback.s, + fallback.cmax, + fallback.input_pos, + ), + (16, 8, 128, 1, 8192, 8191), + ) + self.assertFalse(fallback.kv_f16) + replays = {seq.name: seq for seq in REPLAY_SEQS} self.assertIn("qwen3_fd", replays) qwen3_fd = replays["qwen3_fd"] @@ -306,61 +369,6 @@ def test_sdpa_export_delegates(self) -> None: found, f"Expected VulkanBackend delegate in {cfg.name}.pte" ) - def test_golden_matches_eager_op(self) -> None: - # Oracle self-validation (mirrors Vulkan test_reference_sdpa): the fp64 - # golden and the shipped fp32 CPU op are independent refs that must agree. - for cfg in CONFIGS: - with self.subTest(config=cfg.name): - q, k, v, kc, vc = _det_inputs(cfg) - eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( - cfg, k, v, kc, vc - ) - eager = SdpaModule(cfg.input_pos)( - q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() - ) - golden = _golden(cfg, q, k, v, kc, vc) - torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) - - def test_replay_golden_matches_eager(self) -> None: - # Pure-torch proof of the threading model BEFORE any GPU run: replay the - # eager llama op with a host-threaded cache and assert each step's output - # equals the accumulated-context golden. Covers the large-S-at-offset mask - # path (small step (5,5), llama3 step (57,114)) absent from CONFIGS. - for seq in REPLAY_SEQS: - with self.subTest(seq=seq.name): - self.assertEqual( - (seq.d, seq.hq, seq.hkv), - VULKAN_PARAMS[seq.name], - f"{seq.name}: (d,hq,hkv) diverges from the Vulkan param set", - ) - self.assertLessEqual(sum(seq.seq_lens), seq.cmax) - kc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) - vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) - input_pos = 0 - for t, s in enumerate(seq.seq_lens): - cfg = SdpaConfig( - f"{seq.name}_step{t}", - seq.hq, - seq.hkv, - seq.d, - s, - seq.cmax, - input_pos, - kv_f16=seq.kv_f16, - ) - q, k, v = _step_inputs(seq, t, s) - golden = _golden(cfg, q, k, v, kc, vc) - eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( - cfg, k, v, kc, vc - ) - eager = SdpaModule(input_pos)( - q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() - ) - torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) - kc[0, input_pos : input_pos + s] = k[0] - vc[0, input_pos : input_pos + s] = v[0] - input_pos += s - def test_replay_export_delegates(self) -> None: # Every step .pte (incl. llama3-scale) must delegate to VulkanBackend. for seq in REPLAY_SEQS: diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index 5b647e1ae7c..d2326b5bfee 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -1316,6 +1316,92 @@ const SdpaConfig kSdpaConfigs[] = { /*required=*/true, /*expect_reject=*/false, /*kv_f16=*/true}, + {"qwen3_fd_splits_1", + 16, + 8, + 128, + 1, + 1, + 0, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_2", + 16, + 8, + 128, + 1, + 129, + 128, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_63", + 16, + 8, + 128, + 1, + 7937, + 7936, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_64", + 16, + 8, + 128, + 1, + 8192, + 8191, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_65", + 16, + 8, + 128, + 1, + 8193, + 8192, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_70", + 16, + 8, + 128, + 1, + 8960, + 8959, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_splits_128", + 16, + 8, + 128, + 1, + 16385, + 16384, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_fd_fallback_fp32", + 16, + 8, + 128, + 1, + 8192, + 8191, + 10.0f, + /*required=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). @@ -2036,26 +2122,31 @@ void test_symint_input_narrowing() { ::flatbuffers::FlatBufferBuilder fbb; const std::vector dims = {1u}; std::vector<::flatbuffers::Offset> values; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::INT32, - &dims, - /*constant_id=*/-1, - /*mem_obj_id=*/0) - .Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::INT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); const std::vector args = {0, 1, 2, 3}; std::vector<::flatbuffers::Offset> chain; - chain.push_back(vk::CreateOperatorCallDirect( - fbb, 0, "et_vk.select_as_symint.default", &args)); + chain.push_back( + vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.select_as_symint.default", &args)); const std::vector input_ids = {0}; const std::vector output_ids = {0}; const auto root = vk::CreateVkGraphDirect( @@ -2108,23 +2199,25 @@ void finish_inline_constant_graph( for (int i = 0; i < tensor_count; i++) { const bool is_cache = mark_as_kv_cache && i >= 3; const bool is_constant = !mark_as_kv_cache || is_cache; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &dims, - is_constant ? (is_cache ? i - 3 : 0) : -1, - is_constant ? -1 : i) - .Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + is_constant ? (is_cache ? i - 3 : 0) : -1, + is_constant ? -1 : i) + .Union())); } std::vector<::flatbuffers::Offset> chain; if (mark_as_kv_cache) { const std::vector args = {0, 1, 2, 3, 4}; - chain.push_back(vk::CreateOperatorCallDirect( - fbb, 0, "sdpa_with_kv_cache.default", &args)); + chain.push_back( + vk::CreateOperatorCallDirect( + fbb, 0, "sdpa_with_kv_cache.default", &args)); } std::vector<::flatbuffers::Offset> constants; constants.push_back( @@ -2605,39 +2698,45 @@ static bool test_slice_double_start_case(double start_d, int out_len) { std::vector out_dims = {1u, static_cast(out_len)}; std::vector<::flatbuffers::Offset> values; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &in_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/0) - .Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); // The value under test: `start` serialized as a Double, not an Int. - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Double, - vk::CreateDouble(fbb, /*double_val=*/start_d).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Int, - vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/start_d).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::Int, + vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &out_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/1) - .Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); std::vector args = {0, 1, 2, 3, 4, 5}; std::vector<::flatbuffers::Offset> chain; @@ -2705,38 +2804,44 @@ static bool test_slice_double_start_rejects(double bad_start) { std::vector out_dims = {1u, static_cast(kInLen)}; std::vector<::flatbuffers::Offset> values; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &in_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/0) - .Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Double, - vk::CreateDouble(fbb, /*double_val=*/bad_start).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Int, - vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &out_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/1) - .Union())); + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/bad_start).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::Int, + vk::CreateInt(fbb, /*int_val=*/kInLen).Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, /*int_val=*/1).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); std::vector args = {0, 1, 2, 3, 4, 5}; std::vector<::flatbuffers::Offset> chain; @@ -2789,41 +2894,46 @@ static void finish_select_scalar_graph( std::vector out_dims = {out_len}; std::vector<::flatbuffers::Offset> values; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &in_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/0) - .Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); if (symint_dim) { - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::SymInt, - vk::CreateSymInt(fbb, /*value=*/0).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::SymInt, + vk::CreateSymInt(fbb, /*value=*/0).Union())); } else { - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Double, - vk::CreateDouble(fbb, /*double_val=*/dim).Union())); - } - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Double, - vk::CreateDouble(fbb, /*double_val=*/index).Union())); - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/dim).Union())); + } + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &out_dims, - /*constant_id=*/-1, - /*mem_obj_id=*/1) - .Union())); + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/index).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); std::vector args = {0, 1, 2, 3}; std::vector<::flatbuffers::Offset> chain; @@ -2969,35 +3079,39 @@ void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { const std::vector freqs_dims = {2u, 2u}; std::vector<::flatbuffers::Offset> values; const auto add_tensor = [&](const std::vector& dims, int mem_id) { - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::VkTensor, - vk::CreateVkTensorDirect( + values.push_back( + vk::CreateVkValue( fbb, - vk::VkDataType::FLOAT32, - &dims, - /*constant_id=*/-1, - /*mem_obj_id=*/mem_id) - .Union())); + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/mem_id) + .Union())); }; add_tensor(q_dims, 0); add_tensor(k_dims, 1); add_tensor(freqs_dims, 2); add_tensor(freqs_dims, 3); - values.push_back(vk::CreateVkValue( - fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back( + vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); add_tensor(q_dims, 4); add_tensor(k_dims, 5); const std::vector output_items = {5, 6}; - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::ValueList, - vk::CreateValueListDirect(fbb, &output_items).Union())); + values.push_back( + vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_items).Union())); const std::vector args = {0, 1, 2, 3, 4, 7}; std::vector<::flatbuffers::Offset> chain; - chain.push_back(vk::CreateOperatorCallDirect( - fbb, 0, "et_vk.apply_rotary_emb_hf.default", &args)); + chain.push_back( + vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb_hf.default", &args)); const std::vector input_ids = {0, 1, 2, 3}; const std::vector output_ids = {5, 6}; const auto root = vk::CreateVkGraphDirect( @@ -3319,6 +3433,42 @@ TEST(WebGPUNative, Qwen3SdpaFixtureContract) { std::vector({16, 8, 128, 17, 64, 31})); EXPECT_TRUE(prefill->kv_f16 && boundary->kv_f16); + struct ExpectedFdContext { + const char* name; + int context; + }; + constexpr ExpectedFdContext kExpectedFdContexts[] = { + {"qwen3_fd_splits_1", 1}, + {"qwen3_fd_splits_2", 129}, + {"qwen3_fd_splits_63", 7937}, + {"qwen3_fd_splits_64", 8192}, + {"qwen3_fd_splits_65", 8193}, + {"qwen3_fd_splits_70", 8960}, + {"qwen3_fd_splits_128", 16385}, + }; + for (const auto& expected : kExpectedFdContexts) { + const auto cfg = find_config(expected.name); + ASSERT_NE(cfg, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {cfg->hq, cfg->hkv, cfg->d, cfg->s, cfg->cmax, cfg->input_pos}), + std::vector( + {16, 8, 128, 1, expected.context, expected.context - 1})); + EXPECT_TRUE(cfg->kv_f16); + } + const auto fallback = find_config("qwen3_fd_fallback_fp32"); + ASSERT_NE(fallback, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {fallback->hq, + fallback->hkv, + fallback->d, + fallback->s, + fallback->cmax, + fallback->input_pos}), + std::vector({16, 8, 128, 1, 8192, 8191})); + EXPECT_FALSE(fallback->kv_f16); + const auto replay = std::find_if( std::begin(kSdpaSequences), std::end(kSdpaSequences), @@ -3363,7 +3513,7 @@ TEST(WebGPUNative, Qwen3SdpaRoutes) { // the partial final workgroup's row mask is covered. Unsupported Q32 devices // intentionally fall back to the already-qualified Q16 route. for (const auto& cfg : kSdpaConfigs) { - if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0 || cfg.s == 1) { continue; } const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 98812c0f144..92a6ca57c14 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -240,14 +240,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 138) + self.assertEqual(len(outputs), 140) self.assertEqual( digest.hexdigest(), - "fee848cd069b4c09d3d2e9a7920331f46d5646b74bec30259542dde8f287e504", + "3ded33f0d5abdb62fe2c631c3248bb63e396712d014b834fb04fd19120c2917a", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "477721998b3cd8f3f0fdd485fa797c71035a20cbc10a8b4bf44893e37fa435b8", + "53c9486a181384305c4d024b7c10e8373ebe16d95efe1c84d65dd3dd493f04c4", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: diff --git a/examples/models/qwen3/README.md b/examples/models/qwen3/README.md index 123e65f16c5..c8c55c35f5f 100644 --- a/examples/models/qwen3/README.md +++ b/examples/models/qwen3/README.md @@ -68,5 +68,33 @@ Note that you have to apply the chat template manually for the C++ runner. To run the model on an example iOS or Android app, see the Llama README's [Step 5: Build Mobile apps](../llama/README.md#step-5-build-mobile-apps) section. +### WebGPU export (Qwen3-0.6B) + +`config/qwen3_webgpu_q4gsw.yaml` exports Qwen3-0.6B with group-symmetric 4-bit weights at `max_seq_length` 512 and `max_context_length` 8960, for the WebGPU delegate. + +**Runtime selection is declared, not inferred.** `LlmConfig` exposes no WebGPU backend field, so the config sets `backend.vulkan.enabled` — that is the *serialization* mechanism that produces the program the WebGPU delegate consumes (the delegate registers under the Vulkan backend id). Vulkan serialization on its own is not a WebGPU selection. The export contract `manifests/qwen3_0_6b_webgpu.json` therefore records `target_runtime: webgpu` alongside `serialization_backend: vulkan`, and `webgpu_artifact_manifest.py` rejects any artifact set whose declared runtime target is missing or different. + +Export: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml +``` + +The contract pins the checkpoint and tokenizer to an exact Hugging Face commit, with the published SHA-256 and byte count for each. Build a manifest over an output directory and validate it: +``` +python -m executorch.examples.models.qwen3.webgpu_artifact_manifest create \ + --root \ + --output /manifest.json \ + --role pte=qwen3_0_6b_webgpu_q4gsw.pte \ + --role javascript=runner.js \ + --role wasm=runner.wasm + +python -m executorch.examples.models.qwen3.webgpu_artifact_manifest validate \ + --root \ + --manifest /manifest.json +``` + +Validation fails closed on a missing, extra, symlinked, wrong-size or wrong-hash artifact; on a role whose file extension contradicts it; on a method set other than `forward`; on a graph carrying portable operators or a delegate other than the WebGPU one; and on any acquisition pin that disagrees with the checked contract. + ### FAQ For more help with exporting or running this model, feel free to ask in our [discord channel](https://discord.gg/UEjkY9Zs). diff --git a/examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml b/examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml new file mode 100644 index 00000000000..a217962388d --- /dev/null +++ b/examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml @@ -0,0 +1,32 @@ +# Qwen3-0.6B group-symmetric 4-bit weight export for the WebGPU delegate. +# +# The WebGPU delegate consumes the Vulkan-partitioner program and registers +# under the Vulkan backend id, so `backend.vulkan` is how a WebGPU-consumable +# program is selected today; `LlmConfig.BackendConfig` exposes no `webgpu` +# field. `webgpu_artifact_manifest.py` enforces the WebGPU contract on the +# resulting artifact (single VulkanBackend delegate, no portable operators). + +base: + model_class: "qwen3_0_6b" + params: "examples/models/qwen3/config/0_6b_config.json" + metadata: '{"get_bos_id": 151644, "get_eos_ids":[151645]}' + +model: + use_kv_cache: true + use_sdpa_with_kv_cache: true + enable_dynamic_shape: true + dtype_override: fp32 + +quantization: + qmode: "4w" + group_size: 64 + embedding_quantize: "4,64" + +export: + output_name: qwen3_0_6b_webgpu_q4gsw.pte + max_seq_length: 512 + max_context_length: 8960 + +backend: + vulkan: + enabled: true diff --git a/examples/models/qwen3/manifests/qwen3_0_6b_webgpu.json b/examples/models/qwen3/manifests/qwen3_0_6b_webgpu.json new file mode 100644 index 00000000000..1f4da3d64c2 --- /dev/null +++ b/examples/models/qwen3/manifests/qwen3_0_6b_webgpu.json @@ -0,0 +1,35 @@ +{ + "acquisition": { + "checkpoint": { + "bytes": 1503300328, + "filename": "model.safetensors", + "repo": "Qwen/Qwen3-0.6B", + "revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "sha256": "f47f71177f32bcd101b7573ec9171e6a57f4f4d31148d38e382306f42996874b" + }, + "tokenizer": { + "bytes": 11422654, + "filename": "tokenizer.json", + "repo": "Qwen/Qwen3-0.6B", + "revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "sha256": "aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4" + } + }, + "export": { + "config": "examples/models/qwen3/config/qwen3_webgpu_q4gsw.yaml", + "max_context_len": 8960, + "max_input_len": 512, + "serialization_backend": "vulkan", + "target_runtime": "webgpu" + }, + "methods": [ + "forward" + ], + "model": "qwen3_0_6b", + "role_suffixes": { + "javascript": ".js", + "pte": ".pte", + "wasm": ".wasm" + }, + "schema_version": 1 +} diff --git a/examples/models/qwen3/targets.bzl b/examples/models/qwen3/targets.bzl index 074fe77d101..5ecbdfb91ce 100644 --- a/examples/models/qwen3/targets.bzl +++ b/examples/models/qwen3/targets.bzl @@ -17,6 +17,7 @@ def define_common_targets(is_fbcode = False): "config/0_6b_config.json": "config/0_6b_config.json", "config/1_7b_config.json": "config/1_7b_config.json", "config/4b_config.json": "config/4b_config.json", + "config/qwen3_webgpu_q4gsw.yaml": "config/qwen3_webgpu_q4gsw.yaml", "config/qwen3_xnnpack_q8da4w.yaml": "config/qwen3_xnnpack_q8da4w.yaml", }, deps = [ @@ -27,3 +28,18 @@ def define_common_targets(is_fbcode = False): ], visibility = ["PUBLIC"], ) + + runtime.python_library( + name = "webgpu_artifact_manifest", + srcs = ["webgpu_artifact_manifest.py"], + _is_external_target = True, + base_module = "executorch.examples.models.qwen3", + resources = { + "manifests/qwen3_0_6b_webgpu.json": "manifests/qwen3_0_6b_webgpu.json", + }, + typing = True, + deps = [ + "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", + ], + visibility = ["PUBLIC"], + ) diff --git a/examples/models/qwen3/tests/BUCK b/examples/models/qwen3/tests/BUCK new file mode 100644 index 00000000000..9123059dfb8 --- /dev/null +++ b/examples/models/qwen3/tests/BUCK @@ -0,0 +1,15 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. + +load("@fbsource//tools/build_defs:fbsource_utils.bzl", "is_fbcode") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +define_common_targets(is_fbcode = is_fbcode()) diff --git a/examples/models/qwen3/tests/targets.bzl b/examples/models/qwen3/tests/targets.bzl new file mode 100644 index 00000000000..b0ef4395ab1 --- /dev/null +++ b/examples/models/qwen3/tests/targets.bzl @@ -0,0 +1,22 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(is_fbcode = False): + if not is_fbcode: + return + + fbcode_target( + _kind = runtime.python_test, + name = "test_webgpu_export_manifest", + srcs = ["test_webgpu_export_manifest.py"], + typing = True, + deps = [ + "//executorch/examples/models/qwen3:webgpu_artifact_manifest", + ], + ) diff --git a/examples/models/qwen3/tests/test_webgpu_export_manifest.py b/examples/models/qwen3/tests/test_webgpu_export_manifest.py new file mode 100644 index 00000000000..4e2c72f65df --- /dev/null +++ b/examples/models/qwen3/tests/test_webgpu_export_manifest.py @@ -0,0 +1,352 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import json +import tempfile +import unittest +from pathlib import Path + +from executorch.examples.models.qwen3.webgpu_artifact_manifest import ( + check_contract_agreement, + check_delegation, + check_export_config, + check_methods, + check_no_extra_files, + check_role_suffixes, + create_qwen3_manifest, + expected_export_config, + load_contract, + ManifestError, + MAX_CONTEXT_LEN, + MAX_INPUT_LEN, + REQUIRED_METHODS, + SERIALIZATION_BACKEND, + TARGET_RUNTIME, + validate_qwen3_manifest, + WEBGPU_BACKEND_ID, +) + +_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca" + +_ACQUISITION = { + "checkpoint": { + "repo": "Qwen/Qwen3-0.6B", + "revision": _REVISION, + "filename": "model.safetensors", + "sha256": "f47f71177f32bcd101b7573ec9171e6a57f4f4d31148d38e382306f42996874b", + "bytes": 1503300328, + }, + "tokenizer": { + "repo": "Qwen/Qwen3-0.6B", + "revision": _REVISION, + "filename": "tokenizer.json", + "sha256": "aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4", + "bytes": 11422654, + }, +} + +_ROLES = { + "pte": Path("qwen3_0_6b.pte"), + "javascript": Path("runner.js"), + "wasm": Path("runner.wasm"), +} + + +def _acquisition(**overrides) -> dict: + result = copy.deepcopy(_ACQUISITION) + for dotted, value in overrides.items(): + kind, _, field = dotted.partition("__") + result[kind][field] = value + return result + + +def _populate(root: Path) -> None: + for name, data in { + "qwen3_0_6b.pte": b"pte-bytes", + "runner.js": b"js-bytes", + "runner.wasm": b"wasm-bytes", + }.items(): + (root / name).write_bytes(data) + + +def _build(root: Path, acquisition=None) -> dict: + return create_qwen3_manifest(root, dict(_ROLES), (), acquisition or _ACQUISITION) + + +class ExportContractTest(unittest.TestCase): + def _config(self, **overrides) -> dict: + config = expected_export_config() + config.update(overrides) + return config + + def test_accepts_the_declared_webgpu_runtime_target(self) -> None: + check_export_config(self._config()) + self.assertEqual(expected_export_config()["target_runtime"], TARGET_RUNTIME) + + def test_rejects_a_wrong_runtime_target(self) -> None: + for wrong in ("vulkan", "xnnpack", "", None): + with self.subTest(target_runtime=wrong): + with self.assertRaisesRegex(ManifestError, "target_runtime"): + check_export_config(self._config(target_runtime=wrong)) + + def test_rejects_a_wrong_serialization_backend(self) -> None: + with self.assertRaisesRegex(ManifestError, "serialization_backend"): + check_export_config(self._config(serialization_backend="xnnpack")) + + def test_vulkan_serialization_alone_is_not_a_webgpu_claim(self) -> None: + config = self._config() + del config["target_runtime"] + self.assertEqual(config["serialization_backend"], SERIALIZATION_BACKEND) + with self.assertRaisesRegex(ManifestError, "target_runtime"): + check_export_config(config) + + def test_rejects_a_wrong_capacity(self) -> None: + for field, value in ( + ("max_input_len", MAX_INPUT_LEN + 1), + ("max_context_len", MAX_CONTEXT_LEN - 1), + ): + with self.subTest(field=field): + with self.assertRaisesRegex(ManifestError, field): + check_export_config(self._config(**{field: value})) + + +class CheckedContractTest(unittest.TestCase): + def test_the_checked_contract_is_self_consistent(self) -> None: + contract = load_contract() + self.assertEqual(contract["model"], "qwen3_0_6b") + self.assertEqual(contract["export"]["target_runtime"], TARGET_RUNTIME) + self.assertEqual( + contract["export"]["serialization_backend"], SERIALIZATION_BACKEND + ) + self.assertEqual(contract["export"]["max_input_len"], MAX_INPUT_LEN) + self.assertEqual(contract["export"]["max_context_len"], MAX_CONTEXT_LEN) + self.assertEqual(contract["acquisition"]["checkpoint"]["revision"], _REVISION) + + def test_a_built_manifest_must_agree_with_the_contract(self) -> None: + contract = load_contract() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root, contract["acquisition"]) + validate_qwen3_manifest(root, manifest, contract["acquisition"], contract) + + def test_contract_agreement_rejects_a_swapped_revision(self) -> None: + contract = load_contract() + manifest = { + "model": "qwen3_0_6b", + "acquisition": _acquisition(checkpoint__revision="0" * 40), + "export": expected_export_config(), + } + with self.assertRaisesRegex(ManifestError, "revision"): + check_contract_agreement(manifest, contract) + + def test_contract_agreement_rejects_a_downgraded_runtime_target(self) -> None: + contract = load_contract() + export = expected_export_config() + export["target_runtime"] = "vulkan" + manifest = { + "model": "qwen3_0_6b", + "acquisition": copy.deepcopy(contract["acquisition"]), + "export": export, + } + with self.assertRaisesRegex(ManifestError, "target_runtime"): + check_contract_agreement(manifest, contract) + + +class DelegationTest(unittest.TestCase): + def test_accepts_a_fully_delegated_webgpu_graph(self) -> None: + check_delegation([WEBGPU_BACKEND_ID], []) + + def test_rejects_a_portable_operator(self) -> None: + for ops in (["aten::add.out"], ["aten::mm.out", "aten::view_copy.out"]): + with self.subTest(ops=ops): + with self.assertRaisesRegex(ManifestError, "portable"): + check_delegation([WEBGPU_BACKEND_ID], ops) + + def test_rejects_a_foreign_or_mixed_delegate_census(self) -> None: + for ids in (["XnnpackBackend"], [WEBGPU_BACKEND_ID, "XnnpackBackend"], []): + with self.subTest(ids=ids): + with self.assertRaisesRegex(ManifestError, "backend"): + check_delegation(ids, []) + + def test_method_census_must_match_exactly(self) -> None: + check_methods(list(REQUIRED_METHODS)) + for observed in ([*REQUIRED_METHODS, "extra"], [], ["not_forward"]): + with self.subTest(observed=observed): + with self.assertRaisesRegex(ManifestError, "method"): + check_methods(observed) + + +class AcquisitionPinTest(unittest.TestCase): + def _expect(self, pattern, **overrides) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + with self.assertRaisesRegex(ManifestError, pattern): + _build(root, _acquisition(**overrides)) + + def test_rejects_a_wrong_checkpoint_digest_shape(self) -> None: + self._expect("sha256", checkpoint__sha256="not-a-digest") + self._expect("sha256", checkpoint__sha256="F" * 64) + + def test_rejects_a_wrong_checkpoint_size(self) -> None: + for bad in (0, -1, "1503300328", None): + with self.subTest(bytes=bad): + self._expect("bytes", checkpoint__bytes=bad) + + def test_rejects_a_non_commit_revision(self) -> None: + for bad in ("main", "v1.0", "c1899de", ""): + with self.subTest(revision=bad): + self._expect("revision", checkpoint__revision=bad) + + def test_rejects_a_missing_acquisition_pin(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + for kind in ("checkpoint", "tokenizer"): + acquisition = copy.deepcopy(_ACQUISITION) + del acquisition[kind] + with self.subTest(kind=kind): + with self.assertRaisesRegex(ManifestError, kind): + _build(root, acquisition) + + def test_rejects_a_tampered_pin_at_validation(self) -> None: + for field, value in ( + ("sha256", "c" * 64), + ("bytes", 1), + ("revision", "d" * 40), + ): + with self.subTest(field=field): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + manifest["acquisition"]["checkpoint"][field] = value + with self.assertRaisesRegex(ManifestError, field): + validate_qwen3_manifest(root, manifest, _ACQUISITION) + + +class ArtifactRoleTest(unittest.TestCase): + def test_rejects_role_confusion_between_pte_and_wasm(self) -> None: + artifacts = [ + {"role": "pte", "path": "runner.wasm"}, + {"role": "wasm", "path": "runner.wasm"}, + ] + with self.assertRaisesRegex(ManifestError, "role confusion"): + check_role_suffixes(artifacts) + + def test_rejects_javascript_role_pointing_at_the_pte(self) -> None: + with self.assertRaisesRegex(ManifestError, "role confusion"): + check_role_suffixes([{"role": "javascript", "path": "qwen3_0_6b.pte"}]) + + def test_accepts_correctly_typed_roles(self) -> None: + check_role_suffixes( + [ + {"role": "pte", "path": "qwen3_0_6b.pte"}, + {"role": "javascript", "path": "runner.js"}, + {"role": "wasm", "path": "runner.wasm"}, + {"role": "ptd", "path": "weights.ptd"}, + ] + ) + + def test_create_rejects_a_role_confused_artifact_set(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + roles = dict(_ROLES) + roles["pte"] = Path("runner.wasm") + with self.assertRaisesRegex(ManifestError, "role confusion"): + create_qwen3_manifest(root, roles, (), _ACQUISITION) + + +class ManifestRoundTripTest(unittest.TestCase): + def test_round_trip_validates_and_records_no_local_paths(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + validate_qwen3_manifest(root, manifest) + self.assertEqual(manifest["model"], "qwen3_0_6b") + self.assertEqual(manifest["ptd_order"], []) + self.assertNotIn(str(root), json.dumps(manifest)) + + def test_rejects_a_missing_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + (root / "runner.wasm").unlink() + with self.assertRaises(ManifestError): + validate_qwen3_manifest(root, manifest) + + def test_rejects_a_missing_required_role(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + manifest["artifacts"] = [ + a for a in manifest["artifacts"] if a["role"] != "wasm" + ] + (root / "runner.wasm").unlink() + with self.assertRaisesRegex(ManifestError, "wasm"): + validate_qwen3_manifest(root, manifest) + + def test_rejects_an_extra_untracked_file(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + (root / "stowaway.bin").write_bytes(b"extra") + with self.assertRaisesRegex(ManifestError, "stowaway.bin"): + validate_qwen3_manifest(root, manifest) + + def test_rejects_a_wrong_hash(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + (root / "qwen3_0_6b.pte").write_bytes(b"tampered!") + with self.assertRaises(ManifestError): + validate_qwen3_manifest(root, manifest) + + def test_rejects_a_wrong_recorded_size(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + manifest = _build(root) + for artifact in manifest["artifacts"]: + if artifact["role"] == "pte": + artifact["bytes"] += 1 + with self.assertRaises(ManifestError): + validate_qwen3_manifest(root, manifest) + + def test_rejects_a_symlinked_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + _populate(root) + (root / "aliased.pte").symlink_to(root / "qwen3_0_6b.pte") + roles = dict(_ROLES) + roles["pte"] = Path("aliased.pte") + with self.assertRaises(ManifestError): + create_qwen3_manifest(root, roles, (), _ACQUISITION) + + +class ExtraFileScanTest(unittest.TestCase): + def test_reports_the_offending_relative_path(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "kept.bin").write_bytes(b"kept") + nested = root / "nested" + nested.mkdir() + (nested / "sneaky.bin").write_bytes(b"sneaky") + check_no_extra_files(root, {"kept.bin", "nested/sneaky.bin"}) + with self.assertRaisesRegex(ManifestError, "nested/sneaky.bin"): + check_no_extra_files(root, {"kept.bin"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/models/qwen3/webgpu_artifact_manifest.py b/examples/models/qwen3/webgpu_artifact_manifest.py new file mode 100644 index 00000000000..10a3b267b14 --- /dev/null +++ b/examples/models/qwen3/webgpu_artifact_manifest.py @@ -0,0 +1,309 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Qwen3-0.6B WebGPU acquisition/artifact manifest over the generic validator. + +The generic byte-role layer (hashing, containment, symlink rejection, ordered +PTD bookkeeping) lives in +``executorch.backends.webgpu.scripts.webgpu_artifact_manifest`` and is reused +here rather than reimplemented. This module adds only what is Qwen3-specific: +pinned checkpoint/tokenizer acquisition, the required artifact roles and method +set, the export capacities, and a fail-closed scan for undeclared files. + +Runtime selection is declared explicitly. ``LlmConfig`` exposes no WebGPU +backend field, so ``backend.vulkan`` remains the *serialization* mechanism that +produces the program; it is not by itself a WebGPU selection. The export +contract in ``manifests/qwen3_0_6b_webgpu.json`` therefore records +``target_runtime: webgpu`` alongside ``serialization_backend: vulkan``, and +both are validated here. +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +from executorch.backends.webgpu.scripts.webgpu_artifact_manifest import ( + create_manifest, + validate_manifest, +) + +MODEL_ID = "qwen3_0_6b" +TARGET_RUNTIME = "webgpu" +# LlmConfig has no WebGPU backend field; Vulkan serialization is the mechanism +# that produces a WebGPU-consumable program, not the runtime claim itself. +SERIALIZATION_BACKEND = "vulkan" +# The WebGPU delegate consumes the Vulkan-partitioner program and registers +# under the Vulkan backend id; a graph carrying any other id is not WebGPU. +WEBGPU_BACKEND_ID = "VulkanBackend" +REQUIRED_METHODS = ("forward",) +REQUIRED_ROLES = ("javascript", "pte", "wasm") +ROLE_SUFFIXES = {"javascript": ".js", "pte": ".pte", "wasm": ".wasm"} +ACQUISITION_KINDS = ("checkpoint", "tokenizer") +MAX_INPUT_LEN = 512 +MAX_CONTEXT_LEN = 8960 + +CONTRACT_PATH = Path(__file__).parent / "manifests" / "qwen3_0_6b_webgpu.json" + +_SHA256 = re.compile(r"\A[0-9a-f]{64}\Z") +_GIT_SHA = re.compile(r"\A[0-9a-f]{40}\Z") + + +class ManifestError(ValueError): + """Raised when a Qwen3 WebGPU manifest or its artifacts fail validation.""" + + +def _wrap(action: str, fn, *args): + try: + return fn(*args) + except ManifestError: + raise + except (OSError, ValueError) as error: + raise ManifestError(f"{action}: {error}") from error + + +def load_contract(path: Path | None = None) -> dict[str, Any]: + contract = json.loads((path or CONTRACT_PATH).read_text(encoding="utf-8")) + _check_acquisition(contract["acquisition"]) + check_export_config(contract["export"]) + check_methods(contract["methods"]) + return contract + + +def expected_export_config() -> dict[str, Any]: + return { + "target_runtime": TARGET_RUNTIME, + "serialization_backend": SERIALIZATION_BACKEND, + "max_input_len": MAX_INPUT_LEN, + "max_context_len": MAX_CONTEXT_LEN, + } + + +def check_export_config(config: Mapping[str, Any]) -> None: + if config.get("target_runtime") != TARGET_RUNTIME: + raise ManifestError( + f"target_runtime must be {TARGET_RUNTIME}: {config.get('target_runtime')}" + ) + if config.get("serialization_backend") != SERIALIZATION_BACKEND: + raise ManifestError( + "serialization_backend must be " + f"{SERIALIZATION_BACKEND}: {config.get('serialization_backend')}" + ) + for field, expected in ( + ("max_input_len", MAX_INPUT_LEN), + ("max_context_len", MAX_CONTEXT_LEN), + ): + if config.get(field) != expected: + raise ManifestError(f"{field} must be {expected}: {config.get(field)}") + + +def check_methods(observed: Sequence[str]) -> None: + if sorted(observed) != sorted(REQUIRED_METHODS): + raise ManifestError( + f"method set must be {sorted(REQUIRED_METHODS)}: {sorted(observed)}" + ) + + +def check_delegation( + backend_ids: Sequence[str], + operator_names: Sequence[str], +) -> None: + if operator_names: + raise ManifestError( + f"graph retains portable operators: {sorted(set(operator_names))}" + ) + unique = sorted(set(backend_ids)) + if unique != [WEBGPU_BACKEND_ID]: + raise ManifestError( + f"graph backend ids must be [{WEBGPU_BACKEND_ID}]: {unique}" + ) + + +def check_role_suffixes(artifacts: Sequence[Mapping[str, Any]]) -> None: + for artifact in artifacts: + role = artifact.get("role") + path = str(artifact.get("path", "")) + expected = ROLE_SUFFIXES.get(str(role)) + if expected is not None and not path.endswith(expected): + raise ManifestError( + f"artifact role confusion: role {role} expects {expected}: {path}" + ) + + +def check_no_extra_files(root: Path, declared: Iterable[str]) -> None: + allowed = set(declared) + for path in sorted(root.rglob("*")): + if not path.is_file() and not path.is_symlink(): + continue + relative = path.relative_to(root).as_posix() + if relative not in allowed: + raise ManifestError(f"undeclared artifact under manifest root: {relative}") + + +def _check_acquisition(acquisition: Mapping[str, Any]) -> None: + for kind in ACQUISITION_KINDS: + entry = acquisition.get(kind) + if not isinstance(entry, Mapping): + raise ManifestError(f"missing {kind} acquisition pin") + for field in ("repo", "filename"): + if not isinstance(entry.get(field), str) or not entry[field]: + raise ManifestError(f"{kind} acquisition {field} must be a string") + revision = entry.get("revision") + if not isinstance(revision, str) or not _GIT_SHA.match(revision): + raise ManifestError( + f"{kind} acquisition revision must be a 40-hex commit: {revision}" + ) + digest = entry.get("sha256") + if not isinstance(digest, str) or not _SHA256.match(digest): + raise ManifestError( + f"{kind} acquisition sha256 must be 64 lowercase hex digits" + ) + size = entry.get("bytes") + if not isinstance(size, int) or isinstance(size, bool) or size <= 0: + raise ManifestError(f"{kind} acquisition bytes must be a positive integer") + unknown = sorted(set(acquisition) - set(ACQUISITION_KINDS)) + if unknown: + raise ManifestError(f"unsupported acquisition kinds: {unknown}") + + +def check_contract_agreement( + manifest: Mapping[str, Any], + contract: Mapping[str, Any], +) -> None: + if manifest.get("model") != contract.get("model"): + raise ManifestError("manifest model disagrees with the export contract") + for kind in ACQUISITION_KINDS: + recorded = manifest["acquisition"][kind] + pinned = contract["acquisition"][kind] + for field in ("repo", "revision", "filename", "sha256", "bytes"): + if recorded.get(field) != pinned.get(field): + raise ManifestError( + f"acquisition {field} mismatch for {kind}: " + f"{recorded.get(field)} != {pinned.get(field)}" + ) + for field in ("target_runtime", "serialization_backend"): + if manifest["export"].get(field) != contract["export"].get(field): + raise ManifestError(f"export {field} disagrees with the export contract") + + +def create_qwen3_manifest( + root: Path, + role_paths: Mapping[str, Path], + ptd_paths: Sequence[Path], + acquisition: Mapping[str, Any], +) -> dict[str, Any]: + _check_acquisition(acquisition) + missing = sorted(set(REQUIRED_ROLES) - set(role_paths)) + if missing: + raise ManifestError(f"missing required artifact roles: {missing}") + manifest: dict[str, Any] = _wrap( + "artifact", create_manifest, root, role_paths, ptd_paths + ) + check_role_suffixes(manifest["artifacts"]) + manifest["model"] = MODEL_ID + manifest["export"] = expected_export_config() + manifest["methods"] = list(REQUIRED_METHODS) + manifest["acquisition"] = { + kind: dict(acquisition[kind]) for kind in ACQUISITION_KINDS + } + return manifest + + +def validate_qwen3_manifest( + root: Path, + manifest: Mapping[str, Any], + expected_acquisition: Mapping[str, Any] | None = None, + contract: Mapping[str, Any] | None = None, +) -> None: + if manifest.get("model") != MODEL_ID: + raise ManifestError( + f"manifest model must be {MODEL_ID}: {manifest.get('model')}" + ) + export = manifest.get("export") + if not isinstance(export, Mapping): + raise ManifestError("manifest export config must be an object") + check_export_config(export) + methods = manifest.get("methods") + if not isinstance(methods, list): + raise ManifestError("manifest method list must be a list") + check_methods(methods) + + acquisition = manifest.get("acquisition") + if not isinstance(acquisition, Mapping): + raise ManifestError("manifest acquisition block must be an object") + _check_acquisition(acquisition) + if expected_acquisition is not None: + for kind in ACQUISITION_KINDS: + for field in ("sha256", "bytes", "revision"): + recorded = acquisition[kind].get(field) + pinned = expected_acquisition[kind].get(field) + if recorded != pinned: + raise ManifestError( + f"acquisition pin mismatch for {kind}.{field}: " + f"{recorded} != {pinned}" + ) + if contract is not None: + check_contract_agreement(manifest, contract) + + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise ManifestError("manifest artifacts must be a list") + observed_roles = {a.get("role") for a in artifacts if isinstance(a, Mapping)} + missing = sorted(set(REQUIRED_ROLES) - observed_roles) + if missing: + raise ManifestError(f"missing required artifact roles: {missing}") + check_role_suffixes(artifacts) + + _wrap("artifact", validate_manifest, root, manifest) + declared = {a["path"] for a in artifacts if isinstance(a, Mapping) and "path" in a} + check_no_extra_files(root, declared) + + +def _parse_roles(values: Iterable[str]) -> dict[str, Path]: + result: dict[str, Path] = {} + for value in values: + role, separator, path = value.partition("=") + if not separator or not role or not path or role in result: + raise ManifestError(f"invalid or duplicate ROLE=PATH: {value}") + result[role] = Path(path) + return result + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + create = commands.add_parser("create") + create.add_argument("--root", type=Path, required=True) + create.add_argument("--output", type=Path, required=True) + create.add_argument("--contract", type=Path, default=None) + create.add_argument("--role", action="append", default=[]) + create.add_argument("--ptd", action="append", type=Path, default=[]) + validate = commands.add_parser("validate") + validate.add_argument("--root", type=Path, required=True) + validate.add_argument("--manifest", type=Path, required=True) + validate.add_argument("--contract", type=Path, default=None) + args = parser.parse_args(argv) + + contract = load_contract(args.contract) + if args.command == "create": + manifest = create_qwen3_manifest( + args.root, _parse_roles(args.role), args.ptd, contract["acquisition"] + ) + args.output.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return 0 + + manifest = json.loads(args.manifest.read_text(encoding="utf-8")) + validate_qwen3_manifest(args.root, manifest, contract["acquisition"], contract) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())