@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 152 lines 4.25 KB
Raw
1 #pragma once
2
3 #ifndef _marlin_cuh
4 #define _marlin_cuh
5 #include <cuda.h>
6 #include <cuda_fp16.h>
7 #include <cuda_runtime.h>
8 #include <iostream>
9
10 #ifndef MARLIN_NAMESPACE_NAME
11 #define MARLIN_NAMESPACE_NAME marlin
12 #endif
13
14 namespace MARLIN_NAMESPACE_NAME {
15
16 // Marlin params
17
18 // 8 warps are a good choice since every SM has 4 schedulers and having more
19 // than 1 warp per schedule allows some more latency hiding. At the same time,
20 // we want relatively few warps to have many registers per warp and small tiles.
21 static constexpr int default_threads = 256;
22
23 static constexpr int pipe_stages =
24 4; // 4 pipeline stages fit into shared memory
25
26 static constexpr int min_thread_n = 64;
27 static constexpr int min_thread_k = 64;
28 static constexpr int max_thread_n = 256;
29
30 static constexpr int tile_size = 16;
31 static constexpr int max_par = 16;
32
33 // Repack params
34 static constexpr int repack_stages = 8;
35
36 static constexpr int repack_threads = 256;
37
38 static constexpr int tile_k_size = tile_size;
39 static constexpr int tile_n_size = tile_k_size * 4;
40
41 // Helpers
42 template <typename T, int n>
43 struct Vec {
44 T elems[n];
45 __device__ T& operator[](int i) { return elems[i]; }
46 };
47
48 using I4 = Vec<int, 4>;
49
50 constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
51
52 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
53
54 __device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
55 bool pred = true) {
56 if (pred) {
57 reinterpret_cast<int32_t*>(smem_ptr)[0] =
58 reinterpret_cast<const int32_t*>(glob_ptr)[0];
59 }
60 }
61
62 __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
63 bool pred = true) {
64 if (pred) {
65 reinterpret_cast<int64_t*>(smem_ptr)[0] =
66 reinterpret_cast<const int64_t*>(glob_ptr)[0];
67 }
68 }
69
70 __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
71 bool pred = true) {
72 if (pred) {
73 reinterpret_cast<int4*>(smem_ptr)[0] =
74 reinterpret_cast<const int4*>(glob_ptr)[0];
75 }
76 }
77
78 __device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
79 reinterpret_cast<int4*>(smem_ptr)[0] =
80 reinterpret_cast<const int4*>(glob_ptr)[0];
81 }
82
83 __device__ inline void cp_async_fence() {}
84
85 template <int n>
86 __device__ inline void cp_async_wait() {}
87
88 #else
89
90 __device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
91 bool pred = true) {
92 const int BYTES = 4;
93 uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
94 asm volatile(
95 "{\n"
96 " .reg .pred p;\n"
97 " setp.ne.b32 p, %0, 0;\n"
98 " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
99 "}\n" ::"r"((int)pred),
100 "r"(smem), "l"(glob_ptr), "n"(BYTES));
101 }
102
103 __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
104 bool pred = true) {
105 const int BYTES = 8;
106 uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
107 asm volatile(
108 "{\n"
109 " .reg .pred p;\n"
110 " setp.ne.b32 p, %0, 0;\n"
111 " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
112 "}\n" ::"r"((int)pred),
113 "r"(smem), "l"(glob_ptr), "n"(BYTES));
114 }
115
116 __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
117 bool pred = true) {
118 const int BYTES = 16;
119 uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
120 asm volatile(
121 "{\n"
122 " .reg .pred p;\n"
123 " setp.ne.b32 p, %0, 0;\n"
124 " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
125 "}\n" ::"r"((int)pred),
126 "r"(smem), "l"(glob_ptr), "n"(BYTES));
127 }
128
129 __device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
130 const int BYTES = 16;
131 uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
132 asm volatile(
133 "{\n"
134 " cp.async.cg.shared.global [%0], [%1], %2;\n"
135 "}\n" ::"r"(smem),
136 "l"(glob_ptr), "n"(BYTES));
137 }
138
139 __device__ inline void cp_async_fence() {
140 asm volatile("cp.async.commit_group;\n" ::);
141 }
142
143 template <int n>
144 __device__ inline void cp_async_wait() {
145 asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
146 }
147
148 #endif
149
150 } // namespace MARLIN_NAMESPACE_NAME
151
152 #endif