[ExecuTorch][WebGPU] Add Qwen3 K16 attention and fp16 KV cache

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -90,7 +90,7 @@ Result<DelegateHandle*> WebGPUBackend::init(
 
   // Parse header to locate flatbuffer and constant data
   Result<WebGPUDelegateHeader> 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<DelegateHandle*> 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<DelegateHandle*> WebGPUBackend::init(
       config.f16_accumulate_gemm = spec.get();
     }
   }
+  {
+    Result<int> spec = context.get_runtime_spec<int>("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<int64_t> 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<std::pair<void*, size_t>> outputs;
+    std::vector<OutputData> 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) {

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> WebGPUDelegateHeader::parse(const void* data) {
+Result<WebGPUDelegateHeader> 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> 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;
 }
 

backends/webgpu/runtime/WebGPUDelegateHeader.h

@@ -8,6 +8,8 @@
 
 #pragma once
 
+#include <cstddef>
+
 #include <executorch/runtime/core/result.h>
 
 namespace executorch {
@@ -18,7 +20,8 @@ struct WebGPUDelegateHeader {
   bool is_valid() const;
 
   static executorch::runtime::Result<WebGPUDelegateHeader> parse(
-      const void* data);
+      const void* data,
+      size_t buffer_size);
 
   uint32_t header_size;
   uint32_t flatbuffer_offset;