bugfix: fix dispatch fp16 type when enable fp8 (#430)

Fix #402 (comment)
flashinfer-ai · Aug 8, 2024 · daa5566 · daa5566
1 parent d52f2da
commit daa5566
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/csrc/pytorch_extension_utils.h b/python/csrc/pytorch_extension_utils.h
@@ -146,6 +146,10 @@ using namespace flashinfer;
 #define DISPATCH_PYTORCH_DTYPE_TO_CTYPE(pytorch_dtype, c_type, ...)                          \
   [&]() -> bool {                                                                            \
     switch (pytorch_dtype) {                                                                 \
+      case at::ScalarType::Half: {                                                           \
+        using c_type = nv_half;                                                              \
+        return __VA_ARGS__();                                                                \
+      }                                                                                      \
       case at::ScalarType::Float8_e4m3fn: {                                                  \
         using c_type = __nv_fp8_e4m3;                                                        \
         return __VA_ARGS__();                                                                \