[DCP] Enable decode context parallel for Kimi K2.5 NVFP4 (#31514)
Co-authored-by: kpham-sgl <kpham-sgl@users.noreply.github.com>
This commit is contained in:
@@ -181,7 +181,7 @@ def all_gather_kv_cache_for_mha_extend(
|
|||||||
k_pe: torch.Tensor,
|
k_pe: torch.Tensor,
|
||||||
):
|
):
|
||||||
prefix_kv_a, prefix_k_pe = token_to_kv_pool.get_mla_kv_buffer(
|
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)
|
extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu)
|
||||||
gathered_kv_cache = all_gather_kv_cache_for_dcp(
|
gathered_kv_cache = all_gather_kv_cache_for_dcp(
|
||||||
|
|||||||
@@ -786,6 +786,35 @@ class KimiK25ForConditionalGeneration(nn.Module):
|
|||||||
else {}
|
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(
|
def forward(
|
||||||
self,
|
self,
|
||||||
input_ids: torch.Tensor,
|
input_ids: torch.Tensor,
|
||||||
|
|||||||
Reference in New Issue
Block a user