Torch Compile Llama Slow Down - Unnecessary Copies

jamesdborin: <p>Hi there,</p>
<p>I was trying to recreate some of the speed ups for Llama models using torch compile like i have seen in a few blog posts, <a href="https://blog.fireworks.ai/speed-python-pick-two-how-cuda-graphs-enable-fast-python-code-for-deep-learning-353bf6241248" rel="noopener nofollow ugc">here</a> and <a href="https://pytorch.org/blog/pytorch-compile-to-speed-up-inference/" rel="noopener nofollow ugc">here</a>.</p>
<p>I see speed ups for small batch size and sequence lengths but when I get to large batch sizes and sequence lengths it slows down compared to torch eager.</p>
<p>I looked at the profiler and saw a large amount of time spent copying data after the model was finished executing:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0d9c41a88b2762947b8fbad7f81af17216e9a4b.png" data-download-href="https://discuss.pytorch.org/uploads/default/a0d9c41a88b2762947b8fbad7f81af17216e9a4b" title="torchscreenshot2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0d9c41a88b2762947b8fbad7f81af17216e9a4b_2_690x296.png" alt="torchscreenshot2" data-base62-sha1="mWWXk5y23DNeMQhsEZp9Hr50Qu7" width="690" height="296" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0d9c41a88b2762947b8fbad7f81af17216e9a4b_2_690x296.png, https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0d9c41a88b2762947b8fbad7f81af17216e9a4b.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0d9c41a88b2762947b8fbad7f81af17216e9a4b.png 2x" data-dominant-color="AABEDE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">torchscreenshot2</span><span class="informations">979×421 53.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In <code>fx_graph_readable.py</code>, these copies are at the end:</p>
<pre data-code-wrap="python"><code class="lang-python">copy_: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg204_1, slice_scatter_2);  arg204_1 = slice_scatter_2 = None
copy__1: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg205_1, slice_scatter_5);  arg205_1 = slice_scatter_5 = None
copy__2: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg206_1, slice_scatter_8);  arg206_1 = slice_scatter_8 = None
copy__3: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg207_1, slice_scatter_11);  arg207_1 = slice_scatter_11 = None
copy__4: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg208_1, slice_scatter_14);  arg208_1 = slice_scatter_14 = None
copy__5: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg209_1, slice_scatter_17);  arg209_1 = slice_scatter_17 = None
copy__6: "bf16[s0, s1, 4, 64]" = torch.ops.aten.copy_.default(arg210_1, slice_scatter_20);  arg210_1 = slice_scatter_20 = None
...
</code></pre>
<p>And it looks like they have something to do with the kv-cache tensors.</p>
<p>When we add the newly generated k and v states to the kv-cache torch compile generates this code:</p>
<pre data-code-wrap="python"><code class="lang-python"># File: /home/llama.py:311, code: k_cache[:bsz, current_seq_pos, :, :] = key_states
slice_5: bf16[s0, s1, 4, 64] = torch.ops.aten.slice.Tensor(arg206_1, 2, 0, 9223372036854775807)
slice_6: bf16[s0, s1, 4, 64] = torch.ops.aten.slice.Tensor(slice_5, 3, 0, 9223372036854775807);  slice_5 = None
index_put: bf16[s0, s1, 4, 64] = torch.ops.aten.index_put.default(slice_6, [None, arg250_1], permute_6);  slice_6 = permute_6 = None
slice_7: bf16[s0, s1, 4, 64] = torch.ops.aten.slice.Tensor(arg206_1, 2, 0, 9223372036854775807)
slice_scatter: bf16[s0, s1, 4, 64] = torch.ops.aten.slice_scatter.default(slice_7, index_put, 3, 0, 9223372036854775807);  slice_7 = index_put = None
slice_scatter_1: bf16[s0, s1, 4, 64] = torch.ops.aten.slice_scatter.default(arg206_1, slice_scatter, 2, 0, 9223372036854775807);  slice_scatter = None
</code></pre>
<p>slice_scatter_1 (and all following slice scatters) are the tensors at the end getting copied back into the original kv cache buffers (arg206_1, arg207_1, etc).</p>
<p>It looks like adding the values from <code>key_states</code> into the slice <code>k_cache[:bsz, current_seq_pos, :, :]</code> doesn’t do this in place, but creates a separate buffer, the contents of which have to be copied back into k_cache in order to be used again for the next iteration.</p>
<p>Attempts to make this operation happen in place using <code>k_cache[slice].copy_(...)</code> didn’t change the result.</p>
<p>Is there a way to get torch.compile to see that this is in an place operation and edit the buffer directly, so that <code>slice_scatter_1</code> is replaced with <code>arg206_1</code> and we can avoid these copies? It looks like torch eager does this correctly according to its profile.</p>
<p>Thanks in advance!</p>

bdhirsh: <p>Do you have a self contained repro of the code you’re using that shows the slowdown?</p>
<p>Also - can you try running on a nightly and see any improved perf? (there have been some recent changes to inductor’s reinplacing logic that may or may not make this faster)</p>