[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 #20773
• -> [ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF) #20772
• [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 the generated f16 (_half) kernel variants that bind the cache as f16 and widen to f32 on read — compute and accumulation are unchanged.

Implementation:

• Retires the 4 SDPA/decode kernels (sdpa_compute_attn_weights, sdpa_compute_out, sdpa_fd_split, update_cache) into DTYPE-templated variants via the landed WGSL shader-variant codegen: each <kernel>.wgsl + a <kernel>.yaml (DTYPE: float/half) generates BOTH the existing f32 header (regenerated byte-identical) and a new _half header — the f16 variant is generated, not a hand-duplicated file. The f16 delta is storage-only: $if DTYPE == "half" enables f16, the K/V cache binding becomes array<${buffer_gvec_type(DTYPE, 4)}> (QK/AV) or array<${buffer_scalar_type(DTYPE)}> (fd_split/update_cache), widened to f32 on read (vec4<f32>(...) / f32(...)); update_cache writes f16(...). Mirrors the codegen usage of the landed rms_norm (an existing kernel retired into a template) and steel_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 generated _half 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, the templated f32 headers regenerate byte-identical (drift-check verified), no ABI change, no KV-cache-layout migration; 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.
@exported-using-ghexport

Differential Revision: D110919974

Differential Revision: D110919974

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Cancelled Job

As of commit 6e03bac with merge base f8c8334 ():

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / unittest-nxp-neutron / linux-job (gh)
  ##[error]The operation was canceled.

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.