mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 05:34:57 +08:00
[Core] Zero-copy asdict for InputMetadata (#3475)
This commit is contained in:
parent
9fdf3de346
commit
49eedea373
@ -1,5 +1,5 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, fields
|
||||
from typing import Optional, Any, Dict
|
||||
|
||||
import torch
|
||||
|
||||
@ -31,3 +31,12 @@ class InputMetadata:
|
||||
def __post_init__(self):
|
||||
# will not appear in the __repr__ and __init__
|
||||
self.attn_bias = None
|
||||
|
||||
def asdict_zerocopy(self) -> Dict[str, Any]:
|
||||
"""Similar to dataclasses.asdict, but avoids deepcopying."""
|
||||
# Note that if we add dataclasses as fields, they will need
|
||||
# similar handling.
|
||||
return {
|
||||
field.name: getattr(self, field.name)
|
||||
for field in fields(self)
|
||||
}
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import time
|
||||
from typing import Dict, List, Optional, Tuple, Set, Union
|
||||
|
||||
@ -527,7 +526,7 @@ class ModelRunner:
|
||||
"lora_requests": lora_requests,
|
||||
"lora_mapping": lora_mapping,
|
||||
}
|
||||
metadata_dict.update(dataclasses.asdict(input_metadata))
|
||||
metadata_dict.update(input_metadata.asdict_zerocopy())
|
||||
broadcast_tensor_dict(metadata_dict, src=0)
|
||||
else:
|
||||
metadata_dict = broadcast_tensor_dict(src=0)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user