A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 19:00 UTC
[platform] Support additional forward context for OOT (#31674)
Signed-off-by: zzzzwwjj <1183291235@qq.com> Signed-off-by: zzzzwwjj <34335947+zzzzwwjj@users.noreply.github.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
zzzzwwjj committed
Jan 5, 2026 at 18:25 UTC
caaa482aca0117ae95e96a34c8df862a5a467a69
2 files changed
+25
-3
vllm/forward_context.py
+18
-3
@@ -4,7 +4,7 @@
4
import time
5
from collections import defaultdict
6
from contextlib import contextmanager
7
-from dataclasses import dataclass
7
+from dataclasses import dataclass, field
8
from typing import Any, NamedTuple
9
10
import torch
@@ -13,6 +13,7 @@ import vllm.envs as envs
13
from vllm.attention.backends.abstract import AttentionMetadata
14
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
15
from vllm.logger import init_logger
16
+from vllm.platforms import current_platform
17
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
18
from vllm.v1.worker.ubatch_utils import UBatchSlices
19
@@ -206,6 +207,8 @@ class ForwardContext:
207
208
ubatch_slices: UBatchSlices | None = None
209
210
+ additional_kwargs: dict[str, Any] = field(default_factory=dict)
211
+
212
def __post_init__(self):
213
assert self.cudagraph_runtime_mode.valid_runtime_modes(), (
214
f"Invalid cudagraph runtime mode: {self.cudagraph_runtime_mode}"
@@ -236,6 +239,7 @@ def create_forward_context(
239
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
240
batch_descriptor: BatchDescriptor | None = None,
241
ubatch_slices: UBatchSlices | None = None,
242
+ additional_kwargs: dict[str, Any] | None = None,
243
):
244
return ForwardContext(
245
no_compile_layers=vllm_config.compilation_config.static_forward_context,
@@ -245,6 +249,7 @@ def create_forward_context(
249
cudagraph_runtime_mode=cudagraph_runtime_mode,
250
batch_descriptor=batch_descriptor,
251
ubatch_slices=ubatch_slices,
252
+ additional_kwargs=additional_kwargs or {},
253
)
254
255
@@ -310,6 +315,17 @@ def set_forward_context(
315
if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None:
316
batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens)
317
318
+ additional_kwargs = current_platform.set_additional_forward_context(
319
+ attn_metadata=attn_metadata,
320
+ vllm_config=vllm_config,
321
+ virtual_engine=virtual_engine,
322
+ num_tokens=num_tokens,
323
+ num_tokens_across_dp=num_tokens_across_dp,
324
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
325
+ batch_descriptor=batch_descriptor,
326
+ ubatch_slices=ubatch_slices,
327
+ )
328
+
329
forward_context = create_forward_context(
330
attn_metadata,
331
vllm_config,
@@ -318,6 +334,7 @@ def set_forward_context(
334
cudagraph_runtime_mode,
335
batch_descriptor,
336
ubatch_slices,
337
+ additional_kwargs,
338
)
339
340
try:
@@ -330,8 +347,6 @@ def set_forward_context(
347
# we use synchronous scheduling right now,
348
# adding a sync point here should not affect
349
# scheduling of the next batch
333
- from vllm.platforms import current_platform
334
-
350
synchronize = current_platform.synchronize
351
if synchronize is not None:
352
synchronize()
vllm/platforms/interface.py
+7
@@ -693,6 +693,13 @@ class Platform:
693
"""
694
return max_model_len
695
696
+ @classmethod
697
+ def set_additional_forward_context(cls, *args, **kwargs) -> dict[str, Any]:
698
+ """
699
+ Set some additional forward context for the current platform if needs.
700
+ """
701
+ return {}
702
+
703
704
class UnspecifiedPlatform(Platform):
705
_enum = PlatformEnum.UNSPECIFIED