[Perf][DSpark] Add KV-only context insertion across V4.1 cache formats (#56441)
Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Co-authored-by: 0z5a <0z5a@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
0z5a committed
Sep 16, 2026 at 12:34 UTC
6ecd97f1b7fb13991065370f6d6be3dffe9fe1db
5 files changed
+234
-61
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+86
-1
@@ -657,7 +657,12 @@ static void launchFusedDeepseekV4Templated(
657
config.attrs = attrs;
658
config.numAttrs = (sm_version >= 90) ? 1 : 0;
659
660
- if (num_tokens_full < NUM_TOKEN_CUTOFF) {
660
+ // The reduced grid pays off by iterating a token's Q heads inside one CTA.
661
+ // With zero query heads there is nothing to iterate, so one warp per block
662
+ // would do all the work and the other seven idle; keep the warp-per-slot
663
+ // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single
664
+ // insert, where it is worth ~2x.
665
+ if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) {
666
cudaLaunchKernelEx(&config,
667
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
668
scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
@@ -993,6 +998,86 @@ static void launchFullCacheKernel(
998
// ────────────────────────────────────────────────────────────────────────────
999
// Torch op wrapper
1000
// ────────────────────────────────────────────────────────────────────────────
1001
+void fused_deepseek_v4_kv_rope_insert(
1002
+ torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
1003
+ torch::stable::Tensor const& slot_mapping,
1004
+ torch::stable::Tensor const& position_ids,
1005
+ torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
1006
+ std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8) {
1007
+ using torch::headeronly::ScalarType;
1008
+ STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() &&
1009
+ kv.dim() == 2 && kv.size(1) == 512 &&
1010
+ kv.scalar_type() == ScalarType::BFloat16,
1011
+ "kv must be contiguous CUDA bfloat16 [tokens, 512]");
1012
+ auto const device = kv.get_device_index();
1013
+ STD_TORCH_CHECK(k_cache.device().is_cuda() &&
1014
+ k_cache.get_device_index() == device &&
1015
+ k_cache.dim() == 3 && cache_block_size > 0 &&
1016
+ k_cache.size(1) == cache_block_size &&
1017
+ k_cache.stride(2) == 1,
1018
+ "k_cache must be a paged cache on the KV device");
1019
+ for (auto const* tensor : {&slot_mapping, &position_ids}) {
1020
+ STD_TORCH_CHECK(tensor->device().is_cuda() &&
1021
+ tensor->get_device_index() == device &&
1022
+ tensor->scalar_type() == ScalarType::Long &&
1023
+ tensor->dim() == 1 && tensor->is_contiguous(),
1024
+ "slots and positions must be contiguous int64 on the KV device");
1025
+ }
1026
+ STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) &&
1027
+ slot_mapping.size(0) <= kv.size(0),
1028
+ "positions must match KV rows; slots may omit padded rows");
1029
+ STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
1030
+ cos_sin_cache.get_device_index() == device &&
1031
+ cos_sin_cache.scalar_type() == ScalarType::Float &&
1032
+ cos_sin_cache.dim() == 2 &&
1033
+ cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(),
1034
+ "cos_sin_cache must be contiguous float32 [max_pos, 64]");
1035
+ bool const packed = k_cache.scalar_type() == ScalarType::Byte;
1036
+ bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn;
1037
+ STD_TORCH_CHECK(packed || fp8 ||
1038
+ k_cache.scalar_type() == ScalarType::BFloat16,
1039
+ "cache must be uint8, bfloat16 or float8_e4m3fn");
1040
+ STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512),
1041
+ "cache row width does not match its dtype");
1042
+ if (fp8) {
1043
+ STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() &&
1044
+ fp8_scale->get_device_index() == device &&
1045
+ fp8_scale->scalar_type() == ScalarType::Float &&
1046
+ fp8_scale->numel() == 1,
1047
+ "fp8 cache requires one float32 KV scale on the KV device");
1048
+ }
1049
+ int const num_tokens = static_cast<int>(slot_mapping.size(0));
1050
+ if (num_tokens == 0) return;
1051
+ const torch::stable::accelerator::DeviceGuard device_guard(device);
1052
+ const cudaStream_t stream = get_current_cuda_stream(device);
1053
+ using scalar_t = c10::BFloat16;
1054
+ auto const* input = reinterpret_cast<scalar_t const*>(kv.const_data_ptr());
1055
+ auto* cache = reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr());
1056
+ auto const* slots = slot_mapping.const_data_ptr<int64_t>();
1057
+ auto const* positions = position_ids.const_data_ptr<int64_t>();
1058
+ auto const* cos_sin = cos_sin_cache.const_data_ptr<float>();
1059
+ // Zero query heads schedules only the existing KV branch, preserving its
1060
+ // RoPE rounding and quantization without allocating any query tensors.
1061
+ if (packed) {
1062
+ vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated<scalar_t, 0, false>(
1063
+ nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f,
1064
+ num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
1065
+ static_cast<int>(k_cache.stride(0)), kv_mxfp8, stream);
1066
+ } else if (fp8) {
1067
+ vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, true>(
1068
+ nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin,
1069
+ fp8_scale->const_data_ptr<float>(), nullptr, 0.0f, num_tokens, num_tokens,
1070
+ 0, static_cast<int>(cache_block_size), k_cache.stride(0),
1071
+ k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream);
1072
+ } else {
1073
+ vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, false>(
1074
+ nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr,
1075
+ nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
1076
+ k_cache.stride(0) * 2, k_cache.stride(1) * 2, false,
1077
+ "fused_deepseek_v4_kv_rope_insert", stream);
1078
+ }
1079
+}
1080
+
1081
torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1082
torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16
1083
torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only)
csrc/libtorch_stable/ops.h
+7
@@ -259,6 +259,13 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
259
torch::stable::Tensor& position_ids,
260
int64_t forced_token_heads_per_warp);
261
262
+void fused_deepseek_v4_kv_rope_insert(
263
+ torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
264
+ torch::stable::Tensor const& slot_mapping,
265
+ torch::stable::Tensor const& position_ids,
266
+ torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
267
+ std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8);
268
+
269
torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
270
torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
271
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
csrc/libtorch_stable/torch_bindings.cpp
+8
@@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
425
"bool is_neox, Tensor position_ids, "
426
"int forced_token_heads_per_warp=-1) -> ()");
427
428
+ ops.def(
429
+ "fused_deepseek_v4_kv_rope_insert("
430
+ "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, "
431
+ "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, "
432
+ "bool kv_mxfp8=False) -> "
433
+ "()");
434
ops.def(
435
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
436
"Tensor q_in, Tensor kv, Tensor! k_cache, "
@@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
813
// Positional encoding kernels (shared CUDA/ROCm)
814
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
815
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
816
+ ops.impl("fused_deepseek_v4_kv_rope_insert",
817
+ TORCH_BOX(&fused_deepseek_v4_kv_rope_insert));
818
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
819
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
820
ops.impl(
tests/kernels/test_compressor_kv_cache.py
+116
@@ -42,6 +42,122 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
42
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
43
44
45
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
46
+@pytest.mark.parametrize(
47
+ "cache_dtype,kv_mxfp8",
48
+ [
49
+ (torch.uint8, False),
50
+ (torch.uint8, True),
51
+ (torch.bfloat16, False),
52
+ (torch.float8_e4m3fn, False),
53
+ ],
54
+)
55
+@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024])
56
+@pytest.mark.parametrize("use_graph", [False, True])
57
+def test_dspark_context_kv_matches_query_insert(
58
+ cache_dtype, kv_mxfp8, num_tokens, use_graph
59
+):
60
+ """KV-only insertion must preserve every cache byte, including graph replay."""
61
+ from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv
62
+
63
+ torch.manual_seed(42)
64
+ block_size = 256
65
+ num_blocks = math.ceil((num_tokens + 7) / block_size)
66
+ row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512
67
+ page_stride = math.ceil(block_size * row_size / 576) * 576 + 576
68
+ backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype)
69
+ cache = backing.as_strided(
70
+ (num_blocks, block_size, row_size), (page_stride, row_size, 1)
71
+ )
72
+ reference_backing = backing.clone()
73
+ reference = reference_backing.as_strided(cache.shape, cache.stride())
74
+ kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16)
75
+ positions = torch.arange(num_tokens + 3, device="cuda") + 7
76
+ angles = torch.randn(num_tokens + 16, 32, device="cuda")
77
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
78
+ slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens]
79
+ slots[::5] = -1
80
+ scale = torch.tensor([0.7], device="cuda")
81
+ # Only what _insert_context_kv actually reads: the record width comes off
82
+ # the cache tensor, so no per-record width attribute is stubbed here.
83
+ attn = SimpleNamespace(
84
+ swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size),
85
+ kv_mxfp8=kv_mxfp8,
86
+ rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin),
87
+ _flashinfer_fp8_kv_scale=scale,
88
+ )
89
+
90
+ def legacy_insert():
91
+ q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda")
92
+ if cache_dtype == torch.uint8:
93
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
94
+ q,
95
+ kv,
96
+ reference.view(num_blocks, -1),
97
+ slots,
98
+ positions,
99
+ cos_sin,
100
+ 8,
101
+ 1e-20,
102
+ block_size,
103
+ True,
104
+ kv_mxfp8,
105
+ )
106
+ elif cache_dtype == torch.bfloat16:
107
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
108
+ q,
109
+ kv,
110
+ reference,
111
+ slots,
112
+ positions,
113
+ cos_sin,
114
+ 1e-20,
115
+ block_size,
116
+ )
117
+ else:
118
+ q_fp8 = torch.empty_like(q, dtype=cache_dtype)
119
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
120
+ q,
121
+ kv,
122
+ q_fp8,
123
+ reference,
124
+ slots,
125
+ positions,
126
+ cos_sin,
127
+ scale,
128
+ scale,
129
+ 1e-20,
130
+ block_size,
131
+ )
132
+
133
+ def insert():
134
+ _insert_context_kv(attn, kv, positions, slots)
135
+
136
+ insert()
137
+ graph = None
138
+ if use_graph:
139
+ graph = torch.cuda.CUDAGraph()
140
+ with torch.cuda.graph(graph):
141
+ insert()
142
+ for iteration in range(3):
143
+ kv.normal_()
144
+ positions.add_(1)
145
+ if iteration == 2:
146
+ slots.fill_(-1)
147
+ backing.view(torch.uint8).fill_(165)
148
+ reference_backing.copy_(backing)
149
+ before = kv.clone()
150
+ legacy_insert()
151
+ insert() if graph is None else graph.replay()
152
+ torch.testing.assert_close(
153
+ backing.view(torch.uint8),
154
+ reference_backing.view(torch.uint8),
155
+ rtol=0,
156
+ atol=0,
157
+ )
158
+ torch.testing.assert_close(kv, before, rtol=0, atol=0)
159
+
160
+
161
@pytest.mark.skipif(
162
not current_platform.is_cuda_alike(), reason="graph capture coverage"
163
)
vllm/models/deepseek_v41/nvidia/dspark.py
+17
-60
@@ -234,67 +234,24 @@ def _insert_context_kv(
234
positions: torch.Tensor,
235
slot_mapping: torch.Tensor,
236
) -> None:
237
- """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
238
-
239
- Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
240
- query and discard it, since context tokens have no query). Mirrors
241
- ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
242
- """
243
- swa_cache = attn.swa_cache_layer.kv_cache
244
- block_size = attn.swa_cache_layer.block_size
245
- cos_sin_cache = attn.rotary_emb.cos_sin_cache
246
- cache_dtype = swa_cache.dtype
247
- n_ctx = kv.shape[0]
248
- dummy_q = torch.zeros(
249
- (n_ctx, attn.n_local_heads, attn.head_dim),
250
- dtype=kv.dtype,
251
- device=kv.device,
237
+ """Insert normalized context KV without constructing an unused query."""
238
+ # The bound cache already is [num_blocks, block_size, bytes_per_token], so
239
+ # the record width needs no per-layer attribute -- one whose name the
240
+ # attention layer is free to change (kv_bytes_per_token has to become
241
+ # swa_bytes_per_token once the sliding-window and compressed records can
242
+ # differ). block_size stays the layer's, so the op's size(1) check still
243
+ # cross-checks the two rather than restating the tensor.
244
+ cache = attn.swa_cache_layer.kv_cache
245
+ torch.ops._C.fused_deepseek_v4_kv_rope_insert(
246
+ kv,
247
+ cache,
248
+ slot_mapping,
249
+ positions,
250
+ attn.rotary_emb.cos_sin_cache,
251
+ attn.swa_cache_layer.block_size,
252
+ attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None,
253
+ attn.kv_mxfp8,
254
)
253
- if cache_dtype == torch.uint8:
254
- # fp8_ds_mla UE8M0 paged layout
255
- swa_2d = swa_cache.view(swa_cache.shape[0], -1)
256
- torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
257
- dummy_q,
258
- kv,
259
- swa_2d,
260
- slot_mapping,
261
- positions,
262
- cos_sin_cache,
263
- attn.padded_heads,
264
- attn.eps,
265
- block_size,
266
- True, # apply_q_norm (unused: the query is a discarded dummy)
267
- attn.kv_mxfp8,
268
- )
269
- elif cache_dtype == torch.bfloat16:
270
- swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
271
- torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
272
- dummy_q,
273
- kv,
274
- swa_3d,
275
- slot_mapping,
276
- positions,
277
- cos_sin_cache,
278
- attn.eps,
279
- block_size,
280
- )
281
- else: # per-tensor fp8 (torch.float8_e4m3fn)
282
- # TODO(ben): double-check if this is being dispatched correctly for FI backend
283
- swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
284
- dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
285
- torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
286
- dummy_q,
287
- kv,
288
- dummy_q_fp8,
289
- swa_3d,
290
- slot_mapping,
291
- positions,
292
- cos_sin_cache,
293
- attn._flashinfer_fp8_kv_scale,
294
- attn._flashinfer_fp8_q_scale_inv,
295
- attn.eps,
296
- block_size,
297
- )
255
256
257
class DSparkDeepseekV4ForCausalLM(nn.Module):