diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 254b9783bc..344b82fbda 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -657,7 +657,12 @@ static void launchFusedDeepseekV4Templated( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - if (num_tokens_full < NUM_TOKEN_CUTOFF) { + // The reduced grid pays off by iterating a token's Q heads inside one CTA. + // With zero query heads there is nothing to iterate, so one warp per block + // would do all the work and the other seven idle; keep the warp-per-slot + // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single + // insert, where it is worth ~2x. + if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { cudaLaunchKernelEx(&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, @@ -993,6 +998,86 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() && + kv.dim() == 2 && kv.size(1) == 512 && + kv.scalar_type() == ScalarType::BFloat16, + "kv must be contiguous CUDA bfloat16 [tokens, 512]"); + auto const device = kv.get_device_index(); + STD_TORCH_CHECK(k_cache.device().is_cuda() && + k_cache.get_device_index() == device && + k_cache.dim() == 3 && cache_block_size > 0 && + k_cache.size(1) == cache_block_size && + k_cache.stride(2) == 1, + "k_cache must be a paged cache on the KV device"); + for (auto const* tensor : {&slot_mapping, &position_ids}) { + STD_TORCH_CHECK(tensor->device().is_cuda() && + tensor->get_device_index() == device && + tensor->scalar_type() == ScalarType::Long && + tensor->dim() == 1 && tensor->is_contiguous(), + "slots and positions must be contiguous int64 on the KV device"); + } + STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) && + slot_mapping.size(0) <= kv.size(0), + "positions must match KV rows; slots may omit padded rows"); + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() && + cos_sin_cache.get_device_index() == device && + cos_sin_cache.scalar_type() == ScalarType::Float && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(), + "cos_sin_cache must be contiguous float32 [max_pos, 64]"); + bool const packed = k_cache.scalar_type() == ScalarType::Byte; + bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK(packed || fp8 || + k_cache.scalar_type() == ScalarType::BFloat16, + "cache must be uint8, bfloat16 or float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512), + "cache row width does not match its dtype"); + if (fp8) { + STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() && + fp8_scale->get_device_index() == device && + fp8_scale->scalar_type() == ScalarType::Float && + fp8_scale->numel() == 1, + "fp8 cache requires one float32 KV scale on the KV device"); + } + int const num_tokens = static_cast(slot_mapping.size(0)); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard(device); + const cudaStream_t stream = get_current_cuda_stream(device); + using scalar_t = c10::BFloat16; + auto const* input = reinterpret_cast(kv.const_data_ptr()); + auto* cache = reinterpret_cast(k_cache.mutable_data_ptr()); + auto const* slots = slot_mapping.const_data_ptr(); + auto const* positions = position_ids.const_data_ptr(); + auto const* cos_sin = cos_sin_cache.const_data_ptr(); + // Zero query heads schedules only the existing KV branch, preserving its + // RoPE rounding and quantization without allocating any query tensors. + if (packed) { + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, + num_tokens, num_tokens, 0, static_cast(cache_block_size), + static_cast(k_cache.stride(0)), kv_mxfp8, stream); + } else if (fp8) { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, + fp8_scale->const_data_ptr(), nullptr, 0.0f, num_tokens, num_tokens, + 0, static_cast(cache_block_size), k_cache.stride(0), + k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream); + } else { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr, + nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), + k_cache.stride(0) * 2, k_cache.stride(1) * 2, false, + "fused_deepseek_v4_kv_rope_insert", stream); + } +} + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index e19f49d603..bc08d90a48 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -259,6 +259,13 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 54548dfd05..3918ab8498 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); ops.def( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" "Tensor q_in, Tensor kv, Tensor! k_cache, " @@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Positional encoding kernels (shared CUDA/ROCm) ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); ops.impl( diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 7e092229fe..c1de94be92 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -42,6 +42,122 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize( + "cache_dtype,kv_mxfp8", + [ + (torch.uint8, False), + (torch.uint8, True), + (torch.bfloat16, False), + (torch.float8_e4m3fn, False), + ], +) +@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024]) +@pytest.mark.parametrize("use_graph", [False, True]) +def test_dspark_context_kv_matches_query_insert( + cache_dtype, kv_mxfp8, num_tokens, use_graph +): + """KV-only insertion must preserve every cache byte, including graph replay.""" + from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv + + torch.manual_seed(42) + block_size = 256 + num_blocks = math.ceil((num_tokens + 7) / block_size) + row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(num_tokens + 3, device="cuda") + 7 + angles = torch.randn(num_tokens + 16, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens] + slots[::5] = -1 + scale = torch.tensor([0.7], device="cuda") + # Only what _insert_context_kv actually reads: the record width comes off + # the cache tensor, so no per-record width attribute is stubbed here. + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size), + kv_mxfp8=kv_mxfp8, + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=scale, + ) + + def legacy_insert(): + q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda") + if cache_dtype == torch.uint8: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + reference.view(num_blocks, -1), + slots, + positions, + cos_sin, + 8, + 1e-20, + block_size, + True, + kv_mxfp8, + ) + elif cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + reference, + slots, + positions, + cos_sin, + 1e-20, + block_size, + ) + else: + q_fp8 = torch.empty_like(q, dtype=cache_dtype) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + reference, + slots, + positions, + cos_sin, + scale, + scale, + 1e-20, + block_size, + ) + + def insert(): + _insert_context_kv(attn, kv, positions, slots) + + insert() + graph = None + if use_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + insert() + for iteration in range(3): + kv.normal_() + positions.add_(1) + if iteration == 2: + slots.fill_(-1) + backing.view(torch.uint8).fill_(165) + reference_backing.copy_(backing) + before = kv.clone() + legacy_insert() + insert() if graph is None else graph.replay() + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close(kv, before, rtol=0, atol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index c19dca8005..214681f3a9 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -234,67 +234,24 @@ def _insert_context_kv( positions: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: - """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. - - Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy - query and discard it, since context tokens have no query). Mirrors - ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. - """ - swa_cache = attn.swa_cache_layer.kv_cache - block_size = attn.swa_cache_layer.block_size - cos_sin_cache = attn.rotary_emb.cos_sin_cache - cache_dtype = swa_cache.dtype - n_ctx = kv.shape[0] - dummy_q = torch.zeros( - (n_ctx, attn.n_local_heads, attn.head_dim), - dtype=kv.dtype, - device=kv.device, + """Insert normalized context KV without constructing an unused query.""" + # The bound cache already is [num_blocks, block_size, bytes_per_token], so + # the record width needs no per-layer attribute -- one whose name the + # attention layer is free to change (kv_bytes_per_token has to become + # swa_bytes_per_token once the sliding-window and compressed records can + # differ). block_size stays the layer's, so the op's size(1) check still + # cross-checks the two rather than restating the tensor. + cache = attn.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + slot_mapping, + positions, + attn.rotary_emb.cos_sin_cache, + attn.swa_cache_layer.block_size, + attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None, + attn.kv_mxfp8, ) - if cache_dtype == torch.uint8: - # fp8_ds_mla UE8M0 paged layout - swa_2d = swa_cache.view(swa_cache.shape[0], -1) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - dummy_q, - kv, - swa_2d, - slot_mapping, - positions, - cos_sin_cache, - attn.padded_heads, - attn.eps, - block_size, - True, # apply_q_norm (unused: the query is a discarded dummy) - attn.kv_mxfp8, - ) - elif cache_dtype == torch.bfloat16: - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - dummy_q, - kv, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn.eps, - block_size, - ) - else: # per-tensor fp8 (torch.float8_e4m3fn) - # TODO(ben): double-check if this is being dispatched correctly for FI backend - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - dummy_q, - kv, - dummy_q_fp8, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn._flashinfer_fp8_kv_scale, - attn._flashinfer_fp8_q_scale_inv, - attn.eps, - block_size, - ) class DSparkDeepseekV4ForCausalLM(nn.Module):