Fix dispatch utils for AMD

Signed-off-by: luka <[email protected]>
neuralmagic · Dec 12, 2024 · 21abcff · 21abcff
1 parent 720d537
commit 21abcff
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h
@@ -14,9 +14,16 @@
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
+// TODO(luka/varun): use FP8_TYPE macro after refactoring
+#ifndef USE_ROCM
 #define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                    \
   AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
   AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+#else
+#define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                      \
+  AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+#endif
 
 #define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))