@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 60 lines 2.23 KB
Raw
1 #pragma once
2
3 #include <cuda_fp8.h>
4 #include <torch/headeronly/core/ScalarType.h>
5 #include <torch/headeronly/util/Exception.h>
6
7 #define MOE_SWITCH(TYPE, ...) \
8 const auto _st = (TYPE); \
9 switch (_st) { \
10 __VA_ARGS__ \
11 default: \
12 STD_TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \
13 }
14
15 #define MOE_DISPATCH_CASE(enum_type, ...) \
16 case enum_type: { \
17 using scalar_t = ScalarType2CudaType<enum_type>::type; \
18 __VA_ARGS__(); \
19 break; \
20 }
21
22 #define MOE_DISPATCH_FLOAT_CASE(...) \
23 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
24 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
25 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) \
26 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e5m2, __VA_ARGS__) \
27 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
28 MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__)
29
30 #define MOE_DISPATCH(TYPE, ...) \
31 MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
32
33 template <torch::headeronly::ScalarType type>
34 struct ScalarType2CudaType;
35
36 template <>
37 struct ScalarType2CudaType<torch::headeronly::ScalarType::Float> {
38 using type = float;
39 };
40 template <>
41 struct ScalarType2CudaType<torch::headeronly::ScalarType::Half> {
42 using type = half;
43 };
44 template <>
45 struct ScalarType2CudaType<torch::headeronly::ScalarType::BFloat16> {
46 using type = __nv_bfloat16;
47 };
48 // uint8 for packed fp4
49 template <>
50 struct ScalarType2CudaType<torch::headeronly::ScalarType::Byte> {
51 using type = uint8_t;
52 };
53 template <>
54 struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e5m2> {
55 using type = __nv_fp8_e5m2;
56 };
57 template <>
58 struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e4m3fn> {
59 using type = __nv_fp8_e4m3;
60 };