diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py
index 1655fcd3..86e55a8b 100644
--- a/src/liger_kernel/ops/fused_linear_cross_entropy.py
+++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py
@@ -40,6 +40,25 @@ _TORCH_VERSION = Version(torch.__version__.split("+")[0])
 _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
 
 
+def _can_use_mm_out(_input, weight):
+    # out= bypasses autocast. Only use it when @ would use these exact operands
+    # and return their dtype; keep AMP casts/rounding on the original path.
+    if (
+        _input.device.type != "cuda"
+        or is_hip()
+        or torch.is_grad_enabled()
+        or _input.dtype != weight.dtype
+        or _input.dtype not in (torch.float16, torch.bfloat16, torch.float32)
+    ):
+        return False
+    if torch.is_autocast_enabled():
+        autocast_dtype = (
+            torch.get_autocast_dtype("cuda") if _TORCH_VERSION >= Version("2.4.0") else torch.get_autocast_gpu_dtype()
+        )
+        return _input.dtype == autocast_dtype
+    return True
+
+
 def fused_linear_cross_entropy_forward(
     _input,
     weight,
@@ -133,15 +152,30 @@ def fused_linear_cross_entropy_forward(
         if ce_weight.stride(-1) != 1:
             ce_weight = ce_weight.contiguous()
 
+    use_mm_out = _can_use_mm_out(_input, weight)
+    # A promoted bias needs the original out-of-place add. Do not reserve an
+    # unused low-precision buffer alongside those promoted logits.
+    logits_storage = (
+        torch.empty((chunk_size, V), dtype=_input.dtype, device=device)
+        if use_mm_out and (bias is None or bias.dtype == _input.dtype)
+        else None
+    )
+
     for chunk_id in range(num_chunks):
         start_idx = chunk_id * chunk_size
         end_idx = min((chunk_id + 1) * chunk_size, BT)
         _input_chunk = _input[start_idx:end_idx]  # chunk_size x H
 
         # when doing matmul, use the original precision
-        logits_chunk = _input_chunk @ weight.t()  # chunk_size x V
-        if bias is not None:
-            logits_chunk = logits_chunk + bias
+        if logits_storage is not None:
+            logits_chunk = logits_storage[: end_idx - start_idx]
+            torch.mm(_input_chunk, weight.t(), out=logits_chunk)
+            if bias is not None:
+                logits_chunk.add_(bias)
+        else:
+            logits_chunk = _input_chunk @ weight.t()  # chunk_size x V
+            if bias is not None:
+                logits_chunk = logits_chunk + bias
 
         target_chunk = target[start_idx:end_idx]  # chunk_size,
 
@@ -155,6 +189,7 @@ def fused_linear_cross_entropy_forward(
 
             # Compute softmax to get predicted probabilities
             probs = torch.softmax(logits_for_softmax, dim=-1)
+            del logits_for_softmax
 
             # Get predicted probabilities for token scaling, handling ignored targets
             valid_target_mask = target_chunk != ignore_index
@@ -164,16 +199,19 @@ def fused_linear_cross_entropy_forward(
                 # Gather probabilities only for valid targets
                 valid_probs = probs[valid_target_mask]
                 pred_probs_valid = torch.gather(valid_probs, -1, valid_targets.unsqueeze(-1)).squeeze(-1)
+                del valid_probs
 
                 # Create full tensor with zeros for ignored targets
                 pred_probs = torch.zeros_like(target_chunk, dtype=probs.dtype, device=probs.device)
                 pred_probs[valid_target_mask] = pred_probs_valid
+                del pred_probs_valid
             else:
                 # All targets are ignored
                 pred_probs = torch.zeros_like(target_chunk, dtype=probs.dtype, device=probs.device)
 
             # Store the scaling factors
             scaling_factors = pred_probs.detach()  # Detach to ensure no gradient flow
+            del probs, pred_probs, valid_target_mask, valid_targets
 
         # ensure _input and target are contiguous
         logits_chunk = logits_chunk.contiguous()
@@ -248,6 +286,11 @@ def fused_linear_cross_entropy_forward(
                 impl=ce_impl,
                 mode=ce_mode,
             )
+            del ce_args
+
+        # CE may return an alias of logits. Keep only the gradient reference so
+        # replacing it with scaled gradients also releases fallback logits.
+        del logits_chunk
 
         # Apply token scaling if requested
         if use_token_scaling:
@@ -268,6 +311,7 @@ def fused_linear_cross_entropy_forward(
             # Expand scaling factors to match gradient dimensions
             scaling_factors_expanded = scaling_factors.unsqueeze(-1)  # chunk_size x 1
             grad_logits_chunk = grad_logits_chunk * scaling_factors_expanded
+            del scaling_factors, scaling_factors_expanded
 
         # reduction="none": fold the per-token upstream gradient in HERE, before the
         # projections below sum over the token dimension. grad_weight is
@@ -280,7 +324,10 @@ def fused_linear_cross_entropy_forward(
             )
 
         if input_requires_grad:
-            grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
+            if use_mm_out and grad_logits_chunk.dtype == grad_input.dtype:
+                torch.mm(grad_logits_chunk, weight, out=grad_input[start_idx:end_idx])
+            else:
+                grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
 
         if grad_weight is not None and input_requires_grad:
             grad_logits_t = grad_logits_chunk.t()
@@ -306,6 +353,7 @@ def fused_linear_cross_entropy_forward(
                     out_dtype=torch.float32,
                     out=grad_weight,
                 )
+                del input_chunk
             elif (
                 grad_weight.device.type == "cuda"
                 and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
@@ -327,11 +375,13 @@ def fused_linear_cross_entropy_forward(
                     input_chunk,
                     out=grad_weight,
                 )
+                del input_chunk
             else:
                 # Legacy fallback: unsupported torch/device (no out_dtype), fp64 accumulators,
                 # or a dtype mismatch (e.g. fp32 grad_logits promoted under AMP). Correct but
                 # allocates a parameter-sized fp32 temporary before summing into grad_weight.
                 grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
+            del grad_logits_t
 
         if bias is not None and input_requires_grad:
             torch.add(
@@ -341,6 +391,13 @@ def fused_linear_cross_entropy_forward(
                 alpha=1.0,
             )
 
+        # In particular, no transposed/scaled gradient or CE argument tuple may
+        # keep the previous chunk alive when the fallback allocates new logits.
+        del grad_logits_chunk, _input_chunk, target_chunk
+
+    # Release even the persistent buffer before allocating a dtype-converted dW.
+    del logits_storage
+
     # Need extra calculations for backward if reduction=='none'. Not supporting reduction='none' now.
     # if reduction == "none":
     #     loss = loss_1d
diff --git a/test/ops/test_fused_linear_cross_entropy.py b/test/ops/test_fused_linear_cross_entropy.py
index fb915f5f..5ff8a73b 100644
--- a/test/ops/test_fused_linear_cross_entropy.py
+++ b/test/ops/test_fused_linear_cross_entropy.py
@@ -19,6 +19,7 @@ import torch
 import torch.nn.functional as F
 import triton
 
+from torch.multiprocessing.reductions import StorageWeakRef
 from torch.utils._python_dispatch import TorchDispatchMode
 
 import liger_kernel
@@ -317,6 +318,11 @@ requires_addmm_out_dtype = pytest.mark.skipif(
     reason="addmm out_dtype overload (fp32 accumulator fast path) requires torch >= 2.8",
 )
 
+requires_nvidia_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available() or flce_ops.is_hip(),
+    reason="reusable mm(out=) buffers require NVIDIA CUDA",
+)
+
 
 def _production_chunk_size(BT: int, V: int, H: int) -> int:
     """The chunk size the production loop will pick for these dims, using the
@@ -337,13 +343,17 @@ _ADDMM_OPS = ("aten.addmm.out", "aten.addmm.dtype_out")
 
 class _OpRecorder(TorchDispatchMode):
     """Records the aten ops we care about: addmm (with the out-tensor shape and
-    whether the out_dtype overload was used) and ``_to_copy`` (with in/out dtype
-    and shape so we can detect parameter-sized upcasts)."""
+    whether the out_dtype overload was used), mm outputs/copies, and dtype casts.
+    Store metadata only, so the recorder cannot extend a buffer's lifetime."""
 
     def __init__(self):
         self.addmm_out_shapes = []  # list of torch.Size for aten.addmm.out
         self.addmm_dtype_out_shapes = []  # list of torch.Size for aten.addmm.dtype_out
         self.mm_default = 0  # count of aten.mm.default
+        self.mm_default_shapes = []
+        self.mm_out = []  # (shape, storage pointer, data pointer)
+        self.copies = []  # (destination shape, destination storage pointer)
+        self.empty_shapes = []
         self.to_copies = []  # list of (in_dtype, out_dtype, shape)
 
     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
@@ -358,6 +368,15 @@ class _OpRecorder(TorchDispatchMode):
                 self.addmm_dtype_out_shapes.append(shape)
         elif name == "aten.mm.default":
             self.mm_default += 1
+            self.mm_default_shapes.append((args[0].shape[0], args[1].shape[1]))
+        elif name == "aten.mm.out":
+            out = kwargs["out"]
+            self.mm_out.append((tuple(out.shape), out.untyped_storage().data_ptr(), out.data_ptr()))
+        elif name == "aten.copy_.default":
+            dst = args[0]
+            self.copies.append((tuple(dst.shape), dst.untyped_storage().data_ptr()))
+        elif name == "aten.empty.memory_format":
+            self.empty_shapes.append(tuple(args[0]))
         elif name == "aten._to_copy.default":
             src = args[0] if args else None
             if isinstance(src, torch.Tensor):
@@ -383,6 +402,55 @@ class _OpRecorder(TorchDispatchMode):
         return any(src in low_dtypes and dst == torch.float32 and sh == tgt for (src, dst, sh) in self.to_copies)
 
 
+class _BufferRecorder(_OpRecorder):
+    """Check storage lifetime, including views that outlive their base tensor.
+
+    StorageWeakRef observes the C++ storage without keeping it alive. Python
+    tensor weakrefs alone would miss aliases such as grad_logits_t and ce_args.
+    """
+
+    def __init__(self, BT, V, H):
+        super().__init__()
+        self.V, self.H = V, H
+        self.chunk_size = _production_chunk_size(BT, V, H)
+        self.chunk_storages = []
+        self.logits_gemms = 0
+        self.final_casts = 0
+
+    def assert_released(self, reusable_ptr=None):
+        assert all(ref.expired() or ptr == reusable_ptr for ptr, ref in self.chunk_storages), (
+            "a logits/probability/scaled-gradient storage is still live"
+        )
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        name = str(func)
+        if name in ("aten.mm.default", "aten.mm.out") and tuple(args[1].shape) == (self.H, self.V):
+            out = kwargs.get("out")
+            reusable_ptr = out.untyped_storage().data_ptr() if out is not None else None
+            self.assert_released(reusable_ptr)
+            self.logits_gemms += 1
+        if (
+            name == "aten._to_copy.default"
+            and tuple(args[0].shape) == (self.V, self.H)
+            and args[0].dtype == torch.float32
+            and kwargs.get("dtype") in (torch.bfloat16, torch.float16)
+        ):
+            self.assert_released()
+            self.final_casts += 1
+
+        result = super().__torch_dispatch__(func, types, args, kwargs)
+        if (
+            isinstance(result, torch.Tensor)
+            and result.ndim == 2
+            and result.shape[1] == self.V
+            and 0 < result.shape[0] <= self.chunk_size
+        ):
+            storage = result.untyped_storage()
+            self.chunk_storages.append((storage.data_ptr(), StorageWeakRef(storage)))
+        return result
+
+
 # ---------------------------------------------------------------------------
 # Reference + driver helpers
 # ---------------------------------------------------------------------------
@@ -447,6 +515,7 @@ def _run_flce(
     token_grad_output=None,
     recorder: Optional[_OpRecorder] = None,
     autocast_dtype=None,
+    ce_impl=None,
 ):
     """Invoke the forward the same way the autograd Function does: under
     ``no_grad`` (Function.forward runs with grad disabled), optionally under
@@ -466,6 +535,7 @@ def _run_flce(
             accum_dtype=accum_dtype,
             use_token_scaling=use_token_scaling,
             token_grad_output=token_grad_output,
+            ce_impl=ce_impl,
         )
 
     import contextlib
@@ -487,6 +557,8 @@ _ADDMM_TOLS = {
     torch.float32: dict(atol=1e-4, rtol=1e-4),
 }
 
+_SHARED_CE_IMPLS = [None] + [b for b in _REGISTERED_BACKENDS if b in ("nvidia-triton", "nvidia-cutedsl")]
+
 
 def _mk_inputs(
     BT,
@@ -525,7 +597,7 @@ def _mk_inputs(
 
 
 class TestFusedLinearCrossEntropyAddmm:
-    """Regression suite for the low-precision ``grad_weight`` addmm path.
+    """Regression suite for chunk buffers and the ``grad_weight`` addmm path.
 
     Grouped in a class so the ``_force_multichunk`` autouse fixture (which pins
     ``_CHUNK_MEM_CONST=1``) is scoped to *only* these migrated tests and never
@@ -763,7 +835,7 @@ class TestFusedLinearCrossEntropyAddmm:
         if accum_dtype is torch.float32:
             # fp32 accumulator falls back to legacy torch.mm(...).float(): mm + a
             # parameter-sized fp32 upcast, and no addmm at all into the weight buffer.
-            assert rec.mm_default >= 1
+            assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
             assert rec.has_param_upcast(weight.shape)
             assert not rec.has_addmm_out_into(weight.shape)
         else:
@@ -775,6 +847,257 @@ class TestFusedLinearCrossEntropyAddmm:
         torch.testing.assert_close(gw.float(), ref_gw, **_ADDMM_TOLS[torch.bfloat16])
         torch.testing.assert_close(gi.float(), ref_gi, **_ADDMM_TOLS[torch.bfloat16])
 
+    # -----------------------------------------------------------------------
+    # Logits storage reuse and direct dX output, including ragged/strided views
+    # -----------------------------------------------------------------------
+
+    @requires_nvidia_cuda
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("bias", [False, True])
+    @pytest.mark.parametrize("layout", ["contiguous", "transposed", "strided"])
+    def test_reuses_logits_and_writes_dx_out(self, monkeypatch, dtype, bias, layout):
+        BT, V, H = 40, 128, 32
+        _input, weight, b, _, target = _mk_inputs(BT, V, H, dtype, bias=bias, ce_weight=False, seed=15)
+        if layout == "transposed":
+            _input = _input.t().contiguous().t().detach().requires_grad_(True)
+            weight = weight.t().contiguous().t().detach().requires_grad_(True)
+        elif layout == "strided":
+            _input = _input.repeat_interleave(2, dim=1)[:, ::2].detach().requires_grad_(True)
+            weight = weight.repeat_interleave(2, dim=1)[:, ::2].detach().requires_grad_(True)
+        if layout != "contiguous":
+            assert not _input.is_contiguous() and not weight.is_contiguous()
+        originals = tuple(t.detach().clone() if t is not None else None for t in (_input, weight, b))
+
+        rec = _BufferRecorder(BT, V, H)
+        actual = _run_flce(_input, weight, target, bias=b, recorder=rec)
+        gi = actual[1]
+        cs = _production_chunk_size(BT, V, H)
+        rows = [min(cs, BT - start) for start in range(0, BT, cs)]
+        assert len(rows) > 1 and rows[-1] < cs
+        logits_out = [entry for entry in rec.mm_out if entry[0][1] == V]
+        dx_out = [entry for entry in rec.mm_out if entry[0][1] == H]
+        assert [shape for shape, _, _ in logits_out] == [(n, V) for n in rows]
+        assert len({ptr for _, ptr, _ in logits_out}) == 1
+        assert len({ptr for _, _, ptr in logits_out}) == 1
+        assert rec.empty_shapes.count((cs, V)) == 1
+        assert [shape for shape, _, _ in dx_out] == [(n, H) for n in rows]
+        assert {ptr for _, ptr, _ in dx_out} == {gi.untyped_storage().data_ptr()}
+        assert [ptr for _, _, ptr in dx_out] == [
+            gi.data_ptr() + start * gi.stride(0) * gi.element_size() for start in range(0, BT, cs)
+        ]
+        assert all(ptr != gi.untyped_storage().data_ptr() for _, ptr in rec.copies)
+        # Any remaining allocating mm must be the unchanged dW fallback, never logits/dX.
+        assert all(shape == (V, H) for shape in rec.mm_default_shapes)
+        assert rec.logits_gemms == len(rows)
+        rec.assert_released()
+        torch.testing.assert_close((_input, weight, b), originals, atol=0, rtol=0)
+
+        monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
+        expected = _run_flce(_input, weight, target, bias=b)
+        torch.testing.assert_close(actual, expected)
+
+    @pytest.mark.parametrize("ce_impl", _SHARED_CE_IMPLS)
+    @pytest.mark.parametrize("buffer_path", ["reuse", "forced-fallback", "promoted-bias"])
+    @pytest.mark.parametrize("use_token_scaling", [False, True])
+    @pytest.mark.parametrize("forward_only", [False, True])
+    def test_chunk_storage_released_before_next_chunk_and_final_cast(
+        self, monkeypatch, ce_impl, buffer_path, use_token_scaling, forward_only
+    ):
+        BT, V, H = 40, 128, 32
+        cs = _production_chunk_size(BT, V, H)
+        _input, weight, b, cw, target = _mk_inputs(
+            BT,
+            V,
+            H,
+            torch.bfloat16,
+            bias=True,
+            ce_weight=True,
+            seed=16,
+            first_chunk_size=cs,
+        )
+        reuse = buffer_path == "reuse" and not flce_ops.is_hip()
+        promoted_bias = buffer_path == "promoted-bias"
+        if buffer_path == "forced-fallback":
+            monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
+        if promoted_bias:
+            b = b.detach().float().requires_grad_(True)
+        kwargs = dict(
+            bias=b,
+            ce_weight=cw,
+            accum_dtype=torch.float32,
+            use_token_scaling=use_token_scaling,
+            compute_gradients=not forward_only,
+            softcap=1.5,
+            label_smoothing=0.1,
+            lse_square_scale=1e-4,
+            return_z_loss=True,
+            return_token_accuracy=True,
+            return_predicted_tokens=True,
+            ce_impl=ce_impl,
+        )
+        rec = _BufferRecorder(BT, V, H)
+        with torch.no_grad(), torch.autocast("cuda", enabled=promoted_bias, dtype=torch.bfloat16), rec:
+            actual = fused_linear_cross_entropy_forward(_input, weight, target, **kwargs)
+        assert rec.logits_gemms == triton.cdiv(BT, cs)
+        assert rec.final_casts == (0 if forward_only else 1)
+        assert bool(rec.mm_out) == reuse
+        rec.assert_released()
+        if forward_only:
+            assert actual[5] is None and actual[6] is None
+            assert torch.count_nonzero(actual[4]) == 0
+            assert len(rec.mm_out) == (triton.cdiv(BT, cs) if reuse else 0)
+
+        monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
+        with torch.no_grad(), torch.autocast("cuda", enabled=promoted_bias, dtype=torch.bfloat16):
+            expected = fused_linear_cross_entropy_forward(_input, weight, target, **kwargs)
+        torch.testing.assert_close(actual, expected)
+
+    @pytest.mark.parametrize(
+        "input_dtype,weight_dtype,bias_dtype,amp_dtype,reuse,dx_out",
+        [
+            (torch.bfloat16, torch.bfloat16, None, torch.bfloat16, True, True),
+            (torch.float16, torch.float16, torch.float16, torch.float16, True, True),
+            (torch.float32, torch.float32, None, torch.bfloat16, False, False),
+            (torch.float32, torch.bfloat16, None, torch.bfloat16, False, False),
+            (torch.bfloat16, torch.float32, None, torch.bfloat16, False, False),
+            (torch.bfloat16, torch.bfloat16, torch.float32, torch.bfloat16, False, False),
+            (torch.float32, torch.float32, torch.float32, torch.bfloat16, False, False),
+            (torch.float16, torch.float16, None, torch.bfloat16, False, False),
+            (torch.float32, torch.float32, torch.bfloat16, None, False, True),
+        ],
+    )
+    def test_buffer_amp_gates_preserve_matmul_and_bias_semantics(
+        self, monkeypatch, input_dtype, weight_dtype, bias_dtype, amp_dtype, reuse, dx_out
+    ):
+        reuse = reuse and not flce_ops.is_hip()
+        dx_out = dx_out and not flce_ops.is_hip()
+        BT, V, H = 40, 128, 32
+        cs = _production_chunk_size(BT, V, H)
+        _input, weight, b, _, target = _mk_inputs(
+            BT, V, H, weight_dtype, input_dtype=input_dtype, bias=bias_dtype is not None, ce_weight=False, seed=17
+        )
+        if b is not None:
+            b = b.detach().to(bias_dtype).requires_grad_(True)
+
+        expected_logits = []
+        with torch.no_grad(), torch.autocast("cuda", enabled=amp_dtype is not None, dtype=amp_dtype):
+            for start in range(0, BT, cs):
+                logits = _input[start : start + cs] @ weight.t()
+                expected_logits.append(logits if b is None else logits + b)
+
+        grad_logits = []
+        real_dispatch = flce_ops.dispatch
+
+        def check_ce_logits(op, *args, **kwargs):
+            torch.testing.assert_close(args[0], expected_logits[len(grad_logits)], atol=0, rtol=0)
+            result = real_dispatch(op, *args, **kwargs)
+            grad_logits.append(result[-1].clone())
+            return result
+
+        monkeypatch.setattr(flce_ops, "dispatch", check_ce_logits)
+        rec = _OpRecorder()
+        actual = _run_flce(
+            _input, weight, target, bias=b, autocast_dtype=amp_dtype, ce_impl="nvidia-triton", recorder=rec
+        )
+        n_chunks = triton.cdiv(BT, cs)
+        assert sum(shape[1] == V for shape, _, _ in rec.mm_out) == (n_chunks if reuse else 0)
+        assert sum(shape[1] == H for shape, _, _ in rec.mm_out) == (n_chunks if dx_out else 0)
+        assert rec.empty_shapes.count((cs, V)) == (1 if reuse else 0)
+        assert len(grad_logits) == n_chunks
+        with torch.no_grad(), torch.autocast("cuda", enabled=amp_dtype is not None, dtype=amp_dtype):
+            expected_dx = torch.cat([g @ weight for g in grad_logits]).to(_input.dtype)
+        torch.testing.assert_close(actual[1], expected_dx, atol=0, rtol=0)
+
+        monkeypatch.setattr(flce_ops, "dispatch", real_dispatch)
+        monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
+        expected = _run_flce(_input, weight, target, bias=b, autocast_dtype=amp_dtype, ce_impl="nvidia-triton")
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+    @requires_nvidia_cuda
+    def test_mm_out_gate_uses_legacy_autocast_query_on_old_torch(self, monkeypatch):
+        _input = torch.empty(1, 1, device="cuda", dtype=torch.bfloat16)
+        monkeypatch.setattr(flce_ops, "_TORCH_VERSION", flce_ops.Version("2.1.2"))
+        monkeypatch.setattr(torch, "is_autocast_enabled", lambda: True)
+        monkeypatch.setattr(torch, "get_autocast_gpu_dtype", lambda: torch.bfloat16)
+
+        def unexpected_modern_query(*args):
+            pytest.fail("torch < 2.4 must not use get_autocast_dtype(device_type)")
+
+        monkeypatch.setattr(torch, "get_autocast_dtype", unexpected_modern_query, raising=False)
+        with torch.no_grad():
+            assert flce_ops._can_use_mm_out(_input, _input)
+            monkeypatch.setattr(torch, "get_autocast_gpu_dtype", lambda: torch.float16)
+            assert not flce_ops._can_use_mm_out(_input, _input)
+        assert not flce_ops._can_use_mm_out(_input, _input)
+
+    def test_mm_out_gate_preserves_hip_fallback(self, monkeypatch):
+        _input = torch.empty(1, 1, device="cuda", dtype=torch.float32)
+        monkeypatch.setattr(flce_ops, "is_hip", lambda: True)
+        with torch.no_grad():
+            assert not flce_ops._can_use_mm_out(_input, _input)
+
+    @pytest.mark.parametrize("ce_impl", _SHARED_CE_IMPLS)
+    @pytest.mark.parametrize("use_token_scaling", [False, True])
+    @pytest.mark.parametrize("weight_requires_grad", [False, True])
+    def test_reduction_none_reenters_with_reusable_buffers(
+        self, monkeypatch, ce_impl, use_token_scaling, weight_requires_grad
+    ):
+        BT, V, H = 40, 128, 32
+        cs = _production_chunk_size(BT, V, H)
+        _input, weight, b, cw, target = _mk_inputs(
+            BT,
+            V,
+            H,
+            torch.bfloat16,
+            bias=True,
+            ce_weight=True,
+            seed=18,
+            first_chunk_size=cs,
+            weight_requires_grad=weight_requires_grad,
+        )
+        upstream = torch.linspace(-0.5, 1.5, BT, device="cuda")
+
+        def run():
+            outputs = flce_ops.LigerFusedLinearCrossEntropyFunction.apply(
+                _input,
+                weight,
+                target,
+                b,
+                cw,
+                -100,
+                1e-4,
+                0.1,
+                "none",
+                1.5,
+                True,
+                torch.float32,
+                use_token_scaling,
+                True,
+                True,
+                ce_impl,
+            )
+            outputs[0].backward(upstream)
+            return outputs, _input.grad, weight.grad, b.grad
+
+        rec = _BufferRecorder(BT, V, H)
+        with rec:
+            actual = run()
+        assert rec.logits_gemms == 2 * triton.cdiv(BT, cs)
+        assert rec.final_casts == int(weight_requires_grad)
+        expected_mm_out = 0 if flce_ops.is_hip() else 3 * triton.cdiv(BT, cs)
+        assert len(rec.mm_out) == expected_mm_out  # logits twice, dX only during backward
+        rec.assert_released()
+        assert torch.count_nonzero(actual[1][:cs]) == 0
+        if not weight_requires_grad:
+            assert actual[2] is None
+            assert not rec.has_addmm_out_into(weight.shape)
+            assert not rec.has_addmm_dtype_out_into(weight.shape)
+
+        _input.grad = weight.grad = b.grad = None
+        monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
+        expected = run()
+        torch.testing.assert_close(actual, expected)
+
     # -----------------------------------------------------------------------
     # 8. AMP: in-place out ops are not autocast, so _input_chunk must be aligned
     # -----------------------------------------------------------------------
@@ -839,7 +1162,7 @@ class TestFusedLinearCrossEntropyAddmm:
         # grad_logits promoted to fp32 -> no fp32-out fast path, no low-prec out path.
         assert not rec.has_addmm_dtype_out_into(weight.shape)
         assert not rec.has_addmm_out_into(weight.shape)
-        assert rec.mm_default >= 1
+        assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
         _, ref_gi, ref_gw, ref_gb = _addmm_reference_grads(_input, weight, target, bias=bias)
         torch.testing.assert_close(gw.float(), ref_gw, atol=1e-2, rtol=1e-2)
         torch.testing.assert_close(gi.float(), ref_gi, atol=1e-2, rtol=1e-2)
@@ -877,7 +1200,7 @@ class TestFusedLinearCrossEntropyAddmm:
         # Sub-SM80 capability -> neither addmm fast path; legacy mm + fp32 upcast instead.
         assert not rec.has_addmm_out_into(weight.shape)
         assert not rec.has_addmm_dtype_out_into(weight.shape)
-        assert rec.mm_default >= 1
+        assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
         assert rec.has_param_upcast(weight.shape)
         _, ref_gi, ref_gw, _ = _addmm_reference_grads(_input, weight, target)
         torch.testing.assert_close(gw.float(), ref_gw, **_ADDMM_TOLS[torch.bfloat16])