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

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -98,14 +98,22 @@ Result<DelegateHandle*> WebGPUBackend::init(
       enable_f16_accumulate_gemm = spec.get();
     }
   }
+  int sdpa_query_tile = 0;
+  {
+    Result<int> spec = context.get_runtime_spec<int>("sdpa_query_tile");
+    if (spec.ok()) {
+      sdpa_query_tile = spec.get();
+    }
+  }
 
   try {
     graph->build(
         flatbuffer_data,
         constant_data,
         context.get_named_data_map(),
         enable_f16_kv_cache,
-        enable_f16_accumulate_gemm);
+        enable_f16_accumulate_gemm,
+        sdpa_query_tile);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU graph build failed: %s", e.what());
     graph->~WebGPUGraph();
@@ -140,8 +148,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);
@@ -177,12 +190,15 @@ Error WebGPUBackend::execute(
     graph->execute(graph_options);
 
     // 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, graph_options);
   } catch (const std::exception& e) {

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -11,6 +11,7 @@
 
 #include <executorch/backends/vulkan/serialization/schema_generated.h>
 #include <executorch/runtime/core/named_data_map.h>
+#include <executorch/runtime/core/portable_type/half.h>
 
 #include <executorch/backends/webgpu/runtime/WebGPUCompat.h>
 #include <executorch/backends/webgpu/runtime/WebGPUDevice.h>
@@ -813,7 +814,8 @@ void WebGPUGraph::build(
     const uint8_t* constant_data,
     const executorch::runtime::NamedDataMap* named_data_map,
     bool f16_kv_cache,
-    bool f16_accumulate_gemm) {
+    bool f16_accumulate_gemm,
+    int sdpa_query_tile) {
   if (!device_) {
     auto* ctx = get_default_webgpu_context();
     if (ctx) {
@@ -843,6 +845,10 @@ void WebGPUGraph::build(
   // additionally gates the kernel on the negotiated shader-f16 feature.
   f16_accumulate_gemm_ = f16_accumulate_gemm;
 
+  // SDPA query-tile selector (runtime opt-in); 0 = geometry default (Q16),
+  // 32 = Q32 candidate. Read at the SDPA op-lowering selection site.
+  sdpa_query_tile_ = sdpa_query_tile;
+
   // Phase 1: Create all values
   const auto* values = graph->values();
   const int num_vals = values ? values->size() : 0;
@@ -1725,6 +1731,22 @@ void WebGPUGraph::copy_inputs(const std::vector<InputData>& inputs) {
       continue;
     }
 
+    const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2;
+    // 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<const float*>(in.data);
+      std::vector<executorch::runtime::etensor::Half> 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" +
@@ -1775,7 +1797,11 @@ 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.rfind("sdpa_streaming_attention_", 0) == 0 &&
+      kernel_name.find("k16_causal_bound") != std::string::npos) {
+    // llama + qwen3 (Q16/Q32) streaming causal-bound kernels; the
+    // sdpa_streaming_attention_ prefix guards against an unrelated future
+    // kernel whose label merely contains the k16_causal_bound substring.
     bits = kRoutePrefill | kRouteK16CausalBound;
   } else if (kernel_name == "sdpa_streaming_attention_k16") {
     bits = kRoutePrefill | kRouteK16;
@@ -2030,7 +2056,7 @@ void buffer_map_callback(
 } // namespace
 
 void WebGPUGraph::copy_outputs(
-    std::vector<std::pair<void*, size_t>>& outputs,
+    std::vector<OutputData>& outputs,
     const WebGPUGraphExecutionOptions& options) {
   const size_t count = std::min(outputs.size(), output_staging_buffers_.size());
   const WebGPUExecutionPlan plan = plan_webgpu_execution(
@@ -2043,38 +2069,71 @@ void WebGPUGraph::copy_outputs(
   std::vector<MapCallbackData> cb_data(count);
   std::vector<WGPUFuture> map_futures(count, WGPUFuture{});
 
+  // Single source of truth for the fp16-widen predicate + host map size, shared
+  // by the map-request and map-result loops so the two cannot drift apart.
+  auto output_map_size = [&](size_t i) -> std::pair<bool, size_t> {
+    const auto& tensor = tensors_[output_ids_[i]];
+    const bool widen_fp16 = !tensor.is_int && tensor.elem_size == 2 &&
+        outputs[i].nbytes == tensor.cur_nbytes * 2;
+    return {widen_fp16, widen_fp16 ? tensor.cur_nbytes : outputs[i].nbytes};
+  };
+
+  // Validate dtypes up front, before any wgpuBufferMapAsync is issued: require
+  // an explicit fp32 host dtype for the fp16->fp32 widen (mirrors the
+  // copy_inputs narrow guard) so a same-2:1-ratio non-fp32 host is not silently
+  // reinterpreted as fp32. Throwing here — rather than mid map-request loop —
+  // keeps in-flight async maps (whose callbacks point at cb_data) from being
+  // left dangling when the exception unwinds this frame.
   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;
+    }
+    if (output_map_size(i).first && !outputs[i].host_is_fp32) {
+      throw std::runtime_error(
+          "WebGPU: fp16 device output requires an fp32 host tensor");
+    }
+  }
+
+  for (size_t i = 0; i < count; i++) {
+    if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) {
       cb_data[i].status = WGPUMapAsyncStatus_Success;
       continue;
     }
     WGPUBufferMapCallbackInfo cb_info = {};
     cb_info.mode = WGPUCallbackMode_WaitAnyOnly;
     cb_info.callback = buffer_map_callback;
     cb_info.userdata1 = &cb_data[i];
+    const size_t map_nbytes = output_map_size(i).second;
     map_futures[i] = wgpuBufferMapAsync(
-        output_staging_buffers_[i],
-        WGPUMapMode_Read,
-        0,
-        outputs[i].second,
-        cb_info);
+        output_staging_buffers_[i], WGPUMapMode_Read, 0, map_nbytes, cb_info);
   }
 
   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 &&
         webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) {
       throw std::runtime_error("WebGPU: WaitAny failed for output 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;
     }
     if (cb_data[i].status == WGPUMapAsyncStatus_Success) {
+      const auto [widen_fp16, map_nbytes] = output_map_size(i);
       const void* mapped = wgpuBufferGetConstMappedRange(
-          output_staging_buffers_[i], 0, outputs[i].second);
-      std::memcpy(outputs[i].first, mapped, outputs[i].second);
+          output_staging_buffers_[i], 0, map_nbytes);
+      if (widen_fp16) {
+        const auto* src =
+            static_cast<const executorch::runtime::etensor::Half*>(mapped);
+        auto* dst = static_cast<float*>(outputs[i].data);
+        const size_t numel = map_nbytes / sizeof(*src);
+        for (size_t e = 0; e < numel; e++) {
+          dst[e] = static_cast<float>(src[e]);
+        }
+      } else {
+        std::memcpy(outputs[i].data, mapped, outputs[i].nbytes);
+      }
       wgpuBufferUnmap(output_staging_buffers_[i]);
     } else {
       throw std::runtime_error("WebGPU buffer map failed for output");

backends/webgpu/runtime/WebGPUGraph.h

@@ -43,6 +43,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 {
@@ -103,7 +112,8 @@ class WebGPUGraph {
       const uint8_t* constant_data,
       const executorch::runtime::NamedDataMap* named_data_map = nullptr,
       bool f16_kv_cache = false,
-      bool f16_accumulate_gemm = false);
+      bool f16_accumulate_gemm = false,
+      int sdpa_query_tile = 0);
 
   // Copy input tensor data from host pointers into GPU buffers.
   void copy_inputs(const std::vector<InputData>& inputs);
@@ -114,7 +124,7 @@ class WebGPUGraph {
   // Copy output tensor data from GPU buffers back to host pointers.
   // Uses mapAsync + ASYNCIFY in Wasm.
   void copy_outputs(
-      std::vector<std::pair<void*, size_t>>& outputs,
+      std::vector<OutputData>& outputs,
       const WebGPUGraphExecutionOptions& options);
 
   const std::vector<int>& input_ids() const {
@@ -383,6 +393,12 @@ class WebGPUGraph {
     return f16_accumulate_gemm_;
   }
 
+  // Runtime-selected SDPA query-tile candidate; 0 = geometry default (Q16),
+  // 32 = Q32 candidate.
+  int sdpa_query_tile() const {
+    return sdpa_query_tile_;
+  }
+
  private:
 #ifdef WGPU_BACKEND_ENABLE_PROFILING
   void record_active_route(const std::string& kernel_name);
@@ -391,6 +407,7 @@ class WebGPUGraph {
   bool kv_f16_ = false;
   std::unordered_set<int> kv_cache_ids_;
   bool f16_accumulate_gemm_ = false;
+  int sdpa_query_tile_ = 0;
 
  private:
   WGPUInstance instance_ = nullptr;