diff --git a/python/sglang/srt/layers/dcp/layout.py b/python/sglang/srt/layers/dcp/layout.py index 4157e2ea5..56f9a1939 100644 --- a/python/sglang/srt/layers/dcp/layout.py +++ b/python/sglang/srt/layers/dcp/layout.py @@ -51,6 +51,25 @@ def filter_dcp_local_kv_indices(kv_indices: torch.Tensor): return kv_indices +def filter_dcp_local_chunk_kv_indices( + kv_indices: torch.Tensor, + chunk_starts_cpu: torch.Tensor, + chunk_seq_lens_cpu: torch.Tensor, +) -> torch.Tensor: + parallel = get_parallel() + if not parallel.dcp_enabled: + return kv_indices + + dcp_size = parallel.dcp_size + parts = [] + offset = 0 + for start, length in zip(chunk_starts_cpu.tolist(), chunk_seq_lens_cpu.tolist()): + first = (parallel.dcp_rank - start) % dcp_size + parts.append(kv_indices[offset + first : offset + length : dcp_size]) + offset += length + return torch.cat(parts) // dcp_size + + def update_local_kv_lens_for_dcp(kv_len_arr): """In-place per-rank KV length: the start=0 case of get_dcp_lens. diff --git a/python/sglang/srt/layers/dcp/planner.py b/python/sglang/srt/layers/dcp/planner.py index d41d23630..7a8a25e5c 100644 --- a/python/sglang/srt/layers/dcp/planner.py +++ b/python/sglang/srt/layers/dcp/planner.py @@ -109,10 +109,9 @@ def prepare_decode_context_parallel_metadata( extend_prefix_lens_sum, parallel.dcp_size, ) + # Prefix lengths are dcp_size-aligned (widened allocator page), so no nonzero(). dcp_local_prefix_kv_indices = ( - dcp_prefix_kv_indices[ - dcp_prefix_kv_indices % parallel.dcp_size == parallel.dcp_rank - ] + dcp_prefix_kv_indices[parallel.dcp_rank :: parallel.dcp_size] // parallel.dcp_size ) dcp_kv_buffer = torch.empty( diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 770958c76..5125cc7f0 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -10,6 +10,7 @@ from sglang.kernels.ops.kvcache.kv_indices import ( create_flashinfer_kv_indices_triton, ) from sglang.srt.environ import envs +from sglang.srt.layers.dcp.layout import filter_dcp_local_chunk_kv_indices from sglang.srt.model_executor.forward_context import ( get_req_to_token_pool, get_token_to_kv_pool, @@ -84,6 +85,11 @@ class ForwardBatchDeepSeekMHAMixin: chunk_kv_indices, req_to_token.shape[1], ) + chunk_kv_indices = filter_dcp_local_chunk_kv_indices( + chunk_kv_indices, + self.prefix_chunk_starts_cpu[idx], + self.prefix_chunk_seq_lens_cpu[idx], + ) self.prefix_chunk_kv_indices.append(chunk_kv_indices) # Here we suppose the length of each chunk is equal diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 4eeb8bf49..01e66499a 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -450,7 +450,6 @@ class DeepseekMHAForwardMixin: forward_batch: ForwardBatch, ): if _is_cuda: - kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices) kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer( self.attn_mha, kv_indices, dst_dtype ) diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index 8c7f5731d..53242ec0f 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/test/registered/dcp/test_dcp_layout_unit.py @@ -22,7 +22,10 @@ import torch from sglang.srt import runtime_context as rc from sglang.srt.configs.model_config import ModelConfig from sglang.srt.layers.attention.triton_backend import TritonAttnBackend -from sglang.srt.layers.dcp.layout import get_dcp_lens +from sglang.srt.layers.dcp.layout import ( + filter_dcp_local_chunk_kv_indices, + get_dcp_lens, +) from sglang.srt.layers.linear import QKVParallelLinear from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator @@ -47,6 +50,97 @@ def _legacy_inplace_formula(length: int, n: int, rank: int) -> int: return (length - rank - 1) // n + 1 +class TestFilterDcpLocalChunkKvIndices(CustomTestCase): + PAGE = 64 + + def _build_chunk(self, starts, lens, dcp_size, seed=0): + g = torch.Generator().manual_seed(seed) + widened = self.PAGE * dcp_size + runs = [] + for start, length in zip(starts, lens): + if length == 0: + runs.append(torch.empty(0, dtype=torch.int64)) + continue + pos = torch.arange(start, start + length) + page_of = pos // widened + bases = ( + torch.randint(0, 64, (int(page_of.max()) + 1,), generator=g) * widened + ) + runs.append(bases[page_of] + pos % widened) + return torch.cat(runs) if runs else torch.empty(0, dtype=torch.int64) + + def _owner_rule(self, kv, dcp_size, dcp_rank): + return kv[kv % dcp_size == dcp_rank] // dcp_size + + def _run(self, starts, lens, dcp_size, dcp_rank, seed=0): + kv = self._build_chunk(starts, lens, dcp_size, seed) + with rc.get_parallel().override( + dcp_enabled=dcp_size > 1, dcp_size=dcp_size, dcp_rank=dcp_rank + ): + got = filter_dcp_local_chunk_kv_indices( + kv, torch.tensor(starts), torch.tensor(lens) + ) + return kv, got + + def test_matches_owner_rule_on_unaligned_runs(self): + cases = [ + ([0], [1]), + ([0, 0], [2048, 2048]), + ([8192], [3]), + ([5, 13, 31], [7, 0, 19]), + ([1365, 1365, 1365], [1365, 1365, 1365]), + ([43690, 43690], [43690, 17]), + ([7, 7, 7, 7, 7], [11, 13, 0, 1, 40]), + ] + for dcp_size in [2, 3, 4, 8]: + for dcp_rank in range(dcp_size): + for i, (starts, lens) in enumerate(cases): + kv, got = self._run(starts, lens, dcp_size, dcp_rank, seed=i) + torch.testing.assert_close( + got, + self._owner_rule(kv, dcp_size, dcp_rank), + msg=f"size={dcp_size} rank={dcp_rank} " + f"starts={starts} lens={lens}", + ) + + def test_row_count_matches_get_dcp_lens(self): + starts, lens = [1365, 2730, 0], [1365, 900, 37] + for dcp_size in [2, 3, 4, 8]: + for dcp_rank in range(dcp_size): + _, got = self._run(starts, lens, dcp_size, dcp_rank) + expected = get_dcp_lens( + torch.tensor(lens), dcp_size, dcp_rank, start=torch.tensor(starts) + ) + self.assertEqual( + got.numel(), + int(expected.sum()), + f"size={dcp_size} rank={dcp_rank}", + ) + + def test_second_chunk_at_batch_three(self): + starts, lens = [2730, 2730, 2730], [1365, 1365, 1365] + for dcp_rank in range(8): + kv, got = self._run(starts, lens, 8, dcp_rank) + torch.testing.assert_close(got, self._owner_rule(kv, 8, dcp_rank)) + naive = kv[dcp_rank::8] // 8 + self.assertNotEqual( + got.numel(), + naive.numel(), + f"rank={dcp_rank}: phase-free stride happens to match here, " + f"so this case no longer guards the phase term", + ) + + def test_identity_without_dcp(self): + kv = torch.arange(37) + with rc.get_parallel().override(dcp_enabled=False, dcp_size=1, dcp_rank=0): + self.assertIs( + filter_dcp_local_chunk_kv_indices( + kv, torch.tensor([0]), torch.tensor([37]) + ), + kv, + ) + + class TestGetDcpLens(CustomTestCase): def test_start_none_matches_owner_count(self): for n in DCP_SIZES: