@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 64 lines 2.25 KB
Raw
1 #pragma once
2
3 #include <cuda_runtime_api.h>
4 #include <algorithm>
5
6 // maximum blocks per SM cap
7 #ifndef VLLM_LAUNCH_BLOCKS_CAP
8 #define VLLM_LAUNCH_BLOCKS_CAP 4
9 #endif
10
11 // Compile-time estimate of max threads per SM for launch bounds.
12 // Families: 1024, 1536, 2048 threads/SM.
13 #ifndef VLLM_MAX_THREADS_PER_SM
14 #ifdef __CUDA_ARCH__
15
16 /* 1024 thr/SM: Turing (sm_75) */
17 #if (__CUDA_ARCH__ == 750)
18 #define VLLM_MAX_THREADS_PER_SM 1024
19
20 /* 1536 thr/SM: Ampere GA10x (sm_86/87), Ada (sm_89),
21 GB20x consumer (sm_120/121), Thor (sm_101 or sm_110) */
22 #elif (__CUDA_ARCH__ == 860) || (__CUDA_ARCH__ == 870) || \
23 (__CUDA_ARCH__ == 890) || (__CUDA_ARCH__ == 1010) || \
24 (__CUDA_ARCH__ == 1100) || (__CUDA_ARCH__ == 1200) || \
25 (__CUDA_ARCH__ == 1210)
26 #define VLLM_MAX_THREADS_PER_SM 1536
27
28 /* 2048 thr/SM: Volta (sm_70/72), Ampere GA100 (sm_80),
29 Hopper (sm_90), Blackwell (sm_100/103) */
30 #elif (__CUDA_ARCH__ == 700) || (__CUDA_ARCH__ == 720) || \
31 (__CUDA_ARCH__ == 800) || (__CUDA_ARCH__ == 900) || \
32 (__CUDA_ARCH__ == 1000) || (__CUDA_ARCH__ == 1030)
33 #define VLLM_MAX_THREADS_PER_SM 2048
34
35 /* Fallback: use 2048 for unknown future CCs */
36 #else
37 #define VLLM_MAX_THREADS_PER_SM 2048
38 #endif
39
40 #else
41 /* Host pass (no __CUDA_ARCH__): neutral default */
42 #define VLLM_MAX_THREADS_PER_SM 2048
43 #endif
44 #endif
45
46 // compute the number of blocks per SM to request in __launch_bounds__
47 #define VLLM_BLOCKS_DIV(VAL) (VLLM_MAX_THREADS_PER_SM / (VAL))
48 #define VLLM_CLAMP_BLOCKS_PER_SM(VAL) \
49 (((VAL) <= 0) \
50 ? 1 \
51 : (((VAL) < VLLM_LAUNCH_BLOCKS_CAP) ? (VAL) : VLLM_LAUNCH_BLOCKS_CAP))
52 #define VLLM_BLOCKS_PER_SM(BLOCK_THREADS) \
53 VLLM_CLAMP_BLOCKS_PER_SM(VLLM_BLOCKS_DIV(BLOCK_THREADS))
54
55 // runtime-time helper to compute blocks/SM
56 static inline int vllm_runtime_blocks_per_sm(int block_threads) {
57 int device = -1;
58 cudaGetDevice(&device);
59 int max_threads_per_sm = VLLM_MAX_THREADS_PER_SM;
60 cudaDeviceGetAttribute(&max_threads_per_sm,
61 cudaDevAttrMaxThreadsPerMultiProcessor, device);
62 int blocks = (block_threads > 0) ? (max_threads_per_sm / block_threads) : 1;
63 return VLLM_CLAMP_BLOCKS_PER_SM(blocks);
64 }