Support Ampere architecture (#204)

* Update README

* Update `setup.py`

* Fix headers

* Add `DISABLE_NVSHMEM` for APIs

* Fix launch

* Fix TMA settings

* Fix TMA usages

* Fix dlink

* Separate layout kernels

* Update version

* Add `is_sm90_compiled`

* Fix tests

* Add NVLink connection checks

* Update README

* Fix tests

* Add some comments

* Minor fix

* Minor fix

* Fix bugs

csrc/kernels/launch.cuh

@@ -1,8 +1,10 @@
 #pragma once
 
 #include "configs.cuh"
+#include "exception.cuh"
 
 #ifndef SETUP_LAUNCH_CONFIG
+#ifndef DISABLE_SM90_FEATURES
 #define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \
     cudaLaunchConfig_t cfg = {(num_sms), (num_threads), 0, stream, nullptr, 0}; \
     cudaLaunchAttribute attr[1]; \
@@ -10,10 +12,39 @@
     attr[0].val.cooperative = 1; \
     cfg.attrs = attr; \
     cfg.numAttrs = 1
+#else
+#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \
+    int __num_sms = (sms); \
+    int __num_threads = (threads); \
+    auto __stream = (stream)
+#endif
 #endif
 
 #ifndef LAUNCH_KERNEL
+#ifndef DISABLE_SM90_FEATURES
 #define LAUNCH_KERNEL(config, kernel, ...) CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__))
+#else
+#define LAUNCH_KERNEL(config, kernel, ...) \
+do { \
+    kernel<<<__num_sms, __num_threads, 0, __stream>>>(__VA_ARGS__); \
+    cudaError_t e = cudaGetLastError(); \
+    if (e != cudaSuccess) { \
+        EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \
+        fprintf(stderr, "%s\n", cuda_exception.what()); \
+        throw cuda_exception; \
+    } \
+} while (0)
+#endif
+#endif
+
+#ifndef SET_SHARED_MEMORY_FOR_TMA
+#ifndef DISABLE_SM90_FEATURES
+#define SET_SHARED_MEMORY_FOR_TMA(kernel) \
+EP_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess); \
+cfg.dynamicSmemBytes = smem_size;
+#else
+#define SET_SHARED_MEMORY_FOR_TMA(kernel) void()
+#endif
 #endif
 
 #define SWITCH_RANKS(case_macro) \