[Fix] Disable use_cache during export dummy forward for custom configs

The export step's collect_shared_input_modules() runs a dummy forward pass to
trace shared input modules for resmoothing/requantization. This forward call
hits the same AttributeError as PR #1324 on configs that don't assign
use_cache (e.g., stepfun-ai/Step-3.5-Flash's Step3p5Config):

  AttributeError: 'Step3p5Config' object has no attribute 'use_cache'

Wrap the dummy forward with the _disable_use_cache context manager so it
sets/restores config.use_cache around the call, mirroring the calibration
fix from #1324. Covers all callsites since both unified_export_hf and
plugins/vllm_fakequant_hf funnel through collect_shared_input_modules.

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

Author: meenchen

modelopt/torch/export/unified_export_hf.py

@@ -55,6 +55,7 @@
 from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names
+from modelopt.torch.utils.dataset_utils import _disable_use_cache
 
 try:
     from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config
@@ -213,11 +214,14 @@ def _output_hook(module, input, output):
     if not handles:
         return input_to_linear, output_to_layernorm
 
-    # Run dummy forward pass to collect modules sharing same input
+    # Run dummy forward pass to collect modules sharing same input.
+    # `_disable_use_cache` keeps the probe forward working on configs that don't
+    # set `use_cache` (e.g., stepfun-ai/Step-3.5-Flash's Step3p5Config).
     try:
         with (
             torch.no_grad(),
             set_quantizer_by_cfg_context(model, [{"quantizer_name": "*", "enable": False}]),
+            _disable_use_cache(model),
         ):
             dummy_forward_fn()
     finally: