iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][DSpark] Add KV-only context insertion across V4.1 cache formats (#56441)

Browse Source 
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>
This commit is contained in:
0z5a
2026-09-15 21:34:55 -07:00
committed by GitHub
co-authored by 0z5a Codex Yongye Zhu Claude Opus 5
parent a31ec3a68b
commit 6ecd97f1b7
5 changed files with 234 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+86 -1
View File
@@ -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<torch::stable::Tensor> 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<int>(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<scalar_t const*>(kv.const_data_ptr());
auto* cache = reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr());
auto const* slots = slot_mapping.const_data_ptr<int64_t>();
auto const* positions = position_ids.const_data_ptr<int64_t>();
auto const* cos_sin = cos_sin_cache.const_data_ptr<float>();
// 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<scalar_t, 0, false>(
nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f,
num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
static_cast<int>(k_cache.stride(0)), kv_mxfp8, stream);
} else if (fp8) {
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, true>(
nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin,
fp8_scale->const_data_ptr<float>(), nullptr, 0.0f, num_tokens, num_tokens,
0, static_cast<int>(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<scalar_t, false, false>(
nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr,
nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast<int>(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)
csrc/libtorch_stable/ops.h
+7
View File
@@ -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<torch::stable::Tensor> 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,
csrc/libtorch_stable/torch_bindings.cpp
+8
View File
@@ -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(
tests/kernels/test_compressor_kv_cache.py
+116
View File
@@ -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"
)
vllm/models/deepseek_v41/nvidia/dspark.py
+17 -60
View File
@@ -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):