@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 142 lines 3.79 KB
Raw
1 #pragma once
2
3 #include "cutlass/cutlass.h"
4 #include <climits>
5 #include "cuda_runtime.h"
6 #include <cstdio>
7 #include <cstdlib>
8
9 #include <torch/headeronly/util/shim_utils.h>
10
11 /**
12 * Helper function for checking CUTLASS errors
13 */
14 #define CUTLASS_CHECK(status) \
15 { \
16 cutlass::Status error = status; \
17 STD_TORCH_CHECK(error == cutlass::Status::kSuccess, \
18 cutlassGetStatusString(error)); \
19 }
20
21 inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) {
22 int max_shared_mem_per_block_opt_in = 0;
23 cudaDeviceGetAttribute(&max_shared_mem_per_block_opt_in,
24 cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
25 return max_shared_mem_per_block_opt_in;
26 }
27
28 int32_t get_sm_version_num();
29
30 /**
31 * A wrapper for a kernel that is used to guard against compilation on
32 * architectures that will never use the kernel. The purpose of this is to
33 * reduce the size of the compiled binary.
34 * __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
35 * into code that will be executed on the device where it is defined.
36 */
37
38 template <typename Kernel>
39 struct enable_sm75_to_sm80 : Kernel {
40 template <typename... Args>
41 CUTLASS_DEVICE static void invoke(Args&&... args) {
42 #if defined __CUDA_ARCH__
43 #if __CUDA_ARCH__ >= 750 && __CUDA_ARCH__ < 800
44 Kernel::invoke(std::forward<Args>(args)...);
45 #else
46 printf("This kernel only supports sm[75, 80).\n");
47 asm("trap;");
48 #endif
49 #endif
50 }
51 };
52
53 template <typename Kernel>
54 struct enable_sm80_to_sm89 : Kernel {
55 template <typename... Args>
56 CUTLASS_DEVICE static void invoke(Args&&... args) {
57 #if defined __CUDA_ARCH__
58 #if __CUDA_ARCH__ >= 800 && __CUDA_ARCH__ < 890
59 Kernel::invoke(std::forward<Args>(args)...);
60 #else
61 printf("This kernel only supports sm[80, 89).\n");
62 asm("trap;");
63 #endif
64 #endif
65 }
66 };
67
68 template <typename Kernel>
69 struct enable_sm89_to_sm90 : Kernel {
70 template <typename... Args>
71 CUTLASS_DEVICE static void invoke(Args&&... args) {
72 #if defined __CUDA_ARCH__
73 #if __CUDA_ARCH__ >= 890 && __CUDA_ARCH__ < 900
74 Kernel::invoke(std::forward<Args>(args)...);
75 #else
76 printf("This kernel only supports sm[89, 90).\n");
77 asm("trap;");
78 #endif
79 #endif
80 }
81 };
82
83 template <typename Kernel>
84 struct enable_sm90_or_later : Kernel {
85 template <typename... Args>
86 CUTLASS_DEVICE void operator()(Args&&... args) {
87 #if defined __CUDA_ARCH__
88 #if __CUDA_ARCH__ >= 900
89 Kernel::operator()(std::forward<Args>(args)...);
90 #else
91 printf("This kernel only supports sm >= 90.\n");
92 asm("trap;");
93 #endif
94 #endif
95 }
96 };
97
98 template <typename Kernel>
99 struct enable_sm100_to_sm120 : Kernel {
100 template <typename... Args>
101 CUTLASS_DEVICE void operator()(Args&&... args) {
102 #if defined __CUDA_ARCH__
103 #if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
104 Kernel::operator()(std::forward<Args>(args)...);
105 #else
106 printf("This kernel only supports sm[100, 120).\n");
107 asm("trap;");
108 #endif
109 #endif
110 }
111 };
112
113 template <typename Kernel>
114 struct enable_sm120_only : Kernel {
115 template <typename... Args>
116 CUTLASS_DEVICE void operator()(Args&&... args) {
117 #if defined __CUDA_ARCH__
118 #if __CUDA_ARCH__ == 1200
119 Kernel::operator()(std::forward<Args>(args)...);
120 #else
121 printf("This kernel only supports sm120a.\n");
122 asm("trap;");
123 #endif
124 #endif
125 }
126 };
127
128 // SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10)
129 template <typename Kernel>
130 struct enable_sm120_family : Kernel {
131 template <typename... Args>
132 CUTLASS_DEVICE void operator()(Args&&... args) {
133 #if defined __CUDA_ARCH__
134 #if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
135 Kernel::operator()(std::forward<Args>(args)...);
136 #else
137 printf("This kernel only supports sm120f.\n");
138 asm("trap;");
139 #endif
140 #endif
141 }
142 };