From 2310ca96b9a2ba16d05f5eb970f06ac05674f959 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Wed, 5 Mar 2025 19:03:46 +0000 Subject: [PATCH 1/2] [Bugfix] Fix num_tokens_across_dp Signed-off-by: Tyler Michael Smith --- vllm/forward_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 540a35e1ecb9..ecc6aff3c79c 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -88,6 +88,7 @@ def set_forward_context(attn_metadata: Any, dtype=torch.int32) from vllm.distributed.parallel_state import get_dp_group dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group) + num_tokens_across_dp = num_tokens_tensor.tolist() cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0) dp_metadata = DPMetadata(num_tokens_across_dp, cu_tokens_across_dp_cpu) From 3cf2244b61b56c388dc8944ac13f792e5ae7003b Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Wed, 5 Mar 2025 21:07:13 +0000 Subject: [PATCH 2/2] Remove it Signed-off-by: Tyler Michael Smith --- vllm/forward_context.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/forward_context.py b/vllm/forward_context.py index ecc6aff3c79c..e195a03c5cac 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -27,7 +27,6 @@ @dataclass class DPMetadata: - num_tokens_across_dp: list[int] cu_tokens_across_dp_cpu: torch.Tensor @@ -88,9 +87,8 @@ def set_forward_context(attn_metadata: Any, dtype=torch.int32) from vllm.distributed.parallel_state import get_dp_group dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group) - num_tokens_across_dp = num_tokens_tensor.tolist() cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0) - dp_metadata = DPMetadata(num_tokens_across_dp, cu_tokens_across_dp_cpu) + dp_metadata = DPMetadata(cu_tokens_across_dp_cpu) global _forward_context prev_context = _forward_context