From 294e7507f8a1f588715e78f409a4df21d158c4a8 Mon Sep 17 00:00:00 2001
From: Przemek Tredak <ptredak@nvidia.com>
Date: Wed, 29 Jul 2026 07:16:54 -0700
Subject: [PATCH 1/2] First version, not reviewed yet

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
---
 .../benchmark_parallel_cross_entropy.py       | 204 ++++++++++++++
 docs/api/pytorch.rst                          |   2 +
 tests/pytorch/test_parallel_cross_entropy.py  | 256 +++++++++++++++++-
 .../common/triton/cross_entropy.py            | 222 +++++++++++++++
 transformer_engine/pytorch/__init__.py        |   5 +-
 transformer_engine/pytorch/cross_entropy.py   | 149 ++++++++++
 .../pytorch/triton/cross_entropy.py           | 136 ++++++++++
 7 files changed, 972 insertions(+), 2 deletions(-)
 create mode 100644 benchmarks/benchmark_parallel_cross_entropy.py

diff --git a/benchmarks/benchmark_parallel_cross_entropy.py b/benchmarks/benchmark_parallel_cross_entropy.py
new file mode 100644
index 0000000000..baf0dae38c
--- /dev/null
+++ b/benchmarks/benchmark_parallel_cross_entropy.py
@@ -0,0 +1,204 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Benchmark fused cross entropy saved-state and peak-memory tradeoffs."""
+
+import argparse
+import gc
+from statistics import mean
+
+import torch
+
+from transformer_engine.pytorch import (
+    parallel_cross_entropy,
+    parallel_cross_entropy_recompute,
+)
+
+MIB = 1024**2
+
+
+def parse_args():
+    """Parse command-line arguments."""
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--sequence-length", type=int, default=2048)
+    parser.add_argument("--vocab-size", type=int, default=128000)
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--trials", type=int, default=10)
+    parser.add_argument("--label-smoothing", type=float, default=0.0)
+    return parser.parse_args()
+
+
+def _operator(implementation, logits, target, label_smoothing):
+    if implementation == "existing":
+        return parallel_cross_entropy(
+            logits,
+            target,
+            label_smoothing=label_smoothing,
+            reduce_loss=True,
+        )
+    return parallel_cross_entropy_recompute(
+        logits,
+        target,
+        label_smoothing=label_smoothing,
+        reduce_loss=True,
+        overwrite_input=implementation == "destructive",
+    )
+
+
+def _make_inputs(shape, dtype):
+    gc.collect()
+    torch.cuda.empty_cache()
+    caller_input = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+    target = torch.randint(0, shape[-1], shape[:-1], device="cuda")
+    torch.cuda.synchronize()
+    return caller_input, target
+
+
+def _measure_memory(implementation, shape, dtype, label_smoothing):
+    """Measure memory while keeping the caller-owned input alive throughout."""
+
+    caller_input, target = _make_inputs(shape, dtype)
+    baseline = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+    loss = _operator(implementation, caller_input, target, label_smoothing)
+    torch.cuda.synchronize()
+    forward_peak = torch.cuda.max_memory_allocated()
+    forward_end_live = torch.cuda.memory_allocated()
+
+    backward_entry = forward_end_live
+    torch.cuda.reset_peak_memory_stats()
+    loss.backward()
+    torch.cuda.synchronize()
+    backward_peak = torch.cuda.max_memory_allocated()
+    backward_end_live = torch.cuda.memory_allocated()
+
+    # Referencing the input here makes its consistent lifetime explicit. In
+    # destructive mode its storage now contains the derivative.
+    assert caller_input.data_ptr() != 0
+    return {
+        "baseline": baseline,
+        "forward_peak": forward_peak,
+        "forward_end": forward_end_live,
+        "backward_entry": backward_entry,
+        "backward_peak": backward_peak,
+        "backward_end": backward_end_live,
+        "e2e_peak": max(forward_peak, backward_peak),
+    }
+
+
+def _measure_latency(implementation, shape, dtype, label_smoothing):
+    """Measure uninterrupted forward and backward device latency."""
+
+    caller_input, target = _make_inputs(shape, dtype)
+    forward_start = torch.cuda.Event(enable_timing=True)
+    forward_end = torch.cuda.Event(enable_timing=True)
+    backward_end = torch.cuda.Event(enable_timing=True)
+
+    forward_start.record()
+    loss = _operator(implementation, caller_input, target, label_smoothing)
+    forward_end.record()
+    loss.backward()
+    backward_end.record()
+    torch.cuda.synchronize()
+
+    assert caller_input.data_ptr() != 0
+    return {
+        "forward_ms": forward_start.elapsed_time(forward_end),
+        "backward_ms": forward_end.elapsed_time(backward_end),
+        "total_ms": forward_start.elapsed_time(backward_end),
+    }
+
+
+def _format_memory(value, baseline):
+    return f"{value / MIB:9.1f} / +{(value - baseline) / MIB:7.1f}"
+
+
+def _print_results(results):
+    memory_fields = (
+        ("fwd peak", "forward_peak"),
+        ("fwd end", "forward_end"),
+        ("bwd entry", "backward_entry"),
+        ("bwd peak", "backward_peak"),
+        ("e2e peak", "e2e_peak"),
+        ("bwd end", "backward_end"),
+    )
+    header = ["implementation", "baseline MiB"]
+    header.extend(f"{label} abs/+inc MiB" for label, _ in memory_fields)
+    header.extend(("fwd ms", "bwd ms", "total ms"))
+    rows = []
+    for implementation, measurements in results.items():
+        memory_sample = measurements["memory"]
+        timings = measurements["timings"]
+        row = [implementation, f"{memory_sample['baseline'] / MIB:.1f}"]
+        row.extend(
+            _format_memory(memory_sample[field], memory_sample["baseline"])
+            for _, field in memory_fields
+        )
+        row.extend(
+            f"{mean(sample[field] for sample in timings):.3f}"
+            for field in ("forward_ms", "backward_ms", "total_ms")
+        )
+        rows.append(row)
+
+    widths = [max(len(header[idx]), *(len(row[idx]) for row in rows)) for idx in range(len(header))]
+    print("  ".join(value.ljust(widths[idx]) for idx, value in enumerate(header)))
+    print("  ".join("-" * width for width in widths))
+    for row in rows:
+        print("  ".join(value.ljust(widths[idx]) for idx, value in enumerate(row)))
+
+
+def main():
+    """Run the benchmark."""
+
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("This benchmark requires CUDA")
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    shape = (args.batch_size, args.sequence_length, args.vocab_size)
+    n_logits = args.batch_size * args.sequence_length * args.vocab_size
+    n_rows = args.batch_size * args.sequence_length
+    input_bytes = n_logits * dtype.itemsize
+
+    expected_multipliers = {
+        "existing": dtype.itemsize + 4,
+        "safe": 2 * dtype.itemsize,
+        "destructive": dtype.itemsize,
+    }
+    print(f"shape={shape}, dtype={dtype}, live input={input_bytes / MIB:.1f} MiB")
+    print("Expected logits-sized absolute forward peak (includes live caller input):")
+    for name, bytes_per_logit in expected_multipliers.items():
+        print(f"  {name:11s}: {bytes_per_logit}N = {bytes_per_logit * n_logits / MIB:.1f} MiB")
+    print(
+        "Recompute saved row metadata: "
+        f"{(n_rows * (2 * 4 + 8) + 8) / MIB:.3f} MiB "
+        "(FP32 max/denominator, int64 targets/count)"
+    )
+    print(f"Existing forward row temporaries: {(n_rows * (3 * 4 + 4) + 8) / MIB:.3f} MiB")
+    print(f"Recompute forward loss temporary: {n_rows * 4 / MIB:.3f} MiB")
+    print("Tensor-parallel communication buffers: 0 MiB (single-GPU benchmark)")
+    print("Memory cells below are 'absolute / +incremental-from-pre-forward-baseline'.")
+
+    implementations = ("existing", "safe", "destructive")
+    for implementation in implementations:
+        for _ in range(args.warmup):
+            _measure_latency(implementation, shape, dtype, args.label_smoothing)
+
+    results = {}
+    for implementation in implementations:
+        results[implementation] = {
+            "memory": _measure_memory(implementation, shape, dtype, args.label_smoothing),
+            "timings": [
+                _measure_latency(implementation, shape, dtype, args.label_smoothing)
+                for _ in range(args.trials)
+            ],
+        }
+    _print_results(results)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index 5fac0a89a6..ec76417646 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -55,6 +55,8 @@ PyTorch
 
 .. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy
 
+.. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy_recompute
+
 .. autoapifunction:: transformer_engine.pytorch.interleave_glu_tensor
 
 .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor
diff --git a/tests/pytorch/test_parallel_cross_entropy.py b/tests/pytorch/test_parallel_cross_entropy.py
index 5e0762f4aa..19505ee529 100644
--- a/tests/pytorch/test_parallel_cross_entropy.py
+++ b/tests/pytorch/test_parallel_cross_entropy.py
@@ -2,9 +2,17 @@
 #
 # See LICENSE for license information.
 
+import os
 import random
+import tempfile
+import pytest
 import torch
-from transformer_engine.pytorch import parallel_cross_entropy
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from transformer_engine.pytorch import (
+    parallel_cross_entropy,
+    parallel_cross_entropy_recompute,
+)
 
 from utils import dtype_tols
 
@@ -233,3 +241,249 @@ def test_bfloat16_loss_matches_float32_input():
     # differ by a small number of FP32 rounding steps.
     fp32_eps = torch.finfo(torch.float32).eps
     torch.testing.assert_close(bf16_loss, fp32_loss, rtol=2 * fp32_eps, atol=2 * fp32_eps)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"])
+@pytest.mark.parametrize("reduce_loss", [False, True], ids=["none", "mean"])
+@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"])
+@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"])
+def test_recompute_matches_pytorch_and_existing(
+    dtype,
+    reduce_loss,
+    label_smoothing,
+    overwrite_input,
+):
+    """Check loss and externally-scaled gradients against both reference paths."""
+
+    torch.manual_seed(1234)
+    shape = (2, 5, 37)
+    target = torch.randint(0, shape[-1], shape[:-1], device="cuda")
+    target[0, 1] = -100
+    values = torch.randn(shape, dtype=dtype, device="cuda")
+
+    logits = values.clone().requires_grad_()
+    existing_logits = values.clone().requires_grad_()
+    ref_logits = values.float().clone().requires_grad_()
+
+    loss = parallel_cross_entropy_recompute(
+        logits,
+        target,
+        label_smoothing,
+        reduce_loss,
+        overwrite_input=overwrite_input,
+    )
+    existing_loss = parallel_cross_entropy(
+        existing_logits,
+        target,
+        label_smoothing,
+        reduce_loss,
+    )
+    ref_loss = torch.nn.functional.cross_entropy(
+        ref_logits.reshape(-1, shape[-1]),
+        target.reshape(-1),
+        label_smoothing=label_smoothing,
+        reduction="mean" if reduce_loss else "none",
+    )
+    if reduce_loss:
+        ref_loss = ref_loss.reshape_as(loss)
+    else:
+        ref_loss = ref_loss.reshape_as(target)
+
+    external_grad = (
+        torch.full_like(loss, 0.37) if reduce_loss else torch.randn_like(loss, dtype=torch.float32)
+    )
+    loss.backward(external_grad)
+    existing_loss.backward(external_grad)
+    ref_loss.backward(external_grad)
+
+    tols = dtype_tols(dtype)
+    torch.testing.assert_close(loss, ref_loss, **tols)
+    torch.testing.assert_close(loss, existing_loss, **tols)
+    expected_grad = ref_logits.grad.to(dtype)
+    torch.testing.assert_close(logits.grad, expected_grad, **tols)
+    torch.testing.assert_close(logits.grad, existing_logits.grad, **tols)
+
+
+@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"])
+def test_recompute_saved_state_and_buffer_reuse(overwrite_input):
+    """The saved logits buffer is input-typed and becomes the returned derivative."""
+
+    torch.manual_seed(42)
+    logits = torch.randn(2, 3, 11, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    target = torch.tensor([[0, -100, 4], [7, 2, 10]], device="cuda")
+    before = logits.detach().clone()
+    saved_tensors = []
+
+    def pack_hook(tensor):
+        saved_tensors.append(tensor)
+        return tensor
+
+    with torch.autograd.graph.saved_tensors_hooks(pack_hook, lambda tensor: tensor):
+        loss = parallel_cross_entropy_recompute(
+            logits,
+            target,
+            label_smoothing=0.1,
+            overwrite_input=overwrite_input,
+        )
+
+    assert len(saved_tensors) == 4
+    saved_logits, stats, saved_target, n_non_ignore = saved_tensors
+    assert saved_logits.shape == logits.shape
+    assert saved_logits.dtype == logits.dtype
+    assert stats.shape == (target.numel(), 2)
+    assert stats.dtype == torch.float32
+    assert saved_target.numel() == target.numel()
+    assert saved_target.dtype == torch.int64
+    assert n_non_ignore.shape == (1,)
+    assert n_non_ignore.dtype == torch.int64
+    assert n_non_ignore.item() == target.numel() - 1
+
+    if overwrite_input:
+        assert saved_logits.data_ptr() == logits.data_ptr()
+    else:
+        assert saved_logits.data_ptr() != logits.data_ptr()
+    torch.testing.assert_close(logits, before, rtol=0.0, atol=0.0)
+
+    expected_max = before.float().amax(dim=-1).reshape(-1)
+    expected_denominator = (
+        torch.exp(before.float() - expected_max.reshape(logits.shape[0], logits.shape[1], 1))
+        .sum(dim=-1)
+        .reshape(-1)
+    )
+    torch.testing.assert_close(stats[:, 0], expected_max, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(stats[:, 1], expected_denominator, rtol=1e-6, atol=1e-6)
+
+    external_grad = torch.randn_like(loss)
+    loss.backward(external_grad)
+
+    # Backward writes directly into the tensor saved by forward.
+    torch.testing.assert_close(saved_logits, logits.grad, rtol=0.0, atol=0.0)
+    if overwrite_input:
+        assert not torch.equal(logits, before)
+    else:
+        torch.testing.assert_close(logits, before, rtol=0.0, atol=0.0)
+
+
+@pytest.mark.parametrize("layout", ["transpose", "strided_vocab"])
+def test_recompute_safe_mode_supported_layouts(layout):
+    """Safe mode copies non-contiguous logical layouts directly into its work buffer."""
+
+    torch.manual_seed(17)
+    if layout == "transpose":
+        values = torch.randn(3, 2, 13, device="cuda").transpose(0, 1)
+    else:
+        values = torch.randn(2, 3, 26, device="cuda")[..., ::2]
+    logits = values.detach().requires_grad_()
+    reference = values.detach().clone().requires_grad_()
+    target = torch.randint(0, values.shape[-1], values.shape[:-1], device="cuda")
+    external_grad = torch.randn(values.shape[:-1], device="cuda")
+
+    loss = parallel_cross_entropy_recompute(logits, target)
+    ref_loss = torch.nn.functional.cross_entropy(
+        reference.reshape(-1, reference.shape[-1]),
+        target.reshape(-1),
+        reduction="none",
+    ).reshape_as(target)
+    loss.backward(external_grad)
+    ref_loss.backward(external_grad)
+
+    torch.testing.assert_close(loss, ref_loss, **dtype_tols(torch.float32))
+    torch.testing.assert_close(logits.grad, reference.grad, **dtype_tols(torch.float32))
+
+
+def test_recompute_validation_errors():
+    """Reject inputs that violate dtype, shape, device, or overwrite assumptions."""
+
+    target = torch.zeros((2, 3), dtype=torch.int64, device="cuda")
+    logits = torch.randn(2, 3, 5, device="cuda")
+
+    with pytest.raises(ValueError, match="3D"):
+        parallel_cross_entropy_recompute(logits.reshape(6, 5), target)
+    with pytest.raises(TypeError, match="BF16 or FP32"):
+        parallel_cross_entropy_recompute(logits.half(), target)
+    with pytest.raises(TypeError, match="int64"):
+        parallel_cross_entropy_recompute(logits, target.int())
+    with pytest.raises(ValueError, match="one target"):
+        parallel_cross_entropy_recompute(logits, target[:, :2])
+    with pytest.raises(ValueError, match="\\[0, 1\\]"):
+        parallel_cross_entropy_recompute(logits, target, label_smoothing=1.1)
+    with pytest.raises(ValueError, match="contiguous"):
+        parallel_cross_entropy_recompute(
+            logits.transpose(0, 1),
+            target.transpose(0, 1),
+            overwrite_input=True,
+        )
+    with pytest.raises(ValueError, match="CUDA"):
+        parallel_cross_entropy_recompute(logits.cpu(), target.cpu())
+
+
+def _run_recompute_tensor_parallel(rank, world_size, init_file):
+    """Two-rank correctness worker for the tensor-parallel pre/post kernels."""
+
+    torch.cuda.set_device(rank)
+    device = torch.device("cuda", rank)
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"file://{init_file}",
+        rank=rank,
+        world_size=world_size,
+    )
+    try:
+        generator = torch.Generator().manual_seed(2025)
+        shape = (2, 3, 22)
+        local_vocab = shape[-1] // world_size
+        target = torch.randint(
+            0,
+            shape[-1],
+            shape[:-1],
+            generator=generator,
+        ).to(device)
+        target[0, 2] = -100
+        external_grad = torch.randn(shape[:-1], generator=generator).to(device)
+
+        for dtype in (torch.float32, torch.bfloat16):
+            global_values = torch.randn(shape, generator=generator).to(
+                device=device,
+                dtype=dtype,
+            )
+            vocab_start = rank * local_vocab
+            local_values = global_values[..., vocab_start : vocab_start + local_vocab]
+            local_logits = local_values.clone().requires_grad_()
+            ref_logits = global_values.float().clone().requires_grad_()
+
+            loss = parallel_cross_entropy_recompute(
+                local_logits,
+                target,
+                label_smoothing=0.1,
+                dist_process_group=dist.group.WORLD,
+            )
+            ref_loss = torch.nn.functional.cross_entropy(
+                ref_logits.reshape(-1, shape[-1]),
+                target.reshape(-1),
+                label_smoothing=0.1,
+                reduction="none",
+            ).reshape_as(target)
+            loss.backward(external_grad)
+            ref_loss.backward(external_grad)
+
+            tols = dtype_tols(dtype)
+            torch.testing.assert_close(loss, ref_loss, **tols)
+            expected_grad = ref_logits.grad[..., vocab_start : vocab_start + local_vocab].to(dtype)
+            torch.testing.assert_close(local_logits.grad, expected_grad, **tols)
+    finally:
+        dist.destroy_process_group()
+
+
+def test_recompute_tensor_parallel():
+    """Validate global statistics, smoothed loss, and local gradients on two ranks."""
+
+    if torch.cuda.device_count() < 2:
+        pytest.skip("tensor-parallel cross entropy test requires two CUDA devices")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        init_file = os.path.join(temp_dir, "distributed_init")
+        mp.spawn(
+            _run_recompute_tensor_parallel,
+            args=(2, init_file),
+            nprocs=2,
+            join=True,
+        )
diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py
index afc2ad4854..05b855db9b 100644
--- a/transformer_engine/common/triton/cross_entropy.py
+++ b/transformer_engine/common/triton/cross_entropy.py
@@ -264,3 +264,225 @@ def element_mul_kernel(
         X_offsets = i + tl.arange(0, BLOCK_SIZE)
         X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols)
         tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols)
+
+
+@triton.jit
+def cross_entropy_recompute_forward_kernel(
+    X_ptr,
+    X_stride_0,
+    X_stride_1,
+    X_stride_2,
+    logits_ptr,
+    Y_ptr,
+    loss_ptr,
+    stats_ptr,
+    n_non_ignore,
+    n_cols,
+    n_rows_1,
+    ignore_idx,
+    label_smoothing: tl.constexpr,
+    COPY_LOGITS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Compute single-rank loss/statistics and optionally preserve the logits."""
+
+    row = tl.program_id(0).to(tl.int64)
+    row_0 = row // n_rows_1
+    row_1 = row - row_0 * n_rows_1
+    X_ptr += row_0 * X_stride_0 + row_1 * X_stride_1
+    logits_ptr += row * n_cols
+
+    y = tl.load(Y_ptr + row)
+    if y != ignore_idx:
+        tl.atomic_add(n_non_ignore, 1)
+
+    m = float("-inf")
+    d = 0.0
+    x_sum = 0.0
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_cols
+        x = tl.load(X_ptr + offsets * X_stride_2, mask=mask, other=float("-inf"))
+        if COPY_LOGITS:
+            tl.store(logits_ptr + offsets, x, mask=mask)
+        x = x.to(tl.float32)
+        block_max = tl.max(x)
+        m_new = tl.maximum(m, block_max)
+        d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new))
+        m = m_new
+        if label_smoothing > 0:
+            x_sum += tl.sum(tl.where(mask, x, 0.0))
+
+    tl.store(stats_ptr + row * 2, m)
+    tl.store(stats_ptr + row * 2 + 1, d)
+
+    if y == ignore_idx:
+        tl.store(loss_ptr + row, 0.0)
+        return
+
+    x_y = float("-inf")
+    if y >= 0:
+        if y < n_cols:
+            x_y = tl.load(X_ptr + y * X_stride_2).to(tl.float32)
+
+    loss = -(x_y - m - tl.log(d))
+    if label_smoothing > 0:
+        eps = label_smoothing / n_cols
+        smooth_loss = -eps * x_sum + label_smoothing * (m + tl.log(d))
+        loss = loss * (1 - label_smoothing) + smooth_loss
+    tl.store(loss_ptr + row, loss)
+
+
+@triton.jit
+def cross_entropy_recompute_tp_pre_kernel(
+    X_ptr,
+    X_stride_0,
+    X_stride_1,
+    X_stride_2,
+    logits_ptr,
+    Y_ptr,
+    local_data_ptr,
+    n_non_ignore,
+    rank,
+    n_cols,
+    n_rows_1,
+    ignore_idx,
+    COPY_LOGITS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Compute the local statistics needed by tensor-parallel cross entropy."""
+
+    row = tl.program_id(0).to(tl.int64)
+    row_0 = row // n_rows_1
+    row_1 = row - row_0 * n_rows_1
+    X_ptr += row_0 * X_stride_0 + row_1 * X_stride_1
+    logits_ptr += row * n_cols
+
+    y = tl.load(Y_ptr + row)
+    if y != ignore_idx:
+        tl.atomic_add(n_non_ignore, 1)
+
+    vocab_start = rank * n_cols
+    x_y = float("-inf")
+    if y >= vocab_start:
+        if y < vocab_start + n_cols:
+            x_y = tl.load(X_ptr + (y - vocab_start) * X_stride_2).to(tl.float32)
+
+    m = float("-inf")
+    d = 0.0
+    x_sum = 0.0
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_cols
+        x = tl.load(X_ptr + offsets * X_stride_2, mask=mask, other=float("-inf"))
+        if COPY_LOGITS:
+            tl.store(logits_ptr + offsets, x, mask=mask)
+        x = x.to(tl.float32)
+        block_max = tl.max(x)
+        m_new = tl.maximum(m, block_max)
+        d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new))
+        m = m_new
+        x_sum += tl.sum(tl.where(mask, x, 0.0))
+
+    local_data_ptr += row * 4
+    tl.store(local_data_ptr, m)
+    tl.store(local_data_ptr + 1, d)
+    tl.store(local_data_ptr + 2, x_y)
+    tl.store(local_data_ptr + 3, x_sum)
+
+
+@triton.jit
+def cross_entropy_recompute_tp_post_kernel(
+    gathered_data_ptr,
+    Y_ptr,
+    loss_ptr,
+    stats_ptr,
+    world_size,
+    n_rows,
+    n_cols,
+    ignore_idx,
+    label_smoothing: tl.constexpr,
+):
+    """Combine tensor-parallel statistics and compute loss/global statistics."""
+
+    row = tl.program_id(0).to(tl.int64)
+    data_ptr = gathered_data_ptr + row * 4
+    m = tl.load(data_ptr)
+    d = tl.load(data_ptr + 1)
+    x_y = tl.load(data_ptr + 2)
+    x_sum = tl.load(data_ptr + 3)
+
+    for rank_idx in range(1, world_size):
+        rank_data_ptr = data_ptr + rank_idx * n_rows * 4
+        m_new = tl.load(rank_data_ptr)
+        d_new = tl.load(rank_data_ptr + 1)
+        global_m = tl.maximum(m, m_new)
+        d = d * tl.exp(m - global_m) + d_new * tl.exp(m_new - global_m)
+        m = global_m
+        x_y = tl.maximum(x_y, tl.load(rank_data_ptr + 2))
+        x_sum += tl.load(rank_data_ptr + 3)
+
+    tl.store(stats_ptr + row * 2, m)
+    tl.store(stats_ptr + row * 2 + 1, d)
+
+    y = tl.load(Y_ptr + row)
+    if y == ignore_idx:
+        tl.store(loss_ptr + row, 0.0)
+        return
+
+    loss = -(x_y - m - tl.log(d))
+    if label_smoothing > 0:
+        vocab_size = n_cols * world_size
+        eps = label_smoothing / vocab_size
+        smooth_loss = -eps * x_sum + label_smoothing * (m + tl.log(d))
+        loss = loss * (1 - label_smoothing) + smooth_loss
+    tl.store(loss_ptr + row, loss)
+
+
+@triton.jit
+def cross_entropy_recompute_backward_kernel(
+    logits_ptr,
+    Y_ptr,
+    stats_ptr,
+    n_non_ignore_ptr,
+    grad_output_ptr,
+    grad_output_stride,
+    rank,
+    world_size,
+    n_cols,
+    ignore_idx,
+    reduce_loss: tl.constexpr,
+    label_smoothing: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Reconstruct the FP32 derivative and store it in the logits buffer."""
+
+    row = tl.program_id(0).to(tl.int64)
+    logits_ptr += row * n_cols
+    y = tl.load(Y_ptr + row)
+
+    if y == ignore_idx:
+        for i in range(0, n_cols, BLOCK_SIZE):
+            offsets = i + tl.arange(0, BLOCK_SIZE)
+            tl.store(logits_ptr + offsets, 0.0, mask=offsets < n_cols)
+        return
+
+    m = tl.load(stats_ptr + row * 2)
+    d = tl.load(stats_ptr + row * 2 + 1)
+    grad_output = tl.load(grad_output_ptr + row * grad_output_stride).to(tl.float32)
+    if reduce_loss:
+        grad_output /= tl.load(n_non_ignore_ptr)
+
+    eps = label_smoothing / (n_cols * world_size)
+    vocab_start = rank * n_cols
+    target_col = y - vocab_start
+    target_is_local = (y >= vocab_start) & (y < vocab_start + n_cols)
+
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_cols
+        x = tl.load(logits_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32)
+        grad = tl.exp(x - m) / d - eps
+        is_target = target_is_local & (offsets == target_col)
+        grad -= tl.where(is_target, 1 - label_smoothing, 0.0)
+        tl.store(logits_ptr + offsets, grad * grad_output, mask=mask)
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 06db28ee27..17e19b659b 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -69,7 +69,10 @@
 from transformer_engine.pytorch import ops
 from transformer_engine.pytorch import optimizers
 from transformer_engine.pytorch.export import onnx_export
-from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy
+from transformer_engine.pytorch.cross_entropy import (
+    parallel_cross_entropy,
+    parallel_cross_entropy_recompute,
+)
 from transformer_engine.pytorch.newton_schulz import (
     CusolverMpCtx,
     newton_schulz,
diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py
index bb6e74fdc0..1a58a6592e 100644
--- a/transformer_engine/pytorch/cross_entropy.py
+++ b/transformer_engine/pytorch/cross_entropy.py
@@ -13,6 +13,7 @@
 
 __all__ = [
     "parallel_cross_entropy",
+    "parallel_cross_entropy_recompute",
 ]
 
 
@@ -90,6 +91,98 @@ def backward(ctx, grad_output):
         )
 
 
+class CrossEntropyRecomputeFunction(torch.autograd.Function):
+    """Cross entropy autograd function that recomputes its derivative."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        inp,
+        target,
+        label_smoothing=0.0,
+        reduce_loss=False,
+        dist_process_group=None,
+        ignore_idx=-100,
+        is_cg_capturable=False,
+        overwrite_input=False,
+    ):
+        loss, logits, stats, target, n_non_ignore = (
+            triton_cross_entropy.cross_entropy_recompute_forward(
+                inp,
+                target,
+                label_smoothing,
+                reduce_loss,
+                dist_process_group,
+                ignore_idx,
+                overwrite_input,
+            )
+        )
+        ctx.save_for_backward(logits.detach(), stats, target, n_non_ignore)
+        ctx.label_smoothing = label_smoothing
+        ctx.reduce_loss = reduce_loss
+        ctx.dist_process_group = dist_process_group
+        ctx.ignore_idx = ignore_idx
+        ctx.is_cg_capturable = is_cg_capturable
+        ctx.did_backward = False
+        return loss
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.did_backward:
+            raise RuntimeError(
+                "parallel_cross_entropy_recompute does not support repeated backward passes "
+                "because backward reuses its saved logits buffer"
+            )
+        ctx.did_backward = True
+        logits, stats, target, n_non_ignore = ctx.saved_tensors
+        grad_input = triton_cross_entropy.cross_entropy_recompute_backward(
+            logits,
+            stats,
+            target,
+            n_non_ignore,
+            grad_output,
+            ctx.label_smoothing,
+            ctx.reduce_loss,
+            ctx.dist_process_group,
+            ctx.ignore_idx,
+            ctx.is_cg_capturable,
+        )
+        return grad_input, None, None, None, None, None, None, None
+
+
+def _validate_recompute_inputs(
+    inp: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float,
+    overwrite_input: bool,
+) -> None:
+    """Validate assumptions made by the recompute Triton kernels."""
+
+    if inp.ndim != 3:
+        raise ValueError(f"inp must be a 3D tensor, but got shape {tuple(inp.shape)}")
+    if any(size == 0 for size in inp.shape):
+        raise ValueError(f"inp dimensions must be non-zero, but got shape {tuple(inp.shape)}")
+    if inp.dtype not in (torch.bfloat16, torch.float32):
+        raise TypeError(f"inp must have BF16 or FP32 dtype, but got {inp.dtype}")
+    if not inp.is_cuda:
+        raise ValueError("inp must be a CUDA tensor")
+    if target.device != inp.device:
+        raise ValueError("target must be on the same CUDA device as inp")
+    if target.dtype != torch.int64:
+        raise TypeError(f"target must have torch.int64 dtype, but got {target.dtype}")
+    if target.numel() != inp.shape[0] * inp.shape[1]:
+        raise ValueError(
+            "Each input row needs one target token ID: "
+            f"expected {inp.shape[0] * inp.shape[1]}, got {target.numel()}"
+        )
+    if not 0.0 <= label_smoothing <= 1.0:
+        raise ValueError(f"label_smoothing must be in [0, 1], but got {label_smoothing}")
+    if overwrite_input and (not inp.is_contiguous() or torch._debug_has_internal_overlap(inp) != 0):
+        raise ValueError(
+            "overwrite_input=True requires a contiguous input with no internal overlap"
+        )
+
+
 def parallel_cross_entropy(
     inp: torch.Tensor,
     target: torch.Tensor,
@@ -151,3 +244,59 @@ def parallel_cross_entropy(
         ignore_idx,
         is_cg_capturable,
     )
+
+
+def parallel_cross_entropy_recompute(
+    inp: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    reduce_loss: bool = False,
+    dist_process_group: Optional[torch.distributed.ProcessGroup] = None,
+    ignore_idx: int = -100,
+    is_cg_capturable: bool = False,
+    *,
+    overwrite_input: bool = False,
+) -> torch.Tensor:
+    """Memory-efficient cross entropy with derivative recomputation.
+
+    Loss and derivative calculations use FP32 arithmetic for BF16 and FP32
+    inputs. Instead of saving a full FP32 derivative, this function saves an
+    input-typed logits buffer and per-row FP32 softmax maximum/denominator
+    statistics. The returned loss is FP32.
+
+    By default, the saved logits are a private contiguous copy and ``inp`` is
+    preserved. This safe mode necessarily has both the caller's input and the
+    copy live during forward. With ``overwrite_input=True``, the original
+    contiguous, non-overlapping input is used as the saved buffer and is
+    overwritten with its gradient during backward. Callers must not read or
+    otherwise reuse that input after starting backward.
+
+    The saved logits buffer is consumed by backward, so repeated backward
+    passes on the same result are not supported.
+
+    Parameters are the same as :func:`parallel_cross_entropy`, with the
+    addition of ``overwrite_input``.
+
+    Parameters
+    ----------
+    overwrite_input : bool, default = False
+        Reuse and overwrite ``inp`` rather than allocating a private logits
+        copy. The input must be contiguous and have no internal overlap.
+
+    Returns
+    -------
+    torch.Tensor
+        The computed loss.
+    """
+
+    _validate_recompute_inputs(inp, target, label_smoothing, overwrite_input)
+    return CrossEntropyRecomputeFunction.apply(
+        inp,
+        target,
+        label_smoothing,
+        reduce_loss,
+        dist_process_group,
+        ignore_idx,
+        is_cg_capturable,
+        overwrite_input,
+    )
diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py
index 7ca7cfed2f..bcc913b976 100644
--- a/transformer_engine/pytorch/triton/cross_entropy.py
+++ b/transformer_engine/pytorch/triton/cross_entropy.py
@@ -17,6 +17,10 @@
     online_softmax_kernel,
     cross_entropy_kernel,
     element_mul_kernel,
+    cross_entropy_recompute_forward_kernel,
+    cross_entropy_recompute_tp_pre_kernel,
+    cross_entropy_recompute_tp_post_kernel,
+    cross_entropy_recompute_backward_kernel,
 )
 
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
@@ -143,3 +147,135 @@ def cross_entropy_backward(
         )
 
     return grad_input
+
+
+def cross_entropy_recompute_forward(
+    _input: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float,
+    reduce_loss: bool,
+    dist_process_group: Union[dist.ProcessGroup, None],
+    ignore_idx: int,
+    overwrite_input: bool,
+):
+    """Forward implementation that saves logits and compact softmax statistics."""
+
+    B, SQ, V = _input.shape
+    n_rows = B * SQ
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+
+    target = target.contiguous().reshape(-1)
+    logits = (
+        _input
+        if overwrite_input
+        else torch.empty(_input.shape, dtype=_input.dtype, device=_input.device)
+    )
+    loss_1d = torch.empty(n_rows, dtype=torch.float32, device=_input.device)
+    stats = torch.empty((n_rows, 2), dtype=torch.float32, device=_input.device)
+    n_non_ignore = torch.zeros(1, dtype=torch.int64, device=_input.device)
+
+    rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group)
+    world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)
+
+    if world_size == 1:
+        cross_entropy_recompute_forward_kernel[(n_rows,)](
+            X_ptr=_input,
+            X_stride_0=_input.stride(0),
+            X_stride_1=_input.stride(1),
+            X_stride_2=_input.stride(2),
+            logits_ptr=logits,
+            Y_ptr=target,
+            loss_ptr=loss_1d,
+            stats_ptr=stats,
+            n_non_ignore=n_non_ignore,
+            n_cols=V,
+            n_rows_1=SQ,
+            ignore_idx=ignore_idx,
+            label_smoothing=label_smoothing,
+            COPY_LOGITS=not overwrite_input,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=32,
+        )
+    else:
+        local_data = torch.empty((n_rows, 4), dtype=torch.float32, device=_input.device)
+        cross_entropy_recompute_tp_pre_kernel[(n_rows,)](
+            X_ptr=_input,
+            X_stride_0=_input.stride(0),
+            X_stride_1=_input.stride(1),
+            X_stride_2=_input.stride(2),
+            logits_ptr=logits,
+            Y_ptr=target,
+            local_data_ptr=local_data,
+            n_non_ignore=n_non_ignore,
+            rank=rank,
+            n_cols=V,
+            n_rows_1=SQ,
+            ignore_idx=ignore_idx,
+            COPY_LOGITS=not overwrite_input,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=32,
+        )
+        gathered_data = torch.empty(
+            (world_size * n_rows, 4), dtype=torch.float32, device=_input.device
+        )
+        dist.all_gather_into_tensor(gathered_data, local_data, group=dist_process_group)
+        cross_entropy_recompute_tp_post_kernel[(n_rows,)](
+            gathered_data_ptr=gathered_data,
+            Y_ptr=target,
+            loss_ptr=loss_1d,
+            stats_ptr=stats,
+            world_size=world_size,
+            n_rows=n_rows,
+            n_cols=V,
+            ignore_idx=ignore_idx,
+            label_smoothing=label_smoothing,
+            num_warps=1,
+        )
+
+    n_non_ignore.clamp_(min=1)
+    loss = loss_1d.reshape(B, SQ)
+    if reduce_loss:
+        loss = loss_1d.sum() / n_non_ignore
+
+    return loss, logits, stats, target, n_non_ignore
+
+
+def cross_entropy_recompute_backward(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    target: torch.Tensor,
+    n_non_ignore: torch.Tensor,
+    grad_output: torch.Tensor,
+    label_smoothing: float,
+    reduce_loss: bool,
+    dist_process_group: Union[dist.ProcessGroup, None],
+    ignore_idx: int,
+    is_cg_capturable: bool = False,
+):
+    """Reconstruct the derivative in FP32 and overwrite the saved logits buffer."""
+
+    del is_cg_capturable
+    B, SQ, V = logits.shape
+    n_rows = B * SQ
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group)
+    world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)
+    grad_output = grad_output.contiguous()
+
+    cross_entropy_recompute_backward_kernel[(n_rows,)](
+        logits_ptr=logits,
+        Y_ptr=target,
+        stats_ptr=stats,
+        n_non_ignore_ptr=n_non_ignore,
+        grad_output_ptr=grad_output,
+        grad_output_stride=0 if reduce_loss else 1,
+        rank=rank,
+        world_size=world_size,
+        n_cols=V,
+        ignore_idx=ignore_idx,
+        reduce_loss=reduce_loss,
+        label_smoothing=label_smoothing,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=32,
+    )
+    return logits

From 41157db8e1e1b8aa682438c706967fe8ed199146 Mon Sep 17 00:00:00 2001
From: Przemek Tredak <ptredak@nvidia.com>
Date: Wed, 29 Jul 2026 07:59:28 -0700
Subject: [PATCH 2/2] Removed the original, still not reviewed

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
---
 .../benchmark_parallel_cross_entropy.py       |  24 +-
 docs/api/pytorch.rst                          |   2 -
 tests/pytorch/test_parallel_cross_entropy.py  |  68 +++--
 .../common/triton/cross_entropy.py            | 266 +-----------------
 transformer_engine/pytorch/__init__.py        |   5 +-
 transformer_engine/pytorch/cross_entropy.py   | 191 +++----------
 .../pytorch/triton/cross_entropy.py           | 146 +---------
 7 files changed, 87 insertions(+), 615 deletions(-)

diff --git a/benchmarks/benchmark_parallel_cross_entropy.py b/benchmarks/benchmark_parallel_cross_entropy.py
index baf0dae38c..2af5c1dc14 100644
--- a/benchmarks/benchmark_parallel_cross_entropy.py
+++ b/benchmarks/benchmark_parallel_cross_entropy.py
@@ -2,7 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""Benchmark fused cross entropy saved-state and peak-memory tradeoffs."""
+"""Benchmark fused cross entropy safe-copy and destructive-reuse modes."""
 
 import argparse
 import gc
@@ -10,10 +10,7 @@
 
 import torch
 
-from transformer_engine.pytorch import (
-    parallel_cross_entropy,
-    parallel_cross_entropy_recompute,
-)
+from transformer_engine.pytorch import parallel_cross_entropy
 
 MIB = 1024**2
 
@@ -33,14 +30,7 @@ def parse_args():
 
 
 def _operator(implementation, logits, target, label_smoothing):
-    if implementation == "existing":
-        return parallel_cross_entropy(
-            logits,
-            target,
-            label_smoothing=label_smoothing,
-            reduce_loss=True,
-        )
-    return parallel_cross_entropy_recompute(
+    return parallel_cross_entropy(
         logits,
         target,
         label_smoothing=label_smoothing,
@@ -165,7 +155,6 @@ def main():
     input_bytes = n_logits * dtype.itemsize
 
     expected_multipliers = {
-        "existing": dtype.itemsize + 4,
         "safe": 2 * dtype.itemsize,
         "destructive": dtype.itemsize,
     }
@@ -174,16 +163,15 @@ def main():
     for name, bytes_per_logit in expected_multipliers.items():
         print(f"  {name:11s}: {bytes_per_logit}N = {bytes_per_logit * n_logits / MIB:.1f} MiB")
     print(
-        "Recompute saved row metadata: "
+        "Saved row metadata: "
         f"{(n_rows * (2 * 4 + 8) + 8) / MIB:.3f} MiB "
         "(FP32 max/denominator, int64 targets/count)"
     )
-    print(f"Existing forward row temporaries: {(n_rows * (3 * 4 + 4) + 8) / MIB:.3f} MiB")
-    print(f"Recompute forward loss temporary: {n_rows * 4 / MIB:.3f} MiB")
+    print(f"Forward loss temporary: {n_rows * 4 / MIB:.3f} MiB")
     print("Tensor-parallel communication buffers: 0 MiB (single-GPU benchmark)")
     print("Memory cells below are 'absolute / +incremental-from-pre-forward-baseline'.")
 
-    implementations = ("existing", "safe", "destructive")
+    implementations = ("safe", "destructive")
     for implementation in implementations:
         for _ in range(args.warmup):
             _measure_latency(implementation, shape, dtype, args.label_smoothing)
diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index ec76417646..5fac0a89a6 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -55,8 +55,6 @@ PyTorch
 
 .. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy
 
-.. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy_recompute
-
 .. autoapifunction:: transformer_engine.pytorch.interleave_glu_tensor
 
 .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor
diff --git a/tests/pytorch/test_parallel_cross_entropy.py b/tests/pytorch/test_parallel_cross_entropy.py
index 19505ee529..26ac91e763 100644
--- a/tests/pytorch/test_parallel_cross_entropy.py
+++ b/tests/pytorch/test_parallel_cross_entropy.py
@@ -9,10 +9,7 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from transformer_engine.pytorch import (
-    parallel_cross_entropy,
-    parallel_cross_entropy_recompute,
-)
+from transformer_engine.pytorch import parallel_cross_entropy
 
 from utils import dtype_tols
 
@@ -213,8 +210,8 @@ def pack_hook(tensor):
     with torch.autograd.graph.saved_tensors_hooks(pack_hook, lambda tensor: tensor):
         loss = parallel_cross_entropy(logits, target, 0.0, False, None)
 
-    assert len(saved_tensors) == 1
-    assert saved_tensors[0].dtype == torch.float32
+    assert len(saved_tensors) == 4
+    assert saved_tensors[0].dtype == torch.bfloat16
     torch.testing.assert_close(logits, logits_before, rtol=0.0, atol=0.0)
     loss.backward(external_grad)
 
@@ -247,13 +244,13 @@ def test_bfloat16_loss_matches_float32_input():
 @pytest.mark.parametrize("reduce_loss", [False, True], ids=["none", "mean"])
 @pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"])
 @pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"])
-def test_recompute_matches_pytorch_and_existing(
+def test_parallel_cross_entropy_matches_pytorch(
     dtype,
     reduce_loss,
     label_smoothing,
     overwrite_input,
 ):
-    """Check loss and externally-scaled gradients against both reference paths."""
+    """Check loss and externally-scaled gradients against PyTorch."""
 
     torch.manual_seed(1234)
     shape = (2, 5, 37)
@@ -262,22 +259,15 @@ def test_recompute_matches_pytorch_and_existing(
     values = torch.randn(shape, dtype=dtype, device="cuda")
 
     logits = values.clone().requires_grad_()
-    existing_logits = values.clone().requires_grad_()
     ref_logits = values.float().clone().requires_grad_()
 
-    loss = parallel_cross_entropy_recompute(
+    loss = parallel_cross_entropy(
         logits,
         target,
         label_smoothing,
         reduce_loss,
         overwrite_input=overwrite_input,
     )
-    existing_loss = parallel_cross_entropy(
-        existing_logits,
-        target,
-        label_smoothing,
-        reduce_loss,
-    )
     ref_loss = torch.nn.functional.cross_entropy(
         ref_logits.reshape(-1, shape[-1]),
         target.reshape(-1),
@@ -293,19 +283,16 @@ def test_recompute_matches_pytorch_and_existing(
         torch.full_like(loss, 0.37) if reduce_loss else torch.randn_like(loss, dtype=torch.float32)
     )
     loss.backward(external_grad)
-    existing_loss.backward(external_grad)
     ref_loss.backward(external_grad)
 
     tols = dtype_tols(dtype)
     torch.testing.assert_close(loss, ref_loss, **tols)
-    torch.testing.assert_close(loss, existing_loss, **tols)
     expected_grad = ref_logits.grad.to(dtype)
     torch.testing.assert_close(logits.grad, expected_grad, **tols)
-    torch.testing.assert_close(logits.grad, existing_logits.grad, **tols)
 
 
 @pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"])
-def test_recompute_saved_state_and_buffer_reuse(overwrite_input):
+def test_parallel_cross_entropy_saved_state_and_buffer_reuse(overwrite_input):
     """The saved logits buffer is input-typed and becomes the returned derivative."""
 
     torch.manual_seed(42)
@@ -319,7 +306,7 @@ def pack_hook(tensor):
         return tensor
 
     with torch.autograd.graph.saved_tensors_hooks(pack_hook, lambda tensor: tensor):
-        loss = parallel_cross_entropy_recompute(
+        loss = parallel_cross_entropy(
             logits,
             target,
             label_smoothing=0.1,
@@ -365,7 +352,7 @@ def pack_hook(tensor):
 
 
 @pytest.mark.parametrize("layout", ["transpose", "strided_vocab"])
-def test_recompute_safe_mode_supported_layouts(layout):
+def test_parallel_cross_entropy_safe_mode_supported_layouts(layout):
     """Safe mode copies non-contiguous logical layouts directly into its work buffer."""
 
     torch.manual_seed(17)
@@ -378,7 +365,7 @@ def test_recompute_safe_mode_supported_layouts(layout):
     target = torch.randint(0, values.shape[-1], values.shape[:-1], device="cuda")
     external_grad = torch.randn(values.shape[:-1], device="cuda")
 
-    loss = parallel_cross_entropy_recompute(logits, target)
+    loss = parallel_cross_entropy(logits, target)
     ref_loss = torch.nn.functional.cross_entropy(
         reference.reshape(-1, reference.shape[-1]),
         target.reshape(-1),
@@ -391,33 +378,44 @@ def test_recompute_safe_mode_supported_layouts(layout):
     torch.testing.assert_close(logits.grad, reference.grad, **dtype_tols(torch.float32))
 
 
-def test_recompute_validation_errors():
+def test_parallel_cross_entropy_validation_errors():
     """Reject inputs that violate dtype, shape, device, or overwrite assumptions."""
 
     target = torch.zeros((2, 3), dtype=torch.int64, device="cuda")
     logits = torch.randn(2, 3, 5, device="cuda")
 
     with pytest.raises(ValueError, match="3D"):
-        parallel_cross_entropy_recompute(logits.reshape(6, 5), target)
+        parallel_cross_entropy(logits.reshape(6, 5), target)
     with pytest.raises(TypeError, match="BF16 or FP32"):
-        parallel_cross_entropy_recompute(logits.half(), target)
+        parallel_cross_entropy(logits.half(), target)
     with pytest.raises(TypeError, match="int64"):
-        parallel_cross_entropy_recompute(logits, target.int())
+        parallel_cross_entropy(logits, target.int())
     with pytest.raises(ValueError, match="one target"):
-        parallel_cross_entropy_recompute(logits, target[:, :2])
+        parallel_cross_entropy(logits, target[:, :2])
     with pytest.raises(ValueError, match="\\[0, 1\\]"):
-        parallel_cross_entropy_recompute(logits, target, label_smoothing=1.1)
+        parallel_cross_entropy(logits, target, label_smoothing=1.1)
     with pytest.raises(ValueError, match="contiguous"):
-        parallel_cross_entropy_recompute(
+        parallel_cross_entropy(
             logits.transpose(0, 1),
             target.transpose(0, 1),
             overwrite_input=True,
         )
     with pytest.raises(ValueError, match="CUDA"):
-        parallel_cross_entropy_recompute(logits.cpu(), target.cpu())
+        parallel_cross_entropy(logits.cpu(), target.cpu())
+
+
+def test_parallel_cross_entropy_deprecated_input_alias():
+    """The deprecated _input keyword remains compatible with the replaced path."""
+
+    logits = torch.randn(2, 3, 5, device="cuda")
+    target = torch.zeros((2, 3), dtype=torch.int64, device="cuda")
+    with pytest.warns(FutureWarning, match="_input"):
+        alias_loss = parallel_cross_entropy(logits, target, _input=logits)
+    direct_loss = parallel_cross_entropy(logits, target)
+    torch.testing.assert_close(alias_loss, direct_loss)
 
 
-def _run_recompute_tensor_parallel(rank, world_size, init_file):
+def _run_tensor_parallel(rank, world_size, init_file):
     """Two-rank correctness worker for the tensor-parallel pre/post kernels."""
 
     torch.cuda.set_device(rank)
@@ -451,7 +449,7 @@ def _run_recompute_tensor_parallel(rank, world_size, init_file):
             local_logits = local_values.clone().requires_grad_()
             ref_logits = global_values.float().clone().requires_grad_()
 
-            loss = parallel_cross_entropy_recompute(
+            loss = parallel_cross_entropy(
                 local_logits,
                 target,
                 label_smoothing=0.1,
@@ -474,7 +472,7 @@ def _run_recompute_tensor_parallel(rank, world_size, init_file):
         dist.destroy_process_group()
 
 
-def test_recompute_tensor_parallel():
+def test_parallel_cross_entropy_tensor_parallel():
     """Validate global statistics, smoothed loss, and local gradients on two ranks."""
 
     if torch.cuda.device_count() < 2:
@@ -482,7 +480,7 @@ def test_recompute_tensor_parallel():
     with tempfile.TemporaryDirectory() as temp_dir:
         init_file = os.path.join(temp_dir, "distributed_init")
         mp.spawn(
-            _run_recompute_tensor_parallel,
+            _run_tensor_parallel,
             args=(2, init_file),
             nprocs=2,
             join=True,
diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py
index 05b855db9b..74ca0085de 100644
--- a/transformer_engine/common/triton/cross_entropy.py
+++ b/transformer_engine/common/triton/cross_entropy.py
@@ -9,265 +9,7 @@
 
 
 @triton.jit
-def online_softmax_kernel(
-    X_ptr,
-    X_stride,
-    Y_ptr,
-    Y_stride,
-    m_d_X_y_ptr,
-    m_d_X_y_stride,
-    rank,
-    n_cols,
-    ignore_idx,
-    n_non_ignore,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    This kernel computes the m/d components on this TP rank for the online softmax.
-
-    Parameters:
-    X_ptr: Pointer to input tensor.
-    X_stride (int): The stride of the input tensor.
-    Y_ptr: Pointer to target tensor.
-    Y_stride (int): The stride of the target tensor.
-    m_d_X_y_ptr: Pointer to m/d/X_y tensor.
-    m_d_X_y_stride (int): The stride of the m/d/X_y tensor.
-    rank (int): The rank of this device in the TP group.
-    n_cols (int): The number of columns in the input tensor.
-    ignore_idx (int): The index to ignore for loss calculation.
-    n_non_ignore: The number of non-ignored elements in the batch.
-    BLOCK_SIZE (int): The block size for Triton operations.
-    """
-
-    program_id = tl.program_id(0).to(tl.int64)
-
-    # locate the start index
-    X_ptr += program_id * X_stride
-
-    # Load Y_ptr
-    Y_ptr += program_id * Y_stride
-    y = tl.load(Y_ptr)
-
-    if y != ignore_idx:
-        tl.atomic_add(n_non_ignore, 1)
-
-    vocab_start_idx = rank * n_cols
-    vocab_end_idx = (rank + 1) * n_cols
-    if y >= vocab_start_idx:
-        if y < vocab_end_idx:
-            X_y = tl.load(X_ptr + y - vocab_start_idx).to(tl.float32)
-        else:
-            X_y = float("-inf")
-    else:
-        X_y = float("-inf")
-
-    m_d_X_y_ptr += program_id * m_d_X_y_stride * 3
-
-    # 3. [Online softmax] first pass: find max + sum
-    m = float("-inf")  # m is the max value. use the notation from the paper
-    d = 0.0  # d is the sum. use the notation from the paper
-
-    for i in range(0, n_cols, BLOCK_SIZE):
-        X_offsets = i + tl.arange(0, BLOCK_SIZE)
-        X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf")).to(
-            tl.float32
-        )
-        block_max = tl.max(X_block)
-        m_new = tl.maximum(m, block_max)
-        d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new))
-        m = m_new
-
-    tl.store(m_d_X_y_ptr, m)
-    tl.store(m_d_X_y_ptr + m_d_X_y_stride, d)
-    tl.store(m_d_X_y_ptr + (2 * m_d_X_y_stride), X_y)
-
-
-@triton.jit
-def cross_entropy_kernel(
-    X_ptr,
-    X_stride,
-    grad_input_ptr,
-    grad_input_stride,
-    Y_ptr,
-    Y_stride,
-    loss_ptr,
-    loss_stride,
-    m_d_X_y_ptr,
-    m_d_X_y_stride,
-    rank,
-    world_size,
-    ignore_idx,
-    n_cols,
-    n_rows,
-    n_non_ignore,
-    reduce_loss: tl.constexpr,
-    label_smoothing: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    This kernel computes both cross entropy loss and the gradient of the input.
-
-    Parameters:
-    X_ptr: Pointer to input tensor.
-    X_stride (int): The stride of the input tensor.
-    grad_input_ptr: Pointer to the FP32 tensor that stores the input gradient.
-    grad_input_stride (int): The stride of the input gradient tensor.
-    Y_ptr: Pointer to target tensor.
-    Y_stride (int): The stride of the target tensor.
-    loss_ptr: Pointer to tensor to store the loss.
-    loss_stride (int): The stride of the loss tensor.
-    m_d_X_y_ptr: Pointer to m/d/X_y tensor.
-    m_d_X_y_stride: The stride of m/d/X_y tensor.
-    rank (int): The rank of this device in the TP group.
-    world_size (int): The size of world involved in this distributed loss calculation.
-    ignore_idx (int): Tokens to be ignored for loss and gradient calculation.
-    n_cols (int): The number of columns in the input tensor.
-    n_rows (int): The number of rows in the batch (B * SQ), used for buffer indexing.
-    n_non_ignore: The number of non-ignored elements in the batch.
-    label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
-    BLOCK_SIZE (int): The block size for Triton operations.
-    """
-
-    program_id = tl.program_id(0).to(tl.int64)
-    n_non_ignore = tl.load(n_non_ignore)
-
-    # locate the start index
-    X_ptr += program_id * X_stride
-    grad_input_ptr += program_id * grad_input_stride
-
-    # Load Y_ptr
-    Y_ptr += program_id * Y_stride
-    y = tl.load(Y_ptr)
-
-    if y == ignore_idx:
-        # Set the input gradient to zero.
-        for i in range(0, n_cols, BLOCK_SIZE):
-            X_offsets = i + tl.arange(0, BLOCK_SIZE)
-            tl.store(grad_input_ptr + X_offsets, 0.0, mask=X_offsets < n_cols)
-        return
-
-    loss_ptr += program_id * loss_stride
-    m_d_X_y_ptr += program_id * 3 * m_d_X_y_stride
-
-    # Need to reduce the m/d/X_y values from other TP ranks
-    m = tl.load(m_d_X_y_ptr)
-    d = tl.load(m_d_X_y_ptr + m_d_X_y_stride)
-    ori_X_y = tl.load(m_d_X_y_ptr + (2 * m_d_X_y_stride))
-
-    for i in range(1, world_size):
-        offset = i * 3 * n_rows * m_d_X_y_stride
-        access_ptr = m_d_X_y_ptr + offset
-        m_new = tl.load(access_ptr)
-        d_new = tl.load(access_ptr + m_d_X_y_stride)
-        X_y_new = tl.load(access_ptr + (2 * m_d_X_y_stride))
-
-        d = d * tl.exp(m - tl.maximum(m, m_new)) + d_new * tl.exp(m_new - tl.maximum(m, m_new))
-        m = tl.maximum(m, m_new)
-        ori_X_y = tl.maximum(ori_X_y, X_y_new)
-
-    # Label smoothing is a general case of normal cross entropy
-    scaled_x_sum = 0.0
-    eps = label_smoothing / (n_cols * world_size)
-
-    # 4. [Online softmax] second pass: calculate the gradients
-    # dx_y = (softmax(x_y) - 1) / N
-    # dx_i = softmax(x_i) / N, i != y
-    # N is the number of non ignored elements in the batch
-    # For label smoothing:
-    # dx_i = (softmax(x_y) - label_smoothing / V) / N, V = n_cols, i != y
-    # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N
-    #      = dx_i - (1 - label_smoothing) / N
-    for i in range(0, n_cols, BLOCK_SIZE):
-        X_offsets = i + tl.arange(0, BLOCK_SIZE)
-        X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf"))
-        X_block = X_block.to(tl.float32)
-        if label_smoothing > 0:
-            # scale X beforehand to avoid overflow
-            scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block, 0.0))
-        # Scale gradients based on reduction mode
-        # For reduce_loss=True: PyTorch will scale by 1/n_rows, so we need to scale by n_rows/n_non_ignore
-        # For reduce_loss=False: No additional scaling from PyTorch, so we don't scale here
-        if reduce_loss:
-            X_block = (tl.exp(X_block - m) / d - eps) / (n_non_ignore)
-        else:
-            X_block = tl.exp(X_block - m) / d - eps
-        tl.store(grad_input_ptr + X_offsets, X_block, mask=X_offsets < n_cols)
-
-    # Ensure the gradient is written before updating the target-token element.
-    tl.debug_barrier()
-
-    # 5. Calculate the loss
-
-    # loss = log (softmax(X_y)) = log ((e ^ (X_y - max(X)) / sum(e ^ (X - max(X))))
-    #      = (X_y - max(X)) - log(sum(e ^ (X - max(X))))
-    loss = -(ori_X_y - m - tl.log(d))
-
-    # Orginal loss = H(q, p),  with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps
-    # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p)
-    #          = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i))
-    # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as:
-    #          = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd))
-    # Refer to H(q', p) in section 7 of the paper: https://arxiv.org/pdf/1512.00567
-    if label_smoothing > 0:
-        smooth_loss = scaled_x_sum + label_smoothing * (m + tl.log(d))
-        loss = loss * (1 - label_smoothing) + smooth_loss
-
-    # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N`
-    vocab_start_idx = rank * n_cols
-    vocab_end_idx = (rank + 1) * n_cols
-    if y >= vocab_start_idx:
-        if y < vocab_end_idx:
-            X_y = tl.load(grad_input_ptr + y - vocab_start_idx)
-            # Apply the same conditional scaling logic for the target token
-            if reduce_loss:
-                X_y += -(1 - label_smoothing) / (n_non_ignore)
-            else:
-                X_y += -(1 - label_smoothing)
-            tl.store(grad_input_ptr + y - vocab_start_idx, X_y)
-
-    tl.store(loss_ptr, loss)
-
-
-@triton.jit
-def element_mul_kernel(
-    X_ptr,
-    X_stride,
-    grad_output_ptr,
-    grad_output_stride,
-    n_cols,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr.
-    The multiplication is performed in-place on the tensor pointed by X_ptr.
-
-    Parameters:
-    X_ptr: Pointer to the input tensor.
-    X_stride (int): The stride of the input tensor.
-    grad_output_ptr: Pointer to the gradient output value.
-    n_cols (int): The number of columns in the input tensor.
-    BLOCK_SIZE (int): The block size for Triton operations.
-    """
-
-    # Get the program ID and convert it to int64 to avoid overflow
-    program_id = tl.program_id(0).to(tl.int64)
-
-    # Locate the start index
-    X_ptr += program_id * X_stride
-
-    # Load the gradient output value
-    grad_output_ptr += program_id * grad_output_stride
-    grad_output = tl.load(grad_output_ptr)
-
-    # Perform the element-wise multiplication
-    for i in range(0, n_cols, BLOCK_SIZE):
-        X_offsets = i + tl.arange(0, BLOCK_SIZE)
-        X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols)
-        tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols)
-
-
-@triton.jit
-def cross_entropy_recompute_forward_kernel(
+def cross_entropy_forward_kernel(
     X_ptr,
     X_stride_0,
     X_stride_1,
@@ -334,7 +76,7 @@ def cross_entropy_recompute_forward_kernel(
 
 
 @triton.jit
-def cross_entropy_recompute_tp_pre_kernel(
+def cross_entropy_tp_pre_kernel(
     X_ptr,
     X_stride_0,
     X_stride_1,
@@ -392,7 +134,7 @@ def cross_entropy_recompute_tp_pre_kernel(
 
 
 @triton.jit
-def cross_entropy_recompute_tp_post_kernel(
+def cross_entropy_tp_post_kernel(
     gathered_data_ptr,
     Y_ptr,
     loss_ptr,
@@ -440,7 +182,7 @@ def cross_entropy_recompute_tp_post_kernel(
 
 
 @triton.jit
-def cross_entropy_recompute_backward_kernel(
+def cross_entropy_backward_kernel(
     logits_ptr,
     Y_ptr,
     stats_ptr,
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 17e19b659b..06db28ee27 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -69,10 +69,7 @@
 from transformer_engine.pytorch import ops
 from transformer_engine.pytorch import optimizers
 from transformer_engine.pytorch.export import onnx_export
-from transformer_engine.pytorch.cross_entropy import (
-    parallel_cross_entropy,
-    parallel_cross_entropy_recompute,
-)
+from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy
 from transformer_engine.pytorch.newton_schulz import (
     CusolverMpCtx,
     newton_schulz,
diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py
index 1a58a6592e..c8cddac0fa 100644
--- a/transformer_engine/pytorch/cross_entropy.py
+++ b/transformer_engine/pytorch/cross_entropy.py
@@ -11,18 +11,11 @@
 
 import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy
 
-__all__ = [
-    "parallel_cross_entropy",
-    "parallel_cross_entropy_recompute",
-]
+__all__ = ["parallel_cross_entropy"]
 
 
 class CrossEntropyFunction(torch.autograd.Function):
-    """
-    This class implements a custom autograd function for the Cross Entropy loss. The input
-    tensor can be in BF16/FP32, and loss and gradient calculations happen in FP32. The
-    returned loss is always in FP32.
-    """
+    """Cross entropy autograd function that recomputes its derivative."""
 
     @staticmethod
     def forward(
@@ -34,88 +27,16 @@ def forward(
         dist_process_group=None,
         ignore_idx=-100,
         is_cg_capturable=False,
+        overwrite_input=False,
     ):
-        """
-        The forward pass of the Cross Entropy loss. If dist_process_group is passed for distributed loss calculation, the input to each
-        distributed rank should be (*,V/world_size). Note that each of the ranks should get equal shards along the V dimension.
-
-        Parameters:
-        ctx : The context object.
-        inp (tensor): The input tensor of shape (B, SQ, V) or (SQ, B, V) where B is batch size, SQ is sequence length, V is vocab size.
-        target (tensor): The target tensor of shape (B,SQ) or (SQ, B) where each value is in [0, V-1].
-        label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
-        reduce_loss (bool): If true, returns the averaged loss across the B*SQ dimension.
-        dist_process_group (torch.dist.ProcessGroup): The distributed process group the loss computation is split across, None if on 1 device.
-        ignore_idx (int): The index for which loss and gradients are made to zero
-
-        Returns:
-        tensor: The computed loss.
-        """
-        loss, grad_input = triton_cross_entropy.cross_entropy_forward(
+        loss, logits, stats, target, n_non_ignore = triton_cross_entropy.cross_entropy_forward(
             inp,
             target,
             label_smoothing,
             reduce_loss,
             dist_process_group,
             ignore_idx,
-        )
-
-        ctx.save_for_backward(grad_input.detach())
-        ctx.is_cg_capturable = is_cg_capturable
-        return loss
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        """
-        The backward pass of the Cross Entropy loss.
-
-        Parameters:
-        ctx : The context object with saved tensors.
-        grad_output (tensor): The tensor containing the gradient of the loss with respect to the output.
-
-        Returns:
-        tuple: A tuple with the gradients with respect to the inputs. The elements are tensors or None.
-        """
-        (grad_input,) = ctx.saved_tensors
-        grad_input = triton_cross_entropy.cross_entropy_backward(
-            grad_input, grad_output, ctx.is_cg_capturable
-        )
-        return (
-            grad_input,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
-
-
-class CrossEntropyRecomputeFunction(torch.autograd.Function):
-    """Cross entropy autograd function that recomputes its derivative."""
-
-    @staticmethod
-    def forward(
-        ctx,
-        inp,
-        target,
-        label_smoothing=0.0,
-        reduce_loss=False,
-        dist_process_group=None,
-        ignore_idx=-100,
-        is_cg_capturable=False,
-        overwrite_input=False,
-    ):
-        loss, logits, stats, target, n_non_ignore = (
-            triton_cross_entropy.cross_entropy_recompute_forward(
-                inp,
-                target,
-                label_smoothing,
-                reduce_loss,
-                dist_process_group,
-                ignore_idx,
-                overwrite_input,
-            )
+            overwrite_input,
         )
         ctx.save_for_backward(logits.detach(), stats, target, n_non_ignore)
         ctx.label_smoothing = label_smoothing
@@ -130,12 +51,12 @@ def forward(
     def backward(ctx, grad_output):
         if ctx.did_backward:
             raise RuntimeError(
-                "parallel_cross_entropy_recompute does not support repeated backward passes "
+                "parallel_cross_entropy does not support repeated backward passes "
                 "because backward reuses its saved logits buffer"
             )
         ctx.did_backward = True
         logits, stats, target, n_non_ignore = ctx.saved_tensors
-        grad_input = triton_cross_entropy.cross_entropy_recompute_backward(
+        grad_input = triton_cross_entropy.cross_entropy_backward(
             logits,
             stats,
             target,
@@ -150,13 +71,13 @@ def backward(ctx, grad_output):
         return grad_input, None, None, None, None, None, None, None
 
 
-def _validate_recompute_inputs(
+def _validate_inputs(
     inp: torch.Tensor,
     target: torch.Tensor,
     label_smoothing: float,
     overwrite_input: bool,
 ) -> None:
-    """Validate assumptions made by the recompute Triton kernels."""
+    """Validate assumptions made by the cross entropy Triton kernels."""
 
     if inp.ndim != 3:
         raise ValueError(f"inp must be a 3D tensor, but got shape {tuple(inp.shape)}")
@@ -184,69 +105,6 @@ def _validate_recompute_inputs(
 
 
 def parallel_cross_entropy(
-    inp: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float = 0.0,
-    reduce_loss: bool = False,
-    dist_process_group: Optional[torch.distributed.ProcessGroup] = None,
-    ignore_idx: int = -100,
-    is_cg_capturable: bool = False,
-    *,
-    _input: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """
-    Cross Entropy loss with optional distributed reduction.
-
-    The input tensor can be in BF16/FP32, and loss and gradient calculations happen in
-    FP32. The returned loss is always in FP32.
-
-    If ``dist_process_group`` is passed for distributed loss calculation, the input to each
-    distributed rank should be ``(*, V/world_size)``. Note that each of the ranks should
-    get equal shards along the V dimension.
-
-    Parameters
-    ----------
-    inp : torch.Tensor
-        The input tensor of shape ``(B, SQ, V)`` or ``(SQ, B, V)`` where B is batch size,
-        SQ is sequence length, V is vocab size.
-    target : torch.Tensor
-        The target tensor of shape ``(B, SQ)`` or ``(SQ, B)`` where each value is in ``[0, V-1]``.
-    label_smoothing : float, default = 0.0
-        The amount of smoothing when computing the loss, where 0.0 means no smoothing.
-    reduce_loss : bool, default = False
-        If True, returns the averaged loss across the B*SQ dimension.
-    dist_process_group : torch.distributed.ProcessGroup, default = None
-        The distributed process group the loss computation is split across, None if on 1 device.
-    ignore_idx : int, default = -100
-        The index for which loss and gradients are made to zero.
-    is_cg_capturable : bool, default = False
-        Whether the operation is CUDA graph capturable.
-
-    Returns
-    -------
-    torch.Tensor
-        The computed loss.
-    """
-    # Handle backward compatibility with _input parameter
-    if _input is not None:
-        warnings.warn(
-            "The '_input' parameter is deprecated. Please use 'inp' instead.",
-            FutureWarning,
-        )
-        inp = _input
-
-    return CrossEntropyFunction.apply(
-        inp,
-        target,
-        label_smoothing,
-        reduce_loss,
-        dist_process_group,
-        ignore_idx,
-        is_cg_capturable,
-    )
-
-
-def parallel_cross_entropy_recompute(
     inp: torch.Tensor,
     target: torch.Tensor,
     label_smoothing: float = 0.0,
@@ -256,8 +114,9 @@ def parallel_cross_entropy_recompute(
     is_cg_capturable: bool = False,
     *,
     overwrite_input: bool = False,
+    _input: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """Memory-efficient cross entropy with derivative recomputation.
+    """Cross entropy loss with optional distributed reduction.
 
     Loss and derivative calculations use FP32 arithmetic for BF16 and FP32
     inputs. Instead of saving a full FP32 derivative, this function saves an
@@ -274,11 +133,22 @@ def parallel_cross_entropy_recompute(
     The saved logits buffer is consumed by backward, so repeated backward
     passes on the same result are not supported.
 
-    Parameters are the same as :func:`parallel_cross_entropy`, with the
-    addition of ``overwrite_input``.
-
     Parameters
     ----------
+    inp : torch.Tensor
+        Input logits with shape ``(B, SQ, V)`` or ``(SQ, B, V)``.
+    target : torch.Tensor
+        Target token IDs with one value for each input row.
+    label_smoothing : float, default = 0.0
+        Amount of label smoothing.
+    reduce_loss : bool, default = False
+        Return the mean loss over non-ignored targets when True.
+    dist_process_group : torch.distributed.ProcessGroup, default = None
+        Tensor-parallel process group, or None for a single device.
+    ignore_idx : int, default = -100
+        Target value for ignored rows.
+    is_cg_capturable : bool, default = False
+        Whether the operation is CUDA graph capturable.
     overwrite_input : bool, default = False
         Reuse and overwrite ``inp`` rather than allocating a private logits
         copy. The input must be contiguous and have no internal overlap.
@@ -289,8 +159,15 @@ def parallel_cross_entropy_recompute(
         The computed loss.
     """
 
-    _validate_recompute_inputs(inp, target, label_smoothing, overwrite_input)
-    return CrossEntropyRecomputeFunction.apply(
+    if _input is not None:
+        warnings.warn(
+            "The '_input' parameter is deprecated. Please use 'inp' instead.",
+            FutureWarning,
+        )
+        inp = _input
+
+    _validate_inputs(inp, target, label_smoothing, overwrite_input)
+    return CrossEntropyFunction.apply(
         inp,
         target,
         label_smoothing,
diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py
index bcc913b976..8a4c455c8c 100644
--- a/transformer_engine/pytorch/triton/cross_entropy.py
+++ b/transformer_engine/pytorch/triton/cross_entropy.py
@@ -5,22 +5,16 @@
 """PyTorch wrapper functions for Cross Entropy Triton kernels."""
 
 from typing import Union
-from functools import reduce
-from operator import mul
-
 import torch
 import torch.distributed as dist
 
 import triton
 
 from transformer_engine.common.triton.cross_entropy import (
-    online_softmax_kernel,
-    cross_entropy_kernel,
-    element_mul_kernel,
-    cross_entropy_recompute_forward_kernel,
-    cross_entropy_recompute_tp_pre_kernel,
-    cross_entropy_recompute_tp_post_kernel,
-    cross_entropy_recompute_backward_kernel,
+    cross_entropy_forward_kernel,
+    cross_entropy_tp_pre_kernel,
+    cross_entropy_tp_post_kernel,
+    cross_entropy_backward_kernel,
 )
 
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
@@ -34,128 +28,6 @@ def cross_entropy_forward(
     reduce_loss: bool,
     dist_process_group: Union[dist.ProcessGroup, None],
     ignore_idx: int,
-):
-    """Forward implementation of Cross Entropy kernel"""
-
-    B, SQ, V = _input.shape
-    n_rows = B * SQ
-
-    assert reduce(mul, list(target.size())) == (B * SQ), "Each token needs a target token ID."
-
-    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-
-    # unreduced loss
-    loss_1d = torch.zeros(n_rows, dtype=torch.float32, device=_input.device)
-
-    # tensor to hold this rank's m/d/X_y values
-    m_d_X_y = torch.zeros(n_rows * 3, dtype=torch.float32, device=_input.device)
-
-    n_non_ignore = torch.zeros(1, dtype=torch.int64, device=_input.device)
-
-    # ensure _input and target are contiguous in the last dimension
-    if _input.stride(-1) != 1 or _input.stride(-2) != _input.shape[-1]:
-        _input = _input.contiguous()
-    if target.stride(-1) != 1:
-        target = target.contiguous()
-
-    # Store the input gradient in FP32 so it is not quantized before backward.
-    grad_input = torch.empty_like(_input, dtype=torch.float32)
-
-    rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group)
-
-    online_softmax_kernel[(n_rows,)](
-        X_ptr=_input,
-        X_stride=_input.stride(-2),
-        Y_ptr=target,
-        Y_stride=target.stride(-1),  # always 1
-        m_d_X_y_ptr=m_d_X_y,
-        m_d_X_y_stride=m_d_X_y.stride(-1),
-        rank=rank,
-        n_cols=V,
-        ignore_idx=ignore_idx,
-        n_non_ignore=n_non_ignore,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=32,
-    )
-
-    n_non_ignore = torch.clamp(n_non_ignore, min=1)
-
-    world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)
-
-    if world_size > 1:
-        m_d_X_y_gathered = torch.zeros(
-            n_rows * 3 * world_size, dtype=torch.float32, device=_input.device
-        )
-        dist.all_gather_into_tensor(m_d_X_y_gathered, m_d_X_y, group=dist_process_group)
-    else:
-        m_d_X_y_gathered = m_d_X_y
-
-    cross_entropy_kernel[(n_rows,)](
-        X_ptr=_input,
-        X_stride=_input.stride(-2),
-        grad_input_ptr=grad_input,
-        grad_input_stride=grad_input.stride(-2),
-        Y_ptr=target,
-        Y_stride=target.stride(-1),
-        loss_ptr=loss_1d,
-        loss_stride=loss_1d.stride(-1),
-        m_d_X_y_ptr=m_d_X_y_gathered,
-        m_d_X_y_stride=m_d_X_y_gathered.stride(-1),
-        rank=rank,
-        world_size=world_size,
-        ignore_idx=ignore_idx,
-        n_cols=V,
-        n_rows=n_rows,
-        n_non_ignore=n_non_ignore,
-        reduce_loss=reduce_loss,
-        label_smoothing=label_smoothing,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=32,
-    )
-
-    loss = (
-        torch.reshape(loss_1d, (B, SQ)) if not reduce_loss else (torch.sum(loss_1d) / n_non_ignore)
-    )
-
-    return loss, grad_input
-
-
-def cross_entropy_backward(
-    grad_input: torch.Tensor, grad_output: torch.Tensor, is_cg_capturable: bool = False
-):
-    """Backward implementation of cross entropy loss kernel"""
-
-    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
-    # Only check torch.equal when not in CUDA graph capturable mode
-    if not is_cg_capturable and torch.equal(
-        grad_output, torch.tensor(1.0, device=grad_output.device)
-    ):
-        pass
-    else:
-        B, SQ, V = grad_input.shape
-        n_rows = B * SQ
-        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-
-        element_mul_kernel[(n_rows,)](
-            grad_input,
-            grad_input.stride(-2),
-            grad_output.contiguous(),
-            1 if grad_output.numel() > 1 else 0,
-            V,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=32,
-        )
-
-    return grad_input
-
-
-def cross_entropy_recompute_forward(
-    _input: torch.Tensor,
-    target: torch.Tensor,
-    label_smoothing: float,
-    reduce_loss: bool,
-    dist_process_group: Union[dist.ProcessGroup, None],
-    ignore_idx: int,
     overwrite_input: bool,
 ):
     """Forward implementation that saves logits and compact softmax statistics."""
@@ -178,7 +50,7 @@ def cross_entropy_recompute_forward(
     world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)
 
     if world_size == 1:
-        cross_entropy_recompute_forward_kernel[(n_rows,)](
+        cross_entropy_forward_kernel[(n_rows,)](
             X_ptr=_input,
             X_stride_0=_input.stride(0),
             X_stride_1=_input.stride(1),
@@ -198,7 +70,7 @@ def cross_entropy_recompute_forward(
         )
     else:
         local_data = torch.empty((n_rows, 4), dtype=torch.float32, device=_input.device)
-        cross_entropy_recompute_tp_pre_kernel[(n_rows,)](
+        cross_entropy_tp_pre_kernel[(n_rows,)](
             X_ptr=_input,
             X_stride_0=_input.stride(0),
             X_stride_1=_input.stride(1),
@@ -219,7 +91,7 @@ def cross_entropy_recompute_forward(
             (world_size * n_rows, 4), dtype=torch.float32, device=_input.device
         )
         dist.all_gather_into_tensor(gathered_data, local_data, group=dist_process_group)
-        cross_entropy_recompute_tp_post_kernel[(n_rows,)](
+        cross_entropy_tp_post_kernel[(n_rows,)](
             gathered_data_ptr=gathered_data,
             Y_ptr=target,
             loss_ptr=loss_1d,
@@ -240,7 +112,7 @@ def cross_entropy_recompute_forward(
     return loss, logits, stats, target, n_non_ignore
 
 
-def cross_entropy_recompute_backward(
+def cross_entropy_backward(
     logits: torch.Tensor,
     stats: torch.Tensor,
     target: torch.Tensor,
@@ -262,7 +134,7 @@ def cross_entropy_recompute_backward(
     world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)
     grad_output = grad_output.contiguous()
 
-    cross_entropy_recompute_backward_kernel[(n_rows,)](
+    cross_entropy_backward_kernel[(n_rows,)](
         logits_ptr=logits,
         Y_ptr=target,
         stats_ptr=stats,