From fa826e08b1d0cd27bd3f825954cb86639af8865e Mon Sep 17 00:00:00 2001 From: abing Date: Sat, 19 Sep 2026 10:24:01 -0700 Subject: [PATCH] Reuse shared compressed KV dequantization in DeepSeek V4.1 CP prefill (cherry picked from commit 44eb378e94610e4ccc6b990bf09c105f8c8f9ee9) --- .../layers/attention/deepseek_v4_backend.py | 41 +++-- .../attention/dsv4/sparse_prefill_utils.py | 17 +- .../attention/test_dsv41_shared_prefill_kv.py | 151 ++++++++++++++++++ 3 files changed, 194 insertions(+), 15 deletions(-) create mode 100644 test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 980a7303f..9b95fb127 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -1296,6 +1296,11 @@ class DeepseekV4AttnBackend( ] = None self.online_c128_mtp = OnlineC128MTPController(self) self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device) + # CP V4.1 consumers share compressed KV across layers. Separate ratio + # workspaces keep those prefixes intact while each layer refreshes SWA. + self.shared_compressed_prefill_workspaces = { + ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2) + } spec_alg = model_runner.spec_algorithm self.needs_cpu_seq_lens = not spec_alg.is_dspark() and ( not _is_cuda or self.online_c128_mtp.enabled() @@ -4034,20 +4039,38 @@ class DeepseekV4AttnBackend( compress_ratio, core_attn_metadata, extra_page_size ) n_compressed = flat_token_ids.shape[0] - workspace = self.sparse_prefill_workspace.get( - n_compressed + cache.swa_token_ids.shape[0] + reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch) + workspace_pool = ( + self.shared_compressed_prefill_workspaces[compress_ratio] + if reuse_compressed + else self.sparse_prefill_workspace ) + workspace = workspace_pool.get(n_compressed + cache.swa_token_ids.shape[0]) compressed_slice = workspace[:n_compressed] swa_slice = workspace[n_compressed:] if compressed_slice is not None: - dequantize_k_cache_paged( - extra_k_cache, - flat_token_ids, - page_size=extra_page_size, - out=compressed_slice, - layout=token_to_kv_pool.get_extra_key_layout(layer_id), - ) + source_key = None + if reuse_compressed: + source_layer = token_to_kv_pool.source_layer_of(layer_id) + source_key = (source_layer, workspace.data_ptr()) + gather = cache.compressed[compress_ratio] + # A source layer may have just updated its cache in place. Consumer + # layers only reuse the compressed prefix; their top-k and SWA stay live. + if ( + source_key is None + or layer_id == source_key[0] + or gather.dequantized_source != source_key + ): + dequantize_k_cache_paged( + extra_k_cache, + flat_token_ids, + page_size=extra_page_size, + out=compressed_slice, + layout=token_to_kv_pool.get_extra_key_layout(layer_id), + ) + if source_key is not None: + gather.dequantized_source = source_key dequantize_k_cache_paged( token_to_kv_pool.get_swa_key_buffer_radix(layer_id), cache.swa_token_ids, diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py index 49a8b52dd..bd0a56906 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -78,10 +78,11 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool: class SparsePrefillWorkspace: """Backend-owned scratch storage for sparse prefill KV dequantization. - The workspace contents are fully overwritten before every attention call, - so token buckets and compression ratios can safely share one buffer. Sparse - prefill executes eagerly and serially on the supported paths, which makes it - safe to replace the scratch allocation when a larger extent is needed. + Callers normally overwrite the entire workspace. Shared compressed-KV + callers keep separate workspaces per ratio and track prefix validity in the + per-forward gather cache, including the allocation address. Sparse prefill + executes eagerly and serially on the supported paths, so the allocation can + be replaced when a larger extent is needed. """ def __init__(self, device: torch.device): @@ -275,14 +276,18 @@ class CompressedGather: # chunk-invariant per request; subsequent layers only overwrite that prefix. combined_indices: Optional[torch.Tensor] = None combined_lens: Optional[torch.Tensor] = None + # Valid only for this forward's gather layout. Each ratio has its own + # workspace; its compressed prefix survives consumer layers' SWA writes. + dequantized_source: Optional[tuple[int, int]] = None @dataclass class SparsePrefillChunkCache: """Cache prefill-chunk metadata shared across layers. - Fields depend on request/token mappings and compressed page tables, not - per-layer k_cache; per-layer top-k combinations are recomputed into reused + Gather layouts depend on request/token mappings and compressed page tables. + Shared-source dequantization keys live only for this forward; per-layer + top-k combinations are recomputed into reused buffers. """ diff --git a/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py b/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py new file mode 100644 index 000000000..d8b9de75d --- /dev/null +++ b/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py @@ -0,0 +1,151 @@ +"""Shared compressed-KV workspace validity across layers and forwards.""" +import sys +from types import SimpleNamespace as NS +from unittest.mock import patch + +import pytest +import torch + +from sglang.srt.layers.attention import deepseek_v4_backend as backend +from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( + CompressedGather, + SparsePrefillWorkspace, + WORKSPACE_DIM, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def test_shared_compressed_dequant_lifetime(): + device = "cuda" + obj = object.__new__(backend.DeepseekV4AttnBackend) + obj.sparse_prefill_workspace = SparsePrefillWorkspace(device) + obj.shared_compressed_prefill_workspaces = { + ratio: SparsePrefillWorkspace(device) for ratio in (1, 2) + } + obj.softmax_scale = 0.1 + obj.head_dim_v = WORKSPACE_DIM + sources = {0: 0, 1: 1, 2: 0, 3: 1, 4: 4, 5: 4, 6: 6, 7: 7} + ratios = {0: 1, 1: 2, 2: 1, 3: 2, 4: 1, 5: 1, 6: 0, 7: 4} + compressed = { + source: torch.full( + (128, 1, WORKSPACE_DIM), + float(source + 1), + device=device, + dtype=torch.bfloat16, + ) + for source in (0, 1, 4, 7) + } + swa = torch.empty((16, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16) + pool = NS( + source_layer_of=lambda layer: sources[layer], + get_extra_key_page_size=lambda layer: 1, + get_extra_key_buffer=lambda layer: compressed[sources[layer]], + get_extra_key_layout=lambda layer: None, + get_swa_key_buffer_radix=lambda layer: swa, + get_swa_key_layout=lambda: None, + ) + calls = [] + active = [True] + + def dequant(src, indices, *, out, **kwargs): + calls.append("swa" if src is swa else sources_by_ptr[src.data_ptr()]) + out.copy_(src.index_select(0, indices.long())) + + sources_by_ptr = {v.data_ptr(): k for k, v in compressed.items()} + + def make_cache(n): + gathers = { + ratio: CompressedGather( + flat_token_ids=torch.arange(n // ratio, device=device, dtype=torch.int32), + compressed_base=torch.zeros(1, device=device, dtype=torch.int32), + swa_base=torch.zeros(1, device=device, dtype=torch.int32), + ) + for ratio in (1, 2, 4) + } + indices = torch.zeros((4, 128), device=device, dtype=torch.int32) + lengths = torch.full((4,), 3, device=device, dtype=torch.int32) + cache = NS( + compressed=gathers, + swa_token_ids=torch.arange(3, device=device), + swa_page_size=1, + c0_combined_indices=indices, + c0_combined_lens=lengths, + ) + cache.layer_inputs = lambda ratio, core, page: ( + gathers[ratio].flat_token_ids, + indices, + lengths, + ) + return cache + + def forward(layer): + return obj._forward_prefill_sparse( + torch.empty((4, 1, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16), + layer, + ratios[layer], + NS(), + pool, + NS(), + torch.zeros(1, device=device), + ) + + with ( + patch.object(backend, "is_cp_active", side_effect=lambda _: active[0]), + patch.object(backend, "dequantize_k_cache_paged", side_effect=dequant), + patch( + "sgl_kernel.flash_mla.flash_mla_sparse_fwd", + side_effect=lambda **kw: (kw["kv"].clone(), None, None), + ), + ): + # Same-size replay, then growth and shrink must all read freshly written KV. + for step, n in enumerate((8, 8, 40, 4)): + obj.forward_metadata = NS(sparse_prefill_cache=make_cache(n)) + for source, tensor in compressed.items(): + tensor.fill_(source + 1 + step * 10) + for layer, should_dequant in ( + (0, True), + (1, True), + (2, False), + (3, False), + (4, True), + (5, False), + (2, True), + (3, False), + (6, False), + (7, True), + (7, True), + ): + swa.fill_(100 + layer + step) + before = len(calls) + active[0] = True + actual = forward(layer) + actual_calls = calls[before:] + assert actual_calls.count("swa") == 1 + assert len(actual_calls) == 1 + int(should_dequant) + active[0] = False + expected = forward(layer) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + # Workspace replacement invalidates even an unchanged source identity. + obj.shared_compressed_prefill_workspaces[1].get(256 + step * 256) + active[0] = True + before = len(calls) + actual = forward(2) + assert calls[before:] == [0, "swa"] + active[0] = False + torch.testing.assert_close(actual, forward(2), rtol=0, atol=0) + + # Re-executing a producer may mutate the same cache address in place. + compressed[0].add_(1) + active[0] = True + before = len(calls) + actual = forward(0) + assert calls[before:] == [0, "swa"] + active[0] = False + torch.testing.assert_close(actual, forward(0), rtol=0, atol=0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, *sys.argv[1:]]))