File tree Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Original file line number Diff line number Diff line change @@ -133,6 +133,14 @@ class BatchContents:
133
133
def get_num_tokens (self ):
134
134
return [len (t ) for t in self .token_ids ]
135
135
136
+ def clone (self ):
137
+ return BatchContents (
138
+ req_ids = self .req_ids .copy (),
139
+ token_ids = [t .copy () for t in self .token_ids ],
140
+ context_lens = self .context_lens .copy (),
141
+ blocks = [b .copy () for b in self .blocks ],
142
+ logits_positions = [lp .copy () for lp in self .logits_positions ])
143
+
136
144
137
145
# TODO(kzawora): remove this
138
146
@dataclass
@@ -1600,7 +1608,7 @@ def _create_dummy_prefill_batch_contents(
1600
1608
)
1601
1609
1602
1610
outputs = [
1603
- self ._form_prefill_batch (new_batch_contents )
1611
+ self ._form_prefill_batch (new_batch_contents . clone () )
1604
1612
for _ in range (num_prefills )
1605
1613
]
1606
1614
return outputs
You can’t perform that action at this time.
0 commit comments