Skip to content

Commit 363aaee

Browse files
OthmanMohammadMohammad Othman
andauthored
Fix IntermediateTensors initialization and add type hints (#28743)
Signed-off-by: Mohammad Othman <Mo@MohammadOthman.com> Co-authored-by: Mohammad Othman <Mo@MohammadOthman.com>
1 parent ac86bff commit 363aaee

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

vllm/sequence.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,12 +60,17 @@ class IntermediateTensors:
6060
tensors: dict[str, torch.Tensor]
6161
kv_connector_output: KVConnectorOutput | None
6262

63-
def __init__(self, tensors):
63+
def __init__(
64+
self,
65+
tensors: dict[str, torch.Tensor],
66+
kv_connector_output: KVConnectorOutput | None = None,
67+
) -> None:
6468
# manually define this function, so that
6569
# Dynamo knows `IntermediateTensors()` comes from this file.
6670
# Otherwise, dataclass will generate this function by evaluating
6771
# a string, and we will lose the information about the source file.
6872
self.tensors = tensors
73+
self.kv_connector_output = kv_connector_output
6974

7075
def __getitem__(self, key: str | slice):
7176
if isinstance(key, str):

0 commit comments

Comments
 (0)