[WebGPU] Graph capture support for KV-shared decoder models - #29000
Closed
feich-ms wants to merge 1 commit into
Closed
[WebGPU] Graph capture support for KV-shared decoder models#29000feich-ms wants to merge 1 commit into
feich-ms wants to merge 1 commit into
Conversation
feich-ms
marked this pull request as draft
June 11, 2026 03:44
feich-ms
force-pushed
the
user/feich/webgpu-gc-production
branch
2 times, most recently
from
June 11, 2026 03:52
00f7030 to
a9f679b
Compare
Three changes required for graph capture to work on models with KV-shared layers (e.g. Gemma4), independent of how the model was built: 1. ConstantOfShape WebGPU kernel — KV-shared layers use ConstantOfShape to create zero-filled tensors. Without a WebGPU kernel, these fall back to CPU, flooding the captured graph with MemcpyFromHost. Registered opset-by-opset (9, 20, 21, 23, 24) to match ONNX spec changes and KernelRegistry::VerifyVersion semantics. 2. GraphCacheManager buffer quarantine — during graph capture/replay, released buffers must not be immediately reused because captured bind_groups still reference them. Tracks GraphCaptureState in a member; Default mode keeps PR #27453 immediate-reuse optimization, Capturing/Replaying quarantines to captured_buffers_. 3. PrepareIndirectDispatch for kv_empty layers — KV-shared layers skip CopyKVCache (no new K/V to append), which normally populates the indirect dispatch buffer. Under graph capture, step 2+ crashes with Invalid dispatch group size (0, 1, 1). A single-thread shader now fills the buffer from seqlen_k[0]+1 when CopyKVCache is skipped. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
feich-ms
force-pushed
the
user/feich/webgpu-gc-production
branch
from
June 11, 2026 04:37
a9f679b to
cdcb07a
Compare
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
1 participant
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.
Summary
Two changes required for WebGPU graph capture to work on models with KV-shared layers (e.g. Gemma4), independent of how the model was built (works with both
enable_webgpu_graph=trueand without):seqlen_k[0]+1when CopyKVCache is skipped under graph capture. Also widens the indirect dispatch predicate to cover GPU-residenttotal_sequence_length(sentinel 0).Test plan
enable_webgpu_graph=true(pending genai builder support)🤖 Generated with Claude Code