added ffn blocking and num blocks env variables

Signed-off-by: Onkar Chougule <ochougul@qti.qualcomm.com>

Author: ochougul

QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -47,6 +47,8 @@ def __qeff_init__(self):
 
 class QEffPrefillOnlyGptOssMLP(GptOssMLP):
     def forward(self, hidden: torch.Tensor):
+        if os.environ.get("NUM_FFN_BLOCKS", None) is not None:
+            return self.blocked_ffn_forward(hidden)
         B, S, H = hidden.shape
         T = B * S
         hidden = hidden.view(T, H)
@@ -118,7 +120,7 @@ def blocked_ffn_forward(self, hidden: torch.Tensor):
 
         # ────────────────── allocate the output tensor ─────
         expert_out = hidden.new_zeros((T, H))  # accumulation buffer
-        target_blocks = int(os.environ.get("NUM_BLOCKS", 1))
+        target_blocks = int(os.environ.get("NUM_FFN_BLOCKS", 1))
         block_positions = []
         for j in range(target_blocks):
             block_positions.append(j * (T // target_blocks))

QEfficient/transformers/models/modeling_auto.py

@@ -2314,6 +2314,35 @@ def get_model_config(self) -> dict:
         """
         return self.model.config.__dict__
 
+    def get_seq_len_and_handle_specialized_prefill_model(self, prefill_seq_len: Optional[int] = None) -> int:
+        num_q_blocks = os.environ.get("NUM_Q_BLOCKS", None)
+        if num_q_blocks is None:
+            block_size = 128
+            if prefill_seq_len is None or prefill_seq_len % block_size != 0 or prefill_seq_len < 128:
+                raise ValueError(
+                    f"When prefill_only=True, 'prefill_seq_len' must be explicitly set and divisible by block_size={block_size}. "
+                    f"Or set `NUM_BLOCKS` ENV variable"
+                    f"Received: prefill_seq_len={prefill_seq_len}"
+                )
+
+            num_q_blocks = prefill_seq_len // block_size
+            logger.warning(
+                f"Setting NUM_BLOCKS={num_q_blocks} used in attention Q-blocking for prefill_only model, please set ENV variable `NUM_BLOCKS` to override"
+            )
+            os.environ["NUM_Q_BLOCKS"] = num_q_blocks
+
+        num_ffn_blocks = os.environ.get("NUM_FFN_BLOCKS", None)
+        min_seq_len = int(max(num_q_blocks, num_ffn_blocks)) if num_ffn_blocks else num_q_blocks
+
+        self.prefill(True)
+        self.hash_params["prefill_only"] = True
+        self.hash_params["num_blocks"] = os.environ["NUM_BLOCKS"]
+        return (
+            min_seq_len
+            if min_seq_len > constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN
+            else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN
+        )
+
     def export(
         self,
         export_dir: Optional[str] = None,
@@ -2345,25 +2374,11 @@ def export(
         fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS
         if prefill_only:
             assert not self.continuous_batching, "prefill_only=True is not supported with continuous_batching=True"
-
-            if self.model.config.model_type in SPECIALIZED_PREFILL_ONLY_MODEL_ARCH:
-                block_size = os.environ.get("BLOCK_SIZE", None)
-                if block_size is None:
-                    block_size = 128
-                    logger.warning(
-                        "Setting BLOCK_SIZE=128 for prefill_only model, please set ENV variable `BLOCK_SIZE` to override"
-                    )
-                if prefill_seq_len is None or prefill_seq_len % block_size != 0:
-                    raise ValueError(
-                        f"When prefill_only=True, 'prefill_seq_len' must be explicitly set and divisible by block_size={block_size}. "
-                        f"Received: prefill_seq_len={prefill_seq_len}"
-                    )
-                os.environ["NUM_BLOCKS"] = str(prefill_seq_len // block_size)
-
-                self.prefill(True)
-                self.hash_params["prefill_only"] = True
-                self.hash_params["num_blocks"] = os.environ["NUM_BLOCKS"]
-                seq_len = prefill_seq_len // block_size if (prefill_seq_len // block_size) > seq_len else seq_len
+            seq_len = (
+                self.get_seq_len_and_handle_specialized_prefill_model(prefill_seq_len)
+                if self.model.config.model_type in SPECIALIZED_PREFILL_ONLY_MODEL_ARCH
+                else seq_len
+            )
         else:
             self.prefill(False)
             self.hash_params.pop("prefill_only", None)