diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index d7ba7dbe2cd..4a282f8266e 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -113,6 +113,7 @@ #include #include #include +#include #include #include #include @@ -148,7 +149,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -1003,6 +1004,13 @@ constexpr std::array kShaderRegistry = {{ kSqrtWorkgroupSizeY, kSqrtWorkgroupSizeZ, }, + { + "streaming_attention_k16_causal_bound", + kStreamingAttentionK16CausalBoundWGSL, + kStreamingAttentionK16CausalBoundWorkgroupSizeX, + kStreamingAttentionK16CausalBoundWorkgroupSizeY, + kStreamingAttentionK16CausalBoundWorkgroupSizeZ, + }, { "tanh", kTanhWGSL, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 3c06d60747f..f36a84cc54f 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -6,17 +6,12 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include -#include -#include -#include -#include -#include #include -#include -#include #include @@ -34,6 +29,16 @@ namespace { constexpr int64_t kSdpaTileM = 4; constexpr int64_t kSdpaTileN = 4; +constexpr const char* kUpdateCacheShader = "update_cache"; +constexpr const char* kUpdateCacheHalfShader = "update_cache_half"; +constexpr const char* kAttnWeightsShader = "sdpa_compute_attn_weights"; +constexpr const char* kAttnWeightsHalfShader = "sdpa_compute_attn_weights_half"; +constexpr const char* kSoftmaxShader = "sdpa_softmax"; +constexpr const char* kComputeOutShader = "sdpa_compute_out"; +constexpr const char* kComputeOutHalfShader = "sdpa_compute_out_half"; +constexpr const char* kStreamingK16Shader = + "streaming_attention_k16_causal_bound"; + // Uniform param structs (all 16-byte aligned, matching the WGSL Params). struct UpdateCacheParams { uint32_t numel; @@ -75,6 +80,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; @@ -83,10 +106,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; }; @@ -148,111 +173,102 @@ static ComputeOutParams make_compute_out_params( return p; } -// A buffer + its byte size, for binding. -struct BufferBinding { - WGPUBuffer buffer; - uint64_t size; -}; +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; +} + +static bool streaming_attention_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} -// Build one dispatch (pipeline + bind group) and record it on the graph. -void build_dispatch( +static 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}; +} + +size_t add_sdpa_compute_dispatch( WebGPUGraph& graph, - const char* wgsl_source, - const BufferBinding* storage_bindings, - uint32_t n_storage, // includes the rw output at index 0 + const char* shader_name, + std::vector bindings, WGPUBuffer uniform_buffer, uint64_t uniform_size, - uint32_t workgroup_count_x, - uint32_t workgroup_count_y, + utils::WgCount grid, uint32_t wg_size, - bool retain_uniform = false, const char* kernel_name = "") { - WGPUDevice device = graph.device(); - - // Bind group layout: storage entries then the uniform. - constexpr uint32_t kMaxEntries = 8; - if (n_storage + 1 > kMaxEntries) { - throw std::runtime_error("WebGPU sdpa: n_storage exceeds kMaxEntries"); - } - const uint32_t uniform_binding = n_storage; - std::vector bindings; - bindings.reserve(n_storage + 1u); - for (uint32_t i = 0; i < n_storage; i++) { - bindings.push_back( - {i, - (i == 0) ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage, - storage_bindings[i].buffer, - storage_bindings[i].size}); - } - bindings.push_back( - {uniform_binding, - WGPUBufferBindingType_Uniform, - uniform_buffer, - uniform_size}); - - // All callers pass an override wg_size; a 0 would keep the shader default. - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, - wgsl_source, - bindings, - wg_size != 0 ? &wg_size_constant : nullptr, - wg_size != 0 ? 1u : 0u); - - graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count_x, - kernel_name, - workgroup_count_y}); - - if (retain_uniform) { - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - } else { - // Drop our ref; the bind group keeps the uniform alive. - wgpuBufferRelease(uniform_buffer); + bindings.push_back({uniform_buffer, 0u, uniform_size}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = std::move(bindings); + if (wg_size != 0) { + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; } + descriptor.grid = {grid.x, grid.y}; + return graph.add_compute_dispatch(descriptor); } // Dispatch one update_cache (K or V); returns the retained uniform buffer. static WGPUBuffer record_update_cache_dispatch( WebGPUGraph& graph, - WGPUDevice device, const WebGPUTensor& cache, const WebGPUTensor& src, uint64_t kv_numel, uint32_t kv_dst_offset, uint64_t cache_numel, uint32_t uc_wg, - bool retain_uniform, const char* label) { const uint32_t wgc = utils::compute_1d_workgroup_count( - device, static_cast(kv_numel), uc_wg, label); - UpdateCacheParams uc = + graph.device(), static_cast(kv_numel), uc_wg, label); + const UpdateCacheParams uc = make_update_cache_params(kv_numel, kv_dst_offset, cache_numel); - WGPUBuffer ubuf = graph.make_uniform_buffer(&uc, sizeof(uc)); - BufferBinding bindings[2] = { - {cache.buffer, cache.nbytes}, {src.buffer, src.nbytes}}; - const char* uc_src = kUpdateCacheWGSL; - if (graph.kv_f16()) { - uc_src = kUpdateCacheHalfWGSL; - } - build_dispatch( + WGPUBuffer ubuf = graph.create_params_buffer(uc); + const std::vector bindings = { + {cache.buffer, 0u, cache.nbytes}, {src.buffer, 0u, src.nbytes}}; + add_sdpa_compute_dispatch( graph, - uc_src, + graph.kv_f16() ? kUpdateCacheHalfShader : kUpdateCacheShader, bindings, - 2, ubuf, sizeof(uc), - wgc, - 1, + {wgc, 1u}, uc_wg, - retain_uniform, "update_cache"); return ubuf; } @@ -296,10 +312,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const size_t cn = k_cache.dims.size(); const int64_t Cmax = k_cache.dims[cn - 3]; - // Validate B == 1 (leading dims must all be 1). - for (size_t i = 0; i + 3 < qn; i++) { - if (q.dims[i] != 1) { - throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + // Validate B == 1 for every tensor (leading dims must all be 1). Rank-3 + // tensors are the equivalent squeezed-batch representation. + for (const WebGPUTensor* tensor : {&q, &k, &v, &k_cache, &v_cache, &out}) { + for (size_t i = 0; i + 3 < tensor->dims.size(); i++) { + if (tensor->dims[i] != 1) { + throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + } } } if (S <= 0 || Hq <= 0 || D <= 0 || Hkv <= 0 || Cmax <= 0) { @@ -335,6 +354,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k_cache.dims != v_cache.dims) { throw std::runtime_error("WebGPU sdpa: k_cache and v_cache shape mismatch"); } + if (k_cache.dims[cn - 2] != Hkv) { + throw std::runtime_error( + "WebGPU sdpa: cache num_heads must match projected k/v"); + } + if (out.dims != q.dims) { + throw std::runtime_error("WebGPU sdpa: output shape must match q"); + } // fp32-only: validate byte counts against fp32 element counts. auto numel = [](const WebGPUTensor& t) { @@ -400,14 +426,26 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } const WGPUDevice device = graph.device(); - const uint32_t uc_wg = - utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); + 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, get_webgpu_shader_info(kUpdateCacheShader).workgroup_size_x); const uint32_t qk_wg = utils::clamp_workgroup_size( - device, kSdpaComputeAttnWeightsWorkgroupSizeX); - const uint32_t av_wg = - utils::clamp_workgroup_size(device, kSdpaComputeOutWorkgroupSizeX); - const uint32_t sm_wg = - utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); + device, get_webgpu_shader_info(kAttnWeightsShader).workgroup_size_x); + const uint32_t av_wg = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kComputeOutShader).workgroup_size_x); + const uint32_t sm_wg = utils::clamp_workgroup_size_pow2( + device, get_webgpu_shader_info(kSoftmaxShader).workgroup_size_x); const bool fd_eligible = D <= kSdpaFdMaxHeadDim; const int64_t pos_const = input_pos; @@ -430,7 +468,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]; @@ -467,21 +506,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"); } @@ -492,16 +518,43 @@ 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; // make_sdpa_fd_decode_state requires D % 4 == 0; the op-level guard above // ("head_dim (D) must be a multiple of 4") rejects any other D before this @@ -514,32 +567,29 @@ 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( graph, - device, k_cache, k, initial_state.update_cache.numel, initial_state.update_cache.dst_offset, initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(K)"); WGPUBuffer uc_v_buf = record_update_cache_dispatch( graph, - device, v_cache, v, initial_state.update_cache.numel, initial_state.update_cache.dst_offset, initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(V)"); const size_t uc_k_idx = graph.num_dispatches() - 2; const size_t uc_v_idx = graph.num_dispatches() - 1; @@ -548,7 +598,9 @@ 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; @@ -566,70 +618,81 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { &graph, attn_weights_softmax); materialized_range.begin = graph.num_dispatches(); - qk_buf = graph.make_uniform_buffer( - &initial_state.attn_weights, sizeof(AttnWeightsParams)); - BufferBinding qk_bindings[3] = { - {attn_weights, aw_bytes}, - {q.buffer, q.nbytes}, - {k_cache.buffer, k_cache.nbytes}}; - const char* qk_src = graph.kv_f16() ? kSdpaComputeAttnWeightsHalfWGSL - : kSdpaComputeAttnWeightsWGSL; - build_dispatch( + qk_buf = graph.create_params_buffer(initial_state.attn_weights); + const std::vector qk_bindings = { + {attn_weights, 0u, aw_bytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - qk_src, + graph.kv_f16() ? kAttnWeightsHalfShader : kAttnWeightsShader, qk_bindings, - 3, qk_buf, sizeof(AttnWeightsParams), - initial_state.qk_grid.x, - initial_state.qk_grid.y, + initial_state.qk_grid, qk_wg, - true, "sdpa_compute_attn_weights"); qk_idx = graph.num_dispatches() - 1; - softmax_buf = graph.make_uniform_buffer( - &initial_state.softmax, sizeof(SoftmaxParams)); - BufferBinding softmax_bindings[2] = { - {attn_weights_softmax, aw_bytes}, {attn_weights, aw_bytes}}; - build_dispatch( + softmax_buf = graph.create_params_buffer(initial_state.softmax); + const std::vector softmax_bindings = { + {attn_weights_softmax, 0u, aw_bytes}, {attn_weights, 0u, aw_bytes}}; + add_sdpa_compute_dispatch( graph, - kSdpaSoftmaxWGSL, + kSoftmaxShader, softmax_bindings, - 2, softmax_buf, sizeof(SoftmaxParams), - initial_state.softmax_grid.x, - initial_state.softmax_grid.y, + initial_state.softmax_grid, sm_wg, - true, "sdpa_softmax"); softmax_idx = graph.num_dispatches() - 1; - av_buf = graph.make_uniform_buffer( - &initial_state.compute_out, sizeof(ComputeOutParams)); - BufferBinding av_bindings[3] = { - {out.buffer, out.nbytes}, - {attn_weights_softmax, aw_bytes}, - {v_cache.buffer, v_cache.nbytes}}; - const char* av_src = - graph.kv_f16() ? kSdpaComputeOutHalfWGSL : kSdpaComputeOutWGSL; - build_dispatch( + av_buf = graph.create_params_buffer(initial_state.compute_out); + const std::vector av_bindings = { + {out.buffer, 0u, out.nbytes}, + {attn_weights_softmax, 0u, aw_bytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - av_src, + graph.kv_f16() ? kComputeOutHalfShader : kComputeOutShader, av_bindings, - 3, av_buf, sizeof(ComputeOutParams), - initial_state.av_grid.x, - initial_state.av_grid.y, + initial_state.av_grid, av_wg, - true, "sdpa_compute_out"); av_idx = graph.num_dispatches() - 1; 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.create_params_buffer(initial_state.streaming_k16); + const std::vector k16_bindings = { + {out.buffer, 0u, out.nbytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + const utils::WgCount initial_grid = initial_state.use_fd + ? utils::WgCount{0u, 0u} + : initial_state.streaming_k16_grid; + add_sdpa_compute_dispatch( + graph, + kStreamingK16Shader, + k16_bindings, + k16_buf, + sizeof(StreamingAttentionK16Params), + initial_grid, + 0, + "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) { @@ -637,14 +700,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, @@ -654,11 +720,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); @@ -686,6 +754,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); } @@ -699,8 +775,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) { @@ -714,6 +792,12 @@ 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 59e816798db..9b0cf65a723 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 @@ -37,6 +39,7 @@ #include #include #include +#include #include #include @@ -497,7 +500,8 @@ void run_sdpa_case( int hq, int hkv, int d, - int cmax) { + int cmax, + float max_error_limit = 2e-3f) { 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"); @@ -535,7 +539,8 @@ 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; } void run_sdpa(Module& m, int s) { @@ -548,6 +553,125 @@ 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"; + const std::set tokens = { + 0, std::min(15, s - 1), std::min(16, s - 1), s - 1}; + for (int token : tokens) { + 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 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_reduce"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_softmax"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), + expect_materialized ? 1 : 0); + 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"); @@ -1184,8 +1308,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"; @@ -1301,6 +1425,27 @@ 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"); + } +} + +TEST(DynamicShape, K16CacheHeadMismatchRejectedAtLoad) { + 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); + Module module(g_dir + "/sdpa_k16_bad_cache_heads.pte"); + EXPECT_NE(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + #ifdef WGPU_BACKEND_ENABLE_PROFILING TEST(DynamicShape, SdpaLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); @@ -1325,6 +1470,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 2cba504afac..89a19f1fabc 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 @@ -318,6 +318,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_qkv_routes(out_dir) _export_dynamic_sdpa_wide(out_dir) @@ -771,6 +772,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 ( @@ -810,6 +822,191 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _write_f32_tensors(base, tensors) -> None: + for name, tensor in tensors: + tensor.detach().numpy().astype(" None: + from executorch.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)) + if cache_hkv is not None: + kc = torch.zeros(1, K16_CMAX, cache_hkv, d) + vc = torch.zeros_like(kc) + 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") + if expect_runtime_reject: + return + + 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.") + _write_f32_tensors( + prime_base, + ( + ("q", prime_q), + ("k", prime_k), + ("v", prime_v), + ("control", torch.zeros(1, 1)), + ("golden", prime_golden), + ), + ) + + for s in live_s: + live_cfg = cfg(s) + q, k, v, kc, vc = _det_inputs(live_cfg) + kc.zero_() + vc.zero_() + kc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_k[0] + vc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_v[0] + golden = reference(live_cfg, q, k, v, kc, vc) + base = os.path.join(out_dir, f"{prefix}.S{s}.") + _write_f32_tensors( + base, + ( + ("q", q), + ("k", k), + ("v", v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ), + ) + print(f" golden {prefix} S={s}") + + +def _export_dynamic_k16_sdpa_cases(out_dir: str) -> 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, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_bad_cache_heads", + K16_HQ, + K16_HKV, + (), + cache_hkv=K16_HKV - 1, + expect_runtime_reject=True, + ) + + def _export_combined_routes(out_dir: str) -> None: from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _make_quantized_model, @@ -1249,10 +1446,32 @@ def test_export_dynamic_rms(self) -> None: "dyn_qkv_bk64_bias.pte", "dyn_qkv_bk64_wrong_width.pte", "dyn_qkv_bk64_different_input.pte", + "sdpa_k16_bad_cache_heads.pte", ] 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__":