[llama] Added the fused rotary embedding kernel #719
Merged
Conversation
This file contains 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
rsuderman
force-pushed
the
add_rotary_kernel
branch
from
572a546 to
f1caf4d
Compare
Reworked rotary embedding application to be performed via a custom kernel. This includes dropping `static_table` for the sake of maintenance (it was largely unused). It includes a simple numerical test however under the hood no numerical change should occur. Existing baseline vs hugging face remained unchanged.
rsuderman
force-pushed
the
add_rotary_kernel
branch
from
98580bf to
1ef2f6b
Compare
|
Is it expected that the exported IR becomes much bigger (for 8b f16 bs4 tp1 IR is about 34 GB)? |
Sharding code can cause the prebaked matrix size to explode in size.
Updated how we do causal attention on prefill. It was baking out the attention matrix as a constant which was very beefy. |
aviator19941
approved these changes
Dec 20, 2024
archana-ramalingam
added a commit
that referenced
this pull request
Jan 4, 2025
This reverts commit fc9576b.
monorimet
pushed a commit
that referenced
this pull request
Jan 8, 2025
Reworked rotary embedding application to be performed via a custom kernel. This includes dropping `static_table` for the sake of maintenance (it was largely unused). It includes a simple numerical test however under the hood no numerical change should occur. Existing baseline vs hugging face remained unchanged.
archana-ramalingam
added a commit
that referenced
this pull request
Jan 10, 2025
stbaione
added a commit
to stbaione/SHARK-Platform
that referenced
this pull request
Jan 15, 2025
…kernel (nod-ai#719)" (nod-ai#752)" This reverts commit 63ff841.
stbaione
added a commit
that referenced
this pull request
Jan 16, 2025
…kernel (#719)" (#752)" (#828) This reverts commit 63ff841. # Description Our [user docs](https://github.com/nod-ai/shark-ai/blob/main/docs/shortfin/llm/user/llama_serving.md) started producing repetitive output: ```bash curl http://localhost:8003/generate \ -H "Content-Type: application/json" \ -d '{ "text": "Name the capital of the United States.", "sampling_params": {"max_completion_tokens": 50} }' data: Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington ``` After bisecting, the problematic commit was found to be 63ff841. Without commit reverted: ```bash curl http://localhost:8003/generate \ -H "Content-Type: application/json" \ -d '{ "text": "Name the capital of the United States.", "sampling_params": {"max_completion_tokens": 50} }' data: Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington, D.C.! Washington ``` With commit reverted: ```bash curl http://localhost:8003/generate \ -H "Content-Type: application/json" \ -d '{ "text": "Name the capital of the United States.", "sampling_params": {"max_completion_tokens": 50} }' data: Washington D.C. is the capital of the United States. Washington D.C. is a special district that is not part of any state but is the capital of the United States. Washington D.C. is a federal district that serves as the seat of the ```
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.
Reworked rotary embedding application to be performed via a custom kernel. This includes dropping
static_tablefor the sake of maintenance (it was largely unused). It includes a simple numerical test however under the hood no numerical change should occur.Existing baseline vs hugging face remained unchanged.