WebGPU: Optimize GraphCacheManager by removing pending_buffers - #27453
Merged
Conversation
guschmue
approved these changes
Feb 25, 2026
feich-ms
added a commit
that referenced
this pull request
Jun 11, 2026
Remove pending_buffers_ intermediate queue. Instead, track the current GraphCaptureState in a member and branch directly in ReleaseBuffer: Default mode does immediate bucket reuse (preserving PR #27453's optimization), Capturing/Replaying quarantines to captured_buffers_. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
4 tasks
feich-ms
added a commit
that referenced
this pull request
Jun 11, 2026
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
added a commit
that referenced
this pull request
Jun 11, 2026
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
added a commit
that referenced
this pull request
Jun 11, 2026
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
added a commit
that referenced
this pull request
Jun 11, 2026
Tested end-to-end: graph capture produces correct output ("2+2 is 4")
with PR #27453's immediate buffer reuse. The quarantine was incorrectly
attributed as a fix during the original GC investigation. Phi-4 graph
capture also works without it, confirming immediate reuse is safe.
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
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.
This pull request refactors the buffer management logic in the
GraphCacheManagerclass to simplify buffer handling and improve initialization. The main changes focus on removing the use of a pending buffer queue and ensuring buckets are properly initialized.Benefits:
✅ Simpler, more efficient code
✅ Immediate buffer reuse - no waiting for OnRefresh()
✅ Less memory overhead