Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add sdp support for stablelm 3b #11473

Merged
merged 1 commit into from
Jul 1, 2024
Merged

add sdp support for stablelm 3b #11473

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/llm/src/ipex_llm/transformers/models/stablelm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def stablelm_model_forward(
):
# IPEX-LLM OPT: kv cache and quantize kv cache
use_cache = use_cache if use_cache is not None else self.config.use_cache
use_quantize_kv = (self.layers[0].self_attn.head_dim in [64, 96, 128]
use_quantize_kv = (self.layers[0].self_attn.head_dim in [64, 80, 96, 128]
and use_quantize_kv_cache(self.layers[0].mlp.up_proj, input_ids))
if use_cache:
if use_quantize_kv and not isinstance(past_key_values, DynamicFp8Cache):
Expand Down
4 changes: 2 additions & 2 deletions python/llm/src/ipex_llm/transformers/models/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,7 @@ def use_sdp(q_len, kv_len, head_dim, query_states):
return (
query_states.device.type == "xpu"
and query_states.dtype in [torch.float, torch.half] # fp32/fp16
and head_dim in [64, 96, 128]
and head_dim in [64, 80, 96, 128]
and q_len != kv_len # next token
and q_len <= 32 # lookup
)
Expand All @@ -347,7 +347,7 @@ def use_sdp_fp8(q_len, kv_len, query_states):
def use_sdp_causal(q_len, kv_len, head_dim, query_states, training):
return (
q_len == kv_len # first token
and head_dim in [64, 96, 128] # for now
and head_dim in [64, 80, 96, 128] # for now
and query_states.device.type == "xpu" # GPU
and query_states.dtype in [torch.float, torch.half] # fp32/fp16
and not query_states.requires_grad and not training # not training
Expand Down
Loading