thek0tyara/llama-cpp-turboquant
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 6

Fix small matrices matmuls, imp3 working but slow

Browse source
This commit is contained in:
Corentin 2021-06-28 19:05:39 +09:00
parent a3f7793c17
commit 3962ee70af
5 changed files with 77 additions and 219 deletions
Download patch file Download diff file Expand all files Collapse all files

10
examples/python_naive_matmul/imp2_tiled.py
View file

@ -24,7 +24,7 @@ class MatMulOp:
assert tile_size <= max_workgroup_size[1]
self.tile_size = tile_size
self.shader = f'''
self.shader = '''
#version 450
layout (local_size_x = {tile_size}, local_size_y = {tile_size}) in;
@ -66,7 +66,7 @@ void main()
}}
out_tensor[tensor_size * globalCol + globalRow] = acc;
}}'''
self.compiled_shader = kp.Shader.compile_source(self.shader)
self.compiled_shader = kp.Shader.compile_source(self.shader.format(tile_size=tile_size))
self.tensor_shape: tuple[int, int] = (0, 0)
self.params: list[kp.Tensor] = []
self.algo = None
@ -78,7 +78,9 @@ void main()
if self.algo is None or self.tensor_shape != tensor_shape or self.params != params:
self.tensor_shape = tensor_shape
self.params = params
workgroup = (tensor_shape[0] // self.tile_size, tensor_shape[1] // self.tile_size, 1)
tile_size = min(tensor_shape[0], tensor_shape[1], self.tile_size)
self.compiled_shader = kp.Shader.compile_source(self.shader.format(tile_size=tile_size))
workgroup = [tensor_shape[0] // tile_size, tensor_shape[1] // tile_size, 1]
self.algo = self.mgr.algorithm(
params, # params
self.compiled_shader, # spirv
@ -87,7 +89,7 @@ void main()
[]) # push_consts
(self.mgr.sequence()
.record(kp.OpTensorSyncDevice(self.params))
.record(kp.OpTensorSyncDevice([tensor_in_1, tensor_in_2]))
.record(kp.OpAlgoDispatch(self.algo))
.record(kp.OpTensorSyncLocal([tensor_out]))
.eval())