import torch
import triton

from packaging.version import Version

# Trigger declaration of the ``cross_entropy_loss_and_grad`` op location so the
# dispatcher knows where to discover the Triton + CuTe DSL impls. Without this
# import the first dispatch call would return "no impl registered".
import liger_kernel.functional  # noqa: F401

from liger_kernel.backends import dispatch
from liger_kernel.ops.cross_entropy import liger_cross_entropy_kernel
from liger_kernel.ops.utils import amp_custom_bwd
from liger_kernel.ops.utils import amp_custom_fwd
from liger_kernel.ops.utils import element_mul_kernel
from liger_kernel.ops.utils import is_hip
from liger_kernel.utils import infer_device

# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
# The optimal maximum block size depends on your hardware, your kernel, and your dtype
MAX_FUSED_SIZE = 2048 if infer_device() == "npu" else 65536 // 2

# Public callers default to the historical minimum transient-logits budget:
# approximately BT x H elements regardless of vocabulary size. Integrations
# with stricter reduction-order requirements may explicitly request more.
_CHUNK_MEM_CONST = 1
_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,
    target,
    ce_weight=None,
    bias=None,
    ignore_index=-100,
    lse_square_scale=0.0,
    label_smoothing=0.0,
    reduction="mean",
    softcap=None,
    return_z_loss=False,
    accum_dtype=None,
    use_token_scaling=False,
    return_token_accuracy=False,
    return_predicted_tokens=False,
    ce_impl=None,
    ce_mode=None,
    token_grad_output=None,
    compute_gradients=None,
    weight_requires_grad=None,
    chunk_mem_const: int = _CHUNK_MEM_CONST,
):
    assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}"
    assert isinstance(return_token_accuracy, bool), (
        f"return_token_accuracy must be True or False. Got: {return_token_accuracy}"
    )
    assert isinstance(return_predicted_tokens, bool), (
        f"return_predicted_tokens must be True or False. Got: {return_predicted_tokens}"
    )
    if isinstance(chunk_mem_const, bool) or not isinstance(chunk_mem_const, int) or chunk_mem_const < 1:
        raise ValueError(f"chunk_mem_const must be a positive integer. Got: {chunk_mem_const!r}")
    device = _input.device
    # ``compute_gradients`` / ``weight_requires_grad`` let backward re-enter this
    # function with *detached* saved tensors (whose ``.requires_grad`` is False) and
    # still get gradients out. ``token_grad_output`` is the per-token upstream
    # gradient for reduction="none"; it must be folded into the logits gradient
    # before that gradient is summed into grad_weight/grad_bias.
    input_requires_grad = _input.requires_grad if compute_gradients is None else compute_gradients
    weight_needs_grad = weight.requires_grad if weight_requires_grad is None else weight_requires_grad

    # inputs have shape: BT x H
    # materialized activations will have shape: BT x V
    # the increase in memory = BT x V
    # reduction can be achieved by partitioning the number of tokens BT into smaller chunks.
    # for ex: if we were to achieve the same memory consumption as BT x H, then the chunk size should be:
    # inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor
    # for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048
    BT, H = _input.shape
    V = weight.shape[0]

    # Bound transient logits to C x BT x H.
    inc_factor = triton.cdiv(V, chunk_mem_const * H)
    chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))  # (BT + inc_factor - 1) // inc_factor
    chunk_size = min(chunk_size, BT)  # a single chunk covers BT when the budget allows; never exceed BT
    num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size

    grad_input = torch.zeros_like(_input, device=device)

    # loss_1d is always fp32. The weight/bias grad accumulators are NOT always fp32:
    # with accum_dtype=None they inherit the parameter dtype (e.g. bf16/fp16), and only
    # take accum_dtype when one is explicitly requested. The weight-grad projection below
    # therefore has a dtype-matched low-precision path in addition to the fp32 fast path.
    if input_requires_grad:
        if accum_dtype is None:
            grad_weight = torch.zeros_like(weight, device=device) if weight_needs_grad else None
            grad_bias = torch.zeros_like(bias, device=device) if bias is not None else None
        else:
            grad_weight = torch.zeros_like(weight, dtype=accum_dtype, device=device) if weight_needs_grad else None
            grad_bias = torch.zeros_like(bias, dtype=accum_dtype, device=device) if bias is not None else None
    else:
        grad_weight = None
        grad_bias = None

    loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
    z_loss_1d = torch.zeros(BT, dtype=_input.dtype, device=_input.device) if return_z_loss else None
    token_accuracy_1d = torch.zeros(BT, dtype=torch.float32, device=device) if return_token_accuracy else None
    predicted_tokens_1d = torch.full((BT,), -1, dtype=torch.int64, device=device) if return_predicted_tokens else None

    # TODO: evaluate how CUDA synchronization caused by .item() affects the speed
    target_mask = target != ignore_index
    total_n_non_ignore = target_mask.sum().item()
    total_sum_non_ignore_ce_weight = total_n_non_ignore
    ce_weight_sum = 0.0
    if ce_weight is not None:
        assert ce_weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {ce_weight.shape}"
        assert torch.is_floating_point(ce_weight), (
            f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
        )
        total_sum_non_ignore_ce_weight = (
            torch.gather(ce_weight, dim=0, index=target.masked_select(target_mask)).sum().item()
        )
        ce_weight_sum = ce_weight.sum().item()
        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
        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,

        # Compute predicted probabilities for token scaling if needed
        if use_token_scaling:
            # Compute softmax probabilities for scaling
            # We need to compute this before the cross entropy kernel modifies logits_chunk
            logits_for_softmax = logits_chunk.detach().clone()  # Detach to avoid gradient flow
            if softcap is not None:
                logits_for_softmax = softcap * torch.tanh(logits_for_softmax / softcap)

            # 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
            valid_targets = target_chunk[valid_target_mask]

            if len(valid_targets) > 0:
                # 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()
        target_chunk = target_chunk.contiguous()

        if ce_impl is None:
            loss_1d_slice = loss_1d[start_idx:end_idx]
            z_loss_1d_slice = z_loss_1d[start_idx:end_idx] if return_z_loss else None
            token_accuracy_1d_slice = token_accuracy_1d[start_idx:end_idx] if return_token_accuracy else None
            predicted_tokens_1d_slice = predicted_tokens_1d[start_idx:end_idx] if return_predicted_tokens else None
            block_size = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))

            liger_cross_entropy_kernel[(end_idx - start_idx,)](
                X_ptr=logits_chunk,
                X_stride=logits_chunk.stride(-2),
                Y_ptr=target_chunk,
                Y_stride=target_chunk.stride(-1),
                weight_ptr=ce_weight,
                loss_ptr=loss_1d_slice,
                z_loss_ptr=z_loss_1d_slice,
                loss_stride=loss_1d_slice.stride(-1),
                token_accuracy_ptr=token_accuracy_1d_slice,
                token_accuracy_stride=token_accuracy_1d_slice.stride(-1) if return_token_accuracy else 0,
                predicted_tokens_ptr=predicted_tokens_1d_slice,
                predicted_tokens_stride=predicted_tokens_1d_slice.stride(-1) if return_predicted_tokens else 0,
                n_cols=V,
                n_non_ignore=total_n_non_ignore,
                sum_non_ignore_weight=total_sum_non_ignore_ce_weight,
                weight_sum=ce_weight_sum,
                ignore_index=ignore_index,
                lse_square_scale=lse_square_scale,
                label_smoothing=label_smoothing,
                reduction=reduction,
                softcap=softcap,
                RETURN_Z_LOSS=return_z_loss,
                RETURN_TOKEN_ACCURACY=return_token_accuracy,
                RETURN_PREDICTED_TOKENS=return_predicted_tokens,
                HAS_WEIGHT=ce_weight is not None,
                HAS_SOFTCAPPING=softcap is not None,
                HAS_GRADIENTS=input_requires_grad,
                BLOCK_SIZE=block_size,
                num_warps=32 if not is_hip() else 16,
            )
            grad_logits_chunk = logits_chunk
        else:
            ce_args = (
                logits_chunk,
                target_chunk,
                ce_weight,
                ignore_index,
                lse_square_scale,
                label_smoothing,
                reduction,
                softcap,
                total_n_non_ignore,
                total_sum_non_ignore_ce_weight,
                ce_weight_sum,
                return_z_loss,
                return_token_accuracy,
                return_predicted_tokens,
                input_requires_grad,
            )
            (
                loss_1d_slice,
                z_loss_1d_slice,
                token_accuracy_1d_slice,
                predicted_tokens_1d_slice,
                grad_logits_chunk,
            ) = dispatch(
                "cross_entropy_loss_and_grad",
                *ce_args,
                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:
            loss_1d_slice = loss_1d_slice * scaling_factors
            if return_z_loss:
                z_loss_1d_slice = z_loss_1d_slice * scaling_factors

        loss_1d[start_idx:end_idx] = loss_1d_slice
        if return_z_loss:
            z_loss_1d[start_idx:end_idx] = z_loss_1d_slice
        if return_token_accuracy:
            token_accuracy_1d[start_idx:end_idx] = token_accuracy_1d_slice
        if return_predicted_tokens:
            predicted_tokens_1d[start_idx:end_idx] = predicted_tokens_1d_slice

        # Apply token scaling to gradients if requested
        if use_token_scaling:
            # 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
        # sum_i go_i * (g_i outer x_i); once the sum has happened the per-token
        # weights can no longer be recovered, so scaling grad_weight afterwards
        # (as element_mul_kernel does) is mathematically incapable of being correct.
        if token_grad_output is not None:
            grad_logits_chunk = grad_logits_chunk * token_grad_output[start_idx:end_idx].unsqueeze(-1).to(
                grad_logits_chunk.dtype
            )

        if input_requires_grad:
            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()
            if (
                _ADDMM_SUPPORTS_OUT_DTYPE
                and grad_weight.device.type == "cuda"
                and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
                and grad_weight.dtype == torch.float32
                and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
            ):
                # FP32 accumulator (accum_dtype=torch.float32, or fp32 params under AMP).
                # Unlike torch.mm, torch.addmm's out_dtype path does not participate in
                # autocast operand casting, so under AMP (fp32 params, no bias) _input_chunk
                # can stay fp32 while grad_logits is the autocast dtype. addmm requires mat1
                # and mat2 to share a dtype, so align _input_chunk before accumulating.
                input_chunk = _input_chunk
                if input_chunk.dtype != grad_logits_t.dtype:
                    input_chunk = input_chunk.to(grad_logits_t.dtype)
                torch.addmm(
                    grad_weight,
                    grad_logits_t,
                    input_chunk,
                    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
                and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
                and grad_weight.dtype == grad_logits_t.dtype
            ):
                # Low-precision accumulator (accum_dtype=None with bf16/fp16 params) whose
                # dtype already matches grad_logits. Accumulate straight into grad_weight with
                # addmm(out=grad_weight); this mirrors the CuTe backend's direct bf16 addmm and
                # avoids the legacy path's parameter-sized bf16->fp32 temporary + cast per chunk.
                # In-place out ops are not autocast, so -- as torch.mm's autocast used to do --
                # align _input_chunk (which may be promoted fp32 under bf16 AMP) to grad_logits.
                input_chunk = _input_chunk
                if input_chunk.dtype != grad_logits_t.dtype:
                    input_chunk = input_chunk.to(grad_logits_t.dtype)
                torch.addmm(
                    grad_weight,
                    grad_logits_t,
                    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(
                input=grad_bias,
                other=grad_logits_chunk.sum(dim=0),
                out=grad_bias,
                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
    #     z_loss = z_loss_1d if return_z_loss else None

    if reduction == "none":
        # Return per-token losses
        loss = loss_1d
        z_loss = z_loss_1d if return_z_loss else None
        token_accuracy = token_accuracy_1d if return_token_accuracy else None
    else:
        loss = torch.sum(loss_1d)
        z_loss = torch.sum(z_loss_1d) if return_z_loss else None
        # For accuracy, we compute the mean across all non-ignored tokens
        token_accuracy = torch.sum(token_accuracy_1d) / total_n_non_ignore if return_token_accuracy else None

    predicted_tokens = predicted_tokens_1d if return_predicted_tokens else None

    # Cast back to original dtype
    grad_weight = grad_weight.to(weight.dtype) if grad_weight is not None else None
    grad_bias = grad_bias.to(bias.dtype) if grad_bias is not None else None

    return loss, z_loss, token_accuracy, predicted_tokens, grad_input, grad_weight, grad_bias


def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias):
    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
    if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
        BT, H = grad_input.shape
        n_rows = BT
        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))

        element_mul_kernel[(n_rows,)](
            grad_input,
            grad_input.stride(-2),
            grad_output,
            H,
            BLOCK_SIZE=BLOCK_SIZE,
            X_col_stride=grad_input.stride(-1),
            num_warps=32 if not is_hip() else 16,
        )

        # handle grad_weight
        if grad_weight is not None:
            V, H = grad_weight.shape
            n_rows = V

            element_mul_kernel[(n_rows,)](
                grad_weight,
                grad_weight.stride(-2),
                grad_output,
                H,
                BLOCK_SIZE=BLOCK_SIZE,
                X_col_stride=grad_weight.stride(-1),
                num_warps=32 if not is_hip() else 16,
            )

        if grad_bias is not None:
            V = grad_bias.shape[0]
            n_rows = V

            element_mul_kernel[(n_rows,)](
                grad_bias,
                grad_bias.stride(-1),
                grad_output,
                1,
                BLOCK_SIZE=BLOCK_SIZE,
                num_warps=32 if not is_hip() else 16,
            )
    return grad_input, grad_weight, grad_bias


class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
    supports_inner_impl_dispatch = True

    @staticmethod
    @amp_custom_fwd
    def forward(
        ctx,
        _input,
        weight,
        target,
        bias=None,
        ce_weight=None,
        ignore_index=-100,
        lse_square_scale=0.0,
        label_smoothing=0.0,
        reduction="mean",
        softcap=None,
        return_z_loss: bool = False,
        accum_dtype=None,
        use_token_scaling: bool = False,
        return_token_accuracy: bool = False,
        return_predicted_tokens: bool = False,
        ce_impl=None,
        ce_mode=None,
        chunk_mem_const: int = _CHUNK_MEM_CONST,
    ):
        """
        Fusing the last linear layer with cross-entropy loss
            Reference: https://github.com/mgmalek/efficient_cross_entropy

        Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding
        the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can
        compute the gradient at the forward pass. By doing so, we don't have to store the _input and target
        for the backward pass.

        _input: (B*T, H) where B is batch size, T is sequence length, H is hidden dimension.
        target: (B*T) where each value is in [0, V-1]
        weight: (V, H) where V is the number of classes
        bias: (V) where V is the number of classes
        ce_weight: a manual rescaling weight given to each class. If given, has to be a Tensor of size V and floating point dtype
        ignore_index: the index to ignore in the target
        label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
        reduction: reduction to apply
        accum_dtype (torch.dtype): the dtype of intermediate result buffers for weight and bias gradient accumulations.
            Recommended to set `accum_dtype` to higher precision, e.g. `torch.float32`, if the training is unstable with original dtype. Default: `None`, performing accumulations in original dtype
        use_token_scaling (bool): whether to scale each token's loss by its predicted probability (detached).
            When True, each token's loss is multiplied by the model's predicted probability for that token's true class.
            Default: False.
        return_token_accuracy (bool): When `return_token_accuracy` is `True`, computes and returns per-token accuracy without materializing logits. Default: `False`
        return_predicted_tokens (bool): When `return_predicted_tokens` is `True`, returns per-token predicted class indices (argmax) without materializing logits. Default: `False`
        chunk_mem_const (int): transient-logits budget multiplier. Defaults to
            `1`; larger values trade memory for fewer chunk reductions.
        """

        # With reduction="none" the loss is per-token, so backward receives a
        # per-token grad_output that must be folded into the logits gradient BEFORE
        # it is summed into grad_weight/grad_bias. That weight is not known yet here,
        # so defer the gradient computation to backward (recomputing the chunked
        # logits there) instead of producing an unscaled gradient we could never
        # correctly rescale. mean/sum keep the original compute-in-forward fast path.
        needs_grad = _input.requires_grad or weight.requires_grad or (bias is not None and bias.requires_grad)
        ctx.defer_grads = reduction == "none" and needs_grad

        loss, z_loss, token_accuracy, predicted_tokens, grad_input, grad_weight, grad_bias = (
            fused_linear_cross_entropy_forward(
                _input=_input,
                weight=weight,
                target=target,
                bias=bias,
                ce_weight=ce_weight,
                ignore_index=ignore_index,
                lse_square_scale=lse_square_scale,
                label_smoothing=label_smoothing,
                reduction=reduction,
                softcap=softcap,
                return_z_loss=return_z_loss,
                accum_dtype=accum_dtype,
                use_token_scaling=use_token_scaling,
                return_token_accuracy=return_token_accuracy,
                return_predicted_tokens=return_predicted_tokens,
                ce_impl=ce_impl,
                ce_mode=ce_mode,
                compute_gradients=False if ctx.defer_grads else None,
                chunk_mem_const=chunk_mem_const,
            )
        )

        if ctx.defer_grads:
            ctx.save_for_backward(_input.detach(), weight.detach(), target, bias.detach() if bias is not None else None)
            ctx.fwd_kwargs = dict(
                ce_weight=ce_weight,
                ignore_index=ignore_index,
                lse_square_scale=lse_square_scale,
                label_smoothing=label_smoothing,
                reduction=reduction,
                softcap=softcap,
                accum_dtype=accum_dtype,
                use_token_scaling=use_token_scaling,
                # Preserve the linkedin-managed multi-backend dispatch selection when the
                # reduction="none" path re-enters forward from backward to recompute grads.
                ce_impl=ce_impl,
                ce_mode=ce_mode,
                chunk_mem_const=chunk_mem_const,
            )
            ctx.weight_requires_grad = weight.requires_grad
        else:
            # downcast to dtype and store for backward
            ctx.save_for_backward(
                grad_input.detach(),
                grad_weight.detach() if grad_weight is not None else None,
                grad_bias.detach() if grad_bias is not None else None,
            )
        ctx.return_z_loss = return_z_loss
        ctx.return_token_accuracy = return_token_accuracy
        ctx.return_predicted_tokens = return_predicted_tokens
        return loss, z_loss, token_accuracy, predicted_tokens

    @staticmethod
    @amp_custom_bwd
    def backward(ctx, grad_output, grad_output2, grad_output3, grad_output4):
        if ctx.return_z_loss:
            del grad_output2  # z_loss is only for logging
        if ctx.return_token_accuracy:
            del grad_output3  # token_accuracy is only for metrics
        if ctx.return_predicted_tokens:
            del grad_output4  # predicted_tokens is only for metrics

        if ctx.defer_grads:
            # reduction="none": grad_output is per-token. Recompute the chunked
            # logits gradient and fold grad_output into each row before projecting,
            # which is the only point at which the per-token weights can be applied.
            _input, weight, target, bias = ctx.saved_tensors
            _, _, _, _, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
                _input=_input,
                weight=weight,
                target=target,
                bias=bias,
                token_grad_output=grad_output,
                compute_gradients=True,
                weight_requires_grad=ctx.weight_requires_grad,
                **ctx.fwd_kwargs,
            )
        else:
            (grad_input, grad_weight, grad_bias) = ctx.saved_tensors
            grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_backward(
                grad_output, grad_input, grad_weight, grad_bias
            )
        gradients = (
            grad_input,
            grad_weight,
            None,
            grad_bias,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,  # use_token_scaling
            None,  # return_token_accuracy
            None,  # return_predicted_tokens
            None,  # ce_impl
            None,  # ce_mode
            None,  # chunk_mem_const
        )
        return gradients[: len(ctx.needs_input_grad)]
