From 5f3c4f60e63185cacff4e8eea6c70a43c778d106 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:11:14 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 219 +++++++++++-- .../streaming_attention_k16_causal_bound.wgsl | 268 ++++++++++++++++ ...treaming_attention_k16_causal_bound_wgsl.h | 292 ++++++++++++++++++ .../webgpu/test/native/test_dynamic_shape.cpp | 215 ++++++++++++- .../test_dynamic_shape_export.py | 193 ++++++++++++ 5 files changed, 1155 insertions(+), 32 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 817f96ec19d..a516fda9046 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include #include @@ -14,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -74,6 +76,24 @@ struct ComputeOutParams { }; static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B"); +struct StreamingAttentionK16Params { + uint32_t S; + uint32_t context_len; + uint32_t input_pos; + uint32_t q_token_stride4; + uint32_t q_head_stride4; + uint32_t kv_token_stride4; + uint32_t kv_head_stride4; + uint32_t o_token_stride4; + uint32_t o_head_stride4; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; +}; +static_assert( + sizeof(StreamingAttentionK16Params) == 48, + "StreamingAttentionK16Params must be 48B"); + struct SdpaLiveState { int64_t s; int64_t pos; @@ -82,10 +102,12 @@ struct SdpaLiveState { AttnWeightsParams attn_weights; SoftmaxParams softmax; ComputeOutParams compute_out; + StreamingAttentionK16Params streaming_k16; utils::WgCount update_cache_grid; utils::WgCount qk_grid; utils::WgCount softmax_grid; utils::WgCount av_grid; + utils::WgCount streaming_k16_grid; bool use_fd; SdpaFdDecodeState fd; }; @@ -147,6 +169,56 @@ static ComputeOutParams make_compute_out_params( return p; } +static StreamingAttentionK16Params make_streaming_attention_k16_params( + int64_t S, + int64_t context_len, + int64_t input_pos, + int64_t Hq, + int64_t Hkv, + int64_t D) { + StreamingAttentionK16Params p = {}; + p.S = static_cast(S); + p.context_len = static_cast(context_len); + p.input_pos = static_cast(input_pos); + p.q_token_stride4 = static_cast(Hq * D / 4); + p.q_head_stride4 = static_cast(D / 4); + p.kv_token_stride4 = static_cast(Hkv * D / 4); + p.kv_head_stride4 = static_cast(D / 4); + p.o_token_stride4 = static_cast(Hq * D / 4); + p.o_head_stride4 = static_cast(D / 4); + return p; +} + +bool streaming_attention_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.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} + +utils::WgCount streaming_attention_k16_grid( + WGPUDevice device, + int64_t S, + int64_t Hkv, + int64_t g) { + const uint64_t groups_per_kv = + (static_cast(S) * static_cast(g) + 31u) / 32u; + 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"); + } + if (workgroups > utils::queried_max_workgroups(device)) { + throw std::runtime_error( + "WebGPU sdpa: K16 workgroup count exceeds the 1D dispatch limit"); + } + return {static_cast(workgroups), 1u}; +} + // A buffer + its byte size, for binding. struct BufferBinding { WGPUBuffer buffer; @@ -436,6 +508,23 @@ 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}; + bool k16_buffers_distinct = true; + for (size_t i = 0; i < 6; i++) { + for (size_t j = i + 1; j < 6; j++) { + k16_buffers_distinct = + k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; + } + } + const bool 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); const uint32_t uc_wg = utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); const uint32_t qk_wg = utils::clamp_workgroup_size( @@ -466,7 +555,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { uc_wg, qk_wg, av_wg, - fd_eligible](WebGPUGraph& gr) { + fd_eligible, + k16_eligible](WebGPUGraph& gr) { SdpaLiveState state = {}; const auto& q_live_dims = gr.cur_dims(q_id); state.s = q_live_dims[qn - 3]; @@ -503,21 +593,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(Hkv) * static_cast(D); const uint64_t cache_numel = static_cast(Cmax) * static_cast(Hkv) * static_cast(D); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(state.s) * - 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)); - const uint64_t softmax_rows = - static_cast(Hq) * static_cast(state.s); - const uint64_t av_tiles = static_cast(Hq) * - static_cast(utils::div_up(state.s, kSdpaTileM)) * - static_cast(utils::div_up(D, kSdpaTileN)); if (kv_numel > UINT32_MAX || kv_offset > UINT32_MAX || - cache_numel > UINT32_MAX || aw_floats > UINT32_MAX || - qk_tiles > UINT32_MAX || softmax_rows > UINT32_MAX || - av_tiles > UINT32_MAX) { + cache_numel > UINT32_MAX) { throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32"); } @@ -528,16 +605,44 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { state.softmax = make_softmax_params(Hq, state.s, state.context_len); state.compute_out = make_compute_out_params(state.s, Hq, Hkv, D, state.context_len, g); + 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.update_cache_grid = { utils::compute_1d_workgroup_count( gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"), 1u}; - state.qk_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); - state.softmax_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(softmax_rows), 1, "softmax(resize)"); - state.av_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + if (!k16_eligible) { + const uint64_t aw_floats = static_cast(Hq) * + static_cast(state.s) * + 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)); + const uint64_t softmax_rows = + static_cast(Hq) * static_cast(state.s); + const uint64_t av_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(D, kSdpaTileN)); + if (aw_floats > UINT32_MAX || qk_tiles > UINT32_MAX || + softmax_rows > UINT32_MAX || av_tiles > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: materialized workload exceeds uint32"); + } + state.qk_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + state.softmax_grid = utils::compute_2d_workgroup_count( + gr.device(), + static_cast(softmax_rows), + 1, + "softmax(resize)"); + state.av_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + } state.use_fd = fd_eligible && state.s == 1; if (fd_eligible) { state.fd = make_sdpa_fd_decode_state( @@ -547,9 +652,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { }; const SdpaLiveState initial_state = compute_live_state(graph); - const uint64_t aw_cap_floats = static_cast(Hq) * - static_cast(S) * - static_cast(dynamic_pos ? Cmax : context_len); + const uint64_t aw_cap_floats = k16_eligible + ? 0u + : static_cast(Hq) * static_cast(S) * + static_cast(dynamic_pos ? Cmax : context_len); const uint64_t aw_bytes = aw_cap_floats * sizeof(float); WGPUBuffer uc_k_buf = record_update_cache_dispatch( @@ -581,7 +687,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph.tensor_has_dynamic_dims(v_id); const bool dual_route = utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence); - const bool record_materialized = dual_route || !initial_state.use_fd; + 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; WGPUBuffer qk_buf = nullptr; @@ -663,6 +772,37 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { materialized_range.end = graph.num_dispatches(); } + WGPUBuffer k16_buf = nullptr; + size_t k16_idx = 0; + utils::DispatchRange k16_range = {}; + if (record_k16) { + k16_range.begin = graph.num_dispatches(); + k16_buf = graph.make_uniform_buffer( + &initial_state.streaming_k16, sizeof(StreamingAttentionK16Params)); + BufferBinding k16_bindings[4] = { + {out.buffer, out.nbytes}, + {q.buffer, q.nbytes}, + {k_cache.buffer, k_cache.nbytes}, + {v_cache.buffer, v_cache.nbytes}}; + const utils::WgCount initial_grid = initial_state.use_fd + ? utils::WgCount{0u, 0u} + : initial_state.streaming_k16_grid; + build_dispatch( + graph, + kStreamingAttentionK16CausalBoundWGSL, + k16_bindings, + 4, + k16_buf, + sizeof(StreamingAttentionK16Params), + initial_grid.x, + initial_grid.y, + 0, + true, + "sdpa_streaming_attention_k16_causal_bound"); + k16_idx = graph.num_dispatches() - 1; + k16_range.end = graph.num_dispatches(); + } + SdpaFdDecodeResources fd_resources = {}; size_t route_group = 0; if (record_fd) { @@ -670,14 +810,17 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph, q, k_cache, v_cache, out, initial_state.fd); } if (dual_route) { + const utils::DispatchRange prefill_range = + record_k16 ? k16_range : materialized_range; route_group = graph.register_dispatch_route_group( - {materialized_range, fd_resources.dispatch_range}); + {prefill_range, fd_resources.dispatch_range}); } auto refresh_state = [compute_live_state, q_id, out_id, dual_route, + record_k16, record_materialized, record_fd, fixed_use_fd = initial_state.use_fd, @@ -687,11 +830,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { qk_idx, softmax_idx, av_idx, + k16_idx, uc_k_buf, uc_v_buf, qk_buf, softmax_buf, av_buf, + k16_buf, fd_resources](WebGPUGraph& gr) { const SdpaLiveState state = compute_live_state(gr); @@ -719,6 +864,14 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { wgpuQueueWriteBuffer( gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out)); } + if (record_k16) { + wgpuQueueWriteBuffer( + gr.queue(), + k16_buf, + 0, + &state.streaming_k16, + sizeof(state.streaming_k16)); + } if (record_fd) { write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd); } @@ -732,8 +885,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const std::vector active_grids = state.use_fd ? std::vector< utils::WgCount>{state.fd.split_grid, state.fd.reduce_grid} - : std::vector{ - state.qk_grid, state.softmax_grid, state.av_grid}; + : (record_k16 + ? std::vector{state.streaming_k16_grid} + : std::vector{ + state.qk_grid, state.softmax_grid, state.av_grid}); gr.select_dispatch_route(route_group, active_route, active_grids); } else if (state.use_fd) { if (!fixed_use_fd) { @@ -747,6 +902,14 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { state.fd.reduce_grid.x; gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y = state.fd.reduce_grid.y; + } else if (record_k16) { + 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; } else { if (fixed_use_fd) { throw std::runtime_error("WebGPU sdpa: static route changed"); diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5b8cf1de03c --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl @@ -0,0 +1,268 @@ +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 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +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_k_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 max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 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 * 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: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + 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 * 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_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + 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), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + 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 tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + 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); + 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; + } + } + 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 * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + 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_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + 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_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..255c6afc6e0 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h @@ -0,0 +1,292 @@ +/* + * 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_k16_causal_bound.wgsl - DO NOT EDIT. +// wgsl-sha256: b1435c4f72834cb896eb4248a899aedeabbdcc772818afc419feca03e8957ffd +inline constexpr const char* kStreamingAttentionK16CausalBoundWGSL = 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 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +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_k_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 max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 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 * 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: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + 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 * 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_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + 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), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + 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 tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + 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); + 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; + } + } + 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 * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + 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_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + 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 kStreamingAttentionK16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeY = 4; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 81e04d19877..7e9844d7ff5 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -27,6 +27,8 @@ #include #include #include +#include +#include #include @@ -412,7 +414,9 @@ void run_sdpa_case( int hq, int hkv, int d, - int cmax) { + int cmax, + float max_error_limit = 2e-3f, + bool check_causal_boundaries = false) { const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(s) + "."; auto q = read_bin(b + "q.bin"); auto k = read_bin(b + "k.bin"); @@ -450,7 +454,22 @@ void run_sdpa_case( << s << "," << hq << "," << d << "]"; std::vector got(attn, attn + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 2e-3f) << "sdpa_dyn S=" << s << " max_err=" << e; + EXPECT_LT(e, max_error_limit) + << prefix << " S=" << s << " full-output max_err=" << e; + if (check_causal_boundaries) { + const size_t token_width = static_cast(hq) * d; + for (int token : {0, std::min(15, s - 1), std::min(16, s - 1), s - 1}) { + const size_t begin = static_cast(token) * token_width; + const size_t end = begin + token_width; + float token_error = 0.0f; + for (size_t i = begin; i < end; ++i) { + token_error = std::fmax(token_error, std::fabs(got[i] - golden[i])); + } + EXPECT_LT(token_error, max_error_limit) + << prefix << " S=" << s << " causal token=" << token + << " max_err=" << token_error; + } + } } void run_sdpa(Module& m, int s) { @@ -463,6 +482,121 @@ void check_sdpa(int s) { run_sdpa(m, s); } +constexpr int kK16Hq = 32; +constexpr int kK16Hkv = 8; +constexpr int kK16D = 64; + +bool k16_device_supported() { + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} + +void load_sdpa_module(Module& module, bool f16_kv) { + if (!f16_kv) { + ASSERT_EQ(module.load_forward(), Error::Ok); + return; + } + executorch::runtime::BackendOptions<1> options; + ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); + executorch::runtime::LoadBackendOptionsMap option_map; + ASSERT_EQ(option_map.set_options("VulkanBackend", options.view()), Error::Ok); + ASSERT_EQ(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + +void run_k16_sdpa( + Module& module, + int s, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D, + bool prime = false) { + const std::string base = g_dir + "/" + prefix + + (prime ? ".prime." : ".S" + std::to_string(s) + "."); + auto q = read_bin(base + "q.bin"); + auto k = read_bin(base + "k.bin"); + auto v = read_bin(base + "v.bin"); + auto control = read_bin(base + "control.bin"); + auto golden = read_bin(base + "golden.bin"); + ASSERT_FALSE( + q.empty() || k.empty() || v.empty() || control.empty() || golden.empty()); + auto tq = make_tensor_ptr({1, s, hq, d}, std::move(q)); + auto tk = make_tensor_ptr({1, s, hkv, d}, std::move(k)); + auto tv = make_tensor_ptr({1, s, hkv, d}, std::move(v)); + const int control_size = static_cast(control.size()); + auto tcontrol = make_tensor_ptr({1, control_size}, std::move(control)); + auto result = + module.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tcontrol)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << prefix << " S=" << s << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const size_t token_width = static_cast(hq) * d; + const size_t numel = static_cast(s) * token_width; + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + ASSERT_EQ(got.size(), golden.size()); + constexpr float kMaxError = 3e-3f; + EXPECT_LT(max_err(got, golden), kMaxError) + << prefix << " S=" << s << " full output"; + for (int token : {0, std::min(15, s - 1), std::min(16, s - 1), s - 1}) { + float token_error = 0.0f; + const size_t begin = static_cast(token) * token_width; + for (size_t i = begin; i < begin + token_width; ++i) { + token_error = std::fmax(token_error, std::fabs(got[i] - golden[i])); + } + EXPECT_LT(token_error, kMaxError) + << prefix << " S=" << s << " causal token=" << token; + } +} + +void prime_k16_sdpa( + Module& module, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D) { + run_k16_sdpa(module, 12, prefix, hq, hkv, d, true); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void expect_sdpa_route( + const std::vector& names, + int s, + bool expect_k16) { + 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); + EXPECT_EQ( + std::count( + names.begin(), + names.end(), + "sdpa_streaming_attention_k16_causal_bound"), + expect_k16); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), expect_fd); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), expect_fd); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + expect_materialized); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_softmax"), + expect_materialized); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), + expect_materialized); + EXPECT_EQ( + names.size(), + static_cast(2 + (expect_k16 ? 1 : (expect_fd ? 2 : 3)))); +} +#endif + void run_combined_routes(Module& m, int s) { const std::string b = g_dir + "/combined_routes.S" + std::to_string(s) + "."; auto x = read_bin(b + "x.bin"); @@ -1030,8 +1164,8 @@ TEST(DynamicShape, QuantizedLinearBk64QkvProfileSoleWriter) { TEST(DynamicShape, CombinedLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || - !context->timestamp_supported) { - GTEST_SKIP() << "timestamp queries unavailable"; + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; } Module m(g_dir + "/combined_routes.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; @@ -1147,6 +1281,18 @@ TEST(DynamicShape, SdpaWideMaterializedOnly) { } } +TEST(DynamicShape, K16CausalNumericsReusedGraph) { + if (!k16_device_supported()) { + GTEST_SKIP() << "K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + for (int s : {512, 1, 508, 128, 127, 16, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + } +} + #ifdef WGPU_BACKEND_ENABLE_PROFILING TEST(DynamicShape, SdpaLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); @@ -1171,6 +1317,67 @@ TEST(DynamicShape, SdpaLiveRoutesProfile) { std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); } } + +TEST(DynamicShape, K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, true); + for (int s : {512, 128, 1, 508, 127, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, s > 1); + } +} + +TEST(DynamicShape, K16F32KvFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, false); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, false); + } +} + +TEST(DynamicShape, K16MetadataFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + struct NegativeCase { + const char* prefix; + int hq; + int hkv; + int d; + }; + for (const auto& negative : std::vector{ + {"sdpa_k16_wrong_geometry", 14, 2, 64}, + {"sdpa_k16_wrong_d", 32, 8, 128}, + {"sdpa_k16_wrong_scale", 32, 8, 64}}) { + Module module(g_dir + "/" + negative.prefix + ".pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, s, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), s, false); + } + } +} #endif // K: dynamic embedding (int64 token ids) at several token counts. diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index e712937bc9c..36cbf5b159b 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -284,6 +284,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) + _export_dynamic_k16_sdpa_cases(out_dir) _export_combined_routes(out_dir) _export_dynamic_sdpa_wide(out_dir) _export_static_sdpa(out_dir, 1, "static_sdpa_s1") @@ -724,6 +725,17 @@ def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: SD_CMAX = 64 SD_MAXS = 64 +K16_HQ = 32 +K16_HKV = 8 +K16_D = 64 +K16_CMAX = 525 +K16_MAXS = 512 +K16_INPUT_POS = 13 +K16_PRIME_POS = 1 +K16_PRIME_S = K16_INPUT_POS - K16_PRIME_POS +K16_POS_CONTROL = K16_INPUT_POS +K16_LIVE_S = (K16_MAXS, 508, 128, 127, 16, 1) + def _export_dynamic_sdpa(out_dir: str) -> None: from executorch.backends.webgpu.test.ops.test_sdpa import ( @@ -763,6 +775,166 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _export_dynamic_k16_sdpa_case( + out_dir: str, + prefix: str, + hq: int, + hkv: int, + live_s, + d: int = K16_D, + scale: float | None = None, +) -> None: + from backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + ) + + def cfg(s: int) -> "SdpaConfig": + return SdpaConfig( + prefix, + hq, + hkv, + d, + s, + K16_CMAX, + K16_INPUT_POS, + ) + + q, k, v, kc, vc = _det_inputs(cfg(K16_MAXS)) + kc.zero_() + vc.zero_() + + class K16SdpaModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("k_cache", kc) + self.register_buffer("v_cache", vc) + + def forward(self, q, k, v, pos_control): + input_pos = pos_control.shape[1] + return torch.ops.llama.sdpa_with_kv_cache( + q, + k, + v, + self.k_cache, + self.v_cache, + input_pos, + q.shape[1], + None, + 0.0, + True, + scale, + ) + + model = K16SdpaModule().eval() + inputs = (q, k, v, torch.zeros(1, K16_POS_CONTROL)) + s_dim = torch.export.Dim(f"{prefix}_s", min=1, max=K16_MAXS) + pos_dim = torch.export.Dim(f"{prefix}_pos_control", min=1, max=K16_POS_CONTROL) + ep = torch.export.export( + model, + inputs, + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, {1: pos_dim}), + ) + sym_sizes = [ + node for node in ep.graph.nodes if node.target == torch.ops.aten.sym_size.int + ] + if len(sym_sizes) < 2: + raise RuntimeError(f"{prefix}: dynamic K16 fixture lost S/position symbols") + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + + def reference(live_cfg, q, k, v, kc, vc): + if scale is None: + return _golden(live_cfg, q, k, v, kc, vc) + context_len = live_cfg.s + live_cfg.input_pos + g = hq // hkv + k_full = torch.cat((kc[0, : live_cfg.input_pos].double(), k[0].double()), dim=0) + v_full = torch.cat((vc[0, : live_cfg.input_pos].double(), v[0].double()), dim=0) + q_heads = q[0].double().transpose(0, 1) + k_heads = k_full.repeat_interleave(g, dim=1).transpose(0, 1) + v_heads = v_full.repeat_interleave(g, dim=1).transpose(0, 1) + mask = torch.full((live_cfg.s, context_len), float("-inf"), dtype=torch.float64) + for token in range(live_cfg.s): + mask[token, : live_cfg.input_pos + token + 1] = 0.0 + golden = torch.nn.functional.scaled_dot_product_attention( + q_heads, k_heads, v_heads, attn_mask=mask, scale=scale + ) + return golden.transpose(0, 1).reshape(1, live_cfg.s, hq, d).float().contiguous() + + prime_cfg = SdpaConfig(prefix, hq, hkv, d, K16_PRIME_S, K16_CMAX, K16_PRIME_POS) + prime_q, prime_k, prime_v, prime_kc, prime_vc = _det_inputs(prime_cfg) + prime_kc.zero_() + prime_vc.zero_() + prime_golden = reference(prime_cfg, prime_q, prime_k, prime_v, prime_kc, prime_vc) + prime_base = os.path.join(out_dir, f"{prefix}.prime.") + for name, tensor in ( + ("q", prime_q), + ("k", prime_k), + ("v", prime_v), + ("control", torch.zeros(1, 1)), + ("golden", prime_golden), + ): + tensor.detach().numpy().astype(" None: + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_llama", + K16_HQ, + K16_HKV, + K16_LIVE_S, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_geometry", + 14, + 2, + (128, 1), + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_d", + K16_HQ, + K16_HKV, + (128, 1), + d=128, + scale=0.125, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_scale", + K16_HQ, + K16_HKV, + (128, 1), + scale=0.25, + ) + + def _export_combined_routes(out_dir: str) -> None: from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _make_quantized_model, @@ -1070,6 +1242,27 @@ def test_export_dynamic_rms(self) -> None: for name in expected: with self.subTest(artifact=name): self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for prefix, live_s in ( + ("sdpa_k16_llama", K16_LIVE_S), + ("sdpa_k16_wrong_geometry", (128, 1)), + ("sdpa_k16_wrong_d", (128, 1)), + ("sdpa_k16_wrong_scale", (128, 1)), + ): + with self.subTest(artifact=f"{prefix}.pte"): + self.assertGreater( + os.path.getsize(os.path.join(d, f"{prefix}.pte")), 0 + ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.prime.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for s in live_s: + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.S{s}.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater( + os.path.getsize(os.path.join(d, name)), 0 + ) if __name__ == "__main__":