Fix cache gating

megatron/core/ssm/mamba_mixer.py

@@ -1035,15 +1035,18 @@ 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
         return self._A_neg_exp_cache.view(-1, 1, 1).expand(-1, self.headdim, self.d_state)
 
     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(