diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index bedde5abdb7..35d4225bc7c 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -90,7 +90,7 @@ Result WebGPUBackend::init( // Parse header to locate flatbuffer and constant data Result header = - WebGPUDelegateHeader::parse(processed->data()); + WebGPUDelegateHeader::parse(processed->data(), processed->size()); if (!header.ok()) { ET_LOG(Error, "WebGPUDelegateHeader may be corrupt"); return header.error(); @@ -101,6 +101,17 @@ Result WebGPUBackend::init( const uint8_t* flatbuffer_data = buffer_start + header->flatbuffer_offset; const uint8_t* constant_data = buffer_start + header->bytes_offset; + size_t constant_data_size = header->bytes_size; + if (constant_data_size == 0 && processed->size() > header->bytes_offset) { + constant_data_size = processed->size() - header->bytes_offset; + } + + flatbuffers::Verifier verifier(flatbuffer_data, header->flatbuffer_size); + if (!vkgraph::VerifyVkGraphBuffer(verifier)) { + ET_LOG(Error, "WebGPU delegate FlatBuffer verification failed"); + return Error::DelegateInvalidCompatibility; + } + // Verify FlatBuffer identifier if (!vkgraph::VkGraphBufferHasIdentifier(flatbuffer_data)) { ET_LOG( @@ -125,10 +136,20 @@ Result WebGPUBackend::init( config.f16_accumulate_gemm = spec.get(); } } + { + Result spec = context.get_runtime_spec("sdpa_query_tile"); + if (spec.ok()) { + config.sdpa_query_tile = spec.get(); + } + } try { graph->build( - flatbuffer_data, constant_data, context.get_named_data_map(), config); + flatbuffer_data, + constant_data, + constant_data_size, + context.get_named_data_map(), + config); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); @@ -163,8 +184,13 @@ Error WebGPUBackend::execute( const auto& tensor = args[i]->toTensor(); const bool host_is_int64 = tensor.scalar_type() == executorch::aten::ScalarType::Long; + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; inputs.push_back( - {tensor.const_data_ptr(), tensor.nbytes(), host_is_int64}); + {tensor.const_data_ptr(), + tensor.nbytes(), + host_is_int64, + host_is_fp32}); const auto sizes = tensor.sizes(); std::vector new_dims(sizes.begin(), sizes.end()); graph->resize_input(graph->input_ids()[i], new_dims); @@ -205,12 +231,15 @@ Error WebGPUBackend::execute( graph->execute(plan); // Copy outputs from GPU staging buffers to EValue tensor data pointers - std::vector> outputs; + std::vector outputs; outputs.reserve(num_outputs); for (size_t i = 0; i < num_outputs; i++) { const size_t arg_idx = num_inputs + i; auto& tensor = args[arg_idx]->toTensor(); - outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; + outputs.push_back( + {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32}); } graph->copy_outputs(outputs, plan); } catch (const std::exception& e) { diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp index d1e8b2110a7..69e6dd70536 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp @@ -65,13 +65,19 @@ bool WebGPUDelegateHeader::is_valid() const { if (flatbuffer_size == 0) { return false; } - if (bytes_offset < flatbuffer_offset + flatbuffer_size) { + if (bytes_offset < flatbuffer_offset || + flatbuffer_size > bytes_offset - flatbuffer_offset) { return false; } return true; } -Result WebGPUDelegateHeader::parse(const void* data) { +Result WebGPUDelegateHeader::parse( + const void* data, + size_t buffer_size) { + if (data == nullptr || buffer_size < kExpectedSize) { + return Error::InvalidArgument; + } const uint8_t* header_data = (const uint8_t*)data; const uint8_t* magic_start = header_data + kMagic.offset; @@ -91,6 +97,13 @@ Result WebGPUDelegateHeader::parse(const void* data) { return Error::InvalidArgument; } + if (header.flatbuffer_offset > buffer_size || + header.flatbuffer_size > buffer_size - header.flatbuffer_offset || + header.bytes_offset > buffer_size || + header.bytes_size > buffer_size - header.bytes_offset) { + return Error::InvalidArgument; + } + return header; } diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.h b/backends/webgpu/runtime/WebGPUDelegateHeader.h index 6f2f65130c7..14dcfd14a9a 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.h +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.h @@ -8,6 +8,8 @@ #pragma once +#include + #include namespace executorch { @@ -18,7 +20,8 @@ struct WebGPUDelegateHeader { bool is_valid() const; static executorch::runtime::Result parse( - const void* data); + const void* data, + size_t buffer_size); uint32_t header_size; uint32_t flatbuffer_offset; diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 859e01e7a40..58d7178ddbf 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -34,6 +35,19 @@ namespace executorch::backends::webgpu { namespace { +const uint8_t* checked_inline_constant( + const uint8_t* data, + size_t data_size, + uint64_t offset, + size_t required_size, + const char* error_message) { + if (data == nullptr || offset > data_size || + required_size > data_size - static_cast(offset)) { + throw std::runtime_error(error_message); + } + return data + static_cast(offset); +} + class ScopedBindGroupLayout final { public: explicit ScopedBindGroupLayout(WGPUBindGroupLayout handle) @@ -713,6 +727,7 @@ WebGPUGraph::~WebGPUGraph() { void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map, WebGPUGraphConfig config) { if (!device_) { @@ -733,6 +748,7 @@ void WebGPUGraph::build( // .pte byte sources for prepack-time constant materialization (build-only). constant_data_ = constant_data; + constant_data_size_ = constant_data_size; named_data_map_ = named_data_map; // f16 KV cache (runtime opt-in): store K/V caches as f16 iff the opt-in is @@ -817,11 +833,27 @@ void WebGPUGraph::build( continue; } for (unsigned j = 0; j < a->size(); j++) { - if (kv_cache_ids_.count(static_cast(a->Get(j))) != 0) { + const int id = static_cast(a->Get(j)); + if (kv_cache_ids_.count(id) != 0) { throw std::runtime_error( "WebGPU f16 KV: cache tensor consumed by non-sdpa op '" + nm + "' would misread the f16 buffer"); } + const auto* value = values ? values->Get(id) : nullptr; + if (value && value->value_type() == vkgraph::GraphTypes::ValueList) { + const auto* items = value->value_as_ValueList()->items(); + if (!items) { + continue; + } + for (unsigned k = 0; k < items->size(); k++) { + if (kv_cache_ids_.count(static_cast(items->Get(k))) != 0) { + throw std::runtime_error( + "WebGPU f16 KV: cache tensor consumed through a ValueList " + "by non-sdpa op '" + + nm + "' would misread the f16 buffer"); + } + } + } } } } @@ -843,11 +875,20 @@ void WebGPUGraph::build( size_t numel = 1; if (dims) { for (unsigned j = 0; j < dims->size(); j++) { - tensor.dims.push_back(static_cast(dims->Get(j))); - numel *= dims->Get(j); + const uint32_t dim = dims->Get(j); + tensor.dims.push_back(static_cast(dim)); + if (dim != 0 && numel > std::numeric_limits::max() / dim) { + throw std::runtime_error( + "WebGPU: tensor element count overflows"); + } + numel *= dim; } } tensor.elem_size = vk_datatype_size(vk_tensor->datatype()); + if (tensor.elem_size != 0 && + numel > std::numeric_limits::max() / tensor.elem_size) { + throw std::runtime_error("WebGPU: tensor byte size overflows"); + } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; @@ -860,6 +901,11 @@ void WebGPUGraph::build( // zero-initializes freshly-created buffers, so no explicit clear is // needed. Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && kv_cache_ids_.count(i) != 0) { + if (tensor.is_int || tensor.elem_size != sizeof(float) || + tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU f16 KV: serialized cache tensor must be fp32"); + } tensor.elem_size = 2; tensor.nbytes = numel * 2; tensor.cur_nbytes = tensor.nbytes; @@ -870,6 +916,63 @@ void WebGPUGraph::build( WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; tensor.buffer = wgpuDeviceCreateBuffer(device_, &buf_desc); + + // Mutable caches normally start empty. If the serialized graph owns + // an initialized cache constant, preserve it while changing storage + // representation instead of silently replacing it with zeros. + const int cache_constant_id = vk_tensor->constant_id(); + if (cache_constant_id >= 0) { + const auto* constants = graph->constants(); + if (!constants || + cache_constant_id >= static_cast(constants->size())) { + throw std::runtime_error( + "WebGPU f16 KV: cache constant id is out of range"); + } + const auto* bytes = constants->Get(cache_constant_id); + auto write_fp16_cache = [&](const uint8_t* src) { + std::vector converted(numel); + for (size_t e = 0; e < numel; e++) { + float value = 0.0f; + std::memcpy(&value, src + e * sizeof(float), sizeof(float)); + converted[e] = executorch::runtime::etensor::Half(value); + } + wgpuQueueWriteBuffer( + queue_, + tensor.buffer, + 0, + converted.data(), + converted.size() * sizeof(converted[0])); + }; + if (bytes->offset() != UINT64_MAX) { + write_fp16_cache(checked_inline_constant( + constant_data_, + constant_data_size_, + bytes->offset(), + numel * sizeof(float), + "WebGPU f16 KV: inline cache constant exceeds constant " + "data")); + } else if ( + bytes->named_key() != nullptr && named_data_map_ != nullptr) { + const std::string key = bytes->named_key()->str(); + auto data = named_data_map_->get_data(key.c_str()); + if (!data.ok()) { + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' not found"); + } + if (data->size() < numel * sizeof(float)) { + data->Free(); + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' is undersized"); + } + write_fp16_cache(static_cast(data->data())); + data->Free(); + } else { + throw std::runtime_error( + "WebGPU f16 KV: cache constant has no readable source"); + } + } break; } @@ -1190,6 +1293,7 @@ void WebGPUGraph::build( // The .pte bytes are freed right after build() returns (WebGPUBackend // processed->Free()), so clear the build-only source pointers. constant_data_ = nullptr; + constant_data_size_ = 0; named_data_map_ = nullptr; } @@ -1201,15 +1305,18 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { std::to_string(const_value_id)); } const ConstantSource& cs = it->second; - if (cs.nbytes == 0) { - return; - } if (cs.inline_offset != UINT64_MAX) { - if (constant_data_ == nullptr) { - throw std::runtime_error("WebGPU: inline constant data is null"); - } - wgpuQueueWriteBuffer( - queue_, dst, 0, constant_data_ + cs.inline_offset, cs.nbytes); + const uint8_t* data = checked_inline_constant( + constant_data_, + constant_data_size_, + cs.inline_offset, + cs.nbytes, + "WebGPU: inline constant exceeds constant data"); + if (cs.nbytes != 0) { + wgpuQueueWriteBuffer(queue_, dst, 0, data, cs.nbytes); + } + } else if (cs.nbytes == 0) { + return; } else if (!cs.named_key.empty() && named_data_map_ != nullptr) { auto buf = named_data_map_->get_data(cs.named_key.c_str()); if (!buf.ok()) { @@ -1304,6 +1411,11 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Upload only the live (cur) bytes, not the max allocation; cur_nbytes == // nbytes on a static graph, so this is byte-identical there. const size_t live_nbytes = tensor.cur_nbytes; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !in.host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device input requires an fp32 host tensor"); + } // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { @@ -1332,6 +1444,21 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { continue; } + // Require an explicit fp32 host dtype, not merely "not int64": inferring + // the narrow from the 2:1 byte ratio alone would silently reinterpret a + // same-sized non-fp32 host buffer (e.g. a stale int32) as fp32. + if (in.host_is_fp32 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + wgpuQueueWriteBuffer( + queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + continue; + } + throw std::runtime_error( "WebGPU: unsupported input copy for input " + std::to_string(i) + " (host " + std::to_string(in.nbytes) + " bytes" + @@ -1366,6 +1493,8 @@ constexpr uint32_t kRouteGenericFallback = 1u << 8; constexpr uint32_t kRouteFlashDecoding = 1u << 10; constexpr uint32_t kRouteK16CausalBound = 1u << 11; constexpr uint32_t kRouteBicolSubgroup = 1u << 12; +constexpr uint32_t kRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kRouteQwen3Q32K16 = 1u << 14; #endif // WGPU_BACKEND_ENABLE_PROFILING // Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the @@ -1382,7 +1511,13 @@ bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING void WebGPUGraph::record_active_route(const std::string& kernel_name) { uint32_t bits = 0; - if (kernel_name == "sdpa_streaming_attention_k16_causal_bound") { + if (kernel_name == "sdpa_streaming_attention_qwen3_q32_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q32K16; + } else if (kernel_name == "sdpa_streaming_attention_qwen3_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q16K16; + } else if ( + kernel_name.rfind("sdpa_streaming_attention_", 0) == 0 && + kernel_name.find("k16_causal_bound") != std::string::npos) { bits = kRoutePrefill | kRouteK16CausalBound; } else if (kernel_name == "sdpa_streaming_attention_k16") { bits = kRoutePrefill | kRouteK16; @@ -1665,28 +1800,59 @@ void buffer_map_callback( } // namespace void WebGPUGraph::copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUExecutionPlan& plan) { if (plan.copy_outputs.size() != output_copies_.size()) { throw std::runtime_error("WebGPU: execution plan output count mismatch"); } const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); + // Reject all dtype/size mismatches before issuing an asynchronous map. for (size_t i = 0; i < count; i++) { - if (!plan.copy_outputs[i] || outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { continue; } - const size_t map_nbytes = tensors_[output_ids_[i]].cur_nbytes; + const auto& tensor = tensors_[output_ids_[i]]; + const size_t map_nbytes = tensor.cur_nbytes; if (map_nbytes == 0) { continue; } - const size_t dst_nbytes = outputs[i].second; - const bool widen_int32 = dst_nbytes == 2 * map_nbytes && - tensors_[output_ids_[i]].is_int && - tensors_[output_ids_[i]].elem_size == 4; - if (dst_nbytes != map_nbytes && !widen_int32) { + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !outputs[i].host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device output requires an fp32 host tensor"); + } + if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { + throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); + } + if (dst_nbytes != map_nbytes && !widen_fp16 && !widen_int32) { throw std::runtime_error("WebGPU: output buffer size mismatch"); } + } + + for (size_t i = 0; i < count; i++) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { + continue; + } + const auto& tensor = tensors_[output_ids_[i]]; + const size_t map_nbytes = tensor.cur_nbytes; + if (map_nbytes == 0) { + continue; + } + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; const auto cb_data = std::make_shared(); WGPUBufferMapCallbackInfo cb_info = {}; @@ -1717,16 +1883,24 @@ void WebGPUGraph::copy_outputs( wgpuBufferUnmap(output_staging_buffers_[i]); throw std::runtime_error("WebGPU mapped output range is null"); } - if (!widen_int32) { - std::memcpy(outputs[i].first, mapped, map_nbytes); - } else { + if (widen_fp16) { + const auto* src = + static_cast(mapped); + auto* dst = static_cast(outputs[i].data); + const size_t n = map_nbytes / sizeof(*src); + for (size_t k = 0; k < n; k++) { + dst[k] = static_cast(src[k]); + } + } else if (widen_int32) { // int64 host output backed by an int32 GPU buffer: widen (sign-extend). const int32_t* src = static_cast(mapped); - int64_t* dst = static_cast(outputs[i].first); + int64_t* dst = static_cast(outputs[i].data); const size_t n = map_nbytes / sizeof(int32_t); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } + } else { + std::memcpy(outputs[i].data, mapped, map_nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 292f2ee0cf4..a634ebb4172 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -48,6 +48,15 @@ struct InputData { const void* data = nullptr; size_t nbytes = 0; bool host_is_int64 = false; + bool host_is_fp32 = false; +}; + +// Host destination for a graph output. host_is_fp32 gates the fp16->fp32 widen +// on readback (mirrors InputData's guard on the copy_inputs narrow path). +struct OutputData { + void* data = nullptr; + size_t nbytes = 0; + bool host_is_fp32 = false; }; struct WebGPUDispatch { @@ -130,6 +139,7 @@ struct WebGPUMemoryStats { struct WebGPUGraphConfig { bool f16_kv_cache = false; bool f16_accumulate_gemm = false; + int sdpa_query_tile = 0; bool record_q4gsw_decode_route = false; }; @@ -143,6 +153,7 @@ class WebGPUGraph { void build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map = nullptr, WebGPUGraphConfig config = {}); @@ -159,7 +170,7 @@ class WebGPUGraph { // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. void copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUExecutionPlan& plan); const std::vector& input_ids() const { @@ -541,6 +552,12 @@ class WebGPUGraph { return config_.f16_accumulate_gemm; } + // Runtime-selected SDPA query-tile candidate; 0 = geometry default (Q16), + // 32 = Q32 candidate. + int sdpa_query_tile() const { + return config_.sdpa_query_tile; + } + const WebGPUGraphConfig& config() const { return config_; } @@ -659,6 +676,7 @@ class WebGPUGraph { // materializes these once. constant_data_/named_data_map_ point at the .pte // bytes and are valid only during build(). const uint8_t* constant_data_ = nullptr; + size_t constant_data_size_ = 0; const executorch::runtime::NamedDataMap* named_data_map_ = nullptr; std::unordered_map constant_sources_; diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 4a282f8266e..1cf4b952a5f 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -114,6 +114,8 @@ #include #include #include +#include +#include #include #include #include @@ -149,7 +151,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -1011,6 +1013,20 @@ constexpr std::array kShaderRegistry = {{ kStreamingAttentionK16CausalBoundWorkgroupSizeY, kStreamingAttentionK16CausalBoundWorkgroupSizeZ, }, + { + "streaming_attention_qwen3_k16_causal_bound", + kStreamingAttentionQwen3K16CausalBoundWGSL, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ, + }, + { + "streaming_attention_qwen3_q32_k16_causal_bound", + kStreamingAttentionQwen3Q32K16CausalBoundWGSL, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ, + }, { "tanh", kTanhWGSL, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 9a51f853caf..bcb5deaa770 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -38,6 +38,10 @@ constexpr const char* kComputeOutShader = "sdpa_compute_out"; constexpr const char* kComputeOutHalfShader = "sdpa_compute_out_half"; constexpr const char* kStreamingK16Shader = "streaming_attention_k16_causal_bound"; +constexpr const char* kStreamingQwen3K16Shader = + "streaming_attention_qwen3_k16_causal_bound"; +constexpr const char* kStreamingQwen3Q32K16Shader = + "streaming_attention_qwen3_q32_k16_causal_bound"; // Uniform param structs (all 16-byte aligned, matching the WGSL Params). struct UpdateCacheParams { @@ -205,13 +209,87 @@ static bool streaming_attention_k16_device_supported(WGPUDevice device) { limits.maxComputeWorkgroupStorageSize >= 14720u; } +constexpr uint32_t kLlamaK16QueryTile = 32u; +constexpr uint32_t kQwen3K16QueryTile = 16u; +constexpr uint32_t kQwen3Q32K16QueryTile = 32u; +constexpr uint32_t kQwen3Q16K16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); +constexpr uint32_t kQwen3K16StorageBytes = kQwen3Q16K16StorageBytes; +// Mirrors the Q32 shader's workgroup arrays (t_q_tile vec4x1024, t_kv_tile +// vec4x512, t_scores vec2x256, t_m/t_d/t_alpha f32x32) so the +// device-support gate stays tied to the declared storage, not a literal. +constexpr uint32_t kQwen3Q32K16StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + +constexpr bool streaming_attention_k16_workgroup_count_fits( + int64_t S, + int64_t Hkv, + int64_t g, + uint32_t query_tile, + uint32_t max_workgroups) { + if (S <= 0 || Hkv <= 0 || g <= 0 || query_tile == 0u || + max_workgroups == 0u) { + return false; + } + if (static_cast(S) > UINT64_MAX / static_cast(g)) { + return false; + } + const uint64_t logical_rows = + static_cast(S) * static_cast(g); + if (logical_rows > UINT64_MAX - (query_tile - 1u)) { + return false; + } + const uint64_t groups_per_kv = (logical_rows + query_tile - 1u) / query_tile; + if (groups_per_kv > UINT64_MAX / static_cast(Hkv)) { + return false; + } + const uint64_t workgroups = groups_per_kv * static_cast(Hkv); + return workgroups > 0u && workgroups <= UINT32_MAX && + workgroups <= max_workgroups; +} + +static_assert( + streaming_attention_k16_workgroup_count_fits(65528, 8, 2, 16, 65535)); +static_assert( + !streaming_attention_k16_workgroup_count_fits(65529, 8, 2, 16, 65535)); + +bool qwen3_q16_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQwen3K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQwen3Q32K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + static utils::WgCount streaming_attention_k16_grid( WGPUDevice device, int64_t S, int64_t Hkv, - int64_t g) { + int64_t g, + uint32_t query_tile) { const uint64_t groups_per_kv = - (static_cast(S) * static_cast(g) + 31u) / 32u; + (static_cast(S) * static_cast(g) + query_tile - 1u) / + query_tile; const uint64_t workgroups = static_cast(Hkv) * groups_per_kv; if (workgroups == 0u || workgroups > UINT32_MAX) { throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32"); @@ -362,7 +440,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU sdpa: output shape must match q"); } - // fp32-only: validate byte counts against fp32 element counts. + // q/k/v/out are serialized fp32. KV caches are fp32 by default and use + // dedicated fp16 storage only when the graph-level option is active. auto numel = [](const WebGPUTensor& t) { uint64_t n = 1; for (int64_t d : t.dims) { @@ -370,11 +449,23 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } return n; }; - if (q.nbytes != numel(q) * sizeof(float) || - k.nbytes != numel(k) * sizeof(float) || - v.nbytes != numel(v) * sizeof(float) || - out.nbytes != numel(out) * sizeof(float)) { - throw std::runtime_error("WebGPU sdpa: fp32-only (byte-size mismatch)"); + auto is_fp32 = [&numel](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == sizeof(float) && + t.nbytes == numel(t) * sizeof(float); + }; + if (!is_fp32(q) || !is_fp32(k) || !is_fp32(v) || !is_fp32(out)) { + throw std::runtime_error("WebGPU sdpa: q/k/v/output must be fp32"); + } + const size_t cache_elem_size = + graph.kv_f16() ? sizeof(uint16_t) : sizeof(float); + auto cache_storage_is_valid = [&numel, + cache_elem_size](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == cache_elem_size && + t.nbytes == numel(t) * cache_elem_size; + }; + if (!cache_storage_is_valid(k_cache) || !cache_storage_is_valid(v_cache)) { + throw std::runtime_error( + "WebGPU sdpa: cache dtype does not match the selected storage mode"); } // input_pos: build-time Int (baked) OR runtime SymInt (dynamic decode). @@ -435,9 +526,50 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; } } - const bool k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && - D == 64 && scale == 0.125f && out.dims == q.dims && - k16_buffers_distinct && streaming_attention_k16_device_supported(device); + // The specialized shaders bake the standard Qwen3 scale, so eligibility must + // be exact. A nearby explicit scale has different operator semantics and must + // use the general path. + const float qwen3_expected_scale = 1.0f / std::sqrt(128.0f); + const bool qwen3_k16_geometry = Hq == 16 && Hkv == 8 && g == 2 && D == 128 && + scale == qwen3_expected_scale && out.dims == q.dims; + // Q16 is the default route for exact Qwen3 geometry; Q32 is an explicit + // autotuning candidate requested via the sdpa_query_tile RuntimeSpec. Support + // is evaluated per-tile so an unsupported Q32 request falls back to the Q16 + // streaming route instead of dropping to the materialized path. + const uint32_t device_max_workgroups = utils::queried_max_workgroups(device); + const bool qwen3_q16_supported = + qwen3_k16_geometry && graph.kv_f16() && + qwen3_q16_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3K16QueryTile, device_max_workgroups); + const bool qwen3_q32_requested = qwen3_k16_geometry && + graph.sdpa_query_tile() == static_cast(kQwen3Q32K16QueryTile); + const bool qwen3_q32_supported = + qwen3_q32_requested && graph.kv_f16() && + qwen3_q32_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3Q32K16QueryTile, device_max_workgroups); + const bool qwen3_q32_selected = qwen3_q32_supported; + const bool qwen3_k16_selected = qwen3_q32_selected || qwen3_q16_supported; + const uint32_t qwen3_query_tile = + qwen3_q32_selected ? kQwen3Q32K16QueryTile : kQwen3K16QueryTile; + const bool llama_k16_eligible = + graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && D == 64 && + scale == 0.125f && out.dims == q.dims && + streaming_attention_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kLlamaK16QueryTile, device_max_workgroups); + const bool k16_eligible = + k16_buffers_distinct && (llama_k16_eligible || qwen3_k16_selected); + const uint32_t k16_query_tile = + qwen3_k16_selected ? qwen3_query_tile : kLlamaK16QueryTile; + const char* k16_shader = qwen3_q32_selected ? kStreamingQwen3Q32K16Shader + : qwen3_k16_selected ? kStreamingQwen3K16Shader + : kStreamingK16Shader; + const char* k16_label = qwen3_q32_selected + ? "sdpa_streaming_attention_qwen3_q32_k16_causal_bound" + : qwen3_k16_selected ? "sdpa_streaming_attention_qwen3_k16_causal_bound" + : "sdpa_streaming_attention_k16_causal_bound"; const uint32_t uc_wg = utils::clamp_workgroup_size( device, get_webgpu_shader_info(kUpdateCacheShader).workgroup_size_x); const uint32_t qk_wg = utils::clamp_workgroup_size( @@ -469,7 +601,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { qk_wg, av_wg, fd_eligible, - k16_eligible](WebGPUGraph& gr) { + k16_eligible, + k16_query_tile](WebGPUGraph& gr) { SdpaLiveState state = {}; const auto& q_live_dims = gr.cur_dims(q_id); state.s = q_live_dims[qn - 3]; @@ -521,8 +654,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k16_eligible) { state.streaming_k16 = make_streaming_attention_k16_params( state.s, state.context_len, state.pos, Hq, Hkv, D); - state.streaming_k16_grid = - streaming_attention_k16_grid(gr.device(), state.s, Hkv, g); + state.streaming_k16_grid = streaming_attention_k16_grid( + gr.device(), state.s, Hkv, g, k16_query_tile); } state.update_cache_grid = { utils::compute_1d_workgroup_count( @@ -682,13 +815,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { : initial_state.streaming_k16_grid; add_sdpa_compute_dispatch( graph, - kStreamingK16Shader, + k16_shader, k16_bindings, k16_buf, sizeof(StreamingAttentionK16Params), initial_grid, 0, - "sdpa_streaming_attention_k16_causal_bound"); + k16_label); k16_idx = graph.num_dispatches() - 1; k16_range.end = graph.num_dispatches(); } diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5e27346af6f --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..f69a4d72dae --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from streaming_attention_qwen3_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: b26e81d92ba58c833f7dfbf2da6bce1b9cb468d70d893cef9df3a7fae16ee1fc +inline constexpr const char* kStreamingAttentionQwen3K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX = + 16; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY = + 8; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ = + 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl new file mode 100644 index 00000000000..6135b15b75d --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..e58a9d94bac --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from streaming_attention_qwen3_q32_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: 6d4396945b86cc3445698b1182fc4535ca786ad509ba128a6b6b08977e6a55e0 +inline constexpr const char* kStreamingAttentionQwen3Q32K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY = 8; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index cf044232dde..dd8cc4cc635 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -437,14 +437,17 @@ def embedded_sha256(header_text: str) -> str: def _wg_size_const(base: str, axis: str, val: int) -> str: """One WorkgroupSize constant; wrap to <=80 cols so CLANGFORMAT accepts it. - Long shader names push the single-line form past the 80-col limit (clang-format - then breaks after '=' with a 4-space continuation indent); emit that wrapped - form up front so the generated header matches lintrunner's CLANGFORMAT. + Long shader names push the single-line form past the 80-col limit. Emit the + wrapped form that clang-format selects so generated headers stay byte-stable. """ - decl = f"inline constexpr uint32_t k{base}WorkgroupSize{axis} =" - if len(decl) + len(f" {val};") > 80: - return f"{decl}\n {val};\n" - return f"{decl} {val};\n" + name = f"k{base}WorkgroupSize{axis}" + prefix = f"inline constexpr uint32_t {name} =" + decl = f"{prefix} {val};" + if len(decl) > 85: + return f"inline constexpr uint32_t\n {name} = {val};\n" + if len(decl) > 80: + return f"{prefix}\n {val};\n" + return f"{decl}\n" def render_header( @@ -469,6 +472,14 @@ def render_header( raise ValueError('shader contains )" which would close the R"( literal') base = symbol_base(name) x, y, z = parse_workgroup_size(wgsl_text) + provenance = f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT." + if len(provenance) > 80: + provenance_lines = [ + f"// @generated from {provenance_stem}.wgsl", + "// DO NOT EDIT.", + ] + else: + provenance_lines = [provenance] head = [ _BSD_HEADER, @@ -479,7 +490,7 @@ def render_header( "", "namespace executorch::backends::webgpu {", "", - f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT.", + *provenance_lines, f"// wgsl-sha256: {wgsl_sha256(wgsl_text)}", f'inline constexpr const char* k{base}WGSL = R"(', ] diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index feb6854c59d..b2b0d028567 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -120,7 +120,7 @@ void build_q4_route_graph(WebGPUGraph& graph, Q4RouteSignal signal) { if (signal == Q4RouteSignal::ExplicitOption) { config.record_q4gsw_decode_route = true; } - graph.build(fbb.GetBufferPointer(), nullptr, nullptr, config); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr, config); } std::vector q4_dispatches(WebGPUGraph& graph) { @@ -407,7 +407,7 @@ void build_resize_test_graph(WebGPUGraph& graph) { vk::FinishVkGraphBuffer(fbb, root); graph.set_device(g_device); - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } WebGPUComputeDispatchDescriptor make_dynamic_test_descriptor( @@ -707,7 +707,7 @@ void expect_invalid_rope_graph(const InvalidRopeGraphCase& test_case) { WebGPUGraph graph; std::string error; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& exception) { error = exception.what(); } @@ -801,7 +801,7 @@ TEST(WebGPUExecution, RejectsPlanOutputCountMismatch) { WebGPUGraph graph; WebGPUExecutionPlan plan; plan.copy_outputs = {true}; - std::vector> outputs; + std::vector outputs; EXPECT_THROW(graph.execute(plan), std::runtime_error); EXPECT_THROW(graph.copy_outputs(outputs, plan), std::runtime_error); diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index c9ce9648b43..7c5a5e141f3 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -232,6 +232,7 @@ TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) { EXPECT_FALSE(absent->record_q4gsw_decode_route); EXPECT_FALSE(absent->f16_kv_cache); EXPECT_FALSE(absent->f16_accumulate_gemm); + EXPECT_EQ(absent->sdpa_query_tile, 0); uint8_t false_value = 0; CompileSpec false_spec = { diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 9b0cf65a723..0e72ce6234d 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -556,6 +556,9 @@ void check_sdpa(int s) { constexpr int kK16Hq = 32; constexpr int kK16Hkv = 8; constexpr int kK16D = 64; +constexpr int kQwen3Hq = 16; +constexpr int kQwen3Hkv = 8; +constexpr int kQwen3D = 128; bool k16_device_supported() { const auto* context = get_default_webgpu_context(); @@ -568,6 +571,21 @@ bool k16_device_supported() { limits.maxComputeWorkgroupStorageSize >= 14720u; } +bool qwen3_q16_device_supported() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + void load_sdpa_module(Module& module, bool f16_kv) { if (!f16_kv) { ASSERT_EQ(module.load_forward(), Error::Ok); @@ -587,9 +605,13 @@ void run_k16_sdpa( 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) + "."); + bool prime = false, + float max_error_limit = 3e-3f, + bool initial = false) { + const std::string suffix = initial ? ".initial." + : prime ? ".prime." + : ".S" + std::to_string(s) + "."; + const std::string base = g_dir + "/" + prefix + suffix; auto q = read_bin(base + "q.bin"); auto k = read_bin(base + "k.bin"); auto v = read_bin(base + "v.bin"); @@ -614,8 +636,7 @@ void run_k16_sdpa( 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) + EXPECT_LT(max_err(got, golden), max_error_limit) << prefix << " S=" << s << " full output"; const std::set tokens = { 0, std::min(15, s - 1), std::min(16, s - 1), s - 1}; @@ -625,7 +646,7 @@ void run_k16_sdpa( 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) + EXPECT_LT(token_error, max_error_limit) << prefix << " S=" << s << " causal token=" << token; } } @@ -635,23 +656,22 @@ void prime_k16_sdpa( const char* prefix, int hq = kK16Hq, int hkv = kK16Hkv, - int d = kK16D) { - run_k16_sdpa(module, 12, prefix, hq, hkv, d, true); + int d = kK16D, + float max_error_limit = 3e-3f) { + run_k16_sdpa(module, 12, prefix, hq, hkv, d, true, max_error_limit); } #ifdef WGPU_BACKEND_ENABLE_PROFILING void expect_sdpa_route( const std::vector& names, int s, - bool expect_k16) { + bool expect_k16, + const char* k16_kernel_name = "sdpa_streaming_attention_k16_causal_bound") { 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"), + std::count(names.begin(), names.end(), k16_kernel_name), expect_k16 ? 1 : 0); EXPECT_EQ( std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); @@ -1437,6 +1457,46 @@ TEST(DynamicShape, K16CausalNumericsReusedGraph) { } } +TEST(DynamicShape, Qwen3K16CausalNumericsReusedGraph) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + } +} + +TEST(DynamicShape, Qwen3InitializedConstantCachePreserved) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + run_k16_sdpa( + module, + 17, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + 1e-2f, + true); +} + TEST(DynamicShape, K16CacheHeadMismatchRejectedAtLoad) { executorch::runtime::BackendOptions<1> options; ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); @@ -1487,6 +1547,65 @@ TEST(DynamicShape, K16CausalLiveRoutesProfile) { } } +TEST(DynamicShape, Qwen3K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + constexpr const char* kQwen3Kernel = + "sdpa_streaming_attention_qwen3_k16_causal_bound"; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, true, kQwen3Kernel); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, s > 1, kQwen3Kernel); + } +} + +TEST(DynamicShape, Qwen3NearScaleFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3_near_scale.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, false); + } +} + TEST(DynamicShape, K16F32KvFallsBackToExistingRoutes) { const auto* context = get_default_webgpu_context(); if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || 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 89a19f1fabc..a99c9baf46e 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 @@ -13,6 +13,7 @@ goldens. """ +import math import os import unittest @@ -782,6 +783,10 @@ def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: 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) +QWEN3_HQ = 16 +QWEN3_HKV = 8 +QWEN3_D = 128 +QWEN3_LIVE_S = (128, 17, 1) def _export_dynamic_sdpa(out_dir: str) -> None: @@ -827,6 +832,12 @@ def _write_f32_tensors(base, tensors) -> None: tensor.detach().numpy().astype(" None: + if not initialized_cache: + k_cache.zero_() + v_cache.zero_() + + def _export_dynamic_k16_sdpa_case( out_dir: str, prefix: str, @@ -837,10 +848,14 @@ def _export_dynamic_k16_sdpa_case( scale: float | None = None, cache_hkv: int | None = None, expect_runtime_reject: bool = False, + denom: float = 16.0, + kv_f16_golden: bool = False, + initialized_cache: bool = False, ) -> None: from executorch.backends.webgpu.test.ops.test_sdpa import ( _det_inputs, _golden, + _round_kv_for_storage, SdpaConfig, ) @@ -853,14 +868,15 @@ def cfg(s: int) -> "SdpaConfig": s, K16_CMAX, K16_INPUT_POS, + denom, + kv_f16=kv_f16_golden, ) 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_() + _zero_uninitialized_cache(initialized_cache, kc, vc) class K16SdpaModule(torch.nn.Module): def __init__(self) -> None: @@ -908,6 +924,8 @@ def forward(self, q, k, v, pos_control): def reference(live_cfg, q, k, v, kc, vc): if scale is None: return _golden(live_cfg, q, k, v, kc, vc) + k, v, kc, vc = _round_kv_for_storage(live_cfg, k, v, kc, vc) + runtime_scale = float(torch.tensor(scale, dtype=torch.float32).item()) 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) @@ -919,14 +937,48 @@ def reference(live_cfg, q, k, v, kc, vc): 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 + q_heads, k_heads, v_heads, attn_mask=mask, scale=runtime_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) + if initialized_cache: + initial_cfg = cfg(17) + initial_q, initial_k, initial_v, initial_kc, initial_vc = _det_inputs( + initial_cfg + ) + initial_golden = reference( + initial_cfg, + initial_q, + initial_k, + initial_v, + initial_kc, + initial_vc, + ) + initial_base = os.path.join(out_dir, f"{prefix}.initial.") + _write_f32_tensors( + initial_base, + ( + ("q", initial_q), + ("k", initial_k), + ("v", initial_v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("golden", initial_golden), + ), + ) + + prime_cfg = SdpaConfig( + prefix, + hq, + hkv, + d, + K16_PRIME_S, + K16_CMAX, + K16_PRIME_POS, + denom, + kv_f16=kv_f16_golden, + ) prime_q, prime_k, prime_v, prime_kc, prime_vc = _det_inputs(prime_cfg) - prime_kc.zero_() - prime_vc.zero_() + _zero_uninitialized_cache(initialized_cache, prime_kc, prime_vc) 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( @@ -943,8 +995,7 @@ def reference(live_cfg, q, k, v, kc, vc): for s in live_s: live_cfg = cfg(s) q, k, v, kc, vc = _det_inputs(live_cfg) - kc.zero_() - vc.zero_() + _zero_uninitialized_cache(initialized_cache, kc, vc) 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) @@ -972,6 +1023,28 @@ def _export_dynamic_k16_sdpa_cases(out_dir: str) -> None: K16_HKV, K16_LIVE_S, ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3", + QWEN3_HQ, + QWEN3_HKV, + QWEN3_LIVE_S, + d=QWEN3_D, + denom=10.0, + kv_f16_golden=True, + initialized_cache=True, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3_near_scale", + QWEN3_HQ, + QWEN3_HKV, + (128, 1), + d=QWEN3_D, + scale=1.0 / math.sqrt(float(QWEN3_D)) + 5e-7, + denom=10.0, + kv_f16_golden=True, + ) _export_dynamic_k16_sdpa_case( out_dir, "sdpa_k16_wrong_geometry", @@ -1447,12 +1520,16 @@ def test_export_dynamic_rms(self) -> None: "dyn_qkv_bk64_wrong_width.pte", "dyn_qkv_bk64_different_input.pte", "sdpa_k16_bad_cache_heads.pte", + "sdpa_k16_qwen3.pte", + "sdpa_k16_qwen3_near_scale.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_qwen3", QWEN3_LIVE_S), + ("sdpa_k16_qwen3_near_scale", (128, 1)), ("sdpa_k16_wrong_geometry", (128, 1)), ("sdpa_k16_wrong_d", (128, 1)), ("sdpa_k16_wrong_scale", (128, 1)), @@ -1472,6 +1549,10 @@ def test_export_dynamic_rms(self) -> None: self.assertGreater( os.path.getsize(os.path.join(d, name)), 0 ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"sdpa_k16_qwen3.initial.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) if __name__ == "__main__": diff --git a/backends/webgpu/test/ops/test_sdpa.py b/backends/webgpu/test/ops/test_sdpa.py index 960c62e5203..b2606256e28 100644 --- a/backends/webgpu/test/ops/test_sdpa.py +++ b/backends/webgpu/test/ops/test_sdpa.py @@ -23,7 +23,7 @@ import torch import torch.nn.functional as F -from executorch.backends.vulkan import VulkanPartitioner +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower from executorch.extension.llm.custom_ops import ( # noqa: F401 registers llama ops custom_ops, @@ -40,6 +40,7 @@ class SdpaConfig: cmax: int # kv-cache capacity input_pos: int # number of prior tokens already in the cache (decode) denom: float = 16.0 # ramp divisor; small denom -> large logits (softmax stress) + kv_f16: bool = False # Single source of truth, mirrored by the C++ CONFIGS table in the native test. @@ -64,6 +65,10 @@ class SdpaConfig: # 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV (cap+1). SdpaConfig("llama1b_prefill_512", 32, 8, 64, 512, 512, 0), SdpaConfig("llama1b_prefill_2048", 32, 8, 64, 2048, 2048, 0), + # denom=10 intentionally makes K/V values lossy in fp16, so native + # execution exercises the real fp32->fp16->fp32 cache conversion path. + SdpaConfig("qwen3_prefill", 16, 8, 128, 128, 256, 0, 10.0, kv_f16=True), + SdpaConfig("qwen3_odd_boundary", 16, 8, 128, 17, 64, 31, 10.0, kv_f16=True), ] @@ -83,6 +88,7 @@ class ReplaySeq: d: int # head dim cmax: int # kv-cache capacity (>= sum(seq_lens)) seq_lens: tuple[int, ...] + kv_f16: bool = False # Mirror Vulkan sdpa_test.cpp:856/867/875 (3 param sets); cmax = sum rounded up. @@ -90,11 +96,18 @@ class ReplaySeq: ReplaySeq("small", 8, 4, 4, 16, (3, 1, 1, 5, 1, 1, 2)), ReplaySeq("small_d", 6, 2, 8, 16, (3, 1, 1, 5, 1, 1)), ReplaySeq("llama3", 24, 8, 128, 256, (111, 1, 1, 1, 57, 1, 1)), + ReplaySeq("qwen3_fd", 16, 8, 128, 64, (17, 1), kv_f16=True), ] +DYNAMIC_REPLAY_SEQS = [seq for seq in REPLAY_SEQS if not seq.kv_f16] -# (head_dim, num_heads, num_kv_heads) from sdpa_test.cpp:856/867/875 -- guards a -# transposition of the (hq, hkv, d) field order against the Vulkan source. -VULKAN_PARAMS = {"small": (4, 8, 4), "small_d": (8, 6, 2), "llama3": (128, 24, 8)} +# Guards transposition of the (hq, hkv, d) field order. The first three values +# mirror Vulkan sdpa_test.cpp:856/867/875; Qwen3 extends the same contract. +VULKAN_PARAMS = { + "small": (4, 8, 4), + "small_d": (8, 6, 2), + "llama3": (128, 24, 8), + "qwen3_fd": (128, 16, 8), +} class SdpaModule(torch.nn.Module): @@ -178,6 +191,12 @@ def _det_inputs(cfg: SdpaConfig): return q, k, v, k_cache, v_cache +def _round_kv_for_storage(cfg: SdpaConfig, *tensors: torch.Tensor): + if not cfg.kv_f16: + return tensors + return tuple(tensor.to(torch.float16).to(torch.float32) for tensor in tensors) + + def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """Reference attention output [1,S,Hq,D], computed in fp64 then cast to fp32. @@ -189,6 +208,7 @@ def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """ context_len = cfg.s + cfg.input_pos g = cfg.hq // cfg.hkv + k, v, k_cache, v_cache = _round_kv_for_storage(cfg, k, v, k_cache, v_cache) qd, kd, vd = q.double(), k.double(), v.double() kcd, vcd = k_cache.double(), v_cache.double() @@ -228,6 +248,50 @@ def _export_pte(cfg: SdpaConfig, q, k, v, kc, vc): class TestSdpa(unittest.TestCase): + def test_qwen3_fixture_contract(self) -> None: + configs = {cfg.name: cfg for cfg in CONFIGS} + expected_geometries = { + "qwen3_prefill": (16, 8, 128, 128, 256, 0), + "qwen3_odd_boundary": (16, 8, 128, 17, 64, 31), + } + for name, geometry in expected_geometries.items(): + with self.subTest(config=name): + self.assertIn(name, configs) + cfg = configs[name] + self.assertEqual( + (cfg.hq, cfg.hkv, cfg.d, cfg.s, cfg.cmax, cfg.input_pos), + geometry, + ) + self.assertTrue(cfg.kv_f16) + + replays = {seq.name: seq for seq in REPLAY_SEQS} + self.assertIn("qwen3_fd", replays) + qwen3_fd = replays["qwen3_fd"] + self.assertEqual((qwen3_fd.hq, qwen3_fd.hkv, qwen3_fd.d), (16, 8, 128)) + self.assertEqual(qwen3_fd.seq_lens, (17, 1)) + self.assertTrue(qwen3_fd.kv_f16) + + probe = torch.tensor([0.1], dtype=torch.float32) + (rounded,) = _round_kv_for_storage(configs["qwen3_prefill"], probe) + expected = probe.to(torch.float16).to(torch.float32) + torch.testing.assert_close(rounded, expected, atol=0.0, rtol=0.0) + self.assertFalse(torch.equal(rounded, probe)) + + boundary = configs["qwen3_odd_boundary"] + q, k, v, k_cache, v_cache = _det_inputs(boundary) + self.assertGreater(torch.count_nonzero(k_cache).item(), 0) + self.assertGreater(torch.count_nonzero(v_cache).item(), 0) + initialized = _golden(boundary, q, k, v, k_cache, v_cache) + cleared = _golden( + boundary, + q, + k, + v, + torch.zeros_like(k_cache), + torch.zeros_like(v_cache), + ) + self.assertFalse(torch.equal(initialized, cleared)) + def test_sdpa_export_delegates(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): @@ -248,7 +312,12 @@ def test_golden_matches_eager_op(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): q, k, v, kc, vc = _det_inputs(cfg) - eager = SdpaModule(cfg.input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(cfg.input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) golden = _golden(cfg, q, k, v, kc, vc) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) @@ -277,10 +346,16 @@ def test_replay_golden_matches_eager(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) golden = _golden(cfg, q, k, v, kc, vc) - eager = SdpaModule(input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) kc[0, input_pos : input_pos + s] = k[0] vc[0, input_pos : input_pos + s] = v[0] @@ -302,6 +377,7 @@ def test_replay_export_delegates(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, kc, vc) @@ -347,7 +423,14 @@ def export_replay_sequences(out_dir: str) -> None: input_pos = 0 for t, s in enumerate(seq.seq_lens): cfg = SdpaConfig( - f"{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, s, seq.cmax, input_pos + f"{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + s, + seq.cmax, + input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, ref_kc, ref_vc) @@ -421,7 +504,7 @@ def export_dynamic_decode(out_dir: str) -> None: Mirrors the host accumulation the native test threads: at step t the golden attends over input_pos=t prior tokens plus the new token. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" et = _export_dyn_pte(seq, 1) pte_path = os.path.join(out_dir, f"sdpa_dyn_{seq.name}.pte") @@ -431,7 +514,14 @@ def export_dynamic_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" None: def test_dynamic_decode_golden_matches_eager(self) -> None: # The threaded-cache decode golden must equal the eager op step-by-step. - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: ref_kc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc) @@ -499,7 +596,7 @@ def export_incache_decode(out_dir: str) -> None: """One sdpa_incache_.pte (mutable-buffer KV cache) + per-step decode goldens. forward() feeds only q/k/v + input_pos; the cache persists in-graph. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" m = DecodeCacheModule(seq.hkv, seq.d, seq.cmax) q, k, v = _step_inputs(seq, 0, 1) @@ -517,7 +614,14 @@ def export_incache_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"incache_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"incache_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -309,13 +311,12 @@ bool sdpa_within_tol( const float* golden, int n, float* ma, - float* mr) { + float* mr, + bool kv_f16 = false) { float atol = 1e-4f, rtol = 1e-3f; - // f16 KV (runtime opt-in) reads K/V at reduced precision; loosen the tol on a - // shader-f16 device to cover that rounding. Harmless for f32 KV (looser - // gate). - const WebGPUContext* kv_ctx = get_default_webgpu_context(); - if (kv_ctx != nullptr && kv_ctx->shader_f16_supported) { + // Only fp16-KV cases receive the tolerance needed for storage rounding; + // device capability alone must not weaken unrelated fp32 tests. + if (kv_f16) { atol = 2e-3f; rtol = 1e-2f; } @@ -1250,6 +1251,7 @@ struct SdpaConfig { float denom; // ramp divisor (mirrors Python); small -> large logits bool required = false; // CI (SDPA dir set): absent .pte = FAIL, not skip bool expect_reject = false; // load MUST fail (e.g. D%4 guard), no golden + bool kv_f16 = false; }; const SdpaConfig kSdpaConfigs[] = { @@ -1292,6 +1294,28 @@ const SdpaConfig kSdpaConfigs[] = { 0, 16.0f, /*required=*/true}, + {"qwen3_prefill", + 16, + 8, + 128, + 128, + 256, + 0, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_odd_boundary", + 16, + 8, + 128, + 17, + 64, + 31, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). @@ -1314,9 +1338,8 @@ float sdpa_ramp_t( return static_cast(((i + 31 * t) % mod) - off) / denom; } -// Multi-step replay sequences. Mirror the Python REPLAY_SEQS / Vulkan param -// sets (sdpa_test.cpp:856/867/875). Each seq_lens entry is one step replayed on -// a host-threaded KV cache (big=prefill, mid=multi-token, 1=decode). +// Multi-step replay sequences. The first three mirror Vulkan param sets; Qwen3 +// extends the same Python REPLAY_SEQS contract. struct SdpaSequence { const char* name; int hq; @@ -1324,18 +1347,90 @@ struct SdpaSequence { int d; int cmax; std::vector seq_lens; + bool kv_f16 = false; }; const SdpaSequence kSdpaSequences[] = { {"small", 8, 4, 4, 16, {3, 1, 1, 5, 1, 1, 2}}, {"small_d", 6, 2, 8, 16, {3, 1, 1, 5, 1, 1}}, {"llama3", 24, 8, 128, 256, {111, 1, 1, 1, 57, 1, 1}}, + {"qwen3_fd", 16, 8, 128, 64, {17, 1}, /*kv_f16=*/true}, }; +Error load_sdpa_forward(Module& module, bool kv_f16, int sdpa_query_tile = 0) { + if (!kv_f16 && sdpa_query_tile == 0) { + return module.load_forward(); + } + BackendOptions<2> options; + Error error = Error::Ok; + if (kv_f16) { + error = options.set_option("enable_f16_kv_cache", true); + if (error != Error::Ok) { + return error; + } + } + if (sdpa_query_tile != 0) { + error = options.set_option("sdpa_query_tile", sdpa_query_tile); + if (error != Error::Ok) { + return error; + } + } + LoadBackendOptionsMap option_map; + error = option_map.set_options("VulkanBackend", options.view()); + if (error != Error::Ok) { + return error; + } + return module.load_forward(nullptr, nullptr, &option_map); +} + +bool shader_f16_supported_on_test_device() { + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported; +} + +bool qwen3_q16_supported_on_test_device() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_supported_on_test_device() { + constexpr uint32_t kQ32StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQ32StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kTestRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kTestRouteFlashDecoding = 1u << 10; +constexpr uint32_t kTestRouteK16CausalBound = 1u << 11; +constexpr uint32_t kTestRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kTestRouteQwen3Q32K16 = 1u << 14; +#endif // WGPU_BACKEND_ENABLE_PROFILING + void test_sdpa_config( const SdpaConfig& cfg, const std::string& model_path, - const std::string& golden_path) { + const std::string& golden_path, + int sdpa_query_tile = 0) { // Inputs reconstruct test_sdpa.py::_det_inputs bit-for-bit (/16 exact fp32). printf( "\n--- Test: sdpa_with_kv_cache (%s: Hq=%d,Hkv=%d,D=%d,S=%d,Cmax=%d,pos=%d) ---\n", @@ -1347,8 +1442,13 @@ void test_sdpa_config( cfg.cmax, cfg.input_pos); + if (cfg.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", cfg.name); + return; + } + Module module(model_path); - auto err = module.load_forward(); + auto err = load_sdpa_forward(module, cfg.kv_f16, sdpa_query_tile); if (cfg.expect_reject) { // D not a multiple of 4 must be rejected at load by the head_dim guard. ASSERT_NE(err, Error::Ok) @@ -1393,6 +1493,38 @@ void test_sdpa_config( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() << ")"; + if (cfg.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // Exact Qwen3 geometry + fp16 KV selects the K16 streaming (causal-bound) + // route by default. The sdpa_query_tile RuntimeSpec only swaps the Q16/Q32 + // kernel variant; both map to the K16CausalBound bit. A non-Qwen3 fp16-KV + // shape falls back to the materialized path (or flash-decoding at S==1). + const bool qwen3_geometry = cfg.hq == 16 && cfg.hkv == 8 && cfg.d == 128; + const bool qwen3_streaming = + qwen3_geometry && cfg.s > 1 && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = qwen3_streaming + ? kTestRouteK16CausalBound + : (cfg.s == 1 ? kTestRouteFlashDecoding + : kTestRouteMaterializedAttention); + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route); + EXPECT_EQ(g_last_route_conflict_count, 0u); + const uint32_t qwen3_tile_routes = + g_last_route_mask & (kTestRouteQwen3Q16K16 | kTestRouteQwen3Q32K16); + if (qwen3_streaming) { + const uint32_t expected_tile_route = + sdpa_query_tile == 32 && qwen3_q32_supported_on_test_device() + ? kTestRouteQwen3Q32K16 + : kTestRouteQwen3Q16K16; + EXPECT_EQ(qwen3_tile_routes, expected_tile_route); + } else { + EXPECT_EQ(qwen3_tile_routes, 0u); + } +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outputs = result.get(); // Select the attention output [1,S,Hq,D] by shape; the op returns @@ -1424,8 +1556,8 @@ void test_sdpa_config( ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; float max_abs_err = 0.0f, max_rel_err = 0.0f; - const bool pass = - sdpa_within_tol(out_data, golden.data(), on, &max_abs_err, &max_rel_err); + const bool pass = sdpa_within_tol( + out_data, golden.data(), on, &max_abs_err, &max_rel_err, cfg.kv_f16); printf( "Max abs error: %e Max rel error: %e (checked %d elements)\n", max_abs_err, @@ -1447,6 +1579,10 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { seq.d, seq.cmax, seq.seq_lens.size()); + if (seq.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", seq.name); + return; + } const int cn = seq.cmax * seq.hkv * seq.d; std::vector kc(cn, 0.0f), vc(cn, 0.0f); @@ -1460,7 +1596,7 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { std::to_string(t) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos); Module module(base + ".pte"); - ASSERT_EQ(module.load_forward(), Error::Ok) + ASSERT_EQ(load_sdpa_forward(module, seq.kv_f16), Error::Ok) << "could not load " << base << ".pte"; const int qn = s * seq.hq * seq.d; @@ -1486,6 +1622,26 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward " << base << ".pte (error " << (int)result.error() << ")"; + if (seq.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // S==1 decode -> flash-decoding; a multi-token exact-Qwen3-geometry + // prefill -> the K16 streaming (causal-bound) route by default (no env); + // any other multi-token fp16-KV shape -> materialized. + const bool qwen3_geometry = seq.hq == 16 && seq.hkv == 8 && seq.d == 128; + const bool qwen3_streaming = + qwen3_geometry && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = s == 1 ? kTestRouteFlashDecoding + : qwen3_streaming ? kTestRouteK16CausalBound + : kTestRouteMaterializedAttention; + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route) + << seq.name << " step" << t; + EXPECT_EQ(g_last_route_conflict_count, 0u) << seq.name << " step" << t; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outs = result.get(); // The op returns [k_cache, v_cache, attn_output]: attn has a unique numel; @@ -1533,7 +1689,8 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { ASSERT_FALSE(golden.empty()) << "could not load " << base << ".golden.bin"; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; - const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); + const bool step_ok = + sdpa_within_tol(ad, golden.data(), qn, &ma, &mr, seq.kv_f16); printf( " step%zu (S=%d pos=%d ctx=%d): max abs %e rel %e\n", t, @@ -1906,12 +2063,160 @@ void test_symint_input_narrowing() { vk::FinishVkGraphBuffer(fbb, root); WebGPUGraph graph; - ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); ASSERT_EQ(graph.symint_sources().size(), 1u); const auto& source = graph.symint_sources().front(); exercise_symint_host_inputs(graph, source.symint_id, source.input_tensor_id); } +void write_u16_le(std::vector& data, size_t offset, uint16_t value) { + data.at(offset) = static_cast(value); + data.at(offset + 1) = static_cast(value >> 8); +} + +void write_u32_le(std::vector& data, size_t offset, uint32_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +void write_u64_le(std::vector& data, size_t offset, uint64_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +std::vector make_delegate_header_test_blob() { + std::vector blob(44, 0); + std::memcpy(blob.data() + 4, "VH00", 4); + write_u16_le(blob, 8, 30); + write_u32_le(blob, 10, 32); + write_u32_le(blob, 14, 8); + write_u32_le(blob, 18, 40); + write_u64_le(blob, 22, 4); + return blob; +} + +void finish_inline_constant_graph( + ::flatbuffers::FlatBufferBuilder& fbb, + bool mark_as_kv_cache, + const std::vector& dims, + uint64_t inline_offset = 0) { + namespace vk = vkgraph; + std::vector<::flatbuffers::Offset> values; + const int tensor_count = mark_as_kv_cache ? 5 : 1; + for (int i = 0; i < tensor_count; i++) { + const bool is_cache = mark_as_kv_cache && i >= 3; + const bool is_constant = !mark_as_kv_cache || is_cache; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + is_constant ? (is_cache ? i - 3 : 0) : -1, + is_constant ? -1 : i) + .Union())); + } + + std::vector<::flatbuffers::Offset> chain; + if (mark_as_kv_cache) { + const std::vector args = {0, 1, 2, 3, 4}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "sdpa_with_kv_cache.default", &args)); + } + std::vector<::flatbuffers::Offset> constants; + constants.push_back( + vk::CreateVkBytesDirect(fbb, inline_offset, sizeof(float))); + if (mark_as_kv_cache) { + constants.push_back(vk::CreateVkBytesDirect(fbb, 0, sizeof(float))); + } + const std::vector output_ids = {0}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, nullptr, &output_ids, &constants); + vk::FinishVkGraphBuffer(fbb, root); +} + +TEST(WebGPUNative, DelegateHeaderRejectsTruncatedRanges) { + const auto blob = make_delegate_header_test_blob(); + EXPECT_TRUE(WebGPUDelegateHeader::parse(blob.data(), blob.size()).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), 29).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), blob.size() - 1).ok()); +} + +TEST(WebGPUNative, InlineConstantExtentIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph exact_graph; + EXPECT_NO_THROW(exact_graph.build( + fbb.GetBufferPointer(), data.data(), data.size(), nullptr)); + + WebGPUGraph short_graph; + EXPECT_THROW( + short_graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, ZeroByteInlineConstantOffsetIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {0u}, 1); + const std::array data = {0}; + + WebGPUGraph graph; + EXPECT_THROW( + graph.build(fbb.GetBufferPointer(), data.data(), 0, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, F16KvInlineConstantExtentIsBounded) { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || !context->shader_f16_supported) { + GTEST_SKIP() << "shader-f16 unavailable"; + } + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, true, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph graph; + WebGPUGraphConfig config; + config.f16_kv_cache = true; + try { + graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr, config); + FAIL() << "undersized inline fp16 KV constant was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU f16 KV: inline cache constant exceeds constant data"); + } +} + +void expect_tensor_extent_error( + const std::vector& dims, + const char* expected_error) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, dims); + const std::array data = {0, 0, 0, 0}; + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), data.data(), data.size(), nullptr); + ADD_FAILURE() << "overflowing tensor extent was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), expected_error); + } +} + +TEST(WebGPUNative, TensorExtentOverflowIsRejected) { + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX, 2u}, "WebGPU: tensor element count overflows"); + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX}, "WebGPU: tensor byte size overflows"); +} + struct DelegateBlobView { size_t base_offset; WebGPUDelegateHeader header; @@ -1932,7 +2237,7 @@ std::optional find_delegate_blob( if (std::memcmp(base + kMagicOffset, kMagic, sizeof(kMagic)) != 0) { continue; } - auto header = WebGPUDelegateHeader::parse(base); + auto header = WebGPUDelegateHeader::parse(base, blob.size() - base_offset); if (!header.ok()) { continue; } @@ -1950,6 +2255,43 @@ std::optional find_delegate_blob( return std::nullopt; } +TEST(WebGPUNative, StructurallyInvalidVkGraphIsRejectedAtLoad) { + if (g_symint_blob.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; + } + FILE* input = std::fopen(g_symint_blob.c_str(), "rb"); + ASSERT_NE(input, nullptr); + std::fseek(input, 0, SEEK_END); + const long file_size = std::ftell(input); + std::fseek(input, 0, SEEK_SET); + ASSERT_GT(file_size, 0); + std::vector blob(static_cast(file_size)); + ASSERT_EQ(std::fread(blob.data(), 1, blob.size(), input), blob.size()); + std::fclose(input); + + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()); + ASSERT_GE(delegate->header.flatbuffer_size, sizeof(uint32_t)); + const size_t root_offset = + delegate->base_offset + delegate->header.flatbuffer_offset; + std::fill_n(blob.begin() + root_offset, sizeof(uint32_t), UINT8_MAX); + + const std::string malformed_path = "/tmp/webgpu_invalid_vkgraph_" + + std::to_string(reinterpret_cast(blob.data())) + ".pte"; + FILE* output = std::fopen(malformed_path.c_str(), "wb"); + ASSERT_NE(output, nullptr); + ASSERT_EQ(std::fwrite(blob.data(), 1, blob.size(), output), blob.size()); + std::fclose(output); + + Error load_result = Error::Ok; + { + Module module(malformed_path); + load_result = module.load_forward(); + } + EXPECT_NE(load_result, Error::Ok); + EXPECT_EQ(std::remove(malformed_path.c_str()), 0); +} + // S1 SymInt round-trip: confirm a dynamic input_pos stays live. void test_symint_roundtrip(const std::string& blob_path) { printf("\n--- Test: symint round-trip (%s) ---\n", blob_path.c_str()); @@ -1972,6 +2314,7 @@ void test_symint_roundtrip(const std::string& blob_path) { graph.build( base + delegate->header.flatbuffer_offset, base + delegate->header.bytes_offset, + delegate->header.bytes_size, nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); @@ -2029,6 +2372,7 @@ void test_resize_hook(const std::string& blob_path) { graph.build( base + delegate->header.flatbuffer_offset, base + delegate->header.bytes_offset, + delegate->header.bytes_size, nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); @@ -2178,7 +2522,7 @@ static bool test_slice_double_start_case(double start_d, int out_len) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -2191,8 +2535,8 @@ static bool test_slice_double_start_case(double start_d, int out_len) { std::vector inputs(1); inputs[0] = {in.data(), in.size() * sizeof(float), false}; std::vector out(out_len, -1.0f); - std::vector> outputs(1); - outputs[0] = {out.data(), out.size() * sizeof(float)}; + std::vector outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float), true}; try { graph.copy_inputs(inputs); const WebGPUExecutionPlan plan = graph.make_execution_plan({}); @@ -2276,7 +2620,7 @@ static bool test_slice_double_start_rejects(double bad_start) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("PASS: rejected as expected: %s\n", e.what()); return true; @@ -2374,7 +2718,7 @@ static bool test_select_double_scalar_case( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -2384,8 +2728,8 @@ static bool test_select_double_scalar_case( std::vector inputs = { {in.data(), in.size() * sizeof(float), false}}; std::vector out(expected.size(), -1.0f); - std::vector> outputs = { - {out.data(), out.size() * sizeof(float)}}; + std::vector outputs = { + {out.data(), out.size() * sizeof(float), true}}; try { graph.copy_inputs(inputs); const WebGPUExecutionPlan plan = graph.make_execution_plan({}); @@ -2414,7 +2758,7 @@ static bool test_select_scalar_build_error( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { const std::string error = e.what(); if (error.find(expected_error) != std::string::npos) { @@ -2531,7 +2875,7 @@ void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { vk::FinishVkGraphBuffer(fbb, root); WebGPUGraph graph; - ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); ASSERT_EQ(graph.num_dispatches(), 2u); const uint32_t q_x = graph.dispatch_at(0).workgroup_count_x; const uint32_t q_y = graph.dispatch_at(0).workgroup_count_y; @@ -2812,6 +3156,95 @@ TEST(WebGPUNative, PrepackTied) { } // SDPA sweep: configs self-discover sdpa_.pte; required=FAIL else skip. +TEST(WebGPUNative, Qwen3SdpaFixtureContract) { + const auto find_config = [](const char* name) { + return std::find_if( + std::begin(kSdpaConfigs), + std::end(kSdpaConfigs), + [name](const SdpaConfig& cfg) { + return std::strcmp(cfg.name, name) == 0; + }); + }; + const auto prefill = find_config("qwen3_prefill"); + const auto boundary = find_config("qwen3_odd_boundary"); + ASSERT_NE(prefill, std::end(kSdpaConfigs)); + ASSERT_NE(boundary, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {prefill->hq, + prefill->hkv, + prefill->d, + prefill->s, + prefill->cmax, + prefill->input_pos}), + std::vector({16, 8, 128, 128, 256, 0})); + EXPECT_EQ( + std::vector( + {boundary->hq, + boundary->hkv, + boundary->d, + boundary->s, + boundary->cmax, + boundary->input_pos}), + std::vector({16, 8, 128, 17, 64, 31})); + EXPECT_TRUE(prefill->kv_f16 && boundary->kv_f16); + + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + EXPECT_EQ( + std::vector({replay->hq, replay->hkv, replay->d, replay->cmax}), + std::vector({16, 8, 128, 64})); + EXPECT_EQ(replay->seq_lens, std::vector({17, 1})); + EXPECT_TRUE(replay->kv_f16); +} + +TEST(WebGPUNative, Qwen3SdpaRoutes) { + if (g_sdpa_dir.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SDPA_DIR not set"; + } + if (!qwen3_q16_supported_on_test_device()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + // Default route: exact-Qwen3-geometry fp16-KV configs select the Q16 K16 + // streaming (causal-bound) route by geometry (no runtime config needed) -- + // the per-config assertions live in test_sdpa_config / test_sdpa_replay. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config(cfg, base + ".pte", base + ".golden.bin"); + } + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + test_sdpa_replay(*replay, g_sdpa_dir); + + // Run Q32 over both an aligned prefill and the S=17/nonzero-position case so + // the partial final workgroup's row mask is covered. Unsupported Q32 devices + // intentionally fall back to the already-qualified Q16 route. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config( + cfg, + base + ".pte", + base + ".golden.bin", + /*sdpa_query_tile=*/32); + } +} + TEST(WebGPUNative, SdpaSweep) { const std::string& dir = g_sdpa_dir; bool ran = false; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 0cd110d0ff7..32d59501a3a 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -158,6 +158,22 @@ def test_render_header_embeds_sha256(self) -> None: self.assertEqual(g.embedded_sha256(h), want) self.assertEqual(g.wgsl_sha256(wgsl), want) + def test_render_header_long_name_is_clang_format_stable(self) -> None: + stem = "streaming_attention_qwen3_q32_k16_causal_bound" + wgsl = "@compute @workgroup_size(32, 8, 1)\nfn main(){}\n" + h = g.render_header(Path(f"runtime/ops/sdpa/{stem}.wgsl"), wgsl) + + self.assertIn( + f"// @generated from {stem}.wgsl\n// DO NOT EDIT.", + h, + ) + self.assertIn( + "inline constexpr uint32_t\n" + " kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32;", + h, + ) + self.assertEqual(g.embedded_sha256(h), g.wgsl_sha256(wgsl)) + def test_embedded_sha256_missing_returns_empty(self) -> None: self.assertEqual(g.embedded_sha256("no sha line here\n"), "") @@ -200,6 +216,33 @@ def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: ] self.assertEqual(indices, list(range(8))) + def test_qwen3_runtime_eligibility_is_exact(self) -> None: + sdpa = (g.BACKEND_ROOT / "runtime/ops/sdpa/Sdpa.cpp").read_text() + self.assertIn("q/k/v/output must be fp32", sdpa) + self.assertIn("cache dtype does not match the selected storage mode", sdpa) + self.assertIn("scale == qwen3_expected_scale", sdpa) + self.assertNotIn("std::fabs(scale - qwen3_expected_scale)", sdpa) + + def test_fp16_kv_graph_guards_transfer_and_topology(self) -> None: + graph = (g.BACKEND_ROOT / "runtime/WebGPUGraph.cpp").read_text() + self.assertIn("serialized cache tensor must be fp32", graph) + self.assertIn("consumed through a ValueList", graph) + self.assertIn("preserve it while changing storage", graph) + + copy_inputs = graph.index("void WebGPUGraph::copy_inputs") + input_guard = graph.index( + "fp16 device input requires an fp32 host tensor", copy_inputs + ) + fast_path = graph.index("// Fast path", copy_inputs) + self.assertLess(input_guard, fast_path) + + copy_outputs = graph.index("void WebGPUGraph::copy_outputs") + output_guard = graph.index( + "fp16 device output requires an fp32 host tensor", copy_outputs + ) + map_request = graph.index("wgpuBufferMapAsync", copy_outputs) + self.assertLess(output_guard, map_request) + def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse. self.assertEqual(