@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 188 lines 9.79 KB
Raw
1 /*
2 * Stable ABI compatible dispatch utilities for vLLM.
3 * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros
4 * instead of the ATen (AT_*) macros.
5 *
6 * These macros use:
7 * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH
8 * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE
9 * - torch::headeronly::ScalarType instead of at::ScalarType
10 *
11 * Add more macros here as needed when migrating additional kernels.
12 */
13 #pragma once
14
15 #include <torch/headeronly/core/Dispatch.h>
16 #include <torch/headeronly/core/ScalarType.h>
17 #include <torch/headeronly/util/Exception.h>
18
19 // Need a special dispatch case macro since we will nest the FP8 dispatch.
20 // Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
21 #define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \
22 THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
23
24 // Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids)
25 // nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'.
26 #define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \
27 THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__)
28
29 #define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \
30 VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \
31 __VA_ARGS__) \
32 VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \
33 __VA_ARGS__)
34
35 #define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \
36 THO_DISPATCH_SWITCH(TYPE, NAME, \
37 VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__))
38
39 #define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \
40 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
41 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
42 THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
43
44 #define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
45 THO_DISPATCH_SWITCH(TYPE, NAME, \
46 VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
47
48 #define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \
49 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \
50 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \
51 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \
52 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \
53 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__)
54
55 #define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \
56 VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \
57 THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \
58 THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \
59 THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__)
60
61 #define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
62 THO_DISPATCH_SWITCH(TYPE, NAME, \
63 VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
64
65 #define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
66 THO_DISPATCH_SWITCH( \
67 TYPE, NAME, \
68 VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
69
70 // FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format
71 #ifdef USE_ROCM
72 #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \
73 VLLM_STABLE_DISPATCH_FP8_CASE( \
74 torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
75 VLLM_STABLE_DISPATCH_FP8_CASE( \
76 torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
77 #else
78 #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \
79 VLLM_STABLE_DISPATCH_FP8_CASE( \
80 torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__)
81 #endif
82
83 // When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
84 // See VLLM_STABLE_DISPATCH_FP8_CASE above.
85 #define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
86 THO_DISPATCH_SWITCH(TYPE, NAME, \
87 VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
88
89 // Half types dispatch (Half + BFloat16)
90 #define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \
91 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
92 THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
93
94 #define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
95 THO_DISPATCH_SWITCH(TYPE, NAME, \
96 VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
97
98 // Quant type dispatch (FP8 + INT8)
99 #ifdef USE_ROCM
100 #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \
101 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \
102 __VA_ARGS__) \
103 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \
104 __VA_ARGS__) \
105 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
106 #else
107 #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \
108 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \
109 __VA_ARGS__) \
110 THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
111 #endif
112
113 #define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
114 THO_DISPATCH_SWITCH(TYPE, NAME, \
115 VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
116
117 // Group size dispatch (pure C++ if/else, no ATen dependency)
118 #define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
119 if (group_size == 128) { \
120 constexpr int const_group_size = 128; \
121 __VA_ARGS__(); \
122 } else if (group_size == 64) { \
123 constexpr int const_group_size = 64; \
124 __VA_ARGS__(); \
125 }
126
127 // Boolean dispatch
128 #define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
129 if (expr) { \
130 constexpr bool const_expr = true; \
131 __VA_ARGS__(); \
132 } else { \
133 constexpr bool const_expr = false; \
134 __VA_ARGS__(); \
135 }
136
137 // Vec size dispatch (pure C++ switch, no ATen dependency)
138 #define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \
139 switch (VEC_SIZE) { \
140 case 16: { \
141 constexpr int vec_size = 16; \
142 __VA_ARGS__(); \
143 break; \
144 } \
145 case 8: { \
146 constexpr int vec_size = 8; \
147 __VA_ARGS__(); \
148 break; \
149 } \
150 case 4: { \
151 constexpr int vec_size = 4; \
152 __VA_ARGS__(); \
153 break; \
154 } \
155 case 2: { \
156 constexpr int vec_size = 2; \
157 __VA_ARGS__(); \
158 break; \
159 } \
160 default: { \
161 constexpr int vec_size = 1; \
162 __VA_ARGS__(); \
163 break; \
164 } \
165 }
166
167 // Tensor rank dispatch (2D, 3D, 4D)
168 #define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \
169 switch (NUM_DIMS) { \
170 case 2: { \
171 constexpr int tensor_rank = 2; \
172 __VA_ARGS__(); \
173 break; \
174 } \
175 case 3: { \
176 constexpr int tensor_rank = 3; \
177 __VA_ARGS__(); \
178 break; \
179 } \
180 case 4: { \
181 constexpr int tensor_rank = 4; \
182 __VA_ARGS__(); \
183 break; \
184 } \
185 default: \
186 STD_TORCH_CHECK( \
187 false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \
188 }