Skip to content

PT slice_update in place - #1982

Open
polycrit wants to merge 2 commits into
masterfrom
sheremeta/torch-slice-update-inplace
Open

polycrit wants to merge 2 commits into
masterfrom
sheremeta/torch-slice-update-inplace

Conversation

@polycrit

@polycrit polycrit commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

The torch backend implements slice_update with scatter_, so the frame is written into the target's buffer.

Before, the torch backend used the generic fallback, a where over the whole axis. The bounded KV cache of CausalSelfAttention calls slice_update once per decoding step, so every step copied the whole cache.

start is a scalar or has other dims of the target (one position per batch entry / hypothesis, as the bounded KV cache passes it). A start beyond the axis fails (CPU error, CUDA device-side assert, also under graph capture, without host sync), where the generic fallback silently writes nothing.

Where autograd records the write (grad enabled and target or value requiring grad), it uses scatter and writes into a copy. The index stays on the target's device, so there is no host sync.

Tests:

  • test_slice_update_torch_in_place: the result shares the target's storage, the value lands at start and the rest stays unchanged.
  • test_slice_update_torch_autograd: with a target requiring grad, the target keeps its values and the gradient reaches it.
  • test_slice_update_torch_start_per_entry: per-entry and scalar start, the last position, a start beyond the axis fails.
  • test_causal_self_att_bounded_cache_overflow_fails: the bounded KV cache writes its last free slot, a step beyond fails.

Supersedes #2025.

…ails

From #2025: `start` may have other dims of the target (per-hypothesis lengths of the bounded KV cache).
`scatter_` instead of `index_copy_`, still in place, `scatter` where autograd records the write.
A start beyond the axis fails (CPU error, CUDA device-side assert) instead of writing nothing.
@albertz

albertz commented Oct 17, 2026

Copy link
Copy Markdown
Member

Pushed be10e2d, merging #2025 into this PR (#2025 gets closed as duplicate):

  • start may have other dims of the target, i.e. one position per batch entry / hypothesis.
    The bounded KV cache of CausalSelfAttention passes exactly that (per-hypothesis lengths),
    and the index_copy_ with reshape(1) failed there (shape '[1]' is invalid for input of size 2).
  • scatter_ instead of index_copy_, still in place; scatter where autograd records the write.
  • A start beyond the axis fails (CPU error, CUDA device-side assert, also under graph capture, without host sync),
    where the generic fallback silently writes nothing.
  • rf.slice_update docstring: start scalar or with other dims of target.

Tests added: test_slice_update_torch_start_per_entry (per-entry and scalar start, last position, beyond the axis fails),
test_causal_self_att_bounded_cache_overflow_fails (last free slot written, a step beyond fails).
At RZ (torch 2.7, CPU), test_rf_array.py + test_rf_attention.py -k "slice_update or att": 33 passed, 1 skipped (CUDA only).
Negative control (only the tests on 522c28c): both new tests fail.

@albertz
albertz marked this pull request as ready for review October 17, 2026 22:55
@albertz
albertz requested review from a team and albertz as code owners October 17, 2026 22:55
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.

2 participants