mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-05-11 03:44:42 +08:00
Fix IntermediateTensors initialization and add type hints (#28743)
Signed-off-by: Mohammad Othman <Mo@MohammadOthman.com> Co-authored-by: Mohammad Othman <Mo@MohammadOthman.com>
This commit is contained in:
parent
ac86bff8cb
commit
363aaeef0f
@ -60,12 +60,17 @@ class IntermediateTensors:
|
|||||||
tensors: dict[str, torch.Tensor]
|
tensors: dict[str, torch.Tensor]
|
||||||
kv_connector_output: KVConnectorOutput | None
|
kv_connector_output: KVConnectorOutput | None
|
||||||
|
|
||||||
def __init__(self, tensors):
|
def __init__(
|
||||||
|
self,
|
||||||
|
tensors: dict[str, torch.Tensor],
|
||||||
|
kv_connector_output: KVConnectorOutput | None = None,
|
||||||
|
) -> None:
|
||||||
# manually define this function, so that
|
# manually define this function, so that
|
||||||
# Dynamo knows `IntermediateTensors()` comes from this file.
|
# Dynamo knows `IntermediateTensors()` comes from this file.
|
||||||
# Otherwise, dataclass will generate this function by evaluating
|
# Otherwise, dataclass will generate this function by evaluating
|
||||||
# a string, and we will lose the information about the source file.
|
# a string, and we will lose the information about the source file.
|
||||||
self.tensors = tensors
|
self.tensors = tensors
|
||||||
|
self.kv_connector_output = kv_connector_output
|
||||||
|
|
||||||
def __getitem__(self, key: str | slice):
|
def __getitem__(self, key: str | slice):
|
||||||
if isinstance(key, str):
|
if isinstance(key, str):
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user