File tree Expand file tree Collapse file tree 1 file changed +3
-3
lines changed Expand file tree Collapse file tree 1 file changed +3
-3
lines changed Original file line number Diff line number Diff line change @@ -63,7 +63,7 @@ def get_forward_context() -> ForwardContext:
63
63
def set_forward_context (attn_metadata : Any ,
64
64
vllm_config : VllmConfig ,
65
65
virtual_engine : int = 0 ,
66
- num_tokens : int = 0 ):
66
+ num_tokens : Optional [ int ] = None ):
67
67
"""A context manager that stores the current forward context,
68
68
can be attention metadata, etc.
69
69
Here we can inject common logic for every model forward pass.
@@ -73,8 +73,8 @@ def set_forward_context(attn_metadata: Any,
73
73
if need_to_track_batchsize :
74
74
forward_start_time = time .perf_counter ()
75
75
dp_metadata : Optional [DPMetadata ] = None
76
- if vllm_config .parallel_config .data_parallel_size > 1 :
77
- dp_size = vllm_config . parallel_config . data_parallel_size
76
+ dp_size = vllm_config .parallel_config .data_parallel_size
77
+ if dp_size > 1 and ( attn_metadata is not None or num_tokens is not None ):
78
78
dp_rank = vllm_config .parallel_config .data_parallel_rank
79
79
if attn_metadata is not None and hasattr (attn_metadata ,
80
80
"num_prefill_tokens" ):
You can’t perform that action at this time.
0 commit comments