[ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF)

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Test coverage for the f16 KV cache #20771
• -> [ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF) #20770
• [ExecuTorch][WebGPU] Test coverage for the f16-multiply steel q4gsw GEMM #20753
• [ExecuTorch][WebGPU] Add f16-multiply variant of the steel q4gsw prefill GEMM #20752
• [ExecuTorch][WebGPU] Steel q4gsw prefill GEMM — test coverage #20731
• [ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM #20730
• [ExecuTorch][WebGPU] Request shader-f16 device feature (enablement, fail-open) #20729
• [ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
• [ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727

Add an opt-in native f16 KV cache for WebGPU SDPA — −280 MiB device memory, decode-neutral, default-OFF (byte-identical when off).

Problem: The SDPA K/V cache is stored fp32, so at long context it dominates WebGPU device memory. On a device that negotiated the shader-f16 feature the cache can be stored as native f16 (half the bytes) while attention still accumulates in f32, at no measured decode-speed cost.

Solution: Behind a new opt-in build flag EXECUTORCH_WEBGPU_KV_F16 (defines WGPU_BACKEND_KV_F16, default OFF), and only when the device supports shader-f16 (fail-closed to f32 otherwise):

• Before: the K/V caches are allocated fp32 (numel*4 bytes) and read by the f32 SDPA kernels.
• After: the K/V caches are allocated on a dedicated f16 buffer (numel*2 bytes, zero-init); SDPA and FlashDecoding select _f16 kernel variants that bind the cache as f16 and widen to f32 on read — compute and accumulation are unchanged.

Implementation:

• 4 f16 kernel variants, each a storage-dtype delta of its landed f32 kernel (compute/dispatch/tiling unchanged): sdpa_compute_attn_weights_f16 / sdpa_compute_out_f16 bind array<vec4<f16>> and widen with vec4<f32>(...); sdpa_fd_split_f16 binds scalar array<f16> and widens with f32(...); update_cache_f16 writes f16(...).
• WebGPUGraph::build: collect the sdpa K/V cache value ids (args[3], args[4]), allocate them as a dedicated half-size f16 buffer at the constant-allocation site, plus a defensive guard that throws if a non-sdpa op would misread an f16 cache.
• Sdpa.cpp / SdpaFdDecode.cpp: select the _f16 shader when kv_f16() via a local const char* (mirrors the steel-f16 gate — no function-signature/ABI change).
• Analogous to Vulkan's whole-graph force_fp16 fp32→fp16 storage downcast (backends/vulkan/serialization/vulkan_graph_builder.py get_effective_dtype); here scoped to the K/V cache and gated on the negotiated shader-f16 feature rather than applied graph-wide.

Constraints: the default build (flag OFF) is byte-identical — all f16 code is #ifdef WGPU_BACKEND_KV_F16-gated with no ABI change and no KV-cache-layout migration on the default path; the opt-in build fail-closes to f32 on a non-shader-f16 device; f16 storage requires head_dim % 4 == 0 (already required and guarded).

Co-authored-with: Claude Code.

Differential Revision: D110919974

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20770

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Unrelated Failure

As of commit 3649dd4 with merge base f8c8334 ():

NEW FAILURE - The following job has failed:

• pull / unittest-nxp-neutron / linux-job (gh)
  RuntimeError: Command docker exec -t 10bb0ff012d4ceb591202b13410a85ddd48ceb2d0327f85f8ee12ccc369062fc /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / android / run-emulator (gh) (trunk failure)
  The process '/usr/bin/sh' failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.