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

Pull Request resolved: #21136

Qwen3's attention geometry differs from Llama's, and its KV cache is produced in
fp32 on the host but must be consumed in fp16 on the device. This adds guarded
Qwen3 K16 streaming online-softmax attention schedules — a Q16 schedule that is
the automatic default whenever the exact geometry and capability guards pass,
plus a Q32 candidate that is opt-in through the `sdpa_query_tile` runtime spec
(BackendOption) for future autotuning — together with the exact fp32-host to
fp16-device KV-cache boundary conversion. Selection requires the exact Qwen3
geometry, fp16 KV storage, adapter limits, a valid workgroup count, and an exact
2:1 byte ratio; the established Llama, materialized, and FlashDecoding routes
remain fallbacks. This builds on the HuggingFace rotate-half
RoPE operator. No Vulkan analogue (WebGPU-specific online-softmax attention;
Vulkan has only a materialized attention). It also makes the long generated WGSL
provenance and constant declarations format-stable and covers them with a
generator regression test.

Key changes:
- runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl and
  streaming_attention_qwen3_q32_k16_causal_bound.wgsl (+ generated headers): the
  Q16 and Q32 online-softmax Qwen3 kernels.
- Sdpa.cpp, WebGPUGraph.{cpp,h}: exact Qwen3 geometry and limit guards, Q16
  default route selection, and the fp32-host to fp16-device KV-cache conversion.
- WebGPUBackend.cpp: read the optional `sdpa_query_tile` runtime spec and thread
  it into graph build so the Q32 tile can be requested without a rebuild.
- scripts/gen_wgsl_headers.py (+ test_wgsl_codegen.py): format-stable generated
  headers with a regression test.
ghstack-source-id: 411961459
@exported-using-ghexport

Differential Revision: [D113171744](https://our.internmc.facebook.com/intern/diff/D113171744/)

Author: JCNTH

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;