Skip to content

Commit a8ebc0f

Browse files
committed
added opt swa to hash
Signed-off-by: Onkar Chougule <ochougul@qti.qualcomm.com>
1 parent 276f273 commit a8ebc0f

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

QEfficient/transformers/models/modeling_auto.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2533,8 +2533,9 @@ def get_seq_len_and_handle_specialized_prefill_model(self, prefill_seq_len: Opti
25332533

25342534
self.prefill(True)
25352535
self.hash_params["prefill_only"] = True
2536-
self.hash_params["num_blocks"] = num_q_blocks
2537-
self.hash_params["num_ffn_blocks"] = num_ffn_blocks
2536+
self.hash_params["NUM_Q_BLOCKS"] = num_q_blocks
2537+
self.hash_params["NUM_FFN_BLOCKS"] = num_ffn_blocks
2538+
self.hash_params["ENABLE_OPT_SWA"] = os.environ.get("ENABLE_OPT_SWA", "0")
25382539
return (
25392540
min_seq_len
25402541
if min_seq_len > constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN

0 commit comments

Comments
 (0)