Fix to recent batch-packing prefill refinements

njhill · njhill · commit d1d2bbdbbbb6 · 2023-09-28T15:07:42.000-07:00
Missing pow(2) in padded batch_initial_weight calculation
diff --git a/router/src/batch_types.rs b/router/src/batch_types.rs
@@ -125,7 +125,7 @@ impl BatchType for PaddedBatch {
 
     fn batch_initial_weight(max_in_out_lengths: &Self::Stats, batch_size: usize) -> usize {
         let (max_input_length, _) = max_in_out_lengths;
-        batch_size * max_input_length
+        batch_size * max_input_length.pow(2)
     }
 
     fn prefill_weight(max_in_out_lengths: &Self::Stats, batch_size: usize) -> usize {
diff --git a/server/text_generation_server/models/causal_lm.py b/server/text_generation_server/models/causal_lm.py
@@ -249,7 +249,7 @@ def concatenate(cls, batches: List["CausalLMBatch"]) -> "CausalLMBatch":
             # and to remove unused allocated space
             left_offset = max_sequence_length - batch.max_sequence_length
             batch_left_offset = (
-                    batch.attention_mask.shape[1] - batch.max_sequence_length - batch.padding_right_offset
+                batch.attention_mask.shape[1] - batch.max_sequence_length - batch.padding_right_offset
             )
             attention_mask[
                 start_index:end_index, left_offset:-padding_right_offset,

Original file line number	Diff line number	Diff line change
`@@ -125,7 +125,7 @@ impl BatchType for PaddedBatch {`
`125`	`125`
`126`	`126`	`fn batch_initial_weight(max_in_out_lengths: &Self::Stats, batch_size: usize) -> usize {`
`127`	`127`	`let (max_input_length, _) = max_in_out_lengths;`
`128`		`- batch_size * max_input_length`
	`128`	`+ batch_size * max_input_length.pow(2)`
`129`	`129`	`}`
`130`	`130`
`131`	`131`	`fn prefill_weight(max_in_out_lengths: &Self::Stats, batch_size: usize) -> usize {`
Original file line number	Diff line number	Diff line change
`@@ -249,7 +249,7 @@ def concatenate(cls, batches: List["CausalLMBatch"]) -> "CausalLMBatch":`
`249`	`249`	`# and to remove unused allocated space`
`250`	`250`	`left_offset = max_sequence_length - batch.max_sequence_length`
`251`	`251`	`batch_left_offset = (`
`252`		`- batch.attention_mask.shape[1] - batch.max_sequence_length - batch.padding_right_offset`
	`252`	`+ batch.attention_mask.shape[1] - batch.max_sequence_length - batch.padding_right_offset`
`253`	`253`	`)`
`254`	`254`	`attention_mask[`
`255`	`255`	`start_index:end_index, left_offset:-padding_right_offset,`