ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

update batch kernel condition (#12421)

Browse source
This commit is contained in:
Yishuo Wang 2024-11-21 10:12:46 +08:00 committed by GitHub
parent 7288c759ce
commit 145e8b480f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
python/llm/src/ipex_llm/transformers/low_bit_linear.py
View file

@ -362,14 +362,22 @@ def use_batch_forward(x: torch.Tensor, qtype: int, output_len: int):
batch_size = x.shape[0] batch_size = x.shape[0]
hard_condition = ( hard_condition = (
x.dtype in [torch.float, torch.half] x.dtype in [torch.float, torch.half]
and batch_size <= 48 and x.shape[1] % 128 == 0
and ( and (
( (
qtype in [SYM_INT4, ASYM_INT4, FP8E5, FP8E4] qtype in [SYM_INT4, ASYM_INT4, FP8E5, FP8E4]
and x.shape[1] % 128 == 0 and (
batch_size <= 48
or (
batch_size <= 64
and x.shape[1] % 256 == 0
and output_len % 64 == 0
)
)
) )
or ( or (
qtype in [SYM_INT8, FP4, FP6, Q4_K, Q6_K] qtype in [SYM_INT8, FP4, FP6, Q4_K, Q6_K]
and batch_size <= 48
and device in ["arc", "flex", "pvc", "mtl"] and device in ["arc", "flex", "pvc", "mtl"]
and x.shape[1] % 256 == 0 and x.shape[1] % 256 == 0
and output_len % 32 == 0 and output_len % 32 == 0