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 @@ -127,6 +127,14 @@ class BatchContents:
127
127
def get_num_tokens (self ):
128
128
return [len (t ) for t in self .token_ids ]
129
129
130
+ def clone (self ):
131
+ return BatchContents (
132
+ req_ids = self .req_ids .copy (),
133
+ token_ids = [t .copy () for t in self .token_ids ],
134
+ context_lens = self .context_lens .copy (),
135
+ blocks = [b .copy () for b in self .blocks ],
136
+ logits_positions = [lp .copy () for lp in self .logits_positions ])
137
+
130
138
131
139
# TODO(kzawora): remove this
132
140
@dataclass
@@ -1536,7 +1544,7 @@ def _create_dummy_prefill_batch_contents(
1536
1544
)
1537
1545
1538
1546
outputs = [
1539
- self ._form_prefill_batch (new_batch_contents )
1547
+ self ._form_prefill_batch (new_batch_contents . clone () )
1540
1548
for _ in range (num_prefills )
1541
1549
]
1542
1550
return outputs
You can’t perform that action at this time.
0 commit comments