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.
Stack from ghstack (oldest at bottom):
Summary:
trying to fix the issue with kv_cache update by changing tracing into a
tensor subclass. However it seems we have less success than the fx
tracer. The fx tracer breaks due
k_out[:,:, input_pos] = k_val
getting traced as
new_var = torch.ops.aten.index_put_(k_out, [None, None,
input_pos], k_val)
with new var never being accessed afterward. new_var becomes hte correct
multiInput value, but then is lost.
The subclass ont he other hand, tries to use the func "<slot wrapper 'setitem' of 'torch._C.TensorBase' objects>"
which seems to not want to mutate k_out and so the attempt to make it a
multiTensor fails.
Test Plan: sh run.sh
Reviewers:
Subscribers:
Tasks:
Tags: