Allocate more shared memory to attention kernel (#1154)

2023-09-26 22:27:13 -07:00
parent 03ffd0a022
commit cf5cb1e33e
7 changed files with 87 additions and 3 deletions
--- a/csrc/cuda_utils.cpp
+++ b/csrc/cuda_utils.cpp
@@ -0,0 +1,13 @@
+#include <torch/extension.h>
+
+int get_device_attribute(
+    int attribute,
+    int device_id);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def(
+    "get_device_attribute",
+    &get_device_attribute,
+    "Gets the specified device attribute.");
+}
+