@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 158 lines 7.73 KB
Raw
1 /*
2 * Adapted from
3 * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h
4 */
5 #pragma once
6
7 #include <torch/all.h>
8
9 // Need a special dispatch case macro since we will nest the FP8 dispatch.
10 // Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
11 #define AT_DISPATCH_FP8_CASE(enum_type, ...) \
12 AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
13
14 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
15 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
16 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
17 AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
18
19 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
20 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
21
22 #define VLLM_DISPATCH_CASE_HALF_TYPES(...) \
23 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
24 AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
25
26 #define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
27 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
28
29 // ROCm devices might use either fn or fnuz, so set up dispatch table for both.
30 // A host-based check at runtime will create a preferred FP8 type for ROCm
31 // such that the correct kernel is dispatched.
32 #ifdef USE_ROCM
33 #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
34 AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
35 AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
36
37 #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
38 AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
39 AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
40 AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
41 #else
42 #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
43 AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
44
45 #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
46 AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
47 AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
48 #endif
49
50 // When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
51 // See AT_DISPATCH_FP8_CASE above.
52 #define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
53 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
54
55 #define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
56 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
57
58 #define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \
59 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
60 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
61 AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
62 AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)
63
64 #define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \
65 AT_DISPATCH_SWITCH(TYPE, NAME, \
66 VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__))
67
68 #define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \
69 AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \
70 AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \
71 AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
72 AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
73 AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)
74
75 #define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \
76 AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \
77 AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \
78 AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
79 AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
80 AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \
81 AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \
82 AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \
83 AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__)
84
85 #define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
86 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
87
88 #define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
89 AT_DISPATCH_SWITCH( \
90 TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
91
92 #define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \
93 switch (VEC_SIZE) { \
94 case 16: { \
95 constexpr int vec_size = 16; \
96 __VA_ARGS__(); \
97 break; \
98 } \
99 case 8: { \
100 constexpr int vec_size = 8; \
101 __VA_ARGS__(); \
102 break; \
103 } \
104 case 4: { \
105 constexpr int vec_size = 4; \
106 __VA_ARGS__(); \
107 break; \
108 } \
109 case 2: { \
110 constexpr int vec_size = 2; \
111 __VA_ARGS__(); \
112 break; \
113 } \
114 default: { \
115 constexpr int vec_size = 1; \
116 __VA_ARGS__(); \
117 break; \
118 } \
119 }
120
121 #define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \
122 if (expr) { \
123 constexpr bool const_expr = true; \
124 __VA_ARGS__(); \
125 } else { \
126 constexpr bool const_expr = false; \
127 __VA_ARGS__(); \
128 }
129
130 #define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
131 if (group_size == 128) { \
132 constexpr int const_group_size = 128; \
133 __VA_ARGS__(); \
134 } else if (group_size == 64) { \
135 constexpr int const_group_size = 64; \
136 __VA_ARGS__(); \
137 }
138
139 #define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \
140 switch (NUM_DIMS) { \
141 case 2: { \
142 constexpr int tensor_rank = 2; \
143 __VA_ARGS__(); \
144 break; \
145 } \
146 case 3: { \
147 constexpr int tensor_rank = 3; \
148 __VA_ARGS__(); \
149 break; \
150 } \
151 case 4: { \
152 constexpr int tensor_rank = 4; \
153 __VA_ARGS__(); \
154 break; \
155 } \
156 default: \
157 TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \
158 }