[ExecuTorch] Batch-aware torch.ops.llama.sdpa_with_kv_cache. #4822
Conversation
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
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/4822
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 1 PendingAs of commit 9508947 with merge base 06c0fa3 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
added a commit
to meta-emilian/executorch
that referenced
this pull request
Aug 21, 2024
Summary: Pull Request resolved: pytorch#4822 This is part 1 of a multi-part commit to make torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this code needs to update the cache across the batch dimension. Differential Revision: D61605316
meta-emilian
force-pushed
the
export-D61605316
branch
from
0f8c06a to
54bdb58
Compare
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
added a commit
to meta-emilian/executorch
that referenced
this pull request
Sep 16, 2024
Summary: Pull Request resolved: pytorch#4822 This is part 1 of a multi-part commit to make torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this code needs to update the cache across the batch dimension. Differential Revision: D61605316
meta-emilian
force-pushed
the
export-D61605316
branch
from
54bdb58 to
ae0964a
Compare
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
force-pushed
the
export-D61605316
branch
from
ae0964a to
1e0568b
Compare
meta-emilian
added a commit
to meta-emilian/executorch
that referenced
this pull request
Sep 17, 2024
Summary: Pull Request resolved: pytorch#4822 This changes makes torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. * Makes update_cache update across the batch dimension As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. sdpa_with_kv_cache fuses this update with the flash attention op for tensors that follow a predetermined format [batch, length, heads, dim]. This change removes the assumption that batch == 1. * Makes sdpa_with_kv_cache apply cpu_flash_attention for all batch lines as well. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this is no longer true. Reviewed By: kimishpatel Differential Revision: D61605316
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
force-pushed
the
export-D61605316
branch
from
1e0568b to
a9dae49
Compare
meta-emilian
added a commit
to meta-emilian/executorch
that referenced
this pull request
Sep 17, 2024
Summary: Pull Request resolved: pytorch#4822 This changes makes torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. * Makes update_cache update across the batch dimension As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. sdpa_with_kv_cache fuses this update with the flash attention op for tensors that follow a predetermined format [batch, length, heads, dim]. This change removes the assumption that batch == 1. * Makes sdpa_with_kv_cache apply cpu_flash_attention for all batch lines as well. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this is no longer true. Reviewed By: kimishpatel Differential Revision: D61605316
meta-emilian
changed the title
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
force-pushed
the
export-D61605316
branch
from
a9dae49 to
377e496
Compare
meta-emilian
added a commit
to meta-emilian/executorch
that referenced
this pull request
Sep 17, 2024
Summary: Pull Request resolved: pytorch#4822 This changes makes torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. * Makes update_cache update across the batch dimension As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. sdpa_with_kv_cache fuses this update with the flash attention op for tensors that follow a predetermined format [batch, length, heads, dim]. This change removes the assumption that batch == 1. * Makes sdpa_with_kv_cache apply cpu_flash_attention for all batch lines as well. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this is no longer true. Reviewed By: kimishpatel Differential Revision: D61605316
Summary: Pull Request resolved: pytorch#4822 This changes makes torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search. * Makes update_cache update across the batch dimension As a performance optimization, update_cache implements the following operation ``` k_cache[:, start_pos : start_pos + seq_len, :, :] = k v_cache[:, start_pos : start_pos + seq_len, :, :] = v ``` as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. sdpa_with_kv_cache fuses this update with the flash attention op for tensors that follow a predetermined format [batch, length, heads, dim]. This change removes the assumption that batch == 1. * Makes sdpa_with_kv_cache apply cpu_flash_attention for all batch lines as well. ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this is no longer true. Reviewed By: kimishpatel Differential Revision: D61605316
|
This pull request was exported from Phabricator. Differential Revision: D61605316 |
meta-emilian
force-pushed
the
export-D61605316
branch
from
377e496 to
9508947
Compare
tarun292
approved these changes
Sep 17, 2024
|
This pull request has been merged in 53c1a5f. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
Summary:
This changes makes torch.ops.llama.sdpa_with_kv_cache batch aware. This is needed for batched sdpa cases, for example llm beam search.
As a performance optimization, update_cache implements the following operation
as part of the fused sdpa_with_kv_cache op. A naiive export of this code inserts expensive slice-scatter ops. sdpa_with_kv_cache fuses this update with the flash attention op for tensors that follow a predetermined format [batch, length, heads, dim]. This change removes the assumption that batch == 1.
ExecuTorch-exported Llama models are implemented with a greedy search, so it has not been necessary for this op to be batch-aware. However when working with other models, or when doing LLM beam search, this is no longer true.
Reviewed By: kimishpatel
Differential Revision: D61605316