From afd4ba6138ecd83bda5d7c320be201923604553b Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 16 Aug 2026 21:15:47 +0800 Subject: [PATCH 1/3] [Perf][DSV4] Optimize sparse top-k metadata kernels for higher prefill throughput (#52084) Signed-off-by: chaunceyjiang (cherry picked from commit 836aac92ffdaa337083934181cb6d00b64b2a1a6) --- vllm/models/deepseek_v4/common/ops/cache_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index ec53778e4075..a2b3316c5e25 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -680,7 +680,7 @@ def combine_topk_swa_indices( return combined_indices, combined_lens -_COMBINE_TOPK_SWA_NUM_WORKERS = 128 +_COMBINE_TOPK_SWA_NUM_WORKERS = 256 # Representative pointer alignment variants for Triton pointer specialization. From ae7a9d4062ba11319439538bad98adcf01eeb0a3 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 16 Aug 2026 23:24:14 +0800 Subject: [PATCH 2/3] [Perf][DSV4] Optimize global top-k index kernel with compile-time constants (#51967) Signed-off-by: chaunceyjiang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> (cherry picked from commit 83f591d7f694a3ca3ae3bf22d646e818a1421872) --- vllm/models/deepseek_v4/common/ops/cache_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index a2b3316c5e25..62c8eafdf584 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -513,15 +513,15 @@ def compute_dcp_global_topk_indices_and_lens( @triton.jit def _compute_global_topk_indices_and_lens_kernel( global_topk_indices_ptr, - global_topk_indices_stride, + global_topk_indices_stride: tl.constexpr, topk_lens_ptr, topk_indices_ptr, - topk_indices_stride, - topk, + topk_indices_stride: tl.constexpr, + topk: tl.constexpr, token_to_req_indices_ptr, block_table_ptr, - block_table_stride, - block_size, + block_table_stride: tl.constexpr, + block_size: tl.constexpr, is_valid_token_ptr, TRITON_BLOCK_SIZE: tl.constexpr, ): From 2822782cb76c7bcef952a661f4999992b6527c98 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 18 Aug 2026 22:12:42 +0000 Subject: [PATCH 3/3] test(ds4): cover sparse top-k metadata worker tiles --- tests/models/deepseek_v4/test_cache_utils.py | 62 ++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/models/deepseek_v4/test_cache_utils.py b/tests/models/deepseek_v4/test_cache_utils.py index d1a7fdf54385..2a77de6ba57c 100644 --- a/tests/models/deepseek_v4/test_cache_utils.py +++ b/tests/models/deepseek_v4/test_cache_utils.py @@ -5,6 +5,7 @@ import torch from vllm.models.deepseek_v4.common.ops.cache_utils import ( + combine_topk_swa_indices, compute_dcp_global_topk_indices_and_lens, compute_global_topk_indices_and_lens, ) @@ -64,3 +65,64 @@ def test_dcp_global_topk_ignores_stale_padding_request_index() -> None: assert indices.cpu().tolist() == [[20, 21, -1, -1], [-1, -1, -1, -1]] assert lengths.cpu().tolist() == [2, 0] + + +def test_combine_topk_swa_indices_matches_reference_across_worker_tiles() -> None: + """Sparse and sliding-window metadata cover every query token exactly.""" + device = torch.device("cuda") + query_start = torch.tensor([0, 129, 300], dtype=torch.int32, device=device) + seq_lens = torch.tensor([1024, 2048], dtype=torch.int32, device=device) + gather_lens = torch.tensor([512, 1024], dtype=torch.int32, device=device) + topk = 8 + window_size = 8 + compress_ratio = 4 + req_stride = 4096 + swa_offset = 2048 + topk_indices = torch.arange( + 300 * topk, + dtype=torch.int32, + device=device, + ).reshape(300, topk) + + actual_indices, actual_lens = combine_topk_swa_indices( + topk_indices, + query_start, + seq_lens, + gather_lens, + window_size, + compress_ratio, + topk, + req_stride, + swa_offset, + ) + + expected_indices = torch.full_like(actual_indices, -1) + expected_lens = torch.empty_like(actual_lens) + query_start_cpu = query_start.cpu().tolist() + for req_idx, (start, end) in enumerate( + zip(query_start_cpu[:-1], query_start_cpu[1:], strict=True) + ): + query_len = end - start + seq_len = int(seq_lens[req_idx].item()) + gather_start = seq_len - int(gather_lens[req_idx].item()) + start_pos = seq_len - query_len + for token_idx in range(start, end): + pos = start_pos + token_idx - start + topk_len = min((pos + 1) // compress_ratio, topk) + swa_len = min(pos + 1, window_size) + expected_indices[token_idx, :topk_len] = ( + topk_indices[token_idx, :topk_len] + req_stride * req_idx + ) + expected_indices[token_idx, topk_len : topk_len + swa_len] = ( + torch.arange( + swa_offset + pos - swa_len + 1 - gather_start, + swa_offset + pos + 1 - gather_start, + dtype=torch.int32, + device=device, + ) + + req_stride * req_idx + ) + expected_lens[token_idx] = topk_len + swa_len + + torch.testing.assert_close(actual_indices, expected_indices) + torch.testing.assert_close(actual_lens, expected_lens)