From 3cce033b4bdc948de69ea8b3762373fed08b9e44 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:11:32 -0700 Subject: [PATCH 1/3] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUBackend.cpp | 10 +- backends/webgpu/runtime/WebGPUGraph.cpp | 54 ++- backends/webgpu/runtime/WebGPUGraph.h | 10 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 149 +++++++-- ...ming_attention_qwen3_k16_causal_bound.wgsl | 283 ++++++++++++++++ ...ng_attention_qwen3_k16_causal_bound_wgsl.h | 311 ++++++++++++++++++ ..._attention_qwen3_q32_k16_causal_bound.wgsl | 283 ++++++++++++++++ ...ttention_qwen3_q32_k16_causal_bound_wgsl.h | 311 ++++++++++++++++++ backends/webgpu/scripts/gen_wgsl_headers.py | 27 +- backends/webgpu/test/ops/test_sdpa.py | 111 ++++++- backends/webgpu/test/test_webgpu_native.cpp | 168 +++++++++- backends/webgpu/test/test_wgsl_codegen.py | 16 + 12 files changed, 1668 insertions(+), 65 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 51cca64fd50..27b1ae40ff9 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -98,6 +98,13 @@ Result WebGPUBackend::init( enable_f16_accumulate_gemm = spec.get(); } } + int sdpa_query_tile = 0; + { + Result spec = context.get_runtime_spec("sdpa_query_tile"); + if (spec.ok()) { + sdpa_query_tile = spec.get(); + } + } try { graph->build( @@ -105,7 +112,8 @@ Result WebGPUBackend::init( constant_data, context.get_named_data_map(), enable_f16_kv_cache, - enable_f16_accumulate_gemm); + enable_f16_accumulate_gemm, + sdpa_query_tile); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 6ac261bfdd5..d59c27f49ea 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -808,7 +809,8 @@ void WebGPUGraph::build( const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map, bool f16_kv_cache, - bool f16_accumulate_gemm) { + bool f16_accumulate_gemm, + int sdpa_query_tile) { if (!device_) { auto* ctx = get_default_webgpu_context(); if (ctx) { @@ -838,6 +840,10 @@ void WebGPUGraph::build( // additionally gates the kernel on the negotiated shader-f16 feature. f16_accumulate_gemm_ = f16_accumulate_gemm; + // SDPA query-tile selector (runtime opt-in); 0 = geometry default (Q16), + // 32 = Q32 candidate. Read at the SDPA op-lowering selection site. + sdpa_query_tile_ = sdpa_query_tile; + // Phase 1: Create all values const auto* values = graph->values(); const int num_vals = values ? values->size() : 0; @@ -1718,6 +1724,19 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { continue; } + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (!in.host_is_int64 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + wgpuQueueWriteBuffer( + queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + continue; + } + throw std::runtime_error( "WebGPU: unsupported input copy for input " + std::to_string(i) + " (host " + std::to_string(in.nbytes) + " bytes" + @@ -1761,7 +1780,8 @@ bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING void WebGPUGraph::record_active_route(const std::string& kernel_name) { uint32_t bits = 0; - if (kernel_name == "sdpa_streaming_attention_k16_causal_bound") { + if (kernel_name.find("k16_causal_bound") != std::string::npos) { + // llama + qwen3 (Q16/Q32) streaming causal-bound kernels bits = kRoutePrefill | kRouteK16CausalBound; } else if (kernel_name == "sdpa_streaming_attention_k16") { bits = kRoutePrefill | kRouteK16; @@ -2038,12 +2058,13 @@ void WebGPUGraph::copy_outputs( cb_info.mode = WGPUCallbackMode_WaitAnyOnly; cb_info.callback = buffer_map_callback; cb_info.userdata1 = &cb_data[i]; + const auto& tensor = tensors_[output_ids_[i]]; + const bool widen_fp16 = !tensor.is_int && tensor.elem_size == 2 && + outputs[i].second == tensor.cur_nbytes * 2; + const size_t map_nbytes = + widen_fp16 ? tensor.cur_nbytes : outputs[i].second; map_futures[i] = wgpuBufferMapAsync( - output_staging_buffers_[i], - WGPUMapMode_Read, - 0, - outputs[i].second, - cb_info); + output_staging_buffers_[i], WGPUMapMode_Read, 0, map_nbytes, cb_info); } for (size_t i = 0; i < count; i++) { @@ -2058,9 +2079,24 @@ void WebGPUGraph::copy_outputs( continue; } if (cb_data[i].status == WGPUMapAsyncStatus_Success) { + const auto& tensor = tensors_[output_ids_[i]]; + const bool widen_fp16 = !tensor.is_int && tensor.elem_size == 2 && + outputs[i].second == tensor.cur_nbytes * 2; + const size_t map_nbytes = + widen_fp16 ? tensor.cur_nbytes : outputs[i].second; const void* mapped = wgpuBufferGetConstMappedRange( - output_staging_buffers_[i], 0, outputs[i].second); - std::memcpy(outputs[i].first, mapped, outputs[i].second); + output_staging_buffers_[i], 0, map_nbytes); + if (widen_fp16) { + const auto* src = + static_cast(mapped); + auto* dst = static_cast(outputs[i].first); + const size_t numel = map_nbytes / sizeof(*src); + for (size_t e = 0; e < numel; e++) { + dst[e] = static_cast(src[e]); + } + } else { + std::memcpy(outputs[i].first, mapped, outputs[i].second); + } wgpuBufferUnmap(output_staging_buffers_[i]); } else { throw std::runtime_error("WebGPU buffer map failed for output"); diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 40e459750ca..2ae887b7b33 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -103,7 +103,8 @@ class WebGPUGraph { const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map = nullptr, bool f16_kv_cache = false, - bool f16_accumulate_gemm = false); + bool f16_accumulate_gemm = false, + int sdpa_query_tile = 0); // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); @@ -383,6 +384,12 @@ class WebGPUGraph { return f16_accumulate_gemm_; } + // Runtime-selected SDPA query-tile candidate; 0 = geometry default (Q16), + // 32 = Q32 candidate. + int sdpa_query_tile() const { + return sdpa_query_tile_; + } + private: #ifdef WGPU_BACKEND_ENABLE_PROFILING void record_active_route(const std::string& kernel_name); @@ -391,6 +398,7 @@ class WebGPUGraph { bool kv_f16_ = false; std::unordered_set kv_cache_ids_; bool f16_accumulate_gemm_ = false; + int sdpa_query_tile_ = 0; private: WGPUInstance instance_ = nullptr; diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index b5468854ee3..16c48db86c6 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -6,8 +6,8 @@ * LICENSE file in the root directory of this source tree. */ -#include #include +#include #include #include #include @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -200,17 +202,82 @@ bool streaming_attention_k16_device_supported(WGPUDevice device) { limits.maxComputeWorkgroupStorageSize >= 14720u; } +constexpr uint32_t kQwen3K16QueryTile = 16u; +constexpr uint32_t kQwen3Q32K16QueryTile = 32u; +constexpr uint32_t kQwen3Q16K16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); +constexpr uint32_t kQwen3K16StorageBytes = kQwen3Q16K16StorageBytes; +constexpr uint32_t kQwen3Q32K16StorageBytes = 22912u; + +constexpr bool streaming_attention_k16_workgroup_count_fits( + int64_t S, + int64_t Hkv, + int64_t g, + uint32_t query_tile, + uint32_t max_workgroups) { + if (S <= 0 || Hkv <= 0 || g <= 0 || query_tile == 0u || + max_workgroups == 0u) { + return false; + } + if (static_cast(S) > UINT64_MAX / static_cast(g)) { + return false; + } + const uint64_t logical_rows = + static_cast(S) * static_cast(g); + if (logical_rows > UINT64_MAX - (query_tile - 1u)) { + return false; + } + const uint64_t groups_per_kv = (logical_rows + query_tile - 1u) / query_tile; + if (groups_per_kv > UINT64_MAX / static_cast(Hkv)) { + return false; + } + const uint64_t workgroups = groups_per_kv * static_cast(Hkv); + return workgroups > 0u && workgroups <= UINT32_MAX && + workgroups <= max_workgroups; +} + +static_assert( + streaming_attention_k16_workgroup_count_fits(65528, 8, 2, 16, 65535)); +static_assert( + !streaming_attention_k16_workgroup_count_fits(65529, 8, 2, 16, 65535)); + +bool qwen3_q16_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQwen3K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQwen3Q32K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + utils::WgCount streaming_attention_k16_grid( WGPUDevice device, int64_t S, int64_t Hkv, - int64_t g) { + int64_t g, + uint32_t query_tile) { const uint64_t groups_per_kv = - (static_cast(S) * static_cast(g) + 31u) / 32u; + (static_cast(S) * static_cast(g) + query_tile - 1u) / + query_tile; const uint64_t workgroups = static_cast(Hkv) * groups_per_kv; if (workgroups == 0u || workgroups > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa: K16 workgroup count exceeds uint32"); + throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32"); } if (workgroups > utils::queried_max_workgroups(device)) { throw std::runtime_error( @@ -440,6 +507,9 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (D != k_cache.dims[cn - 1]) { throw std::runtime_error("WebGPU sdpa: q and k_cache head_dim mismatch"); } + if (k_cache.dims[cn - 2] != Hkv) { + throw std::runtime_error("WebGPU sdpa: k and k_cache num_heads mismatch"); + } if (k_cache.dims != v_cache.dims) { throw std::runtime_error("WebGPU sdpa: k_cache and v_cache shape mismatch"); } @@ -509,12 +579,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const WGPUDevice device = graph.device(); const WGPUBuffer k16_buffers[] = { - q.buffer, - k.buffer, - v.buffer, - k_cache.buffer, - v_cache.buffer, - out.buffer}; + q.buffer, k.buffer, v.buffer, k_cache.buffer, v_cache.buffer, out.buffer}; bool k16_buffers_distinct = true; for (size_t i = 0; i < 6; i++) { for (size_t j = i + 1; j < 6; j++) { @@ -522,9 +587,39 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; } } - const bool k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && + const bool qwen3_k16_geometry = Hq == 16 && Hkv == 8 && g == 2 && D == 128 && + scale == 1.0f / std::sqrt(128.0f) && out.dims == q.dims; + // Q16 is the default route for exact Qwen3 geometry; Q32 is an explicit + // autotuning candidate selected only via the sdpa_query_tile RuntimeSpec. + const bool qwen3_query_tile_requested = qwen3_k16_geometry; + const bool qwen3_q32_requested = + graph.sdpa_query_tile() == static_cast(kQwen3Q32K16QueryTile) && + qwen3_query_tile_requested; + const uint32_t qwen3_query_tile = + qwen3_q32_requested ? kQwen3Q32K16QueryTile : kQwen3K16QueryTile; + const bool qwen3_device_supported = + qwen3_query_tile_requested && + (qwen3_q32_requested ? qwen3_q32_k16_device_supported(device) + : qwen3_q16_k16_device_supported(device)) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, qwen3_query_tile, utils::queried_max_workgroups(device)); + const bool qwen3_k16_selected = + qwen3_query_tile_requested && graph.kv_f16() && qwen3_device_supported; + const bool qwen3_q32_selected = qwen3_k16_selected && qwen3_q32_requested; + const bool llama_k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && D == 64 && scale == 0.125f && out.dims == q.dims && - k16_buffers_distinct && streaming_attention_k16_device_supported(device); + streaming_attention_k16_device_supported(device); + const bool k16_eligible = + k16_buffers_distinct && (llama_k16_eligible || qwen3_k16_selected); + const uint32_t k16_query_tile = qwen3_k16_selected ? qwen3_query_tile : 32u; + const char* k16_shader = qwen3_q32_selected + ? kStreamingAttentionQwen3Q32K16CausalBoundWGSL + : qwen3_k16_selected ? kStreamingAttentionQwen3K16CausalBoundWGSL + : kStreamingAttentionK16CausalBoundWGSL; + const char* k16_label = qwen3_q32_selected + ? "sdpa_streaming_attention_qwen3_q32_k16_causal_bound" + : qwen3_k16_selected ? "sdpa_streaming_attention_qwen3_k16_causal_bound" + : "sdpa_streaming_attention_k16_causal_bound"; const uint32_t uc_wg = utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); const uint32_t qk_wg = utils::clamp_workgroup_size( @@ -556,7 +651,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { qk_wg, av_wg, fd_eligible, - k16_eligible](WebGPUGraph& gr) { + k16_eligible, + k16_query_tile](WebGPUGraph& gr) { SdpaLiveState state = {}; const auto& q_live_dims = gr.cur_dims(q_id); state.s = q_live_dims[qn - 3]; @@ -608,8 +704,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k16_eligible) { state.streaming_k16 = make_streaming_attention_k16_params( state.s, state.context_len, state.pos, Hq, Hkv, D); - state.streaming_k16_grid = - streaming_attention_k16_grid(gr.device(), state.s, Hkv, g); + state.streaming_k16_grid = streaming_attention_k16_grid( + gr.device(), state.s, Hkv, g, k16_query_tile); } state.update_cache_grid = { utils::compute_1d_workgroup_count( @@ -621,8 +717,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(state.context_len); const uint64_t qk_tiles = static_cast(Hq) * static_cast(utils::div_up(state.s, kSdpaTileM)) * - static_cast( - utils::div_up(state.context_len, kSdpaTileN)); + static_cast(utils::div_up(state.context_len, kSdpaTileN)); const uint64_t softmax_rows = static_cast(Hq) * static_cast(state.s); const uint64_t av_tiles = static_cast(Hq) * @@ -685,11 +780,9 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const bool dynamic_sequence = graph.tensor_has_dynamic_dims(q_id) || graph.tensor_has_dynamic_dims(k_id) || graph.tensor_has_dynamic_dims(v_id); - const bool dual_route = - utils::should_record_sdpa_dual_route( - fd_eligible, dynamic_sequence, dynamic_pos); - const bool record_k16 = - k16_eligible && (dual_route || !initial_state.use_fd); + const bool dual_route = utils::should_record_sdpa_dual_route( + fd_eligible, dynamic_sequence, dynamic_pos); + const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd); const bool record_materialized = !k16_eligible && (dual_route || !initial_state.use_fd); const bool record_fd = dual_route || initial_state.use_fd; @@ -790,7 +883,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { : initial_state.streaming_k16_grid; build_dispatch( graph, - kStreamingAttentionK16CausalBoundWGSL, + k16_shader, k16_bindings, 4, k16_buf, @@ -799,7 +892,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { initial_grid.y, 0, true, - "sdpa_streaming_attention_k16_causal_bound"); + k16_label); k16_idx = graph.num_dispatches() - 1; k16_range.end = graph.num_dispatches(); } @@ -907,10 +1000,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (fixed_use_fd) { throw std::runtime_error("WebGPU sdpa: static route changed"); } - gr.dispatch_at(k16_idx).workgroup_count_x = - state.streaming_k16_grid.x; - gr.dispatch_at(k16_idx).workgroup_count_y = - state.streaming_k16_grid.y; + gr.dispatch_at(k16_idx).workgroup_count_x = state.streaming_k16_grid.x; + gr.dispatch_at(k16_idx).workgroup_count_y = state.streaming_k16_grid.y; } else { if (fixed_use_fd) { throw std::runtime_error("WebGPU sdpa: static route changed"); diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5e27346af6f --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..f69a4d72dae --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * 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 streaming_attention_qwen3_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: b26e81d92ba58c833f7dfbf2da6bce1b9cb468d70d893cef9df3a7fae16ee1fc +inline constexpr const char* kStreamingAttentionQwen3K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX = + 16; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY = + 8; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ = + 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl new file mode 100644 index 00000000000..6135b15b75d --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..e58a9d94bac --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * 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 streaming_attention_qwen3_q32_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: 6d4396945b86cc3445698b1182fc4535ca786ad509ba128a6b6b08977e6a55e0 +inline constexpr const char* kStreamingAttentionQwen3Q32K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY = 8; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index c66aff2039a..28f3d710979 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -431,14 +431,17 @@ def embedded_sha256(header_text: str) -> str: def _wg_size_const(base: str, axis: str, val: int) -> str: """One WorkgroupSize constant; wrap to <=80 cols so CLANGFORMAT accepts it. - Long shader names push the single-line form past the 80-col limit (clang-format - then breaks after '=' with a 4-space continuation indent); emit that wrapped - form up front so the generated header matches lintrunner's CLANGFORMAT. + Long shader names push the single-line form past the 80-col limit. Emit the + wrapped form that clang-format selects so generated headers stay byte-stable. """ - decl = f"inline constexpr uint32_t k{base}WorkgroupSize{axis} =" - if len(decl) + len(f" {val};") > 80: - return f"{decl}\n {val};\n" - return f"{decl} {val};\n" + name = f"k{base}WorkgroupSize{axis}" + prefix = f"inline constexpr uint32_t {name} =" + decl = f"{prefix} {val};" + if len(decl) > 85: + return f"inline constexpr uint32_t\n {name} = {val};\n" + if len(decl) > 80: + return f"{prefix}\n {val};\n" + return f"{decl}\n" def render_header( @@ -463,6 +466,14 @@ def render_header( raise ValueError('shader contains )" which would close the R"( literal') base = symbol_base(name) x, y, z = parse_workgroup_size(wgsl_text) + provenance = f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT." + if len(provenance) > 80: + provenance_lines = [ + f"// @generated from {provenance_stem}.wgsl", + "// DO NOT EDIT.", + ] + else: + provenance_lines = [provenance] head = [ _BSD_HEADER, @@ -473,7 +484,7 @@ def render_header( "", "namespace executorch::backends::webgpu {", "", - f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT.", + *provenance_lines, f"// wgsl-sha256: {wgsl_sha256(wgsl_text)}", f'inline constexpr const char* k{base}WGSL = R"(', ] diff --git a/backends/webgpu/test/ops/test_sdpa.py b/backends/webgpu/test/ops/test_sdpa.py index 960c62e5203..b595462d5e3 100644 --- a/backends/webgpu/test/ops/test_sdpa.py +++ b/backends/webgpu/test/ops/test_sdpa.py @@ -40,6 +40,7 @@ class SdpaConfig: cmax: int # kv-cache capacity input_pos: int # number of prior tokens already in the cache (decode) denom: float = 16.0 # ramp divisor; small denom -> large logits (softmax stress) + kv_f16: bool = False # Single source of truth, mirrored by the C++ CONFIGS table in the native test. @@ -64,6 +65,8 @@ class SdpaConfig: # 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV (cap+1). SdpaConfig("llama1b_prefill_512", 32, 8, 64, 512, 512, 0), SdpaConfig("llama1b_prefill_2048", 32, 8, 64, 2048, 2048, 0), + SdpaConfig("qwen3_prefill", 16, 8, 128, 128, 256, 0, kv_f16=True), + SdpaConfig("qwen3_odd_boundary", 16, 8, 128, 17, 64, 31, kv_f16=True), ] @@ -83,6 +86,7 @@ class ReplaySeq: d: int # head dim cmax: int # kv-cache capacity (>= sum(seq_lens)) seq_lens: tuple[int, ...] + kv_f16: bool = False # Mirror Vulkan sdpa_test.cpp:856/867/875 (3 param sets); cmax = sum rounded up. @@ -90,11 +94,18 @@ class ReplaySeq: ReplaySeq("small", 8, 4, 4, 16, (3, 1, 1, 5, 1, 1, 2)), ReplaySeq("small_d", 6, 2, 8, 16, (3, 1, 1, 5, 1, 1)), ReplaySeq("llama3", 24, 8, 128, 256, (111, 1, 1, 1, 57, 1, 1)), + ReplaySeq("qwen3_fd", 16, 8, 128, 64, (17, 1), kv_f16=True), ] +DYNAMIC_REPLAY_SEQS = [seq for seq in REPLAY_SEQS if not seq.kv_f16] -# (head_dim, num_heads, num_kv_heads) from sdpa_test.cpp:856/867/875 -- guards a -# transposition of the (hq, hkv, d) field order against the Vulkan source. -VULKAN_PARAMS = {"small": (4, 8, 4), "small_d": (8, 6, 2), "llama3": (128, 24, 8)} +# Guards transposition of the (hq, hkv, d) field order. The first three values +# mirror Vulkan sdpa_test.cpp:856/867/875; Qwen3 extends the same contract. +VULKAN_PARAMS = { + "small": (4, 8, 4), + "small_d": (8, 6, 2), + "llama3": (128, 24, 8), + "qwen3_fd": (128, 16, 8), +} class SdpaModule(torch.nn.Module): @@ -178,6 +189,12 @@ def _det_inputs(cfg: SdpaConfig): return q, k, v, k_cache, v_cache +def _round_kv_for_storage(cfg: SdpaConfig, *tensors: torch.Tensor): + if not cfg.kv_f16: + return tensors + return tuple(tensor.to(torch.float16).to(torch.float32) for tensor in tensors) + + def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """Reference attention output [1,S,Hq,D], computed in fp64 then cast to fp32. @@ -189,6 +206,7 @@ def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """ context_len = cfg.s + cfg.input_pos g = cfg.hq // cfg.hkv + k, v, k_cache, v_cache = _round_kv_for_storage(cfg, k, v, k_cache, v_cache) qd, kd, vd = q.double(), k.double(), v.double() kcd, vcd = k_cache.double(), v_cache.double() @@ -228,6 +246,35 @@ def _export_pte(cfg: SdpaConfig, q, k, v, kc, vc): class TestSdpa(unittest.TestCase): + def test_qwen3_fixture_contract(self) -> None: + configs = {cfg.name: cfg for cfg in CONFIGS} + expected = { + "qwen3_prefill": (16, 8, 128, 128, 256, 0), + "qwen3_odd_boundary": (16, 8, 128, 17, 64, 31), + } + for name, geometry in expected.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), + geometry, + ) + self.assertTrue(cfg.kv_f16) + + replays = {seq.name: seq for seq in REPLAY_SEQS} + self.assertIn("qwen3_fd", replays) + qwen3_fd = replays["qwen3_fd"] + self.assertEqual((qwen3_fd.hq, qwen3_fd.hkv, qwen3_fd.d), (16, 8, 128)) + self.assertEqual(qwen3_fd.seq_lens, (17, 1)) + self.assertTrue(qwen3_fd.kv_f16) + + probe = torch.tensor([0.1], dtype=torch.float32) + (rounded,) = _round_kv_for_storage(configs["qwen3_prefill"], probe) + expected = probe.to(torch.float16).to(torch.float32) + torch.testing.assert_close(rounded, expected, atol=0.0, rtol=0.0) + self.assertFalse(torch.equal(rounded, probe)) + def test_sdpa_export_delegates(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): @@ -248,7 +295,12 @@ def test_golden_matches_eager_op(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): q, k, v, kc, vc = _det_inputs(cfg) - eager = SdpaModule(cfg.input_pos)(q, k, v, kc.clone(), vc.clone()) + 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) @@ -277,10 +329,16 @@ def test_replay_golden_matches_eager(self) -> None: 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 = SdpaModule(input_pos)(q, k, v, kc.clone(), vc.clone()) + 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] @@ -302,6 +360,7 @@ def test_replay_export_delegates(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, kc, vc) @@ -347,7 +406,14 @@ def export_replay_sequences(out_dir: str) -> None: 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 + 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) et = _export_pte(cfg, q, k, v, ref_kc, ref_vc) @@ -421,7 +487,7 @@ def export_dynamic_decode(out_dir: str) -> None: Mirrors the host accumulation the native test threads: at step t the golden attends over input_pos=t prior tokens plus the new token. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" et = _export_dyn_pte(seq, 1) pte_path = os.path.join(out_dir, f"sdpa_dyn_{seq.name}.pte") @@ -431,7 +497,14 @@ def export_dynamic_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" None: def test_dynamic_decode_golden_matches_eager(self) -> None: # The threaded-cache decode golden must equal the eager op step-by-step. - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: ref_kc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc) @@ -499,7 +579,7 @@ def export_incache_decode(out_dir: str) -> None: """One sdpa_incache_.pte (mutable-buffer KV cache) + per-step decode goldens. forward() feeds only q/k/v + input_pos; the cache persists in-graph. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" m = DecodeCacheModule(seq.hkv, seq.d, seq.cmax) q, k, v = _step_inputs(seq, 0, 1) @@ -517,7 +597,14 @@ def export_incache_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"incache_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"incache_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" #include #include +#include #include #include #include @@ -989,6 +990,7 @@ struct SdpaConfig { float denom; // ramp divisor (mirrors Python); small -> large logits bool required = false; // CI (SDPA dir set): absent .pte = FAIL, not skip bool expect_reject = false; // load MUST fail (e.g. D%4 guard), no golden + bool kv_f16 = false; }; const SdpaConfig kSdpaConfigs[] = { @@ -1031,6 +1033,28 @@ const SdpaConfig kSdpaConfigs[] = { 0, 16.0f, /*required=*/true}, + {"qwen3_prefill", + 16, + 8, + 128, + 128, + 256, + 0, + 16.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_odd_boundary", + 16, + 8, + 128, + 17, + 64, + 31, + 16.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). @@ -1053,9 +1077,8 @@ float sdpa_ramp_t( return static_cast(((i + 31 * t) % mod) - off) / denom; } -// Multi-step replay sequences. Mirror the Python REPLAY_SEQS / Vulkan param -// sets (sdpa_test.cpp:856/867/875). Each seq_lens entry is one step replayed on -// a host-threaded KV cache (big=prefill, mid=multi-token, 1=decode). +// Multi-step replay sequences. The first three mirror Vulkan param sets; Qwen3 +// extends the same Python REPLAY_SEQS contract. struct SdpaSequence { const char* name; int hq; @@ -1063,14 +1086,39 @@ struct SdpaSequence { int d; int cmax; std::vector seq_lens; + bool kv_f16 = false; }; const SdpaSequence kSdpaSequences[] = { {"small", 8, 4, 4, 16, {3, 1, 1, 5, 1, 1, 2}}, {"small_d", 6, 2, 8, 16, {3, 1, 1, 5, 1, 1}}, {"llama3", 24, 8, 128, 256, {111, 1, 1, 1, 57, 1, 1}}, + {"qwen3_fd", 16, 8, 128, 64, {17, 1}, /*kv_f16=*/true}, }; +Error load_sdpa_forward(Module& module, bool kv_f16) { + if (!kv_f16) { + return module.load_forward(); + } + BackendOptions<1> options; + auto error = options.set_option("enable_f16_kv_cache", true); + if (error != Error::Ok) { + return error; + } + LoadBackendOptionsMap option_map; + error = option_map.set_options("VulkanBackend", options.view()); + if (error != Error::Ok) { + return error; + } + return module.load_forward(nullptr, nullptr, &option_map); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kTestRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kTestRouteFlashDecoding = 1u << 10; +constexpr uint32_t kTestRouteK16CausalBound = 1u << 11; +#endif // WGPU_BACKEND_ENABLE_PROFILING + void test_sdpa_config( const SdpaConfig& cfg, const std::string& model_path, @@ -1087,7 +1135,7 @@ void test_sdpa_config( cfg.input_pos); Module module(model_path); - auto err = module.load_forward(); + auto err = load_sdpa_forward(module, cfg.kv_f16); if (cfg.expect_reject) { // D not a multiple of 4 must be rejected at load by the head_dim guard. ASSERT_NE(err, Error::Ok) @@ -1132,6 +1180,25 @@ void test_sdpa_config( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() << ")"; + if (cfg.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // Exact Qwen3 geometry + fp16 KV selects the K16 streaming (causal-bound) + // route by default. The sdpa_query_tile RuntimeSpec only swaps the Q16/Q32 + // kernel variant; both map to the K16CausalBound bit. A non-Qwen3 fp16-KV + // shape falls back to the materialized path (or flash-decoding at S==1). + const bool qwen3_geometry = cfg.hq == 16 && cfg.hkv == 8 && cfg.d == 128; + const uint32_t expected_route = qwen3_geometry && cfg.s > 1 + ? kTestRouteK16CausalBound + : (cfg.s == 1 ? kTestRouteFlashDecoding + : kTestRouteMaterializedAttention); + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route); + EXPECT_EQ(g_last_route_conflict_count, 0u); +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outputs = result.get(); // Select the attention output [1,S,Hq,D] by shape; the op returns @@ -1199,7 +1266,7 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { std::to_string(t) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos); Module module(base + ".pte"); - ASSERT_EQ(module.load_forward(), Error::Ok) + ASSERT_EQ(load_sdpa_forward(module, seq.kv_f16), Error::Ok) << "could not load " << base << ".pte"; const int qn = s * seq.hq * seq.d; @@ -1225,6 +1292,24 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward " << base << ".pte (error " << (int)result.error() << ")"; + if (seq.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // S==1 decode -> flash-decoding; a multi-token exact-Qwen3-geometry + // prefill -> the K16 streaming (causal-bound) route by default (no env); + // any other multi-token fp16-KV shape -> materialized. + const bool qwen3_geometry = seq.hq == 16 && seq.hkv == 8 && seq.d == 128; + const uint32_t expected_route = s == 1 ? kTestRouteFlashDecoding + : qwen3_geometry ? kTestRouteK16CausalBound + : kTestRouteMaterializedAttention; + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route) + << seq.name << " step" << t; + EXPECT_EQ(g_last_route_conflict_count, 0u) << seq.name << " step" << t; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outs = result.get(); // The op returns [k_cache, v_cache, attn_output]: attn has a unique numel; @@ -1939,6 +2024,79 @@ TEST(WebGPUNative, PrepackTied) { } // SDPA sweep: configs self-discover sdpa_.pte; required=FAIL else skip. +TEST(WebGPUNative, Qwen3SdpaFixtureContract) { + const auto find_config = [](const char* name) { + return std::find_if( + std::begin(kSdpaConfigs), + std::end(kSdpaConfigs), + [name](const SdpaConfig& cfg) { + return std::strcmp(cfg.name, name) == 0; + }); + }; + const auto prefill = find_config("qwen3_prefill"); + const auto boundary = find_config("qwen3_odd_boundary"); + ASSERT_NE(prefill, std::end(kSdpaConfigs)); + ASSERT_NE(boundary, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {prefill->hq, + prefill->hkv, + prefill->d, + prefill->s, + prefill->cmax, + prefill->input_pos}), + std::vector({16, 8, 128, 128, 256, 0})); + EXPECT_EQ( + std::vector( + {boundary->hq, + boundary->hkv, + boundary->d, + boundary->s, + boundary->cmax, + boundary->input_pos}), + std::vector({16, 8, 128, 17, 64, 31})); + EXPECT_TRUE(prefill->kv_f16 && boundary->kv_f16); + + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + EXPECT_EQ( + std::vector({replay->hq, replay->hkv, replay->d, replay->cmax}), + std::vector({16, 8, 128, 64})); + EXPECT_EQ(replay->seq_lens, std::vector({17, 1})); + EXPECT_TRUE(replay->kv_f16); +} + +TEST(WebGPUNative, Qwen3SdpaRoutes) { + ASSERT_FALSE(g_sdpa_dir.empty()) << "WEBGPU_TEST_SDPA_DIR not set"; + // Default route: exact-Qwen3-geometry fp16-KV configs select the Q16 K16 + // streaming (causal-bound) route by geometry (no runtime config needed) -- + // the per-config assertions live in test_sdpa_config / test_sdpa_replay. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config(cfg, base + ".pte", base + ".golden.bin"); + } + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + test_sdpa_replay(*replay, g_sdpa_dir); + // Q32 is a non-default autotuning candidate selected via the sdpa_query_tile + // RuntimeSpec (=32); it is not route-asserted in default CI because its + // workgroup storage can exceed a device's maxComputeWorkgroupStorageSize + // (then production correctly falls back to materialized). +} + TEST(WebGPUNative, SdpaSweep) { const std::string& dir = g_sdpa_dir; bool ran = false; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 46d285aa60b..c90eff08335 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -133,6 +133,22 @@ def test_render_header_embeds_sha256(self) -> None: self.assertEqual(g.embedded_sha256(h), want) self.assertEqual(g.wgsl_sha256(wgsl), want) + def test_render_header_long_name_is_clang_format_stable(self) -> None: + stem = "streaming_attention_qwen3_q32_k16_causal_bound" + wgsl = "@compute @workgroup_size(32, 8, 1)\nfn main(){}\n" + h = g.render_header(Path(f"runtime/ops/sdpa/{stem}.wgsl"), wgsl) + + self.assertIn( + f"// @generated from {stem}.wgsl\n// DO NOT EDIT.", + h, + ) + self.assertIn( + "inline constexpr uint32_t\n" + " kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32;", + h, + ) + self.assertEqual(g.embedded_sha256(h), g.wgsl_sha256(wgsl)) + def test_embedded_sha256_missing_returns_empty(self) -> None: self.assertEqual(g.embedded_sha256("no sha line here\n"), "") From 54356890e75e5579d5fae2c22f80126bf0e308ef Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 11:38:55 -0700 Subject: [PATCH 2/3] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUBackend.cpp | 7 +++++-- backends/webgpu/runtime/WebGPUGraph.cpp | 24 +++++++++++++++-------- backends/webgpu/runtime/WebGPUGraph.h | 10 +++++++++- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 16 +++++++++++---- 4 files changed, 42 insertions(+), 15 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 292388f9ed4..67aa17f1ab1 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -190,12 +190,15 @@ Error WebGPUBackend::execute( graph->execute(graph_options); // Copy outputs from GPU staging buffers to EValue tensor data pointers - std::vector> outputs; + std::vector outputs; outputs.reserve(num_outputs); for (size_t i = 0; i < num_outputs; i++) { const size_t arg_idx = num_inputs + i; auto& tensor = args[arg_idx]->toTensor(); - outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; + outputs.push_back( + {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32}); } graph->copy_outputs(outputs, graph_options); } catch (const std::exception& e) { diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index c945bea8e74..f92345fdab0 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -2056,7 +2056,7 @@ void buffer_map_callback( } // namespace void WebGPUGraph::copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUGraphExecutionOptions& options) { const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); const WebGPUExecutionPlan plan = plan_webgpu_execution( @@ -2074,12 +2074,20 @@ void WebGPUGraph::copy_outputs( auto output_map_size = [&](size_t i) -> std::pair { const auto& tensor = tensors_[output_ids_[i]]; const bool widen_fp16 = !tensor.is_int && tensor.elem_size == 2 && - outputs[i].second == tensor.cur_nbytes * 2; - return {widen_fp16, widen_fp16 ? tensor.cur_nbytes : outputs[i].second}; + outputs[i].nbytes == tensor.cur_nbytes * 2; + // Require an explicit fp32 host dtype for the widen (mirrors the + // copy_inputs narrow guard): a same-2:1-ratio non-fp32 host must not be + // silently reinterpreted as fp32, and must not fall through to a memcpy + // that would over-read the smaller fp16 staging buffer. + if (widen_fp16 && !outputs[i].host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device output requires an fp32 host tensor"); + } + return {widen_fp16, widen_fp16 ? tensor.cur_nbytes : outputs[i].nbytes}; }; for (size_t i = 0; i < count; i++) { - if (!plan.copy_outputs[i] || outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { cb_data[i].status = WGPUMapAsyncStatus_Success; continue; } @@ -2093,14 +2101,14 @@ void WebGPUGraph::copy_outputs( } for (size_t i = 0; i < count; i++) { - if (plan.copy_outputs[i] && outputs[i].second != 0 && + if (plan.copy_outputs[i] && outputs[i].nbytes != 0 && webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) { throw std::runtime_error("WebGPU: WaitAny failed for output map"); } } for (size_t i = 0; i < count; i++) { - if (!plan.copy_outputs[i] || outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { continue; } if (cb_data[i].status == WGPUMapAsyncStatus_Success) { @@ -2110,13 +2118,13 @@ void WebGPUGraph::copy_outputs( if (widen_fp16) { const auto* src = static_cast(mapped); - auto* dst = static_cast(outputs[i].first); + auto* dst = static_cast(outputs[i].data); const size_t numel = map_nbytes / sizeof(*src); for (size_t e = 0; e < numel; e++) { dst[e] = static_cast(src[e]); } } else { - std::memcpy(outputs[i].first, mapped, outputs[i].second); + std::memcpy(outputs[i].data, mapped, outputs[i].nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } else { diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index c88120fab4d..77f95c41e9e 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -46,6 +46,14 @@ struct InputData { bool host_is_fp32 = false; }; +// Host destination for a graph output. host_is_fp32 gates the fp16->fp32 widen +// on readback (mirrors InputData's guard on the copy_inputs narrow path). +struct OutputData { + void* data = nullptr; + size_t nbytes = 0; + bool host_is_fp32 = false; +}; + struct WebGPUDispatch { WGPUComputePipeline pipeline = nullptr; WGPUBindGroup bind_group = nullptr; @@ -116,7 +124,7 @@ class WebGPUGraph { // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. void copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUGraphExecutionOptions& options); const std::vector& input_ids() const { diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 87ba32e927d..dcef07f41e4 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -209,7 +209,12 @@ constexpr uint32_t kQwen3Q16K16StorageBytes = 512u * 4u * sizeof(float) + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + 3u * 16u * sizeof(float); constexpr uint32_t kQwen3K16StorageBytes = kQwen3Q16K16StorageBytes; -constexpr uint32_t kQwen3Q32K16StorageBytes = 22912u; +// Mirrors the Q32 shader's workgroup arrays (t_q_tile vec4x1024, t_kv_tile +// vec4x512, t_scores vec2x256, t_m/t_d/t_alpha f32x32) so the +// device-support gate stays tied to the declared storage, not a literal. +constexpr uint32_t kQwen3Q32K16StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); constexpr bool streaming_attention_k16_workgroup_count_fits( int64_t S, @@ -615,9 +620,12 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const bool qwen3_k16_selected = qwen3_q32_selected || qwen3_q16_supported; const uint32_t qwen3_query_tile = qwen3_q32_selected ? kQwen3Q32K16QueryTile : kQwen3K16QueryTile; - const bool llama_k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && - g == 4 && D == 64 && scale == 0.125f && out.dims == q.dims && - streaming_attention_k16_device_supported(device); + const bool llama_k16_eligible = + graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && D == 64 && + scale == 0.125f && out.dims == q.dims && + streaming_attention_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kLlamaK16QueryTile, device_max_workgroups); const bool k16_eligible = k16_buffers_distinct && (llama_k16_eligible || qwen3_k16_selected); const uint32_t k16_query_tile = From 2cfbf64984f76015c93669ff9ca6410b9e73e62f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 11:57:31 -0700 Subject: [PATCH 3/3] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUGraph.cpp | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index f92345fdab0..0e4964a3789 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -2075,16 +2075,24 @@ void WebGPUGraph::copy_outputs( const auto& tensor = tensors_[output_ids_[i]]; const bool widen_fp16 = !tensor.is_int && tensor.elem_size == 2 && outputs[i].nbytes == tensor.cur_nbytes * 2; - // Require an explicit fp32 host dtype for the widen (mirrors the - // copy_inputs narrow guard): a same-2:1-ratio non-fp32 host must not be - // silently reinterpreted as fp32, and must not fall through to a memcpy - // that would over-read the smaller fp16 staging buffer. - if (widen_fp16 && !outputs[i].host_is_fp32) { + return {widen_fp16, widen_fp16 ? tensor.cur_nbytes : outputs[i].nbytes}; + }; + + // Validate dtypes up front, before any wgpuBufferMapAsync is issued: require + // an explicit fp32 host dtype for the fp16->fp32 widen (mirrors the + // copy_inputs narrow guard) so a same-2:1-ratio non-fp32 host is not silently + // reinterpreted as fp32. Throwing here — rather than mid map-request loop — + // keeps in-flight async maps (whose callbacks point at cb_data) from being + // left dangling when the exception unwinds this frame. + for (size_t i = 0; i < count; i++) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { + continue; + } + if (output_map_size(i).first && !outputs[i].host_is_fp32) { throw std::runtime_error( "WebGPU: fp16 device output requires an fp32 host tensor"); } - return {widen_fp16, widen_fp16 ? tensor.cur_nbytes : outputs[i].nbytes}; - }; + } for (size_t i = 0; i < count; i++) { if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) {