diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py
index 69c6d89c286..727c6ef5fd6 100644
--- a/megatron/core/ssm/mamba_mixer.py
+++ b/megatron/core/ssm/mamba_mixer.py
@@ -1035,7 +1035,7 @@ def _get_decode_A_neg_exp(self) -> torch.Tensor:
             base = -torch.exp(self.A_log.float())
             return base.view(-1, 1, 1).expand(-1, self.headdim, self.d_state)
         # Inference path. Refill when stale
-        if torch.cuda.is_current_stream_capturing() or self._A_neg_exp_cache_stale:
+        if self._A_neg_exp_cache_stale:
             with torch.no_grad():
                 self._A_neg_exp_cache.copy_(-torch.exp(self.A_log.float()))
             self._A_neg_exp_cache_stale = False
@@ -1043,7 +1043,10 @@ def _get_decode_A_neg_exp(self) -> torch.Tensor:
 
     def train(self, mode: bool = True):
         """Mark the decode cache stale; weights may have updated."""
-        self._A_neg_exp_cache_stale = True
+        if mode:
+            # only mark stale when switching to training mode.
+            # otherwise retain the staleness state.
+            self._A_neg_exp_cache_stale = True
         return super().train(mode)
 
     def _ssm_decode(