update batch kernel condition (#12408)
This commit is contained in:
parent
6c5e8fc70c
commit
3d5fbf2069
1 changed files with 8 additions and 5 deletions
|
|
@ -362,16 +362,19 @@ 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 (
|
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 (
|
or (
|
||||||
x.shape[1] % 256 == 0
|
qtype in [SYM_INT8, FP4, FP6, Q4_K, Q6_K]
|
||||||
and output_len % 32 == 0
|
|
||||||
and device in ["arc", "flex", "pvc", "mtl"]
|
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:
|
if hard_condition:
|
||||||
return (
|
return (
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue