Make key optional for rotary embedding #17566
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
sarckk
force-pushed
the
rotary-emb-key-optional
branch
from
b99fd50 to
b681cca
Compare
| offsets) | ||
| if key is None: | ||
| # XPU kernel doesn't support key=None so fall back to native impl | ||
| # TODO ipex.llm.functional.rotary_embedding_batched support key=None |
There was a problem hiding this comment.
add the github id here? like TODO([sarckk): xxx
|
@sarckk , could you check the failed tests are related to the changes or not? |
these don't seem to be related to my changes, they are failing on another merged PR as well: https://buildkite.com/vllm/ci/builds/19273#01969974-4d6d-46ca-ac55-e52bea52d5b8 |
| return query, key | ||
|
|
||
| def forward_hpu( | ||
| self, | ||
| positions: torch.Tensor, | ||
| query: torch.Tensor, | ||
| key: torch.Tensor, | ||
| key: Optional[torch.Tensor] = None, | ||
| offsets: Optional[torch.Tensor] = None, | ||
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
There was a problem hiding this comment.
nit: Shouldn't we chance the return type?
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
There was a problem hiding this comment.
thanks for the catch, updated! it would affect the typing of downstream call sites though -- is that ok?
Head branch was pushed to by a user without write access
sarckk
force-pushed
the
rotary-emb-key-optional
branch
from
917e289 to
78486cb
Compare
|
@sarckk Can you please merge from main and restart the CI? I'm not sure whether the CI failures are related to the PR. Maybe we should retry. |
Signed-off-by: Yong Hoon Shin <[email protected]>
Signed-off-by: Yong Hoon Shin <[email protected]>
Signed-off-by: Yong Hoon Shin <[email protected]>
sarckk
force-pushed
the
rotary-emb-key-optional
branch
from
0b2e344 to
d12c2a4
Compare
I'm pretty sure failures are not related (e.g. I can reproduce the spec decode test failures locally on trunk) but I've rebased on main and kicked off a new run. |
|
@sarckk Hmm... The tests still failed. Could you please take another look? |
|
may be just run it without PR locally? |
|
none of the highlighted test failures are due to the PR all 3 spec decoding tests fail locally without PR, on commit the examples-test and intel hpu/xpu failures is due to #17426 amd test failures seem to be present before (https://buildkite.com/vllm/ci/builds/19352#0196a33d-129f-4a93-af93-1fb0d1e8c82f) |
Signed-off-by: Yong Hoon Shin <[email protected]>
Head branch was pushed to by a user without write access
|
correction: neuron test failure is actually due to my PR, but it's an issue with the test set up. pushed fix the newly failing |
|
The remaining tests are failing on main, so this should be good to go |
Signed-off-by: Yong Hoon Shin <[email protected]> Signed-off-by: Mu Huai <[email protected]>
Signed-off-by: Yong Hoon Shin <[email protected]>
Make key an optional argument for rotary embedding. This flexibility may be needed in cross-layer KV sharing, e.g. Layer-Condensed KV Cache and Cross-Layer Attention, where there is no K to apply rotary embedding on.
Unit tested with:
E2E tested with offline inference example both with eager and non-eager.
Note: rotary emb kernel in intel-extension-for-pytorch currently does not support key=None, so falling back to native impl for now, followed up in intel/intel-extension-for-pytorch#821