[WebGPU] Support continuous decoding (RewindTo) with graph capture

[qjia7]

This pull request introduces improvements to the handling of attention masks in both the CUDA and WebGPU backends, focusing on more efficient and correct updates of mask buffers during decoding. The main changes are the implementation of a CPU-side update for static attention masks in CUDA and the addition of a reusable staging buffer for efficient mask updates in WebGPU, with logic to avoid redundant work for single-beam cases.

CUDA backend improvements:

• Replaced the previous (commented-out and incorrect) CUDA memory set logic in DefaultPositionInputs::RewindMask with a CPU-side update that correctly sets attended and non-attended positions in the attention mask for each batch/beam, followed by a copy back to the device. This ensures the mask is set with 1s for attended tokens and 0s for future tokens, supporting both int32_t and int64_t types.

WebGPU backend improvements:

• Added a reusable CPU staging buffer (mask_staging_buffer_) to the InterfaceImpl struct for efficient attention mask updates, avoiding repeated allocations and redundant writes.
• Implemented the UpdateAttentionMask method to efficiently update the mask for single-beam cases by only filling new positions with 1s and copying the relevant portion to the device, falling back to CPU for multi-beam cases. This method handles static update path and supports both int32_t and int64_t mask types.

[Copilot]

Pull request overview

This PR improves continuous decoding support when graph capture/static attention masks are used, by updating attention mask buffers more efficiently and correctly across backends (notably CUDA/static mask rewind handling and WebGPU mask updates).

Changes:

• Implement CPU-side rewind/update logic for static attention masks used with graph capture (write attended positions to 1 and future positions to 0, then copy back to device).
• Add a reusable CPU staging buffer and a WebGPU UpdateAttentionMask fast path for single-beam cases to avoid repeated allocations and redundant CPU fills.

Reviewed changes

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

File	Description
src/webgpu/interface.cpp	Adds a reusable staging buffer and implements WebGPU-side UpdateAttentionMask to upload an all-ones prefix for single-beam decoding.
src/models/position_inputs.cpp	Replaces the previous “unsupported”/commented logic with a CPU-written static mask rewind implementation for graph capture.

[qjia7]

Addressed the review comments in 80be21c:

1. RewindMask() now handles both static and dynamic mask paths

• Static path (ShouldUseStaticMaskHandling() — graph capture or NvTensorRtRtx with shared buffers): overwrites the fixed [batch_beam_size, max_length] tensor with 1s for [0, index) and 0s for [index, max_length), then copies to device. Same as before but now gated on ShouldUseStaticMaskHandling() instead of just use_graph_capture, consistent with UpdateAttentionMask().

• Dynamic path (everything else): sets attention_mask_shape_[1] = index. The next Update() call will:

  1. CreateNextAttentionMaskTensor(total_length) — creates a correctly sized tensor using the adjusted shape
  2. CPU UpdateAttentionMask for batch_beam_size == 1 — fills the entire new mask with 1s (ignores old data)
  3. Replaces attention_mask_ with attention_mask_next_

  No tensor recreation or is_first_update_ reset needed — RewindTo already restricts to batch_size == 1 (no padding), so a shape-only adjustment is sufficient.

2. Fixed max_len type mixing — batch_beam_size, max_len, and loop indices now all use size_t consistently.

Tested on WebGPU (NVIDIA RTX):

• Static mask (phi4-graph-prune, enableGraphCapture=1): ✅
• Dynamic mask (phi4-prune, past_present_share_buffer=false): ✅

Both tests: generate 10 tokens -> rewind_to(prompt_len + 5) -> generate 5 more tokens. Output is consistent and coherent after rewind.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 10 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 10 changed files in this pull request and generated 1 comment.