wu-arch/kernels
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Generate S matrix, pull out FA stuff from basic script

Browse Source
This commit is contained in:
Hansung Kim
2024-08-28 16:13:38 -07:00
parent 3f20dd59c0
commit 4260bf7d6e
2 changed files with 12 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/kernel/tensor/flash_attn.py
View File

@@ -94,7 +94,11 @@ if __name__ == "__main__":
def exp2(x): def exp2(x):
return (x**2) / 2.0 + x + 1.0 return (x**2) / 2.0 + x + 1.0
col_to_save = 64 full_S = A_array @ B_array
full_S_T = full_S.transpose([1, 0])
full_S.astype('float32').tofile("full_S.bin")
col_to_save = 128
for col in range(0, seqlen, Bc): for col in range(0, seqlen, Bc):
print(f"tile iteration {col}~{col + Bc} ======================================") print(f"tile iteration {col}~{col + Bc} ======================================")
@@ -137,8 +141,6 @@ if __name__ == "__main__":
rowsum.astype('float32').tofile("rowsum.bin") rowsum.astype('float32').tofile("rowsum.bin")
x = rowmax_prev - rowmax x = rowmax_prev - rowmax
print("haha")
print(exp2(x))
O = O / (exp2(x)[:, np.newaxis]) O = O / (exp2(x)[:, np.newaxis])
if col == col_to_save: if col == col_to_save:
print('O_before_PV:') print('O_before_PV:')

39
tests/kernel/tensor/generate_matrix.py
View File

@@ -46,11 +46,12 @@ def pack_fp16_by_row(array):
if __name__ == "__main__": if __name__ == "__main__":
M, N, K = parse_mnk() M, N, K = parse_mnk()
rand = True rand = False
if not rand: if not rand:
A_array = np.arange(M * K).reshape([M, K]) A_array = np.arange(M * K).reshape([M, K])
B_array = np.arange(K * N).reshape([K, N]) B_array = np.arange(K * N).reshape([K, N])
C_array = np.arange(M * N).reshape([M, N]) # C_array = np.arange(M * N).reshape([M, N])
C_array = np.zeros([M, N])
else: else:
np.random.seed(0) np.random.seed(0)
A_array = np.random.rand(M, K) - 0.5 A_array = np.random.rand(M, K) - 0.5
@@ -76,11 +77,6 @@ if __name__ == "__main__":
np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array) np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array)
D_expected = A_array @ B_array
D_expected.astype('float32').tofile("d_expected.bin")
print('D_expected:')
print(D_expected)
fp16 = False fp16 = False
if fp16: if fp16:
A_packed = pack_fp16_by_row(A_array) A_packed = pack_fp16_by_row(A_array)
@@ -105,29 +101,8 @@ if __name__ == "__main__":
print('B:') print('B:')
print(B_array) print(B_array)
# generate rowmax result in online softmax D_expected = A_array @ B_array
row_max = np.max(D_expected, axis=1) D_expected.astype('float32').tofile("d_expected.bin")
row_max.astype('float32').tofile("rowmax.bin") print('D_expected:')
print(D_expected)
# subtrace row_max from each row by broadcasting
# (placeholder for exp)
x = D_expected - row_max[:, np.newaxis]
P = (x**2) / 2.0 + x + 1.0
# for i in range(3, 4):
# P += (x**i) / np.math.factorial(i)
# P = np.exp(exp)
P.astype('float32').tofile("P_expected.bin")
# print('P error:')
# print(P / np.exp(x))
print('P_expected:')
print(P)
row_sum = np.sum(P, axis=1)
row_sum.astype('float32').tofile("rowsum.bin")
V = C_array
# O = P.transpose([1, 0]) @ V
O = P @ V
O.astype('float32').tofile("O_expected.bin")
print('O_expected:')
print(O)