Skip to content

Commit a70d496

Browse files
committed
Fix dispatch utils for AMD
Signed-off-by: luka <luka@neuralmagic.com>
1 parent 720d537 commit a70d496

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

csrc/dispatch_utils.h

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,16 @@
1414
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
1515
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
1616

17-
#define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
18-
AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
19-
AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
17+
// TODO(luka/varun): use FP8_TYPE macro after refactoring
18+
#ifndef USE_ROCM
19+
#define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
20+
AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
21+
AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
22+
#else
23+
#define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
24+
AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
25+
AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
26+
#endif
2027

2128
#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
2229
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))

0 commit comments

Comments
 (0)