Skip to content

vllm.forward_context

_forward_context module-attribute 

_forward_context: Optional[ForwardContext] = None

batchsize_forward_time module-attribute 

batchsize_forward_time: defaultdict = defaultdict(list)

batchsize_logging_interval module-attribute 

batchsize_logging_interval: float = (
    VLLM_LOG_BATCHSIZE_INTERVAL
)

forward_start_time module-attribute 

forward_start_time: float = 0

last_logging_time module-attribute 

last_logging_time: float = 0

logger module-attribute 

logger = init_logger(__name__)

track_batchsize module-attribute 

track_batchsize: bool = VLLM_LOG_BATCHSIZE_INTERVAL >= 0

DPMetadata dataclass

Source code in vllm/forward_context.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
@dataclass
class DPMetadata:
    max_tokens_across_dp_cpu: torch.Tensor
    cu_tokens_across_dp_cpu: torch.Tensor
    local_sizes: Optional[list[int]] = None

    @staticmethod
    def num_tokens_across_dp(num_tokens: int, dp_size: int,
                             dp_rank: int) -> torch.Tensor:
        """
        Gather the num_tokens across all DP ranks and return results in a
        CPU tensor of size dp_size.
        """
        num_tokens_across_dp = [0] * dp_size
        num_tokens_across_dp[dp_rank] = num_tokens
        num_tokens_tensor = torch.tensor(num_tokens_across_dp,
                                         device="cpu",
                                         dtype=torch.int32)
        from vllm.distributed.parallel_state import get_dp_group
        dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
        return num_tokens_tensor

    @staticmethod
    def make(
            parallel_config: ParallelConfig,
            attn_metadata: Any,
            num_tokens: int,
            num_tokens_across_dp: Optional[torch.Tensor] = None
    ) -> "DPMetadata":

        assert parallel_config.data_parallel_size > 1
        dp_size = parallel_config.data_parallel_size
        dp_rank = parallel_config.data_parallel_rank
        if attn_metadata is not None and hasattr(attn_metadata,
                                                 "num_prefill_tokens"):
            # for v0 attention backends
            batchsize = attn_metadata.num_prefill_tokens + \
                attn_metadata.num_decode_tokens
        else:
            # for v1 attention backends or no attn_metadata
            batchsize = num_tokens

        # If num_tokens_across_dp is None, it will be computed by all_reduce
        # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
        assert (num_tokens_across_dp is None
                or num_tokens_across_dp[dp_rank] == batchsize)
        if num_tokens_across_dp is None:
            num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
                batchsize, dp_size, dp_rank)
        max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)
        cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)
        return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu)

    @contextmanager
    def chunked_sizes(self, max_chunk_size_per_rank: int, chunk_idx: int):
        """
        Context manager to compute and temporarily set the per-rank local token
        sizes for a specific chunk during chunked forward execution.

        This is necessary to ensure each DP (data parallel) rank processes its
        designated portion of tokens in lockstep with others, even when the
        token counts are uneven or some ranks have completed their input early.

        For chunked execution, we break up the total tokens on each rank into
        multiple chunks (of at most `max_chunk_size_per_rank`), and for a given
        `chunk_idx`, this context manager sets `self.local_sizes` to the number
        of tokens to process in that chunk on each rank.

        It uses cumulative sizes (`cu_tokens_across_dp_cpu`) to derive the
        number of tokens per rank, and calls `_compute_chunked_local_num_tokens`
        to determine the chunk-wise split.

        `self.local_sizes` is only valid inside the context.

        Args:
            max_chunk_size_per_rank: The max number of tokens each rank is 
                                     allowed to process in this chunk.
            chunk_idx: The index of the chunk to compute sizes for.
        """
        cu_sizes = self.cu_tokens_across_dp_cpu
        num_tokens_across_dp_cpu = [
            (cu_sizes[i] -
             cu_sizes[i - 1]).item() if i > 0 else cu_sizes[0].item()
            for i in range(len(cu_sizes))
        ]
        self.local_sizes = _compute_chunked_local_num_tokens(
            num_tokens_across_dp_cpu, max_chunk_size_per_rank, chunk_idx)
        try:
            yield self.local_sizes
        finally:
            self.local_sizes = None

    def get_chunk_sizes_across_dp_rank(self) -> Optional[list[int]]:
        return self.local_sizes

cu_tokens_across_dp_cpu instance-attribute 

cu_tokens_across_dp_cpu: Tensor

local_sizes class-attribute instance-attribute 

local_sizes: Optional[list[int]] = None

max_tokens_across_dp_cpu instance-attribute 

max_tokens_across_dp_cpu: Tensor

__init__ 

__init__(
    max_tokens_across_dp_cpu: Tensor,
    cu_tokens_across_dp_cpu: Tensor,
    local_sizes: Optional[list[int]] = None,
) -> None

chunked_sizes 

chunked_sizes(max_chunk_size_per_rank: int, chunk_idx: int)

Context manager to compute and temporarily set the per-rank local token sizes for a specific chunk during chunked forward execution.

This is necessary to ensure each DP (data parallel) rank processes its designated portion of tokens in lockstep with others, even when the token counts are uneven or some ranks have completed their input early.

For chunked execution, we break up the total tokens on each rank into multiple chunks (of at most max_chunk_size_per_rank), and for a given chunk_idx, this context manager sets self.local_sizes to the number of tokens to process in that chunk on each rank.

It uses cumulative sizes (cu_tokens_across_dp_cpu) to derive the number of tokens per rank, and calls _compute_chunked_local_num_tokens to determine the chunk-wise split.

self.local_sizes is only valid inside the context.

Parameters:

Name Type Description Default
max_chunk_size_per_rank int

The max number of tokens each rank is allowed to process in this chunk.

 required
chunk_idx int

The index of the chunk to compute sizes for.

 required
Source code in vllm/forward_context.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
@contextmanager
def chunked_sizes(self, max_chunk_size_per_rank: int, chunk_idx: int):
    """
    Context manager to compute and temporarily set the per-rank local token
    sizes for a specific chunk during chunked forward execution.

    This is necessary to ensure each DP (data parallel) rank processes its
    designated portion of tokens in lockstep with others, even when the
    token counts are uneven or some ranks have completed their input early.

    For chunked execution, we break up the total tokens on each rank into
    multiple chunks (of at most `max_chunk_size_per_rank`), and for a given
    `chunk_idx`, this context manager sets `self.local_sizes` to the number
    of tokens to process in that chunk on each rank.

    It uses cumulative sizes (`cu_tokens_across_dp_cpu`) to derive the
    number of tokens per rank, and calls `_compute_chunked_local_num_tokens`
    to determine the chunk-wise split.

    `self.local_sizes` is only valid inside the context.

    Args:
        max_chunk_size_per_rank: The max number of tokens each rank is 
                                 allowed to process in this chunk.
        chunk_idx: The index of the chunk to compute sizes for.
    """
    cu_sizes = self.cu_tokens_across_dp_cpu
    num_tokens_across_dp_cpu = [
        (cu_sizes[i] -
         cu_sizes[i - 1]).item() if i > 0 else cu_sizes[0].item()
        for i in range(len(cu_sizes))
    ]
    self.local_sizes = _compute_chunked_local_num_tokens(
        num_tokens_across_dp_cpu, max_chunk_size_per_rank, chunk_idx)
    try:
        yield self.local_sizes
    finally:
        self.local_sizes = None

get_chunk_sizes_across_dp_rank 

get_chunk_sizes_across_dp_rank() -> Optional[list[int]]
Source code in vllm/forward_context.py 
136
137
def get_chunk_sizes_across_dp_rank(self) -> Optional[list[int]]:
    return self.local_sizes

make staticmethod 

make(
    parallel_config: ParallelConfig,
    attn_metadata: Any,
    num_tokens: int,
    num_tokens_across_dp: Optional[Tensor] = None,
) -> DPMetadata
Source code in vllm/forward_context.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
@staticmethod
def make(
        parallel_config: ParallelConfig,
        attn_metadata: Any,
        num_tokens: int,
        num_tokens_across_dp: Optional[torch.Tensor] = None
) -> "DPMetadata":

    assert parallel_config.data_parallel_size > 1
    dp_size = parallel_config.data_parallel_size
    dp_rank = parallel_config.data_parallel_rank
    if attn_metadata is not None and hasattr(attn_metadata,
                                             "num_prefill_tokens"):
        # for v0 attention backends
        batchsize = attn_metadata.num_prefill_tokens + \
            attn_metadata.num_decode_tokens
    else:
        # for v1 attention backends or no attn_metadata
        batchsize = num_tokens

    # If num_tokens_across_dp is None, it will be computed by all_reduce
    # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
    assert (num_tokens_across_dp is None
            or num_tokens_across_dp[dp_rank] == batchsize)
    if num_tokens_across_dp is None:
        num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
            batchsize, dp_size, dp_rank)
    max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)
    cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)
    return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu)

num_tokens_across_dp staticmethod 

num_tokens_across_dp(
    num_tokens: int, dp_size: int, dp_rank: int
) -> Tensor

Gather the num_tokens across all DP ranks and return results in a CPU tensor of size dp_size.

Source code in vllm/forward_context.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@staticmethod
def num_tokens_across_dp(num_tokens: int, dp_size: int,
                         dp_rank: int) -> torch.Tensor:
    """
    Gather the num_tokens across all DP ranks and return results in a
    CPU tensor of size dp_size.
    """
    num_tokens_across_dp = [0] * dp_size
    num_tokens_across_dp[dp_rank] = num_tokens
    num_tokens_tensor = torch.tensor(num_tokens_across_dp,
                                     device="cpu",
                                     dtype=torch.int32)
    from vllm.distributed.parallel_state import get_dp_group
    dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
    return num_tokens_tensor

ForwardContext dataclass

Source code in vllm/forward_context.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@dataclass
class ForwardContext:
    # copy from vllm_config.compilation_config.static_forward_context
    no_compile_layers: dict[str, Any]
    """
    Type AttentionMetadata for v0, 
    Type Dict[str, AttentionMetadata] for v1, map from layer_name of each 
    attention layer to its attention metadata
    set dynamically for each forward pass
    """
    attn_metadata: Union["AttentionMetadata", dict[str, "AttentionMetadata"]]
    # TODO: remove after making all virtual_engines share the same kv cache
    virtual_engine: int  # set dynamically for each forward pass
    # set dynamically for each forward pass
    dp_metadata: Optional[DPMetadata] = None
    skip_cuda_graphs: bool = False

attn_metadata instance-attribute 

attn_metadata: Union[
    AttentionMetadata, dict[str, AttentionMetadata]
]

dp_metadata class-attribute instance-attribute 

dp_metadata: Optional[DPMetadata] = None

no_compile_layers instance-attribute 

no_compile_layers: dict[str, Any]

Type AttentionMetadata for v0, Type Dict[str, AttentionMetadata] for v1, map from layer_name of each attention layer to its attention metadata set dynamically for each forward pass

skip_cuda_graphs class-attribute instance-attribute 

skip_cuda_graphs: bool = False

virtual_engine instance-attribute 

virtual_engine: int

__init__ 

__init__(
    no_compile_layers: dict[str, Any],
    attn_metadata: Union[
        AttentionMetadata, dict[str, AttentionMetadata]
    ],
    virtual_engine: int,
    dp_metadata: Optional[DPMetadata] = None,
    skip_cuda_graphs: bool = False,
) -> None

_compute_chunked_local_num_tokens 

_compute_chunked_local_num_tokens(
    num_tokens_across_dp_cpu: list[int],
    max_num_tokens: int,
    chunk_idx: int,
) -> list[int]
Source code in vllm/forward_context.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
def _compute_chunked_local_num_tokens(num_tokens_across_dp_cpu: list[int],
                                      max_num_tokens: int,
                                      chunk_idx: int) -> list[int]:
    dp_size = len(num_tokens_across_dp_cpu)

    local_size = [-1] * dp_size
    for i in range(dp_size):
        dp_tokens = num_tokens_across_dp_cpu[i]
        local_size[i] = min(max_num_tokens,
                            dp_tokens - (max_num_tokens * chunk_idx))
        if local_size[i] <= 0:
            local_size[i] = 1  # ensure lockstep even if done
    return local_size

get_forward_context 

get_forward_context() -> ForwardContext

Get the current forward context.

Source code in vllm/forward_context.py 
161
162
163
164
165
166
def get_forward_context() -> ForwardContext:
    """Get the current forward context."""
    assert _forward_context is not None, (
        "Forward context is not set. "
        "Please use `set_forward_context` to set the forward context.")
    return _forward_context

set_forward_context 

set_forward_context(
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
    num_tokens: Optional[int] = None,
    num_tokens_across_dp: Optional[Tensor] = None,
    skip_cuda_graphs: bool = False,
)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

Source code in vllm/forward_context.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
@contextmanager
def set_forward_context(
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
    num_tokens: Optional[int] = None,
    num_tokens_across_dp: Optional[torch.Tensor] = None,
    skip_cuda_graphs: bool = False,
):
    """A context manager that stores the current forward context,
    can be attention metadata, etc.
    Here we can inject common logic for every model forward pass.
    """
    global forward_start_time
    need_to_track_batchsize = track_batchsize and attn_metadata is not None
    if need_to_track_batchsize:
        forward_start_time = time.perf_counter()
    dp_metadata: Optional[DPMetadata] = None
    if vllm_config.parallel_config.data_parallel_size > 1 and (
            attn_metadata is not None or num_tokens is not None):
        dp_metadata = DPMetadata.make(vllm_config.parallel_config,
                                      attn_metadata, num_tokens or 0,
                                      num_tokens_across_dp)

    global _forward_context
    prev_context = _forward_context
    _forward_context = ForwardContext(
        no_compile_layers=vllm_config.compilation_config.
        static_forward_context,
        virtual_engine=virtual_engine,
        attn_metadata=attn_metadata,
        dp_metadata=dp_metadata,
        skip_cuda_graphs=skip_cuda_graphs,
    )

    try:
        yield
    finally:
        global last_logging_time, batchsize_logging_interval
        if need_to_track_batchsize:
            if hasattr(attn_metadata, "num_prefill_tokens"):
                # for v0 attention backends
                batchsize = attn_metadata.num_prefill_tokens + \
                    attn_metadata.num_decode_tokens
            else:
                # for v1 attention backends
                batchsize = num_tokens
            # we use synchronous scheduling right now,
            # adding a sync point here should not affect
            # scheduling of the next batch
            from vllm.platforms import current_platform
            synchronize = current_platform.synchronize
            if synchronize is not None:
                synchronize()
            now = time.perf_counter()
            # time measurement is in milliseconds
            batchsize_forward_time[batchsize].append(
                (now - forward_start_time) * 1000)
            if now - last_logging_time > batchsize_logging_interval:
                last_logging_time = now
                forward_stats = []
                for bs, times in batchsize_forward_time.items():
                    if len(times) <= 1:
                        # can be cudagraph / profiling run
                        continue
                    medium = torch.quantile(torch.tensor(times), q=0.5).item()
                    medium = round(medium, 2)
                    forward_stats.append((bs, len(times), medium))
                forward_stats.sort(key=lambda x: x[1], reverse=True)
                if forward_stats:
                    logger.info(("Batchsize forward time stats "
                                 "(batchsize, count, median_time(ms)): %s"),
                                forward_stats)

        _forward_context = prev_context