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

flash.py: Write V to file

Browse Source
This commit is contained in:
Hansung Kim
2024-09-01 18:17:05 -07:00
parent 6cc1b5ca37
commit f7603b18d3
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -98,7 +98,7 @@ if __name__ == "__main__":
full_S_T = full_S.transpose([1, 0]) full_S_T = full_S.transpose([1, 0])
full_S.astype('float32').tofile("full_S.bin") full_S.astype('float32').tofile("full_S.bin")
col_to_save = 128 col_to_save = 0
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} ======================================")
@@ -148,6 +148,8 @@ if __name__ == "__main__":
O.astype('float32').tofile("O_before_PV.bin") O.astype('float32').tofile("O_before_PV.bin")
V = C_array[col:col+Bc, :] V = C_array[col:col+Bc, :]
if col == col_to_save:
V.astype('float32').tofile("V_expected.bin")
# O = P.transpose([1, 0]) @ V # O = P.transpose([1, 0]) @ V
O = O + P @ V O = O + P @ V
if col == col_to_save: if col == col_to_save: