Skip to content

Commit

Permalink
update batch kernel condition (#12408)
Browse files Browse the repository at this point in the history
  • Loading branch information
MeouSker77 authored Nov 15, 2024
1 parent 6c5e8fc commit 3d5fbf2
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions python/llm/src/ipex_llm/transformers/low_bit_linear.py
Original file line number Diff line number Diff line change
Expand Up @@ -362,16 +362,19 @@ def use_batch_forward(x: torch.Tensor, qtype: int, output_len: int):
batch_size = x.shape[0]
hard_condition = (
x.dtype in [torch.float, torch.half]
and batch_size <= 48
and (
x.shape[1] % 128 == 0 and qtype in [SYM_INT4]
(
qtype in [SYM_INT4, ASYM_INT4, FP8E5, FP8E4]
and x.shape[1] % 128 == 0
)
or (
x.shape[1] % 256 == 0
and output_len % 32 == 0
qtype in [SYM_INT8, FP4, FP6, Q4_K, Q6_K]
and device in ["arc", "flex", "pvc", "mtl"]
and qtype in [ASYM_INT4, SYM_INT8, FP4, FP8E5, FP6, FP8E4, Q4_K, Q6_K]
and x.shape[1] % 256 == 0
and output_len % 32 == 0
)
)
and batch_size <= 48
)
if hard_condition:
return (
Expand Down

0 comments on commit 3d5fbf2

Please sign in to comment.