Skip to content

Fix webgpu shared buffer issue - #2191

Merged
apsonawane merged 4 commits into
mainfrom
asonawane/webgpu
Jun 2, 2026
Merged

Fix webgpu shared buffer issue#2191
apsonawane merged 4 commits into
mainfrom
asonawane/webgpu

Conversation

@apsonawane

Copy link
Copy Markdown
Contributor

This pull request refactors the management of recurrent state buffers in RecurrentState to better support both graph capture (e.g., TRT-RTX) and general execution providers. The main change is to conditionally use shared or separate buffers for past/present state, improving compatibility and correctness across execution environments.

Recurrent state buffer management:

  • Added a share_buffers_ flag to RecurrentState to determine whether to use the same buffer for both input and output (for graph capture) or separate buffers with swapping (for general EPs like WebGPU). [1] [2]
  • Allocates and manages pasts_ and presents_ vectors conditionally based on share_buffers_, ensuring correct buffer usage and initialization.
  • Refactored the Add and Update methods to handle buffer aliasing and swapping according to the execution mode, maintaining stable buffer addresses for graph capture and proper buffer updates otherwise.
  • Updated the RewindTo method to zero and rebind buffers as needed, supporting both shared and separate buffer modes.

Internal state tracking:

  • Introduced input_index_ and output_index_ members to track the range of recurrent state buffers within the global input/output vectors, enabling efficient buffer rebinding after rewinding or updating. [1] [2]

Copilot AI review requested due to automatic review settings June 1, 2026 04:16
@apsonawane
apsonawane requested a review from a team as a code owner June 1, 2026 04:16

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes a WebGPU shared-buffer issue in RecurrentState by introducing a dual-mode buffer strategy: alias input/output to the same buffer when graph capture is active (TRT-RTX), or use separate past/present buffers with swap on each step for general EPs (e.g., WebGPU, which prohibits aliasing the same buffer as read-only and read-write). The implementation mirrors the established pattern in DefaultKeyValueCache.

Changes:

  • Add share_buffers_ flag (derived from params_->use_graph_capture) and input_index_/output_index_ tracking to RecurrentState.
  • Conditionally allocate/zero a separate pasts_ vector in non-shared mode; Add() binds input to pasts_ vs. presents_ accordingly.
  • Implement swap-and-rebind in Update() and zero+rebind in RewindTo(0) for the non-shared path; shared path is unchanged.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File Description
src/models/recurrent_state.h Adds pasts_, share_buffers_, and index members for non-shared buffer tracking.
src/models/recurrent_state.cpp Allocates separate pasts when not in graph-capture mode; swaps/rebinds buffers in Update() and RewindTo().

@apsonawane
apsonawane enabled auto-merge (squash) June 1, 2026 04:41
Comment thread src/models/recurrent_state.cpp Outdated
Comment thread src/models/recurrent_state.cpp Outdated
Comment thread src/models/recurrent_state.cpp
@apsonawane
apsonawane requested a review from baijumeswani June 1, 2026 23:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants