huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into sliding_window

Browse Source
This commit is contained in:
Wang, Yi A 2024-10-08 10:00:48 -04:00
commit c2a8819edf
2 changed files with 62 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/layers/attention/kv_cache.py
View File

@ -26,8 +26,8 @@ class KVCache:
if (
dtype == torch.float8_e5m2
and ATTENTION != "flashinfer"
and SYSTEM != "cuda"
and (ATTENTION != "flashinfer"
or SYSTEM != "cuda")
):
raise ValueError(
"float8_e5m2 KV cache is currently only supported for flashinfer on CUDA"

38
server/text_generation_server/models/custom_modeling/mllama.py
View File

@ -19,7 +19,12 @@ from typing import Optional, Tuple, List
import torch
import torch.utils.checkpoint
from torch import nn
import flash_attn_2_cuda
from text_generation_server.utils.import_utils import SYSTEM
if SYSTEM == "ipex":
import intel_extension_for_pytorch as ipex
else:
import flash_attn_2_cuda
from transformers.activations import ACT2FN
import torch.nn.functional as F
@ -698,6 +703,37 @@ class MllamaTextCrossAttention(nn.Module):
# logger.info(
# f"Q: {query_states.shape} -K {key_states.shape} - V{value_states.shape}"
# )
if SYSTEM == "ipex":
attn_output = torch.empty_like(query_states)
ipex.llm.functional.varlen_attention(
(
query_states.contiguous()
if query_states.device.type == "xpu"
else query_states
),
(
key_states.contiguous()
if key_states.device.type == "xpu"
else key_states
),
(
value_states.contiguous()
if value_states.device.type == "xpu"
else value_states
),
attn_output,
cu_seqlen_q,
cu_seqlen_k,
max_q,
max_k,
0.0,
self.softmax_scale,
False,
causal,
False,
None,
)
else:
attn_output = flash_attn_2_cuda.varlen_fwd(
query_states,
key_states,