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 8, 12:30 UTC
[Bugfix][DeepSeek-V3.2] fix fp8 kvcache type cast (#33884)
Signed-off-by: Kebe <mail@kebe7jun.com>
Kebe committed
Feb 11, 2026 at 12:31 UTC
5ee5c86eeb00a4d159e2e2cb4c8c85dcc0733e15
1 file changed
+16
-4
csrc/cache_kernels.cu
+16
-4
@@ -1234,8 +1234,13 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1234
"src_cache and seq_lens must be on the same device");
1235
TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1236
"src_cache and workspace_starts must be on the same device");
1237
-
1238
- TORCH_CHECK(src_cache.dtype() == torch::kUInt8, "src_cache must be uint8");
1237
+ auto dtype = src_cache.scalar_type();
1238
+ TORCH_CHECK(
1239
+ dtype == at::ScalarType::Byte || // uint8
1240
+ dtype == at::ScalarType::Float8_e4m3fn || // fp8 e4m3
1241
+ dtype == at::ScalarType::Float8_e5m2, // fp8 e5m2
1242
+ "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
1243
+ src_cache.dtype());
1244
TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");
1245
TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1246
@@ -1244,14 +1249,21 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1249
int64_t cache_entry_stride = src_cache.stride(1);
1250
int64_t dst_entry_stride = dst.stride(0);
1251
1252
+ const uint8_t* src_ptr = nullptr;
1253
+ if (dtype == at::ScalarType::Byte) {
1254
+ src_ptr = src_cache.data_ptr<uint8_t>();
1255
+ } else {
1256
+ // float8_e4m3fn or float8_e5m2
1257
+ src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
1258
+ }
1259
+
1260
// Decide on the number of splits based on the batch size
1261
int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
1262
dim3 grid(batch_size, num_splits);
1263
dim3 block(576);
1264
1265
vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
1253
- src_cache.data_ptr<uint8_t>(),
1254
- reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1266
+ src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1267
block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),
1268
workspace_starts.data_ptr<int32_t>(), block_size, head_dim,
1269
block_table_stride, cache_block_stride, cache_entry_stride,