use new q4_0 batch kernel (#12396)

This commit is contained in:
Yishuo Wang 2024-11-13 18:37:34 +08:00 committed by GitHub
parent d6d63d6b84
commit 00fce5c940
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -362,12 +362,16 @@ 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 x.shape[1] % 256 == 0 and (
x.shape[1] % 128 == 0 and qtype in [SYM_INT4]
or (
x.shape[1] % 256 == 0
and output_len % 32 == 0 and output_len % 32 == 0
and device in ["arc", "flex", "pvc", "mtl"] and device in ["arc", "flex", "pvc", "mtl"]
and qtype in [SYM_INT4, ASYM_INT4, SYM_INT8, FP4, and qtype in [ASYM_INT4, SYM_INT8, FP4, FP8E5, FP6, FP8E4, Q4_K, Q6_K]
FP8E5, FP6, FP8E4, Q4_K, Q6_K] )
and batch_size <= 64 )
and batch_size <= 48
) )
if hard_condition: if hard_condition:
return ( return (