Repository navigation
Conversation
…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.
Member
|
Pushed be10e2d, merging #2025 into this PR (#2025 gets closed as duplicate):
Tests added: |
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
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.
The torch backend implements
slice_updatewithscatter_, so the frame is written into the target's buffer.Before, the torch backend used the generic fallback, a
whereover the whole axis. The bounded KV cache ofCausalSelfAttentioncallsslice_updateonce per decoding step, so every step copied the whole cache.startis a scalar or has other dims of the target (one position per batch entry / hypothesis, as the bounded KV cache passes it). Astartbeyond 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
scatterand 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 atstartand 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.