[ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF)#20772
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20772
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 New Failures, 1 Cancelled Job, 2 PendingAs of commit 5960505 with merge base f4b01a8 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Jul 7, 2026
Merged
meta-cla
Bot
added
the
CLA Signed
This was referenced Jul 7, 2026
This PR needs a
|
This was referenced Jul 9, 2026
SS-JIA
requested changes
Jul 9, 2026
SS-JIA
left a comment
SS-JIA
left a comment
Contributor
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
SS-JIA
approved these changes
Jul 9, 2026
SS-JIA
left a comment
SS-JIA
left a comment
Contributor
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
JCNTH
added a commit
that referenced
this pull request
Jul 9, 2026
…t-OFF)
**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.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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(definesWGPU_BACKEND_KV_F16, default OFF), and only when the device supports shader-f16 (fail-closed to f32 otherwise):numel*4bytes) and read by the f32 SDPA kernels.numel*2bytes, 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:
sdpa_compute_attn_weights,sdpa_compute_out,sdpa_fd_split,update_cache) intoDTYPE-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_halfheader — 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 becomesarray<${buffer_gvec_type(DTYPE, 4)}>(QK/AV) orarray<${buffer_scalar_type(DTYPE)}>(fd_split/update_cache), widened to f32 on read (vec4<f32>(...)/f32(...));update_cachewritesf16(...). Mirrors the codegen usage of the landedrms_norm(an existing kernel retired into a template) andsteel_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_halfshader whenkv_f16()via a localconst char*(mirrors the steel-f16 gate — no function-signature/ABI change).force_fp16fp32→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, 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 requireshead_dim % 4 == 0(already required and guarded).Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D110919974
Differential Revision: D110919974