diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index c540eeb20..577950834 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -181,7 +181,7 @@ def all_gather_kv_cache_for_mha_extend( k_pe: torch.Tensor, ): prefix_kv_a, prefix_k_pe = token_to_kv_pool.get_mla_kv_buffer( - attn_mqa, dcp_local_prefix_kv_indices + attn_mqa, dcp_local_prefix_kv_indices, dst_dtype=kv_a.dtype ) extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu) gathered_kv_cache = all_gather_kv_cache_for_dcp( diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 05bf837e7..c579c7af4 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -786,6 +786,35 @@ class KimiK25ForConditionalGeneration(nn.Module): else {} ) + def prepare_context_parallel_metadata_for_dcp( + self, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens_sum: int, + kv_buffer_shape: torch.Size, + kv_cache_dtype, + kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn, + ): + # DCP metadata is built on the inner DeepSeek-V3 language model. + return self.language_model.prepare_context_parallel_metadata_for_dcp( + seq_lens=seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + seq_lens_sum=seq_lens_sum, + kv_buffer_shape=kv_buffer_shape, + kv_cache_dtype=kv_cache_dtype, + kv_cache_device=kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn, + ) + def forward( self, input_ids: torch.Tensor,