Add flash attention kernel skeleton

2024-08-14 20:46:09 -07:00
parent 014f7cd06f
commit 692d028afd
6 changed files with 4514 additions and 0 deletions
--- a/tests/regression/flash_attention/common.h
+++ b/tests/regression/flash_attention/common.h
@@ -0,0 +1,18 @@
+#ifndef _COMMON_H_
+#define _COMMON_H_
+
+#include <cstdint>
+
+#define KERNEL_ARG_DEV_MEM_ADDR 0x9fff0000
+#define DEV_SMEM_START_ADDR 0xff000000
+
+typedef struct {
+  uint32_t dim_m;
+  uint32_t dim_n;
+  uint32_t dim_k;
+  uint64_t addr_a;
+  uint64_t addr_b;
+  uint64_t addr_c;
+} kernel_arg_t;
+
+#endif