Fix inference mode mismatch in FlashInfer warmup (#33788)

2026-08-06 16:11:41 -07:00
parent b38caebf09
commit bae29f716a
2 changed files with 2 additions and 2 deletions
@@ -621,7 +621,7 @@ class BaseRunner(ABC):
        torch.get_device_module(mr.device).synchronize()
        mr.tp_group.barrier()
        with forward_context(ForwardContext(attn_backend=mr.attn_backend)):
-            with torch.inference_mode(), run_ctx or empty_context():
+            with run_ctx or empty_context():
                run_once()

    def _autotune_buffers(self) -> Tuple[Optional[Any], Optional[int]]:
@@ -190,7 +190,7 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool)

            maybe_skip_logits = autotune_dummy_run_mode()
        skip_ops = get_flashinfer_autotune_skip_ops(mr)
-        with torch.inference_mode(), autotune(
+        with autotune(
            True,
            cache=str(autotune_cache),
            skip_ops=skip_ops,