Fix use_sdpa_with_kv_cache option (#4456)

Summary:
Pull Request resolved: #4456

As titled. In `export_llava.py` `export_text_model()` needs to
respect `use_sdpa_with_kv_cache_op` option.

Reviewed By: cccclai

Differential Revision: D60431561

fbshipit-source-id: 63d49f39339435fb16f0c1c62288fd31c86b3be8

Author: larryliu0820

examples/models/llava/export_llava.py

@@ -83,11 +83,14 @@ def forward(self, input_pos, embeddings):
     )
     quant_transform = get_quant_weight_transform(args, dtype_override, False)
     pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(args)
-
+    source_transforms = []
+    if llava.use_sdpa_with_kv_cache_op:
+        source_transforms.append(replace_sdpa_with_custom_op)
+    source_transforms.append(quant_transform)
     manager = (
         text_model_em.set_output_dir("./")
         .to_dtype(dtype_override)
-        .source_transform([replace_sdpa_with_custom_op, quant_transform])
+        .source_transform(source_transforms)
         .capture_pre_autograd_graph()
         .pt2e_quantize(quantizers)
     )