From 8d044d53e532ff14ea9c7c84b72ec02686476984 Mon Sep 17 00:00:00 2001 From: Hyunwoong Ko Date: Tue, 7 Sep 2021 19:56:55 +0900 Subject: [PATCH] Update matmul.py --- deepspeed/ops/sparse_attention/matmul.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/ops/sparse_attention/matmul.py b/deepspeed/ops/sparse_attention/matmul.py index 518af17b41be..6b5481a71ee6 100755 --- a/deepspeed/ops/sparse_attention/matmul.py +++ b/deepspeed/ops/sparse_attention/matmul.py @@ -173,7 +173,7 @@ def _sdd_matmul(a, 'TN': block * pack, 'TMN': block * block * pack * pack, 'BLOCK': block, - 'TK': 32, + 'TK': TK[0], 'TYPE': dtype, 'STRIDE_AM': '1' if trans_a else 'lda', 'STRIDE_AK': 'lda' if trans_a else '1',