[ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF) by JCNTH · Pull Request #20770 · pytorch/executorch

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

JCNTH · 2026-07-07T17:59:09Z

Stack from ghstack (oldest at bottom):

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.pyget_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

[ghstack-poisoned]

pytorch-bot · 2026-07-07T17:59:13Z

🔗 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 · 2026-07-07T18:00:19Z

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.

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Update
3649dd4
[ghstack-poisoned]

JCNTH requested review from kirklandsign and larryliu0820 as code owners July 7, 2026 17:59

JCNTH temporarily deployed to cadence July 7, 2026 17:59 — with GitHub Actions Inactive

This was referenced Jul 7, 2026
[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727
Merged
[ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
Merged

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

JCNTH temporarily deployed to cadence July 7, 2026 18:23 — with GitHub Actions Inactive

JCNTH closed this Jul 28, 2026

JCNTH had a problem deploying to cherry-pick-bot July 28, 2026 05:16 — with GitHub Actions Failure

Uh oh!

Conversation

JCNTH commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

github-actionsBot commented Jul 7, 2026

This PR needs a release notes: label

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

JCNTH commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

github-actionsBot commented Jul 7, 2026

This PR needs a release notes: label

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

JCNTH commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

github-actionsBot commented Jul 7, 2026

This PR needs a release notes: label

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

JCNTH commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Jul 7, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

github-actionsBot commented Jul 7, 2026

This PR needs a release notes: label

Uh oh!

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading

This PR needs a `release notes:` label

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading

This PR needs a `release notes:` label

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading

This PR needs a `release notes:` label

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading

This PR needs a `release notes:` label

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading

This PR needs a `release notes:` label

JCNTH commented Jul 7, 2026•
edited
Loading

pytorch-botBot commented Jul 7, 2026•
edited
Loading