A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 08:30 UTC
[Bugfix] Fix broken explicit unquantized kv cache dtype support (#38922)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Apr 10, 2026 at 13:27 UTC
58c0a928c99c80ad3c8f16266a2344c98cf05c24
3 files changed
+52
-50
csrc/attention/dtype_fp8.cuh
+16
@@ -17,6 +17,22 @@ enum class Fp8KVCacheDataType {
17
kFp8E5M2 = 2,
18
};
19
20
+inline Fp8KVCacheDataType get_fp8_kv_cache_data_type(
21
+ const std::string& dtype_str) {
22
+ // dtype_str refers to CacheDType at vllm.config.cache.CacheDType
23
+ if (dtype_str == "auto" || dtype_str == "float16" ||
24
+ dtype_str == "bfloat16") {
25
+ // unquantized kv cache
26
+ return Fp8KVCacheDataType::kAuto;
27
+ } else if (dtype_str == "fp8" || dtype_str == "fp8_ds_mla" ||
28
+ dtype_str == "fp8_e4m3") {
29
+ return Fp8KVCacheDataType::kFp8E4M3;
30
+ } else if (dtype_str == "fp8_e5m2") {
31
+ return Fp8KVCacheDataType::kFp8E5M2;
32
+ }
33
+ TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
34
+}
35
+
36
// fp8 vector types for quantization of kv cache
37
template <>
38
struct Vec<uint8_t, 1> {
csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
+13
-14
@@ -639,7 +639,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
639
// function with template<typename scalar_t, typename cache_t,
640
// Fp8KVCacheDataType kv_dt>.
641
#define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \
642
- if (KV_DTYPE == "auto") { \
642
+ vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
643
+ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
644
+ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
645
if (SRC_DTYPE == at::ScalarType::Float) { \
646
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
647
} else if (SRC_DTYPE == at::ScalarType::Half) { \
@@ -649,21 +651,18 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
651
} else { \
652
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
653
} \
652
- } else { \
653
- if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \
654
- if (SRC_DTYPE == at::ScalarType::Float) { \
655
- FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
656
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
657
- FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
658
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
659
- FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
660
- } else { \
661
- TORCH_CHECK(false, \
662
- "Unsupported input type of kv cache: ", SRC_DTYPE); \
663
- } \
654
+ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
655
+ if (SRC_DTYPE == at::ScalarType::Float) { \
656
+ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
657
+ } else if (SRC_DTYPE == at::ScalarType::Half) { \
658
+ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
659
+ } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
660
+ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
661
} else { \
665
- TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
662
+ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
663
} \
664
+ } else { \
665
+ TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
666
}
667
668
} // namespace fp8
csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
+23
-36
@@ -543,7 +543,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
543
// function with template<typename scalar_t, typename cache_t,
544
// Fp8KVCacheDataType kv_dt>.
545
#define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \
546
- if (KV_DTYPE == "auto") { \
546
+ vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
547
+ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
548
+ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
549
if (SRC_DTYPE == at::ScalarType::Float) { \
550
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
551
} else if (SRC_DTYPE == at::ScalarType::Half) { \
@@ -553,43 +555,28 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
555
} else { \
556
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
557
} \
556
- } else { \
557
- if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \
558
- if (SRC_DTYPE == at::ScalarType::Float) { \
559
- FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
560
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
561
- FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
562
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
563
- FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
564
- } else { \
565
- TORCH_CHECK(false, \
566
- "Unsupported input type of kv cache: ", SRC_DTYPE); \
567
- } \
568
- } else if (KV_DTYPE == "fp8_e5m2") { \
569
- if (SRC_DTYPE == at::ScalarType::Float) { \
570
- FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
571
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
572
- FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
573
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
574
- FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
575
- } else { \
576
- TORCH_CHECK(false, \
577
- "Unsupported input type of kv cache: ", SRC_DTYPE); \
578
- } \
579
- } else if (KV_DTYPE == "fp8_ds_mla") { \
580
- if (SRC_DTYPE == at::ScalarType::Float) { \
581
- FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
582
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
583
- FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
584
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
585
- FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
586
- } else { \
587
- TORCH_CHECK(false, \
588
- "Unsupported input type of kv cache: ", SRC_DTYPE); \
589
- } \
558
+ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
559
+ if (SRC_DTYPE == at::ScalarType::Float) { \
560
+ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
561
+ } else if (SRC_DTYPE == at::ScalarType::Half) { \
562
+ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
563
+ } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
564
+ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
565
+ } else { \
566
+ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
567
+ } \
568
+ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \
569
+ if (SRC_DTYPE == at::ScalarType::Float) { \
570
+ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
571
+ } else if (SRC_DTYPE == at::ScalarType::Half) { \
572
+ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
573
+ } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
574
+ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
575
} else { \
591
- TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
576
+ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
577
} \
578
+ } else { \
579
+ TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
580
}
581
582
} // namespace fp8