[Kimi Perf] Group fp8 mla cahche insertion, 4~6x kernel level performance improvement for small batch (#55356)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 11, 2026 at 11:16 UTC
1e1060f9988fa188fd243c093610889594ba18fd
6 files changed
+216
-39
csrc/libtorch_stable/cache_kernels.cu
+92
-25
@@ -456,14 +456,14 @@ __global__ void concat_and_cache_mla_kernel(
456
// Grouped variant of concat_and_cache_mla: inserts the context K/V for every
457
// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
458
// reads its own cache base pointer from kv_cache_ptrs (same pointer-array
459
-// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no
460
-// scaling or quantization; scalar_t is uint16_t for portability.
461
-template <typename scalar_t>
459
+// pattern as copy_blocks_kernel).
460
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
461
__global__ void concat_and_cache_mla_grouped_kernel(
462
const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens,
463
// kv_lora_rank]
464
const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim]
465
const int64_t* __restrict__ kv_cache_ptrs, // [num_layers]
466
+ const float* __restrict__ kv_scales, // [num_layers] or nullptr
467
const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens]
468
const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
469
const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
@@ -481,12 +481,16 @@ __global__ void concat_and_cache_mla_grouped_kernel(
481
const int64_t block_idx = slot_idx / block_size;
482
const int64_t block_offset = slot_idx % block_size;
483
484
- scalar_t* __restrict__ kv_cache =
485
- reinterpret_cast<scalar_t*>(kv_cache_ptrs[layer_idx]);
484
+ cache_t* __restrict__ kv_cache =
485
+ reinterpret_cast<cache_t*>(kv_cache_ptrs[layer_idx]);
486
const scalar_t* __restrict__ kv_c_layer =
487
kv_c + layer_idx * kv_c_layer_stride;
488
const scalar_t* __restrict__ k_pe_layer =
489
k_pe + layer_idx * k_pe_layer_stride;
490
+ float scale = 0.0f;
491
+ if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) {
492
+ scale = kv_scales[layer_idx];
493
+ }
494
495
auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
496
int size, int offset) {
@@ -494,7 +498,12 @@ __global__ void concat_and_cache_mla_grouped_kernel(
498
const int64_t src_idx = token_idx * src_token_stride + i;
499
const int64_t dst_idx =
500
block_idx * block_stride + block_offset * entry_stride + i + offset;
497
- kv_cache[dst_idx] = src[src_idx];
501
+ if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
502
+ kv_cache[dst_idx] = src[src_idx];
503
+ } else {
504
+ kv_cache[dst_idx] =
505
+ fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], scale);
506
+ }
507
}
508
};
509
@@ -1001,20 +1010,55 @@ void concat_and_cache_mla_grouped(
1010
torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim]
1011
torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device
1012
torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64
1004
- int64_t block_size, int64_t block_stride, int64_t entry_stride) {
1005
- int num_layers = kv_c.size(0);
1006
- int num_tokens = kv_c.size(1);
1007
- int kv_lora_rank = kv_c.size(2);
1008
- int pe_dim = k_pe.size(2);
1013
+ int64_t block_size, int64_t block_stride, int64_t entry_stride,
1014
+ std::optional<torch::stable::Tensor> kv_scales, // [num_layers] or None
1015
+ const std::string& kv_cache_dtype) {
1016
+ const bool use_fp8 = kv_cache_dtype == "fp8" ||
1017
+ kv_cache_dtype == "fp8_e4m3" ||
1018
+ kv_cache_dtype == "fp8_e5m2";
1019
+#ifdef USE_ROCM
1020
+ STD_TORCH_CHECK(kv_cache_dtype != "fp8_e5m2",
1021
+ "concat_and_cache_mla_grouped does not support fp8_e5m2 "
1022
+ "KV cache on ROCm");
1023
+#endif
1024
+ STD_TORCH_CHECK(
1025
+ use_fp8 || kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16",
1026
+ "concat_and_cache_mla_grouped only supports BF16 and plain "
1027
+ "FP8 KV cache; got ",
1028
+ kv_cache_dtype);
1029
1030
STD_TORCH_CHECK(
1031
kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
1032
k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1013
- "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=",
1033
+ "concat_and_cache_mla_grouped requires BF16 inputs; got kv_c=",
1034
kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
1035
STD_TORCH_CHECK(
1016
- kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
1017
- "kv_cache_ptrs must be int64");
1036
+ kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long &&
1037
+ slot_mapping.scalar_type() == torch::headeronly::ScalarType::Long,
1038
+ "cache pointers and slot mapping must be int64");
1039
+
1040
+ const int num_layers = kv_c.size(0);
1041
+ const int num_tokens = kv_c.size(1);
1042
+ const int kv_lora_rank = kv_c.size(2);
1043
+ const int pe_dim = k_pe.size(2);
1044
+ const float* kv_scales_ptr = nullptr;
1045
+ if (use_fp8) {
1046
+ STD_TORCH_CHECK(kv_scales.has_value(),
1047
+ "FP8 grouped cache insert requires kv_scales");
1048
+ STD_TORCH_CHECK(
1049
+ kv_scales->scalar_type() == torch::headeronly::ScalarType::Float,
1050
+ "kv_scales must be float32");
1051
+ STD_TORCH_CHECK(kv_scales->numel() == num_layers,
1052
+ "kv_scales must contain one scale per layer");
1053
+ STD_TORCH_CHECK(kv_scales->is_cuda() && kv_scales->get_device_index() ==
1054
+ kv_c.get_device_index(),
1055
+ "kv_scales must be on the same CUDA device as kv_c");
1056
+ STD_TORCH_CHECK(kv_scales->is_contiguous(), "kv_scales must be contiguous");
1057
+ kv_scales_ptr = kv_scales->const_data_ptr<float>();
1058
+ } else {
1059
+ STD_TORCH_CHECK(!kv_scales.has_value(),
1060
+ "BF16 grouped cache insert does not use kv_scales");
1061
+ }
1062
1063
if (num_tokens == 0 || num_layers == 0) {
1064
return;
@@ -1030,17 +1074,40 @@ void concat_and_cache_mla_grouped(
1074
kv_c.get_device_index());
1075
const cudaStream_t stream = get_current_cuda_stream();
1076
1033
- dim3 grid(num_tokens, num_layers);
1034
- dim3 block(std::min(kv_lora_rank, 512));
1035
- vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
1036
- <<<grid, block, 0, stream>>>(
1037
- reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
1038
- reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
1039
- kv_cache_ptrs.const_data_ptr<int64_t>(),
1040
- slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
1041
- kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
1042
- slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
1043
- block_size);
1077
+ const dim3 grid(num_tokens, num_layers);
1078
+ const dim3 block(std::min(kv_lora_rank, 512));
1079
+
1080
+ if (!use_fp8) {
1081
+ vllm::concat_and_cache_mla_grouped_kernel<uint16_t, uint16_t,
1082
+ vllm::Fp8KVCacheDataType::kAuto>
1083
+ <<<grid, block, 0, stream>>>(
1084
+ reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
1085
+ reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
1086
+ kv_cache_ptrs.const_data_ptr<int64_t>(), nullptr,
1087
+ slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
1088
+ kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
1089
+ slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
1090
+ block_size);
1091
+ return;
1092
+ }
1093
+
1094
+#define LAUNCH_GROUPED_FP8(KV_DTYPE) \
1095
+ vllm::concat_and_cache_mla_grouped_kernel<__nv_bfloat16, uint8_t, KV_DTYPE> \
1096
+ <<<grid, block, 0, stream>>>( \
1097
+ reinterpret_cast<const __nv_bfloat16*>(kv_c.data_ptr()), \
1098
+ reinterpret_cast<const __nv_bfloat16*>(k_pe.data_ptr()), \
1099
+ kv_cache_ptrs.const_data_ptr<int64_t>(), kv_scales_ptr, \
1100
+ slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride, \
1101
+ kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, \
1102
+ slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, \
1103
+ block_size)
1104
+
1105
+ if (kv_cache_dtype == "fp8_e5m2") {
1106
+ LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E5M2);
1107
+ } else {
1108
+ LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E4M3);
1109
+ }
1110
+#undef LAUNCH_GROUPED_FP8
1111
}
1112
1113
namespace vllm {
csrc/libtorch_stable/ops.h
+6
-6
@@ -654,12 +654,12 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
654
const std::string& kv_cache_dtype,
655
torch::stable::Tensor& scale);
656
657
-void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c,
658
- torch::stable::Tensor& k_pe,
659
- torch::stable::Tensor& kv_cache_ptrs,
660
- torch::stable::Tensor& slot_mapping,
661
- int64_t block_size, int64_t block_stride,
662
- int64_t entry_stride);
657
+void concat_and_cache_mla_grouped(
658
+ torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe,
659
+ torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping,
660
+ int64_t block_size, int64_t block_stride, int64_t entry_stride,
661
+ std::optional<torch::stable::Tensor> kv_scales,
662
+ const std::string& kv_cache_dtype);
663
664
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
665
void concat_and_cache_mla_rope_fused(
csrc/libtorch_stable/torch_bindings.cpp
+5
-3
@@ -976,14 +976,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
976
" str kv_cache_dtype,"
977
" Tensor scale) -> ()");
978
979
- // Grouped concat_and_cache_mla across all layers (bf16 only). Each
980
- // layer's cache base pointer is read from kv_cache_ptrs.
979
+ // Grouped concat_and_cache_mla across all layers. Each layer's cache base
980
+ // pointer and optional plain-FP8 scale are read from device tensors.
981
ops.def(
982
"concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe,"
983
" Tensor kv_cache_ptrs,"
984
" Tensor slot_mapping,"
985
" int block_size, int block_stride,"
986
- " int entry_stride) -> ()");
986
+ " int entry_stride,"
987
+ " Tensor? kv_scales=None,"
988
+ " str kv_cache_dtype='auto') -> ()");
989
990
// Rotate Q and K, then write to kv cache for MLA
991
ops.def(
tests/kernels/attention/test_cache.py
+88
@@ -805,6 +805,94 @@ def test_concat_and_cache_mla(
805
torch.testing.assert_close(kv_cache, ref_kv_cache)
806
807
808
+@pytest.mark.parametrize("device", CUDA_DEVICES)
809
+@pytest.mark.parametrize("shared_slot_mapping", [False, True])
810
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e5m2"])
811
+@torch.inference_mode()
812
+def test_concat_and_cache_mla_grouped(
813
+ device: str,
814
+ shared_slot_mapping: bool,
815
+ kv_cache_dtype: str,
816
+) -> None:
817
+ if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
818
+ pytest.skip("fp8_e5m2 KV cache is not supported on ROCm/HIP")
819
+
820
+ set_random_seed(0)
821
+ torch.set_default_device(device)
822
+ torch.accelerator.set_device_index(device)
823
+
824
+ num_layers = 5
825
+ num_tokens = 42
826
+ num_blocks = 8
827
+ block_size = 16
828
+ kv_lora_rank = 512
829
+ qk_rope_head_dim = 64
830
+ entry_size = kv_lora_rank + qk_rope_head_dim
831
+ total_slots = num_blocks * block_size
832
+
833
+ kv_c = torch.randn(num_layers, num_tokens, kv_lora_rank, dtype=torch.bfloat16)
834
+ k_pe = torch.randn(num_layers, num_tokens, qk_rope_head_dim, dtype=torch.bfloat16)
835
+ kv_caches = torch.zeros(
836
+ num_layers,
837
+ num_blocks,
838
+ block_size,
839
+ entry_size,
840
+ dtype=torch.bfloat16 if kv_cache_dtype == "auto" else torch.uint8,
841
+ )
842
+ reference = torch.zeros_like(kv_caches)
843
+ scales = torch.linspace(0.05, 0.25, num_layers, dtype=torch.float32)
844
+
845
+ slot_mapping = torch.stack(
846
+ [torch.randperm(total_slots)[:num_tokens] for _ in range(num_layers)]
847
+ )
848
+ slot_mapping[:, -1] = -1
849
+ if shared_slot_mapping:
850
+ slot_mapping = slot_mapping[:1].expand(num_layers, -1)
851
+
852
+ for layer_idx in range(num_layers):
853
+ ops.concat_and_cache_mla(
854
+ kv_c[layer_idx],
855
+ k_pe[layer_idx],
856
+ reference[layer_idx],
857
+ slot_mapping[layer_idx],
858
+ kv_cache_dtype,
859
+ scales[layer_idx],
860
+ )
861
+
862
+ cache_ptrs = torch.tensor(
863
+ [kv_caches[layer_idx].data_ptr() for layer_idx in range(num_layers)],
864
+ dtype=torch.int64,
865
+ )
866
+ ref_cache = kv_caches[0]
867
+
868
+ def run_grouped(kv_scales: torch.Tensor | None) -> None:
869
+ ops.concat_and_cache_mla_grouped(
870
+ kv_c,
871
+ k_pe,
872
+ cache_ptrs,
873
+ slot_mapping,
874
+ ref_cache.size(1),
875
+ ref_cache.stride(0),
876
+ ref_cache.stride(1),
877
+ kv_scales,
878
+ kv_cache_dtype,
879
+ )
880
+
881
+ run_grouped(None if kv_cache_dtype == "auto" else scales)
882
+
883
+ torch.testing.assert_close(kv_caches, reference, rtol=0, atol=0)
884
+
885
+ if kv_cache_dtype == "fp8" and not shared_slot_mapping:
886
+ noncontiguous_scales = torch.ones(num_layers, 2)[:, 0]
887
+ assert not noncontiguous_scales.is_contiguous()
888
+ for invalid_scales, error in (
889
+ (scales.cpu(), "same CUDA device"),
890
+ (noncontiguous_scales, "must be contiguous"),
891
+ ):
892
+ with pytest.raises(RuntimeError, match=error):
893
+ run_grouped(invalid_scales)
894
+
895
+
896
@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
897
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
898
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
vllm/_custom_ops.py
+4
@@ -2828,6 +2828,8 @@ def concat_and_cache_mla_grouped(
2828
block_size: int,
2829
block_stride: int,
2830
entry_stride: int,
2831
+ kv_scales: torch.Tensor | None = None,
2832
+ kv_cache_dtype: str = "auto",
2833
) -> None:
2834
torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
2835
kv_c,
@@ -2837,6 +2839,8 @@ def concat_and_cache_mla_grouped(
2839
block_size,
2840
block_stride,
2841
entry_stride,
2842
+ kv_scales,
2843
+ kv_cache_dtype,
2844
)
2845
2846
vllm/models/kimi_k3/nvidia/dspark_mla.py
+21
-5
@@ -28,6 +28,10 @@ from vllm.models.kimi_k3.nvidia.model import KimiMLP
28
from vllm.utils.torch_utils import is_quantized_kv_cache
29
from vllm.v1.worker.workspace import current_workspace_manager
30
31
+_GROUPED_KV_CACHE_DTYPES = frozenset(
32
+ {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"}
33
+)
34
+
35
36
def _duplicate_context_kv_weights(
37
weights: Iterable[tuple[str, torch.Tensor]], num_layers: int
@@ -233,6 +237,13 @@ class K3DSparkModel(nn.Module):
237
self._context_kv_lora_rank = attn0.kv_lora_rank
238
self._context_rope_dim = attn0.qk_rope_head_dim
239
self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon
240
+ self._context_kv_scales: torch.Tensor | None = None
241
+ if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache(
242
+ attn0.kv_cache_dtype
243
+ ):
244
+ self._context_kv_scales = torch.stack(
245
+ [attn._k_scale.reshape(()) for attn in attentions]
246
+ )
247
248
def _precompute_fused_context_kv(
249
self,
@@ -286,16 +297,18 @@ class K3DSparkModel(nn.Module):
297
return
298
299
cache_layers = [layer.self_attn for layer in self.layers]
300
+ cache_dtype = cache_layers[0].kv_cache_dtype
301
if (
290
- not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype)
302
+ cache_dtype in _GROUPED_KV_CACHE_DTYPES
303
+ and all(cl.kv_cache_dtype == cache_dtype for cl in cache_layers)
304
and self._has_uniform_block_layout(cache_layers)
305
and (
306
isinstance(context_slot_mapping, torch.Tensor)
307
or all(s is not None for s in context_slot_mapping)
308
)
309
):
297
- # Grouped context KV insert only supports unquantized (bf16) KV cache
298
- # and assumes that all layers share the same block layout.
310
+ # Grouped context KV insert assumes that all layers share the same
311
+ # cache dtype and block layout.
312
313
if isinstance(context_slot_mapping, (list, tuple)):
314
per_layer_slot_mappings = [
@@ -321,6 +334,8 @@ class K3DSparkModel(nn.Module):
334
ref_cache.size(1),
335
ref_cache.stride(0),
336
ref_cache.stride(1),
337
+ self._context_kv_scales,
338
+ cache_dtype,
339
)
340
return
341
@@ -483,6 +498,7 @@ class K3DSparkForCausalLM(nn.Module):
498
loader = AutoWeightsLoader(self)
499
# read: 1. all weights. 2. context kv weights
500
weights = _duplicate_context_kv_weights(weights, len(self.model.layers))
486
- loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
501
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
502
+
503
+ def process_weights_after_loading(self) -> None:
504
self.model._build_fused_context_kv_metadata()
488
- return loaded_weights