# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""GroupedLinear API"""

from typing import Union, Optional, Callable, Tuple, List, Sequence
from itertools import chain
import os
import warnings
import weakref

import functools
import torch

import transformer_engine_torch as tex

from transformer_engine.common.recipe import Recipe
from transformer_engine.pytorch.tensor.grouped_tensor import (
    GroupedTensor,
    GroupedTensorStorage,
)
from .base import (
    get_dummy_wgrad,
    quantize_weight,
    TransformerEngineBaseModule,
    _2X_ACC_FPROP,
    _2X_ACC_DGRAD,
    _2X_ACC_WGRAD,
    _attach_high_precision_init_val,
    _clear_high_precision_init_val,
    _get_high_precision_init_val,
)
from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore
from ..quantization import FP8GlobalStateManager, QuantizerRole
from ..utils import (
    divide,
    cast_if_needed,
    clear_tensor_data,
    get_device_compute_capability,
    init_method_constant,
    requires_grad,
    resolve_grouped_linear_single_param_flags,
    get_nvtx_range_context,
)
from ..distributed import (
    set_tensor_model_parallel_attributes,
    get_distributed_world_size,
    is_fp8_activation_recompute_enabled,
    in_fp8_activation_recompute_phase,
)
from ..distributed_weight import (
    is_distributed_weight,
    materialize_weight_for_forward,
    materialize_weight_for_backward,
    finalize_weight_grads,
)
from ..cpp_extensions import (
    general_grouped_gemm,
    general_grouped_gemm_for_grouped_tensor,
)
from ..constants import GemmParallelModes, dist_group_type
from ..jit import no_torch_dynamo
from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload
from ..triton.grouped_dbias_dscales import compute_grouped_dbias

from ..tensor import (
    Float8BlockQuantizer,
    Float8CurrentScalingQuantizer,
    Float8Quantizer,
    HybridQuantizer,
    IdentityQuantizer,
    MXFP8Quantizer,
    NVFP4Quantizer,
)
from ..quantized_tensor import (
    QuantizedTensorStorage,
    Quantizer,
    prepare_for_saving,
    restore_from_func_ctx,
)
from ...debug.pytorch.debug_quantization import DebugQuantizer
from ...debug.pytorch.debug_state import TEDebugState


_NATIVE_SPLIT_QUANTIZER_TYPES = frozenset(
    {
        Float8Quantizer,
        Float8CurrentScalingQuantizer,
        Float8BlockQuantizer,
        MXFP8Quantizer,
        NVFP4Quantizer,
    }
)


def _supports_native_split_quantize(quantizer):
    """Whether ``tex.split_quantize`` has an exact converter for this quantizer."""
    return type(quantizer) in _NATIVE_SPLIT_QUANTIZER_TYPES


def _uses_identity_quantizer(quantizer):
    """Whether a quantizer, including a hybrid sub-quantizer, is Identity-backed."""
    if quantizer is None:
        return False
    if isinstance(quantizer, IdentityQuantizer):
        return True
    if isinstance(quantizer, HybridQuantizer):
        return _uses_identity_quantizer(quantizer.rowwise_quantizer) or _uses_identity_quantizer(
            quantizer.columnwise_quantizer
        )
    return False


def _identity_quantizer_signature(quantizer):
    """Identity usage per GEMM direction: (rowwise, columnwise)."""
    if isinstance(quantizer, HybridQuantizer):
        return (
            _uses_identity_quantizer(quantizer.rowwise_quantizer),
            _uses_identity_quantizer(quantizer.columnwise_quantizer),
        )
    identity = isinstance(quantizer, IdentityQuantizer)
    return (identity, identity)


_DYNAMIC_QUANTIZER_SIGNATURE_FIELDS = frozenset(
    {
        "rowwise_usage",
        "columnwise_usage",
        "internal",
        "optimize_for_gemm",
    }
)


def _backend_quantizer_signature(quantizer):
    """Return backend configuration that grouped kernels require to be uniform."""
    if quantizer is None:
        return None

    # Identity is not registered as a torch.compile value quantizer, but its
    # dtype changes the grouped GEMM input type and therefore must be uniform.
    if isinstance(quantizer, IdentityQuantizer):
        return (type(quantizer), (("dtype", quantizer.dtype),))

    fields = quantizer._value_fields()
    if fields is None:
        # Delayed-scaling Float8Quantizer carries per-expert scale/amax tensors,
        # which are intentionally different, but its emitted FP8 dtype is a
        # group-wide backend choice. Other unregistered/custom quantizers retain
        # the conservative exact-family behavior until they expose value fields.
        fields = ("dtype",) if isinstance(quantizer, Float8Quantizer) else ()

    config = []
    for name in fields:
        if name in _DYNAMIC_QUANTIZER_SIGNATURE_FIELDS:
            continue
        value = getattr(quantizer, name)
        if name == "dtype":
            value = int(value)
        config.append((name, value))
    return (type(quantizer), tuple(config))


def _validate_backend_match(reference, quantizer, operand_name, direction, expert_index):
    """Validate one expert against the group's reference backend."""
    if type(quantizer) is not type(reference):
        raise ValueError(
            f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend"
            f" families across experts: expert 0 uses {type(reference).__name__}, but expert"
            f" {expert_index} uses {type(quantizer).__name__}. Grouped operands require one"
            " quantizer family per direction."
        )
    reference_signature = _backend_quantizer_signature(reference)
    quantizer_signature = _backend_quantizer_signature(quantizer)
    if quantizer_signature != reference_signature:
        raise ValueError(
            f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend"
            f" configurations across experts: expert 0 uses {reference_signature}, but expert"
            f" {expert_index} uses {quantizer_signature}. Grouped operands require the same"
            " backend-relevant configuration per direction."
        )


def _validate_grouped_quantizer_list(quantizers, *, operand_name="operand") -> None:
    """Validate one grouped operand once when its quantizer generation changes."""
    if not quantizers:
        return

    reference = quantizers[0]
    reference_is_hybrid = isinstance(reference, HybridQuantizer)
    reference_identity = _identity_quantizer_signature(reference)

    for expert_index, quantizer in enumerate(quantizers[1:], start=1):
        if (quantizer is None) != (reference is None):
            raise ValueError(
                f"GroupedLinear {operand_name} quantizers mix None and concrete quantizers"
                f" across experts: expert 0 is {type(reference).__name__}, but expert"
                f" {expert_index} is {type(quantizer).__name__}."
            )
        if reference is None:
            continue

        quantizer_is_hybrid = isinstance(quantizer, HybridQuantizer)
        if quantizer_is_hybrid != reference_is_hybrid:
            raise ValueError(
                f"GroupedLinear {operand_name} quantizers mix HybridQuantizer and non-hybrid"
                f" quantizers across experts: expert 0 is {type(reference).__name__}, but expert"
                f" {expert_index} is {type(quantizer).__name__}."
            )

        identity = _identity_quantizer_signature(quantizer)
        if identity != reference_identity:
            raise ValueError(
                f"GroupedLinear {operand_name} quantizers mix Identity-backed and quantized"
                f" directions across experts: expert 0 uses {reference_identity}, but expert"
                f" {expert_index} uses {identity}."
            )

        if reference_is_hybrid:
            _validate_backend_match(
                reference.rowwise_quantizer,
                quantizer.rowwise_quantizer,
                operand_name,
                "rowwise",
                expert_index,
            )
            _validate_backend_match(
                reference.columnwise_quantizer,
                quantizer.columnwise_quantizer,
                operand_name,
                "columnwise",
                expert_index,
            )
            if quantizer.columnwise_source != reference.columnwise_source:
                raise ValueError(
                    f"GroupedLinear {operand_name} HybridQuantizer list has mixed columnwise"
                    " source policies across experts: expert 0 uses"
                    f" {reference.columnwise_source!r}, but expert {expert_index} uses"
                    f" {quantizer.columnwise_source!r}."
                )
        else:
            _validate_backend_match(
                reference,
                quantizer,
                operand_name,
                "plain",
                expert_index,
            )


def _split_quantize_non_hybrid(
    tensor,
    m_splits,
    quantizers,
    activation_dtype,
    *,
    disable_bulk_allocation=False,
    allow_identity_views=True,
):
    """Split and quantize one homogeneous, non-Hybrid quantizer list."""
    reference = quantizers[0]
    if _supports_native_split_quantize(reference):
        return tex.split_quantize(
            tensor,
            m_splits,
            quantizers,
            disable_bulk_allocation=disable_bulk_allocation,
        )

    tensor = cast_if_needed(tensor, activation_dtype)
    if (
        allow_identity_views
        # Only the base IdentityQuantizer can bypass quantization; subclasses
        # may override its behavior and must go through their normal call path.
        and type(reference) is IdentityQuantizer  # pylint: disable=unidiomatic-typecheck
        and (reference.dtype is None or reference.dtype == activation_dtype)
    ):
        return torch.split(tensor, m_splits)

    return [
        quantizer(tensor_part) if quantizer is not None else tensor_part
        for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers)
    ]


def _split_quantize_hybrid(
    tensor,
    m_splits,
    quantizers,
    *,
    disable_bulk_allocation=False,
):
    """Grouped split+quantize for an all-hybrid, generation-validated operand."""
    from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage

    reference = quantizers[0]
    rowwise_enabled = reference.rowwise_usage
    columnwise_enabled = reference.columnwise_usage
    columnwise_source = reference.columnwise_source
    rowwise_quantizers = [quantizer.rowwise_quantizer for quantizer in quantizers]
    columnwise_quantizers = [quantizer.columnwise_quantizer for quantizer in quantizers]

    needs_rowwise_result = rowwise_enabled or (
        columnwise_enabled and columnwise_source == "rowwise_dequantized"
    )
    row_results = (
        _split_quantize_non_hybrid(
            tensor,
            m_splits,
            rowwise_quantizers,
            tensor.dtype,
            disable_bulk_allocation=disable_bulk_allocation,
            allow_identity_views=False,
        )
        if needs_rowwise_result
        else [None] * len(quantizers)
    )

    columnwise_src = tensor
    if columnwise_enabled and columnwise_source == "rowwise_dequantized":
        # Assemble the exact grouped row results in split order. NVFP4 padding
        # and scale layout can differ from independently quantizing each split.
        columnwise_src = torch.cat(
            [result.dequantize(dtype=tensor.dtype) for result in row_results],
            dim=0,
        )
    col_results = (
        _split_quantize_non_hybrid(
            columnwise_src,
            m_splits,
            columnwise_quantizers,
            tensor.dtype,
            disable_bulk_allocation=disable_bulk_allocation,
            allow_identity_views=False,
        )
        if columnwise_enabled
        else [None] * len(quantizers)
    )

    return [
        HybridStorage(
            rowwise_storage=row if rowwise_enabled else None,
            columnwise_storage=col,
            quantizer=q,
            fake_dtype=tensor.dtype,
        )
        for row, col, q in zip(
            row_results,
            col_results,
            quantizers,
        )
    ]


def _split_quantize(
    tensor: torch.Tensor,
    split_sizes: List[int],
    with_quantized_output: bool,
    quantizers: Optional[List[Quantizer]],
    dtype: torch.dtype,
    with_debug_quantizers: bool,
    disable_bulk_allocation: bool,
) -> Sequence[Union[torch.Tensor, QuantizedTensorStorage]]:
    """Split a tensor and quantize each part if needed."""
    if not with_quantized_output:
        return torch.split(cast_if_needed(tensor, dtype), split_sizes)

    if quantizers is None or quantizers[0] is None:
        raise ValueError("Quantizers are required for quantized split output")

    if with_debug_quantizers:
        return DebugQuantizer.multi_tensor_quantize(tensor, quantizers, split_sizes, dtype)

    reference = quantizers[0]
    if isinstance(reference, HybridQuantizer):
        return _split_quantize_hybrid(
            tensor,
            split_sizes,
            quantizers,
            disable_bulk_allocation=disable_bulk_allocation,
        )

    return _split_quantize_non_hybrid(
        tensor,
        split_sizes,
        quantizers,
        dtype,
        disable_bulk_allocation=disable_bulk_allocation,
    )


def _split_quantize_and_bias(
    tensor: torch.Tensor,
    split_sizes: List[int],
    *,
    fp8: bool,
    debug: bool,
    quantizers: Optional[List[Quantizer]],
    dtype: torch.dtype,
    use_bias: bool,
    recipe: Recipe,
    disable_bulk_allocation: bool,
) -> Tuple[
    Sequence[Union[torch.Tensor, QuantizedTensorStorage]],
    List[Optional[torch.Tensor]],
]:
    """Split grad output, quantize if needed, and compute unfused bias gradients."""
    num_splits = len(split_sizes)
    grad_biases = [None] * num_splits
    reference = quantizers[0]
    identity = _uses_identity_quantizer(reference)
    hybrid = isinstance(reference, HybridQuantizer) and not identity

    use_native_bgrad_quantize = (
        fp8
        and not debug
        and not hybrid
        and use_bias
        and not identity
        and (recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8())
    )
    if use_native_bgrad_quantize:
        outputs = [None] * num_splits
        for i, tensor_part in enumerate(torch.split(tensor, split_sizes)):
            grad_biases[i], outputs[i] = tex.bgrad_quantize(tensor_part, quantizers[i])
        return outputs, grad_biases

    with_quantized_output = fp8 or debug
    if with_quantized_output and (use_bias or debug):
        for i, tensor_part in enumerate(torch.split(tensor, split_sizes)):
            grad_biases[i] = tensor_part.sum(dim=0)

    # Preserve the existing CPU-offload policy: only Hybrid split-quantize
    # disables bulk allocation in backward.
    disable_bulk_allocation = disable_bulk_allocation if hybrid else False
    outputs = _split_quantize(
        tensor,
        split_sizes,
        with_quantized_output=with_quantized_output,
        quantizers=quantizers,
        dtype=dtype,
        with_debug_quantizers=debug,
        disable_bulk_allocation=disable_bulk_allocation,
    )
    return outputs, grad_biases


__all__ = ["GroupedLinear"]


class _GroupedLinear(torch.autograd.Function):
    """GroupedLinear semi-top level module
    Calls custom cuda extensions.
    """

    @staticmethod
    def _maybe_dequantize(
        tensor: Union[torch.Tensor, QuantizedTensorStorage],
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Dequantize quantized tensors or cast regular tensors to ``dtype``."""
        if isinstance(tensor, QuantizedTensorStorage):
            return tensor.dequantize(dtype=dtype)
        return cast_if_needed(tensor, dtype)

    @staticmethod
    def _is_grouped_tensor_path_supported(
        *,
        fp8: bool,
        fp8_calibration: bool,
        debug: bool,
        cpu_offloading: bool,
        backward_override: Optional[str],
        save_original_input: bool,
        activation_dtype: torch.dtype,
        input_quantizers: List[Optional[Quantizer]],
        output_quantizers: List[Optional[Quantizer]],
    ) -> bool:
        """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata.

        There are no checks whether split sizes are supported. Splits
        may be in a CUDA tensor, so checking would hurt performance
        and be incompatible with CUDA Graphs.

        Supported Compute Capability (CC) and precisions:
        * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8
          block scaling (1D/2D, including power-of-2 scales).
        * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8
          per-tensor current scaling.
        FP8 delayed scaling is not supported because the corresponding grouped
        quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and
        SM110) raises instead of falling back: the fused path is Hopper-only and has
        no MXFP8-broadcast emulation. Architectures outside the fused-path window
        (e.g. SM120) fall back to the legacy path like every other recipe.
        Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8
        per-tensor current scaling on Hopper); otherwise the legacy path is used.
        Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization
        currently requires RHT.

        Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would
        trigger a fatal error in the cuBLASLt grouped GEMM check.
        """
        # 1. Filter by environment variable
        if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))):
            return False
        # 2. Filter out advanced features
        if (
            debug
            or cpu_offloading
            or fp8_calibration
            or backward_override is not None
            or save_original_input
        ):
            return False
        # 3. Filter by compute capability and cuBLAS version
        device_capability = get_device_compute_capability()
        if not (9, 0) <= device_capability <= (11, 0):
            return False
        cublaslt_version = tex.get_cublasLt_version()
        if cublaslt_version < 130300:
            return False
        if device_capability < (10, 0) and cublaslt_version < 130400:
            return False
        # 4. Output quantization is not supported.
        if any(q is not None for q in output_quantizers):
            return False
        # 5. Filter by quantization recipes.
        if fp8:
            if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers):
                # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.
                if device_capability < (10, 0) and cublaslt_version < 130500:
                    return False
                return True
            if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers):
                # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM
                # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast
                # emulation. On Blackwell (SM100/SM110, the only other arch that reaches
                # this branch) fail loudly rather than silently falling back to the
                # unfused path the user explicitly opted out of.
                if get_device_compute_capability() >= (10, 0):
                    raise RuntimeError(
                        "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the"
                        " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped"
                        " FP8 block-scaling path is Hopper-only. Unset"
                        " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused"
                        " path (emulated via MXFP8 GEMM on Blackwell)."
                    )
                return True
            # MXFP8 and NVFP4 require Blackwell+.
            if not (10, 0) <= device_capability <= (11, 0):
                return False
            return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all(
                isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers
            )
        return activation_dtype in (torch.bfloat16, torch.float16)

    @staticmethod
    def _make_grouped_tensor(
        data: torch.Tensor,
        *,
        num_gemms: int,
        split_sizes: torch.Tensor,
        tensor_offsets: torch.Tensor,
        last_dim: int,
        dtype: torch.dtype,
    ) -> GroupedTensorStorage:
        """Wrap a packed 2D buffer as a varying-first-dimension GroupedTensorStorage."""
        return GroupedTensorStorage(
            shape=(data.size(0), last_dim),
            dtype=dtype,
            num_tensors=num_gemms,
            quantizer=None,
            data=data.reshape(-1),
            first_dims=split_sizes,
            tensor_offsets=tensor_offsets,
        )

    @staticmethod
    def _make_grouped_bias(
        biases: Tuple[torch.Tensor, ...],
        *,
        num_gemms: int,
        out_features: int,
        dtype: torch.dtype,
    ) -> GroupedTensorStorage:
        """Pack per-GEMM biases into the grouped GEMM bias format."""
        bias_data = torch.stack(
            [_GroupedLinear._maybe_dequantize(bias, dtype) for bias in biases],
            dim=0,
        ).contiguous()
        return GroupedTensorStorage(
            shape=(num_gemms, out_features),
            dtype=dtype,
            num_tensors=num_gemms,
            shapes=[(1, out_features)] * num_gemms,
            quantizer=None,
            data=bias_data.reshape(-1),
        )

    @staticmethod
    def _prepare_weights_for_grouped_tensor_gemm(
        weights: Tuple[torch.Tensor, ...],
        weight_quantizers: List[Optional[Quantizer]],
        weight_workspaces: List[Optional[QuantizedTensorStorage]],
        *,
        with_quantized_compute: bool,
        columnwise_usage: bool,
        activation_dtype: torch.dtype,
        is_first_microbatch: Optional[bool],
        skip_fp8_weight_update: Optional[torch.Tensor],
        cache_weight: bool,
    ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]:
        """Prepare discrete weight tensors for GroupedTensor GEMM."""
        weights_for_gemm: List[torch.Tensor] = []
        new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights)
        if not with_quantized_compute:
            return (
                [_GroupedLinear._maybe_dequantize(weight, activation_dtype) for weight in weights],
                new_workspaces,
            )

        update_ws = is_first_microbatch is None or is_first_microbatch
        for idx, weight in enumerate(weights):
            weight_quantizer = weight_quantizers[idx]
            weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
            weight_fp8, new_workspaces[idx] = quantize_weight(
                tensor=weight,
                quantizer=weight_quantizer,
                workspace=weight_workspaces[idx] if weight_workspaces else None,
                update_workspace=update_ws,
                skip_update_flag=skip_fp8_weight_update,
                workspace_dtype=activation_dtype,
                cache=cache_weight,
            )
            weights_for_gemm.append(weight_fp8)
        return weights_for_gemm, new_workspaces

    @staticmethod
    def _validate_or_alloc_output(
        buffer: Optional[torch.Tensor],
        rows: int,
        cols: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> torch.Tensor:
        """Validate and return the caller's output buffer, or allocate one if it is None.

        The buffer must be a 2D, contiguous, non-grad tensor matching the required shape,
        dtype, and device. Validation reads host-side metadata only, with no device sync.
        """
        if buffer is None:
            return torch.empty((rows, cols), dtype=dtype, device=device)
        if buffer.dim() != 2:
            raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.")
        if buffer.size(0) != rows:
            raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.")
        if buffer.size(1) != cols:
            raise ValueError(
                f"Output buffer last dim {buffer.size(1)} does not match required {cols}."
            )
        if buffer.dtype != dtype:
            raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.")
        if buffer.device != device:
            raise ValueError(
                f"Output buffer device {buffer.device} does not match required {device}."
            )
        if not buffer.is_contiguous():
            raise ValueError("Output buffer must be contiguous.")
        if buffer.requires_grad:
            raise ValueError("Output buffer must not require gradient.")
        return buffer

    @staticmethod
    def _forward_grouped_tensor(
        ctx,
        *,
        inp: torch.Tensor,
        m_splits: torch.Tensor,
        use_bias: bool,
        is_first_microbatch: Optional[bool],
        fp8: bool,
        wgrad_store: Optional[WeightGradStore],
        input_quantizers: List[Optional[Quantizer]],
        weight_quantizers: List[Optional[Quantizer]],
        grad_input_quantizers: List[Optional[Quantizer]],
        grad_weight_quantizers: List[Optional[Quantizer]],
        grad_output_quantizers: List[Optional[Quantizer]],
        fuse_wgrad_accumulation: bool,
        activation_dtype: torch.dtype,
        is_grad_enabled: bool,
        weight_workspaces: List[Optional[QuantizedTensorStorage]],
        cache_weight: bool,
        skip_fp8_weight_update: Optional[torch.Tensor],
        weights: Tuple[torch.Tensor, ...],
        biases: Tuple[torch.Tensor, ...],
        out: Optional[torch.Tensor] = None,
        dgrad_out: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, list]:
        """Forward path backed by GroupedTensor + cuBLASLt grouped GEMM."""
        num_gemms = len(m_splits)
        device = inp.device
        in_features = weights[0].size(-1)
        out_features = weights[0].size(0)
        weight_requires_grad = weights[0].requires_grad

        split_sizes, (
            base_split_offsets,
            input_tensor_offsets,
            output_tensor_offsets,
        ) = tex.splits_to_offsets_multi(
            m_splits,
            device,
            strides=[1, in_features, out_features],
            include_leading_zero=[True, True, True],
            dtypes=[torch.int64, torch.int64, torch.int64],
            bulk_allocate=True,
        )

        inp_view = inp.reshape(-1, in_features)
        x = cast_if_needed(inp_view, activation_dtype)
        if fp8:
            input_quantizer = input_quantizers[0]
            input_quantizer.set_usage(
                rowwise=True,
                columnwise=is_grad_enabled and weight_requires_grad,
            )
            input_quantizer.optimize_for_gemm = True
            grouped_x = tex.group_quantize(
                x,
                input_quantizer,
                num_gemms,
                split_sizes,
                tensor_offsets=input_tensor_offsets,
            )
        else:
            grouped_x = _GroupedLinear._make_grouped_tensor(
                x,
                num_gemms=num_gemms,
                split_sizes=split_sizes,
                tensor_offsets=input_tensor_offsets,
                last_dim=in_features,
                dtype=activation_dtype,
            )

        columnwise_usage = is_grad_enabled and inp.requires_grad
        weights_for_gemm, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm(
            weights,
            weight_quantizers,
            weight_workspaces,
            with_quantized_compute=fp8,
            columnwise_usage=columnwise_usage,
            activation_dtype=activation_dtype,
            is_first_microbatch=is_first_microbatch,
            skip_fp8_weight_update=skip_fp8_weight_update,
            cache_weight=cache_weight,
        )

        out = _GroupedLinear._validate_or_alloc_output(
            out,
            x.size(0),
            out_features,
            activation_dtype,
            device,
        )
        grouped_out = _GroupedLinear._make_grouped_tensor(
            out,
            num_gemms=num_gemms,
            split_sizes=split_sizes,
            tensor_offsets=output_tensor_offsets,
            last_dim=out_features,
            dtype=activation_dtype,
        )

        grouped_bias = None
        if use_bias:
            grouped_bias = _GroupedLinear._make_grouped_bias(
                biases,
                num_gemms=num_gemms,
                out_features=out_features,
                dtype=activation_dtype,
            )

        use_split_accumulator = _2X_ACC_FPROP
        if fp8:
            recipe = FP8GlobalStateManager.get_fp8_recipe()
            if hasattr(recipe, "fp8_gemm_fprop"):
                use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator

        general_grouped_gemm_for_grouped_tensor(
            weights_for_gemm,
            grouped_x,
            grouped_out,
            layout="TN",
            bias=grouped_bias,
            use_split_accumulator=use_split_accumulator,
        )

        if is_grad_enabled:
            if weight_requires_grad:
                # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data
                # in backward pass)
                if fp8 and grouped_x.columnwise_data is not None:
                    grouped_x.rowwise_data = None
                    grouped_x.scale_inv = None
            else:
                grouped_x = None

            weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms
            tensors_to_save, tensor_objects = prepare_for_saving(
                grouped_x,
                *weights_to_save,
                split_sizes,
                base_split_offsets,
                input_tensor_offsets,
                output_tensor_offsets,
            )
            ctx.save_for_backward(*tensors_to_save)
            ctx.tensor_objects = tensor_objects

            ctx.use_grouped_tensor_path = True
            ctx.weight_quantizers = weight_quantizers
            ctx.weights_shape_0 = out_features
            ctx.weights_shape_1 = in_features
            ctx.grad_input_quantizers = grad_input_quantizers
            ctx.grad_output_quantizers = grad_output_quantizers
            ctx.grad_weight_quantizers = grad_weight_quantizers
            ctx.weights_requires_grad = weight_requires_grad
            if fuse_wgrad_accumulation and ctx.weights_requires_grad:
                ctx.origin_weight_refs = [weakref.ref(w) for w in weights]
                ctx.origin_weights_overwrite_main_grad = getattr(
                    weights[0], "overwrite_main_grad", False
                )
                if hasattr(weights[0], "__fsdp_param__"):
                    ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)]
                else:
                    ctx.main_grad_funcs = [
                        lambda j=i: weights[j].main_grad for i in range(num_gemms)
                    ]
            ctx.device = device
            ctx.dgrad_out = dgrad_out
            ctx.m_splits = None
            ctx.num_gemms = num_gemms
            ctx.activation_dtype = activation_dtype
            ctx.fp8 = fp8
            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
            ctx.backward_override = None
            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
            ctx.cpu_offloading = False
            ctx.is_first_microbatch = is_first_microbatch
            ctx.use_bias = use_bias
            ctx.inp_shape = inp.shape
            ctx.requires_dgrad = inp.requires_grad
            ctx.reduce_and_update_bwd_fp8_tensors = False
            if ctx.fp8 and requires_grad(inp, weights[0], biases[0]):
                ctx.reduce_and_update_bwd_fp8_tensors = (
                    ctx.reduce_and_update_bwd_fp8_tensors
                    or FP8GlobalStateManager.is_first_fp8_module()
                )
            ctx.wgrad_store = wgrad_store
            ctx.debug = False
            ctx.save_original_input = False
            ctx.input_quantizers = input_quantizers

        return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces

    # pylint: disable=keyword-arg-before-vararg
    @staticmethod
    def forward(
        ctx,
        inp: torch.Tensor,
        m_splits: torch.Tensor,
        non_tensor_args: Tuple,
        out: Optional[torch.Tensor],
        dgrad_out: Optional[torch.Tensor],
        *weights_and_biases,
    ) -> Tuple[torch.Tensor, list]:
        # pylint: disable=missing-function-docstring

        # Reduce number of arguments to autograd function in order
        # to reduce CPU overhead due to pytorch arg checking.
        (
            use_bias,
            is_first_microbatch,
            fp8,
            fp8_calibration,
            wgrad_store,
            input_quantizers,
            weight_quantizers,
            output_quantizers,
            grad_input_quantizers,
            grad_weight_quantizers,
            grad_output_quantizers,
            fuse_wgrad_accumulation,
            cpu_offloading,
            sequence_parallel,
            activation_dtype,
            is_grad_enabled,
            weight_workspaces,
            cache_weight,
            skip_fp8_weight_update,
            save_original_input,
            delayed_scaling_input_quantizer,
            unsafe_requantization_input_quantizer,
            debug,
        ) = non_tensor_args
        if fp8:
            backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override
        else:
            backward_override = None
        if backward_override == "high_precision":
            save_original_input = True
        elif backward_override == "dequantized":
            save_original_input = False

        num_gemms = len(m_splits)
        weights = weights_and_biases[:num_gemms]
        biases = weights_and_biases[num_gemms:]
        device = inp.device
        weight_requires_grad = weights[0].requires_grad

        origin_weights = weights
        is_dist_weight = is_distributed_weight(weights[0])
        if is_dist_weight:
            weights = materialize_weight_for_forward(weights)

        backward_needs_input = is_grad_enabled and weight_requires_grad
        if backward_override is None and save_original_input and backward_needs_input:
            if delayed_scaling_input_quantizer is not None:
                if FP8GlobalStateManager.get_fp8_recipe().custom():
                    warnings.warn(
                        "save_original_input is incompatible with delayed-scaling quantizers "
                        "(Float8Quantizer). Disabling save_original_input for this module.",
                        stacklevel=2,
                    )
                    save_original_input = False
                else:
                    raise ValueError(
                        "DelayedScaling recipe is not supported with save_original_input"
                    )

            # Megatron-Core may enable this automatically to reuse an activation
            # already retained by an upstream operation. The resolved quantizer
            # generation is classified once in ``_validate_quantizer_generation``.
            if save_original_input and unsafe_requantization_input_quantizer is not None:
                warnings.warn(
                    "Ignoring save_original_input=True because the input quantizer cannot "
                    "safely reconstruct the backward operand from the original input "
                    f"({unsafe_requantization_input_quantizer}).",
                    stacklevel=2,
                )
                save_original_input = False

        # Configure quantizers
        if input_quantizers[0] is not None:
            for input_quantizer in input_quantizers:
                input_quantizer.set_usage(
                    rowwise=True,
                    columnwise=(
                        is_grad_enabled
                        and weight_requires_grad
                        and not save_original_input
                        and backward_override is None
                    ),
                )
            columnwise_usage = is_grad_enabled and inp.requires_grad
            if backward_override is not None:
                columnwise_usage = False
            if not columnwise_usage:
                columnwise_usage = (
                    is_fp8_activation_recompute_enabled()
                    and not in_fp8_activation_recompute_phase()
                )
            # No need to set the quantizer states if weight is already quantized
            # for debug mode we create quantizer every iteration, thus we need to set the quantizer states
            if weight_quantizers[0] is not None and (
                not isinstance(weights[0], QuantizedTensorStorage) or debug
            ):
                for weight_quantizer in weight_quantizers:
                    weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
            elif isinstance(weights[0], QuantizedTensorStorage):
                # If weights are already quantized, no need to set quantizer states
                weight_quantizers = [weight._quantizer for weight in weights]
        if output_quantizers[0] is not None:
            for output_quantizer in output_quantizers:
                output_quantizer.set_usage(rowwise=True, columnwise=False)

        # Initialize input tensors
        in_features = weights[0].size(-1)
        if inp.size(-1) != in_features:
            raise ValueError(
                f"Input tensor (shape={tuple(inp.size())}) is not compatible with "
                f"weight tensor (shape={tuple(weights[0].size())})"
            )

        if _GroupedLinear._is_grouped_tensor_path_supported(
            fp8=fp8,
            fp8_calibration=fp8_calibration,
            debug=debug,
            cpu_offloading=cpu_offloading,
            backward_override=backward_override,
            save_original_input=save_original_input,
            activation_dtype=activation_dtype,
            input_quantizers=input_quantizers,
            output_quantizers=output_quantizers,
        ):
            return _GroupedLinear._forward_grouped_tensor(
                ctx,
                inp=inp,
                m_splits=m_splits,
                use_bias=use_bias,
                is_first_microbatch=is_first_microbatch,
                fp8=fp8,
                wgrad_store=wgrad_store,
                input_quantizers=input_quantizers,
                weight_quantizers=weight_quantizers,
                grad_input_quantizers=grad_input_quantizers,
                grad_weight_quantizers=grad_weight_quantizers,
                grad_output_quantizers=grad_output_quantizers,
                fuse_wgrad_accumulation=fuse_wgrad_accumulation,
                activation_dtype=activation_dtype,
                is_grad_enabled=is_grad_enabled,
                weight_workspaces=weight_workspaces,
                cache_weight=cache_weight,
                skip_fp8_weight_update=skip_fp8_weight_update,
                weights=weights,
                biases=biases,
                out=out,
                dgrad_out=dgrad_out,
            )

        # Convert splits to list of ints for compatibility with split functions
        m_splits = m_splits.tolist()

        inp_view = inp.reshape(-1, in_features)
        # Disable bulk allocation when CPU offloading is active: offloading skips small
        # tensors (like scales), but bulk allocation shares storage across all tensors,
        # so if scales can't be offloaded, nothing in the group can be offloaded.
        inputmats = _split_quantize(
            inp_view,
            m_splits,
            with_quantized_output=fp8 or debug,
            quantizers=input_quantizers,
            dtype=activation_dtype,
            with_debug_quantizers=debug,
            disable_bulk_allocation=cpu_offloading,
        )

        if cpu_offloading:
            start_offload(*inputmats)

        # Initialize weights
        weights_fp8: list
        new_workspaces = [None] * num_gemms
        if fp8 or debug:
            weights_fp8 = []
            update_ws = is_first_microbatch is None or is_first_microbatch
            for i in range(num_gemms):
                weight_fp8, new_workspaces[i] = quantize_weight(
                    tensor=weights[i],
                    quantizer=weight_quantizers[i],
                    workspace=weight_workspaces[i] if weight_workspaces else None,
                    update_workspace=update_ws,
                    skip_update_flag=skip_fp8_weight_update,
                    workspace_dtype=activation_dtype,
                    cache=cache_weight,
                )
                weights_fp8.append(weight_fp8)

        else:
            weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]

        # Initialize biases
        bias_dtype = activation_dtype
        if fp8 and activation_dtype == torch.float32:
            bias_dtype = torch.bfloat16  # FP8 GEMM only supports BF16/FP16 bias
        biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases
        # Initialize output tensor
        out = _GroupedLinear._validate_or_alloc_output(
            out,
            sum(m_splits),
            weights_fp8[0].size(0),
            activation_dtype,
            device,
        )

        # Choose whether to use split accumulator
        use_split_accumulator = _2X_ACC_FPROP
        if fp8:
            recipe = FP8GlobalStateManager.get_fp8_recipe()
            if hasattr(recipe, "fp8_gemm_fprop"):
                use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator

        # Perform GEMM
        general_grouped_gemm(
            weights_fp8,
            inputmats,
            [out],
            output_quantizers,
            activation_dtype,
            single_output=True,
            m_splits=m_splits,
            bias=biases,
            use_bias=use_bias,
            use_split_accumulator=use_split_accumulator,
        )

        if fp8_calibration:
            for i in range(num_gemms):
                input_quantizers[i].calibrate(inputmats[i])
                weight_quantizers[i].calibrate(weights[i])

        if cpu_offloading:
            mark_not_offload(*weights_fp8, *weights)

        if is_grad_enabled:
            ctx.use_grouped_tensor_path = False
            ctx.weight_quantizers = weight_quantizers
            ctx.weights_shape_1 = weights[0].shape[1]

            # TODO: update after #1638 is merged. # pylint: disable=fixme
            if weight_requires_grad:
                if save_original_input:
                    inputmats = [None] * num_gemms
                    inputmats[0] = inp
                else:
                    for inputmat in inputmats:
                        if isinstance(inputmat, QuantizedTensorStorage):
                            if backward_override is not None:
                                # In dequantized mode we should dequantize directly from
                                # fprop quantized layouts without retargeting usage.
                                inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
                            else:
                                inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
            else:
                inputmats = [None] * num_gemms

            # Original weights are only needed by high_precision dgrad. The weakrefs
            # used for fused wgrad accumulation serve a different purpose: restoring
            # Python parameter attributes without keeping the parameter alive here.
            saved_weights = (
                weights
                if backward_override == "high_precision" and inp.requires_grad
                else [None] * num_gemms
            )
            if is_dist_weight:
                # GTP: gathered workspace is transient (re-gathered in backward), don't save it.
                weights_fp8 = [None] * num_gemms
                saved_weights = origin_weights
            tensors_to_save, tensor_objects = prepare_for_saving(
                *inputmats,
                *weights_fp8,
                *saved_weights,
                *biases,
            )
            ctx.save_for_backward(*tensors_to_save)
            ctx.tensor_objects = tensor_objects

            ctx.grad_input_quantizers = grad_input_quantizers
            ctx.grad_output_quantizers = grad_output_quantizers
            ctx.grad_weight_quantizers = grad_weight_quantizers

            ctx.weights_requires_grad = weights[0].requires_grad
            if fuse_wgrad_accumulation and ctx.weights_requires_grad:
                # Keep weakrefs to weights to preserve attributes like main_grad
                # when we need to modify the weight python objects
                ctx.origin_weight_refs = [weakref.ref(w) for w in weights]
                ctx.origin_weights_overwrite_main_grad = getattr(
                    weights[0], "overwrite_main_grad", False
                )
                # This check is needed to ensure that main_grad is not created
                # during the forward pass when using MCore FSDP as it creates
                # the main_grad buffer lazily before backprop
                if hasattr(weights[0], "__fsdp_param__"):
                    # MCore FSDP creates main_grad lazily before backward
                    ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)]
                elif is_dist_weight:
                    ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)]
                else:
                    ctx.main_grad_funcs = [
                        lambda j=i: weights[j].main_grad for i in range(num_gemms)
                    ]
            ctx.device = device
            ctx.output_quantizers = output_quantizers
            ctx.m_splits = m_splits
            ctx.num_gemms = num_gemms
            ctx.activation_dtype = activation_dtype
            ctx.fp8 = fp8
            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
            ctx.backward_override = backward_override
            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
            ctx.cpu_offloading = cpu_offloading
            ctx.is_first_microbatch = is_first_microbatch
            ctx.use_bias = use_bias
            ctx.sequence_parallel = sequence_parallel
            ctx.inp_shape = inp.shape
            ctx.requires_dgrad = inp.requires_grad
            ctx.reduce_and_update_bwd_fp8_tensors = False
            if ctx.fp8 and requires_grad(inp, weights[0], biases[0]):
                ctx.reduce_and_update_bwd_fp8_tensors = (
                    ctx.reduce_and_update_bwd_fp8_tensors
                    or FP8GlobalStateManager.is_first_fp8_module()
                )
            ctx.wgrad_store = wgrad_store
            ctx.debug = debug
            ctx.save_original_input = save_original_input
            ctx.input_quantizers = input_quantizers
            ctx.dgrad_out = dgrad_out

            # backward overrides
            if backward_override is not None:
                ctx.fp8 = False
                ctx.debug = False
                ctx.ub_overlap_ag = False
                ctx.ub_overlap_rs_dgrad = False
                ctx.ub_bulk_dgrad = False
                ctx.ub_bulk_wgrad = False
                ctx.grad_input_quantizers = [None] * num_gemms
                ctx.grad_weight_quantizers = [None] * num_gemms
                ctx.grad_output_quantizers = [None] * num_gemms
                ctx.reduce_and_update_bwd_fp8_tensors = False

        # [*, in_features] -> [*, out_features] except first dimension changes for SP
        return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces

    @staticmethod
    def _backward_grouped_tensor(
        ctx,
        grad_output: torch.Tensor,
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        """Backward path paired with ``_forward_grouped_tensor``."""
        saved_tensors = restore_from_func_ctx(ctx)
        N = ctx.num_gemms
        grouped_x = saved_tensors[0]
        weights = saved_tensors[1 : 1 + N]
        split_sizes = saved_tensors[1 + N]
        base_split_offsets = saved_tensors[2 + N]
        input_tensor_offsets = saved_tensors[3 + N]
        output_tensor_offsets = saved_tensors[4 + N]

        origin_weights = [None] * N
        main_grads = [None] * N
        if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad:
            origin_weight_refs = ctx.origin_weight_refs
            ctx.origin_weight_refs = None
            origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs]
            assert all(
                w is not None for w in origin_weights
            ), "weight was removed while fuse_wgrad_accumulation=True"
            main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs]
            for origin_weight, main_grad in zip(origin_weights, main_grads):
                if main_grad is not None:
                    origin_weight.main_grad = main_grad

        grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1])
        dy_2d = cast_if_needed(grad_output_view, ctx.activation_dtype)
        dbias_packed = None
        if ctx.fp8:
            grad_output_quantizer = ctx.grad_output_quantizers[0]
            grad_output_quantizer.set_usage(
                rowwise=ctx.requires_dgrad,
                columnwise=ctx.weights_requires_grad,
            )
            grad_output_quantizer.optimize_for_gemm = True
            # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise
            # pass, so the fusion needs rowwise output (i.e. dgrad required).
            fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
                isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.requires_dgrad
            )
            if ctx.use_bias and fuse_bgrad:
                grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                    dy_2d,
                    grad_output_quantizer,
                    N,
                    split_sizes,
                    tensor_offsets=output_tensor_offsets,
                )
            else:
                grouped_dy = tex.group_quantize(
                    dy_2d,
                    grad_output_quantizer,
                    N,
                    split_sizes,
                    tensor_offsets=output_tensor_offsets,
                )
        else:
            grouped_dy = _GroupedLinear._make_grouped_tensor(
                dy_2d,
                num_gemms=N,
                split_sizes=split_sizes,
                tensor_offsets=output_tensor_offsets,
                last_dim=ctx.weights_shape_0,
                dtype=ctx.activation_dtype,
            )

        grad_biases = [None] * N
        if ctx.use_bias:
            if dbias_packed is None:
                dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N)
            grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)]

        dgrad = None
        if ctx.requires_dgrad:
            dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD
            if ctx.fp8:
                recipe = ctx.fp8_recipe
                if hasattr(recipe, "fp8_gemm_dgrad"):
                    dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
            for weight in weights:
                if isinstance(weight, QuantizedTensorStorage):
                    weight.update_usage(columnwise_usage=True)
            dgrad = _GroupedLinear._validate_or_alloc_output(
                ctx.dgrad_out,
                dy_2d.size(0),
                ctx.weights_shape_1,
                ctx.activation_dtype,
                ctx.device,
            )
            grouped_dgrad = _GroupedLinear._make_grouped_tensor(
                dgrad,
                num_gemms=N,
                split_sizes=split_sizes,
                tensor_offsets=input_tensor_offsets,
                last_dim=ctx.weights_shape_1,
                dtype=ctx.activation_dtype,
            )
            general_grouped_gemm_for_grouped_tensor(
                weights,
                grouped_dy,
                grouped_dgrad,
                layout="NN",
                use_split_accumulator=dgrad_gemm_use_split_accumulator,
            )

        if ctx.is_first_microbatch is not None:
            accumulate_wgrad_into_param_main_grad = (
                ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
            )
        else:
            accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation

        if ctx.weights_requires_grad:
            wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD
            if ctx.fp8:
                recipe = ctx.fp8_recipe
                if hasattr(recipe, "fp8_gemm_wgrad"):
                    wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator
            if ctx.fuse_wgrad_accumulation:
                wgrad_list = main_grads
            else:
                wgrad_packed = torch.empty(
                    N,
                    ctx.weights_shape_0,
                    ctx.weights_shape_1,
                    dtype=ctx.activation_dtype,
                    device=ctx.device,
                )
                wgrad_list = [wgrad_packed[i] for i in range(N)]

            accumulate = (
                accumulate_wgrad_into_param_main_grad
                if not getattr(ctx, "origin_weights_overwrite_main_grad", False)
                else False
            )

            def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights):
                general_grouped_gemm_for_grouped_tensor(
                    inputmats,
                    grad_output_mats,
                    grad_weights,
                    layout="NT",
                    use_split_accumulator=wgrad_gemm_use_split_accumulator,
                    accumulate=accumulate,
                )
                return None, [None] * N, None

            if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
                ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad)
            else:
                grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list)

            def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
                if ctx.weights_requires_grad:
                    if ctx.fuse_wgrad_accumulation and hasattr(weight, "grad_added_to_main_grad"):
                        weight.grad_added_to_main_grad = True
                        if getattr(weight, "zero_out_wgrad", False):
                            wgrad = get_dummy_wgrad(
                                list(main_grad.shape),
                                weight.dtype,
                                zero=True,
                            )
                        else:
                            wgrad = get_dummy_wgrad(
                                list(main_grad.shape),
                                weight.dtype,
                            )
                    elif ctx.fuse_wgrad_accumulation:
                        wgrad = None
                else:
                    wgrad = None
                return wgrad

            wgrad_list = [
                handle_custom_ddp_from_mcore(weight, main_grad, wgrad)
                for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list)
            ]
        else:
            wgrad_list = [None] * N

        if not ctx.use_bias:
            grad_biases = [None] * N

        if ctx.reduce_and_update_bwd_fp8_tensors:
            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
        return (
            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
            None,  # m_splits
            None,  # non_tensor_args
            None,  # out
            None,  # dgrad_out
            *wgrad_list,
            *grad_biases,
        )

    @staticmethod
    def backward(
        ctx, grad_output: torch.Tensor, _grad_workspaces
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        # pylint: disable=missing-function-docstring
        with get_nvtx_range_context("_GroupedLinear_backward"):
            if ctx.use_grouped_tensor_path:
                return _GroupedLinear._backward_grouped_tensor(ctx, grad_output)

            saved_tensors = restore_from_func_ctx(ctx)
            N = ctx.num_gemms
            inputmats = saved_tensors[:N]
            weights = saved_tensors[N : 2 * N]
            saved_weights = saved_tensors[2 * N : 3 * N]
            biases = saved_tensors[3 * N : 4 * N]

            # Restore from weakrefs to get original weight python objects
            # (preserves attributes like main_grad, grad_added_to_main_grad, etc.)
            # Only needed when fuse_wgrad_accumulation is enabled.
            origin_weights = [None] * N
            main_grads = [None] * N
            is_dist_weight = is_distributed_weight(saved_weights[0])
            if is_dist_weight:
                origin_weights = saved_weights
                if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad:
                    main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs]
            elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad:
                origin_weight_refs = ctx.origin_weight_refs
                ctx.origin_weight_refs = None
                origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs]
                assert all(
                    w is not None for w in origin_weights
                ), "weight was removed while fuse_wgrad_accumulation=True"
                main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs]
                for origin_weight, main_grad in zip(origin_weights, main_grads):
                    if main_grad is not None:
                        origin_weight.main_grad = main_grad

            # Preprocess grad output
            grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1])
            grad_output_reference = ctx.grad_output_quantizers[0]
            if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer):
                # Usage is a runtime decision, not part of generation validation.
                # Apply it uniformly so dispatch can read the first parent without
                # rescanning every expert.
                for grad_output_quantizer in ctx.grad_output_quantizers:
                    grad_output_quantizer.set_usage(
                        rowwise=ctx.requires_dgrad,
                        columnwise=ctx.weights_requires_grad,
                    )
            grad_output, grad_biases = _split_quantize_and_bias(
                grad_output_view,
                ctx.m_splits,
                fp8=ctx.fp8,
                debug=ctx.debug,
                quantizers=ctx.grad_output_quantizers,
                dtype=ctx.activation_dtype,
                use_bias=ctx.use_bias,
                recipe=ctx.fp8_recipe,
                disable_bulk_allocation=ctx.cpu_offloading,
            )

            if is_dist_weight:
                accumulate_wgrad_into_param_main_grad = False
            elif ctx.is_first_microbatch is not None:
                accumulate_wgrad_into_param_main_grad = (
                    ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
                )
            else:
                accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation

            if is_dist_weight:
                weights = materialize_weight_for_backward(origin_weights)

            if ctx.requires_dgrad:
                dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD
                if ctx.fp8 or ctx.debug:
                    recipe = ctx.fp8_recipe
                    if hasattr(recipe, "fp8_gemm_dgrad"):
                        dgrad_gemm_use_split_accumulator = (
                            recipe.fp8_gemm_dgrad.use_split_accumulator
                        )
                dgrad = _GroupedLinear._validate_or_alloc_output(
                    ctx.dgrad_out,
                    sum(ctx.m_splits),
                    ctx.weights_shape_1,
                    ctx.activation_dtype,
                    ctx.device,
                )
                weights_for_dgrad = weights
                if ctx.backward_override == "dequantized":
                    weights_for_dgrad = [
                        (
                            weight.dequantize(dtype=ctx.activation_dtype)
                            if isinstance(weight, QuantizedTensorStorage)
                            else cast_if_needed(weight, ctx.activation_dtype)
                        )
                        for weight in weights
                    ]
                elif ctx.backward_override == "high_precision":
                    weights_for_dgrad = [
                        (
                            weight.dequantize(dtype=ctx.activation_dtype)
                            if isinstance(weight, QuantizedTensorStorage)
                            else cast_if_needed(weight, ctx.activation_dtype)
                        )
                        for weight in saved_weights
                    ]
                # Make sure weights are available in column-wise format
                # for dgrad computation.
                for weight in weights_for_dgrad:
                    if isinstance(weight, QuantizedTensorStorage):
                        weight.update_usage(columnwise_usage=True)
                general_grouped_gemm(
                    weights_for_dgrad,
                    grad_output,
                    [dgrad],
                    ctx.grad_input_quantizers,
                    ctx.activation_dtype,
                    single_output=True,
                    layout="NN",
                    m_splits=ctx.m_splits,
                    grad=True,
                    use_split_accumulator=dgrad_gemm_use_split_accumulator,
                )

            if ctx.weights_requires_grad:
                wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD
                if ctx.fp8:
                    recipe = ctx.fp8_recipe
                    if hasattr(recipe, "fp8_gemm_wgrad"):
                        wgrad_gemm_use_split_accumulator = (
                            recipe.fp8_gemm_wgrad.use_split_accumulator
                        )
                if ctx.fuse_wgrad_accumulation:
                    wgrad_list = main_grads
                else:
                    wgrad_packed = torch.empty(
                        ctx.num_gemms,
                        *weights[0].size(),
                        dtype=ctx.activation_dtype,
                        device=ctx.device,
                    )
                    wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)]
                    if is_dist_weight:
                        # Gathered weights are no longer needed after dgrad GEMM.
                        del weights

                if ctx.save_original_input:
                    inp = inputmats[0]
                    in_features = inp.shape[-1]
                    inp_view = inp.reshape(-1, in_features)
                    if ctx.input_quantizers[0] is not None:
                        for input_quantizer in ctx.input_quantizers:
                            if isinstance(
                                input_quantizer,
                                (Float8Quantizer, Float8CurrentScalingQuantizer),
                            ):
                                input_quantizer.set_usage(rowwise=True, columnwise=True)
                            else:
                                input_quantizer.set_usage(rowwise=False, columnwise=True)
                    inputmats = _split_quantize(
                        inp_view,
                        ctx.m_splits,
                        with_quantized_output=ctx.fp8 or ctx.debug,
                        quantizers=ctx.input_quantizers,
                        dtype=ctx.activation_dtype,
                        with_debug_quantizers=ctx.debug,
                        disable_bulk_allocation=ctx.cpu_offloading,
                    )
                elif ctx.backward_override == "dequantized":
                    inputmats_dequant = []
                    for inputmat in inputmats:
                        if isinstance(inputmat, QuantizedTensorStorage):
                            inputmats_dequant.append(
                                inputmat.dequantize(dtype=ctx.activation_dtype)
                            )
                        else:
                            inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype))
                    inputmats = inputmats_dequant
                grouped_gemm_wgrad = functools.partial(
                    general_grouped_gemm,
                    quantization_params=ctx.grad_weight_quantizers,
                    out_dtype=ctx.activation_dtype,
                    layout="NT",
                    grad=True,
                    m_splits=ctx.m_splits,
                    use_bias=ctx.use_bias if grad_biases[0] is None else None,
                    bias=biases,
                    use_split_accumulator=wgrad_gemm_use_split_accumulator,
                    accumulate=(
                        accumulate_wgrad_into_param_main_grad
                        if not is_dist_weight
                        and not getattr(ctx, "origin_weights_overwrite_main_grad", False)
                        else False
                    ),
                )
                # WGRAD
                if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
                    ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad)
                else:
                    _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list)

                    for i in range(ctx.num_gemms):
                        if grad_biases[i] is None:
                            grad_biases[i] = grad_biases_[i]
                    del grad_biases_

                    # Deallocate input tensor
                    clear_tensor_data(*inputmats)

                def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
                    if ctx.weights_requires_grad:
                        # Handle custom DDP from mcore.
                        if ctx.fuse_wgrad_accumulation and hasattr(
                            weight, "grad_added_to_main_grad"
                        ):
                            weight.grad_added_to_main_grad = True
                            if getattr(weight, "zero_out_wgrad", False):
                                wgrad = get_dummy_wgrad(
                                    list(main_grad.shape),
                                    weight.dtype,
                                    zero=True,
                                )
                            else:
                                wgrad = get_dummy_wgrad(
                                    list(main_grad.shape),
                                    weight.dtype,
                                )
                        elif ctx.fuse_wgrad_accumulation:
                            wgrad = None
                    else:
                        wgrad = None
                    return wgrad

                if is_dist_weight:
                    wgrad_list = finalize_weight_grads(origin_weights, wgrad_list)
                else:
                    wgrad_list = [
                        handle_custom_ddp_from_mcore(weight, main_grad, wgrad)
                        for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list)
                    ]
            else:
                wgrad_list = [None] * ctx.num_gemms

            if not ctx.use_bias or (
                ctx.wgrad_store is not None
                and ctx.wgrad_store.delay_wgrad_compute()
                and not ctx.fp8
            ):
                grad_biases = [None] * ctx.num_gemms

        if ctx.reduce_and_update_bwd_fp8_tensors:
            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
        return (
            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
            None,  # m_splits
            None,  # non_tensor_args
            None,  # out
            None,  # dgrad_out
            *wgrad_list,
            *grad_biases,
        )


class GroupedLinear(TransformerEngineBaseModule):
    """Applies linear transformations to the incoming data list
       :math:`y_i = x_iA_i^T + b_i` in a grouped way.

    Parameters
    ----------
    num_gemms : int
                number of GEMMs to be performed simutaneously.
    in_features : int
                 size of each input sample.
    out_features : int
                  size of each output sample.
    bias : bool, default = True
          if set to ``False``, the layer will not learn an additive bias.
    init_method : Callable, default = None
                 used for initializing weights in the following way: ``init_method(weight)``.
                 When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``.
    get_rng_state_tracker : Callable, default = None
                 used to get the random number generator state tracker for initializing weights.
    rng_tracker_name : str, default = None
                 the param passed to get_rng_state_tracker to get the specific rng tracker.
    device : Union[torch.device, str], default = "cuda"
          The device on which the parameters of the model will be allocated. It is the user's
          responsibility to ensure all parameters are moved to the GPU before running the
          forward pass.

    Optimization parameters
    -----------------------
    fuse_wgrad_accumulation : bool, default = False
                             if set to ``True``, enables fusing of creation and accumulation of
                             the weight gradient. When enabled, it is assumed that the weights
                             have an additional ``main_grad`` attribute (used instead of the
                             regular ``grad``) which is a pre-allocated buffer of the correct
                             size to accumulate gradients in. This argument along with
                             weight tensor having attribute 'overwrite_main_grad' set to True
                             will overwrite ``main_grad`` instead of accumulating.
    return_bias : bool, default = False
                 when set to ``True``, this module will not apply the additive bias itself, but
                 instead return the bias value during the forward pass together with the
                 output of the linear transformation :math:`y = xA^T`. This is useful when
                 the bias addition can be fused to subsequent operations.
    params_dtype : torch.dtype, default = torch.get_default_dtype()
                  it controls the type used to allocate the initial parameters. Useful when
                  the model is trained with lower precision and the original FP32 parameters
                  would not fit in GPU memory.
    delay_wgrad_compute : bool, default = False
                         Whether to delay weight gradient computation
    save_original_input : bool, default = False
                       If set to ``True``, always saves the original input tensor rather than the
                       cast tensor. In some scenarios, the input tensor is used by multiple modules,
                       and saving the original input tensor may reduce the memory usage.
                       Requires input quantizers that can safely reproduce their results from the
                       original input. Cannot work with FP8 DelayedScaling recipe.
    single_grouped_weight : bool, default = False
                       If set to ``True``, grouped weights are stored as a single grouped parameter
                       instead of one parameter per GEMM.
                       EXPERIMENTAL and subject to change. Gated by the
                       ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var
                       is not set this argument is forced to ``False`` with a warning.
    single_grouped_bias : bool, default = False
                       If set to ``True``, grouped biases are stored as a single grouped bias
                       instead of one bias per GEMM.
                       EXPERIMENTAL and subject to change. Gated by the
                       ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var
                       is not set this argument is forced to ``False`` with a warning.

    Notes
    -----
    GroupedLinear doesn't really handle the TP communications inside. The ``tp_size`` and
    ``parallel_mode`` are used to determine the shapes of weights and biases.
    The TP communication should be handled in the dispatch and combine stages of MoE models.
    """

    def __init__(
        self,
        num_gemms: int,
        in_features: int,
        out_features: int,
        sequence_parallel: bool = False,
        fuse_wgrad_accumulation: bool = False,
        tp_group: Optional[dist_group_type] = None,
        tp_size: int = 1,
        get_rng_state_tracker: Optional[Callable] = None,
        rng_tracker_name: Optional[str] = None,
        init_method: Optional[Callable] = None,
        bias: bool = True,
        return_bias: bool = False,
        params_dtype: Optional[torch.dtype] = None,
        parallel_mode: Optional[str] = None,
        device: Union[torch.device, str] = "cuda",
        ub_overlap_rs: bool = False,
        ub_overlap_ag: bool = False,
        ub_name: Optional[str] = None,
        delay_wgrad_compute: bool = False,
        save_original_input: bool = False,
        single_grouped_weight: bool = False,
        single_grouped_bias: bool = False,
        name: Optional[str] = None,
    ) -> None:
        super().__init__(name)

        self.params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
        self.num_gemms = num_gemms
        self.in_features = in_features
        self.out_features = out_features
        self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
        self.use_bias = bias
        self.return_bias = return_bias
        self.apply_bias = bias and not return_bias
        self.ub_overlap_rs = ub_overlap_rs
        self.ub_overlap_ag = ub_overlap_ag
        self.ub_name = ub_name
        self.save_original_input = save_original_input
        single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags(
            single_grouped_weight, single_grouped_bias
        )
        self.single_grouped_weight = single_grouped_weight
        self.single_grouped_bias = single_grouped_bias
        if ub_overlap_rs or ub_overlap_ag:
            raise ValueError("GroupedLinear doesn't support Userbuffer overlap.")
        self.init_method = init_method
        self.get_rng_state_tracker = get_rng_state_tracker
        self.rng_tracker_name = rng_tracker_name

        self.wgrad_store = WeightGradStore(delay_wgrad_compute)

        self._offsets = {
            "input": 0,
            "weight": 1,
            "output": 2,
            "grad_output": 0,
            "grad_input": 1,
        }
        self._num_fp8_tensors_per_gemm = {
            "fwd": 3,
            "bwd": 2,
        }
        self._validated_quantizer_generations = {}
        self._delayed_scaling_input_quantizer = None
        self._unsafe_requantization_input_quantizer = None

        if tp_group is None:
            self.tp_size = tp_size
            if tp_size == 1:
                self.set_tensor_parallel_group(tp_group)
        else:
            self.tp_size = get_distributed_world_size(tp_group)
            self.set_tensor_parallel_group(tp_group)
        self.set_nccl_overlap_warning_if_tp()

        if self.tp_size > 1 and bias:
            raise ValueError(
                "GroupedLinear doesn't support bias when TP > 1. "
                "Because the TP communication is handled outside of this module."
            )

        self.parallel_mode = parallel_mode
        if self.parallel_mode not in GemmParallelModes:
            raise ValueError(
                f"parallel_mode {parallel_mode!r} not supported."
                f" Supported modes: {GemmParallelModes}"
            )

        if self.parallel_mode == "column":
            self.out_features = divide(self.out_features, self.tp_size)
        elif self.parallel_mode == "row":
            self.in_features = divide(self.in_features, self.tp_size)

        self.sequence_parallel = (self.tp_size > 1) and sequence_parallel

        for i in range(self.num_gemms):
            # Construct weight parameter
            self.register_parameter(
                f"weight{i}",
                torch.nn.Parameter(
                    torch.empty(
                        self.out_features,
                        self.in_features,
                        device=device,
                        dtype=self.params_dtype,
                    ),
                ),
                init_fn=init_method,
                get_rng_state_tracker=get_rng_state_tracker,
                fp8_meta_index=self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"],
            )

            # Construct bias parameters if needed
            if self.use_bias:
                self.register_parameter(
                    f"bias{i}",
                    torch.nn.Parameter(
                        torch.empty(
                            self.out_features,
                            device=device,
                            dtype=self.params_dtype,
                        ),
                    ),
                    init_fn=init_method_constant(0.0),
                )
            else:
                bias = torch.Tensor().to(dtype=self.params_dtype, device=device)
                setattr(self, f"bias{i}", bias)

        if self.primary_weights_in_fp8:
            self.init_fp8_metadata(num_gemms=self.num_gemms)

        is_meta = torch.device(device).type == "meta"
        self.reset_parameters(defer_init=is_meta)

        if self.wgrad_store.delay_wgrad_compute():
            for name, param in self.named_parameters():
                if name in ("weight", "bias"):
                    param.skip_backward_post_hook = True
                    continue
                for i in range(self.num_gemms):
                    if name in (f"weight{i}", f"bias{i}"):
                        param.skip_backward_post_hook = True

    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
        """Init scales and amaxes for fwd | bwd."""
        super().set_meta_tensor(fwd, recipe)

        # Recipe-specific quantizer configuration
        recipe = FP8GlobalStateManager.get_fp8_recipe()
        if recipe.float8_current_scaling():
            self._customize_quantizers_float8_current_scaling(fwd, recipe)

        self._validate_quantizer_generation(fwd)

    def _validate_quantizer_generation(self, fwd: bool) -> None:
        """Validate grouped-kernel invariants once per quantizer generation."""
        # Recipe state replaces this list object only when it constructs a new
        # quantizer generation. The O(1) identity guard keeps validation off the
        # steady-state forward path. Record a generation only after all of its
        # operand roles pass, so a failed recipe transition is retried.
        meta_key = "scaling_fwd" if fwd else "scaling_bwd"
        generation = self.quantizers.get(meta_key)
        if generation is None:
            return
        if self._validated_quantizer_generations.get(meta_key) is generation:
            return

        if fwd:
            stride = self._num_fp8_tensors_per_gemm["fwd"]
            input_quantizers = tuple(
                generation[self._offsets["input"] + i * stride] for i in range(self.num_gemms)
            )
            weight_quantizers = tuple(
                generation[self._offsets["weight"] + i * stride] for i in range(self.num_gemms)
            )
            _validate_grouped_quantizer_list(input_quantizers, operand_name="input")
            _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight")
            delayed_scaling_input_quantizer = next(
                (q for q in input_quantizers if isinstance(q, Float8Quantizer)),
                None,
            )
            unsafe_requantization_input_quantizer = next(
                (
                    q
                    for q in input_quantizers
                    if q is not None and not can_reconstruct_wgrad_input_from_original(q)
                ),
                None,
            )
            self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer
            self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer
        else:
            stride = self._num_fp8_tensors_per_gemm["bwd"]
            grad_output_quantizers = tuple(
                generation[self._offsets["grad_output"] + i * stride] for i in range(self.num_gemms)
            )
            _validate_grouped_quantizer_list(
                grad_output_quantizers,
                operand_name="grad_output",
            )

        self._validated_quantizer_generations[meta_key] = generation

    def get_quantizer_roles(
        self,
        *,
        fwd: bool,
        num_quantizers: int,
    ) -> Optional[List[QuantizerRole]]:
        """QuantizerRole list for quantizers used by ``GroupedLinear``.

        For grouped GEMMs we repeat the same pattern for each GEMM in
        order.  The output (fwd) and grad-input (bwd) slots default to
        ``None`` (unknown consumer).  Set :attr:`output_quantizer_role` /
        :attr:`grad_input_quantizer_role` to provide consumer identity.
        """
        name = self.name or ""
        if fwd:
            base = [
                QuantizerRole(module_type="grouped_linear", tensor_type="input", name=name),
                QuantizerRole(module_type="grouped_linear", tensor_type="weight", name=name),
                self._output_quantizer_role,
            ]
        else:
            base = [
                QuantizerRole(module_type="grouped_linear", tensor_type="grad_output", name=name),
                self._grad_input_quantizer_role,
            ]
        return [base[i % len(base)] for i in range(num_quantizers)]

    def make_grouped_weights(self, defer_init=False) -> None:
        """
        Convert parameters into a GroupedTensor and re-register them as parameters.
        """

        if defer_init:
            return

        weight_quantizers = self._get_weight_quantizers()
        # TODO(#3158): Support Identity/Hybrid single grouped weights.
        unsupported_quantizers = tuple(
            type(quantizer).__name__
            for quantizer in weight_quantizers
            if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer))
        )
        if unsupported_quantizers:
            quantizer_names = ", ".join(dict.fromkeys(unsupported_quantizers))
            raise NotImplementedError(
                "GroupedLinear(single_grouped_weight=True) does not support "
                f"{quantizer_names} weight quantizers yet. Set "
                "single_grouped_weight=False or unset "
                "NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158."
            )

        recipe = (
            weight_quantizers[0]._get_compatible_recipe()
            if weight_quantizers and weight_quantizers[0] is not None
            else None
        )
        if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()):
            self.set_tensor_parallel_attributes(defer_init=defer_init)
            return

        weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)]

        # TE preserves the original BF16/FP16 initialization on each quantized
        # parameter so distributed optimizers can construct lossless FP32 masters.
        # Packing the parameters must transfer those values to the new registered
        # grouped parameter; otherwise its master is initialized by dequantizing
        # MXFP8 and starts from a different value than the discrete-weight layout.
        high_precision_init_vals = [_get_high_precision_init_val(weight) for weight in weights]
        if any(value is not None for value in high_precision_init_vals) and not all(
            value is not None for value in high_precision_init_vals
        ):
            raise RuntimeError(
                "Grouped weights have inconsistent high-precision initialization state"
            )

        # Create the weight storage.
        grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes(
            num_tensors=self.num_gemms,
            shapes=[(self.out_features, self.in_features)] * self.num_gemms,
            quantizer=weight_quantizers[0],
            dtype=self.params_dtype,
            device=weights[0].device,
        )

        # Copy existing params into storage.
        with torch.no_grad():
            for i in range(self.num_gemms):
                if self.primary_weights_in_fp8:
                    grouped_weights.quantized_tensors[i].copy_from_storage(weights[i])
                else:
                    grouped_weights.quantized_tensors[i].copy_(weights[i])

        # Re-register as a single grouped weight parameter.
        if not (
            isinstance(grouped_weights, torch.Tensor)
            and (weight_quantizers[0] is None or not weight_quantizers[0].internal)
        ):
            raise RuntimeError("Found internal quantizer with `single_grouped_weight=True`.")
        grouped_parameter = torch.nn.Parameter(grouped_weights)
        if all(value is not None for value in high_precision_init_vals):
            _attach_high_precision_init_val(
                grouped_parameter,
                torch.stack(high_precision_init_vals, dim=0),
            )
            for weight in weights:
                _clear_high_precision_init_val(weight)

        self.register_parameter(
            "weight",
            grouped_parameter,
            init_fn=self.init_method,
            get_rng_state_tracker=self.get_rng_state_tracker,
            fp8_meta_index=self._offsets["weight"],
        )
        for i in range(self.num_gemms):
            self.register_parameter(f"weight{i}", None)

        if self.use_bias and self.single_grouped_bias:
            self._make_grouped_biases()

        self.set_tensor_parallel_attributes(defer_init=defer_init)

    def _make_grouped_biases(self) -> None:
        """Pack per-GEMM biases into one ``GroupedTensor`` (``single_grouped_bias``)."""
        biases = [getattr(self, f"bias{i}") for i in range(self.num_gemms)]
        packed = torch.stack([b.detach().clone() for b in biases], dim=0).contiguous()
        grouped_bias = GroupedTensor.make_grouped_tensor_from_rowwise_data(
            num_tensors=self.num_gemms,
            tensor_shape=(self.out_features,),
            rowwise_data=packed,
            dtype=packed.dtype,
        )
        grouped_bias.requires_grad_(True)
        self.register_parameter("bias", torch.nn.Parameter(grouped_bias))
        for i in range(self.num_gemms):
            self.register_parameter(f"bias{i}", None)

    def reset_parameters(self, defer_init=False):
        super().reset_parameters(defer_init=defer_init)
        # Grouped tensor weights / biases are opt-in features.
        if self.single_grouped_weight:
            self.make_grouped_weights(defer_init=defer_init)
        elif self.single_grouped_bias:
            self._make_grouped_biases()

    def set_tensor_parallel_attributes(self, defer_init=False) -> None:
        """Set attributes needed for TP"""

        if not defer_init:
            # Set parallelism attributes for linear weights
            grouped_weight = getattr(self, "weight", None)
            if grouped_weight is not None:
                set_tensor_model_parallel_attributes(
                    tensor=grouped_weight,
                    is_parallel=True,
                    dim=1 if self.parallel_mode == "row" else 0,
                    stride=1,
                )
            else:
                for i in range(self.num_gemms):
                    set_tensor_model_parallel_attributes(
                        tensor=getattr(self, f"weight{i}"),
                        is_parallel=True,
                        dim=1 if self.parallel_mode == "row" else 0,
                        stride=1,
                    )

            # Set parallelism attributes for linear biases
            if self.use_bias:
                grouped_bias = getattr(self, "bias", None)
                if grouped_bias is not None:
                    if self.parallel_mode == "row":
                        setattr(grouped_bias, "sequence_parallel", self.sequence_parallel)
                    elif self.parallel_mode == "column":
                        set_tensor_model_parallel_attributes(grouped_bias, True, 0, 1)
                else:
                    for i in range(self.num_gemms):
                        if self.parallel_mode == "row":
                            setattr(
                                getattr(self, f"bias{i}"),
                                "sequence_parallel",
                                self.sequence_parallel,
                            )
                        elif self.parallel_mode == "column":
                            set_tensor_model_parallel_attributes(
                                getattr(self, f"bias{i}"), True, 0, 1
                            )

    def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None:
        """Remap weight keys between single and per-GEMM checkpoint formats."""
        grouped_weight_key = f"{prefix}weight"
        per_gemm_weight_keys = [f"{prefix}weight{i}" for i in range(self.num_gemms)]
        has_grouped_weight = grouped_weight_key in state_dict
        has_per_gemm_weights = all(key in state_dict for key in per_gemm_weight_keys)

        if self.single_grouped_weight:
            # Backward compatibility: checkpoints saved without single_grouped_weight
            # store one weight tensor per GEMM (weight0..weightN). Convert them into a
            # single stacked grouped weight expected by this module configuration.
            if not has_grouped_weight and has_per_gemm_weights:
                per_gemm_weights = [state_dict.pop(key) for key in per_gemm_weight_keys]
                per_gemm_weights = [
                    weight.dequantize() if isinstance(weight, QuantizedTensorStorage) else weight
                    for weight in per_gemm_weights
                ]
                state_dict[grouped_weight_key] = torch.stack(per_gemm_weights, dim=0)
            elif has_grouped_weight:
                # Drop any redundant per-GEMM keys to avoid strict-load unexpected-key errors.
                for key in per_gemm_weight_keys:
                    state_dict.pop(key, None)
        else:
            # Forward compatibility: checkpoints saved with single_grouped_weight
            # store one grouped `weight`. Convert it back to weight0..weightN.
            if not has_per_gemm_weights and has_grouped_weight:
                grouped_weight = state_dict.pop(grouped_weight_key)
                if hasattr(grouped_weight, "split_into_quantized_tensors"):
                    grouped_members = grouped_weight.quantized_tensors
                    if grouped_members is None:
                        grouped_members = grouped_weight.split_into_quantized_tensors()
                    per_gemm_weights = [
                        (
                            weight.dequantize()
                            if isinstance(weight, QuantizedTensorStorage)
                            else weight
                        )
                        for weight in grouped_members
                    ]
                else:
                    grouped_weight = (
                        grouped_weight.dequantize()
                        if isinstance(grouped_weight, QuantizedTensorStorage)
                        else grouped_weight
                    )
                    per_gemm_weights = list(grouped_weight.unbind(dim=0))
                for i, weight in enumerate(per_gemm_weights):
                    state_dict[f"{prefix}weight{i}"] = weight
            elif has_per_gemm_weights:
                # Drop any redundant grouped key to avoid strict-load unexpected-key errors.
                state_dict.pop(grouped_weight_key, None)

    def _remap_grouped_bias_state_dict_keys(self, state_dict, prefix: str) -> None:
        """Remap bias keys between single grouped and per-GEMM checkpoint formats."""
        if not self.use_bias:
            return
        grouped_bias_key = f"{prefix}bias"
        per_gemm_bias_keys = [f"{prefix}bias{i}" for i in range(self.num_gemms)]
        has_grouped_bias = grouped_bias_key in state_dict
        has_per_gemm_biases = all(key in state_dict for key in per_gemm_bias_keys)

        if self.single_grouped_bias:
            if not has_grouped_bias and has_per_gemm_biases:
                per_gemm = [state_dict.pop(key) for key in per_gemm_bias_keys]
                state_dict[grouped_bias_key] = torch.stack(per_gemm, dim=0)
            elif has_grouped_bias:
                for key in per_gemm_bias_keys:
                    state_dict.pop(key, None)
                val = state_dict[grouped_bias_key]
                if isinstance(val, torch.Tensor) and val.dim() == 3 and val.shape[1] == 1:
                    state_dict[grouped_bias_key] = val.squeeze(1)
        else:
            if not has_per_gemm_biases and has_grouped_bias:
                gb = state_dict.pop(grouped_bias_key)
                if hasattr(gb, "split_into_quantized_tensors"):
                    members = gb.quantized_tensors
                    if members is None:
                        members = gb.split_into_quantized_tensors()
                    per_gemm = [m.reshape(-1) if m.dim() > 1 else m for m in members]
                else:
                    per_gemm = list(gb.unbind(0))
                for i, b in enumerate(per_gemm):
                    state_dict[f"{prefix}bias{i}"] = b.reshape(-1) if b.dim() > 1 else b
            elif has_per_gemm_biases:
                state_dict.pop(grouped_bias_key, None)

    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        """Load state dict with grouped-weight format compatibility."""
        state_dict_copy = state_dict.copy()
        metadata = getattr(state_dict, "_metadata", None)
        if metadata is not None:
            state_dict_copy._metadata = metadata
        self._remap_grouped_weight_state_dict_keys(state_dict_copy, prefix="")
        self._remap_grouped_bias_state_dict_keys(state_dict_copy, prefix="")
        return super().load_state_dict(state_dict_copy, strict=strict, assign=assign)

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        """Load state, including compatibility across grouped-weight checkpoint formats."""
        self._remap_grouped_weight_state_dict_keys(state_dict, prefix)
        self._remap_grouped_bias_state_dict_keys(state_dict, prefix)

        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    @no_torch_dynamo()
    def forward(
        self,
        inp: torch.Tensor,
        m_splits: torch.Tensor,
        is_first_microbatch: Optional[bool] = None,
        out: Optional[torch.Tensor] = None,
        dgrad_out: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
        """
        Apply the linear transformation to the input.

        Parameters
        ----------
        inp : torch.Tensor
             Input tensor.
        m_splits : torch.Tensor
                 Split sizes for the input tensor.
        is_first_microbatch : {True, False, None}, default = None
                             During training using either gradient accumulation or
                             pipeline parallelism a minibatch of data is further split
                             into microbatches. Between the microbatches of the same minibatch
                             the model weights are not updated. Setting this parameter indicates
                             whether the current microbatch is the first in a minibatch or not.
                             When set, this parameter enables additional optimizations:

                             * during FP8 training, it allows caching of the FP8 versions of
                               the weights
                             * it also allows skipping gradient accumulation during the
                               first microbatch (since it is the first gradient being
                               produced)
        out : torch.Tensor, default = None
             Optional preallocated buffer for the forward output; the returned tensor
             aliases it with no copy. Must be a 2D, contiguous, non-grad tensor of shape
             [num_tokens, out_features] in the activation dtype. Only the first
             sum(m_splits) rows are written; any padded trailing rows are left unchanged.
             Can be given independently of dgrad_out. If the buffer is reused across
             iterations, pass ``buffer.detach()`` so autograd does not set its
             ``requires_grad`` (which would trip the non-grad check on the next call).
        dgrad_out : torch.Tensor, default = None
             Optional preallocated buffer for the backward input gradient, of shape
             [num_tokens, in_features] with the same constraints as out. Receives the
             final gradient only when inp has a single consumer in the autograd graph;
             otherwise autograd accumulates into a new tensor.
        """
        debug = self.is_debug_iter()
        is_grad_enabled = torch.is_grad_enabled()
        num_gemms = self.num_gemms

        if FP8GlobalStateManager.fp8_graph_capturing():
            skip_fp8_weight_update = (
                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
            )
        else:
            skip_fp8_weight_update = None
        if skip_fp8_weight_update is not None:
            is_first_microbatch = False

        # Make sure splits are in expected format
        if not isinstance(m_splits, torch.Tensor):
            # Convert list of ints to tensor for backward compatibility
            m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu")
        elif m_splits.dtype != torch.int64:
            m_splits = m_splits.to(dtype=torch.int64)
        if m_splits.size() != (num_gemms,):
            raise ValueError(
                f"Shape of splits tensor ({tuple(m_splits.size())}) "
                f"does not match number of GEMMs ({num_gemms})."
            )

        if FP8GlobalStateManager.fp8_graph_capturing():
            skip_fp8_weight_update = (
                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
            )
        else:
            skip_fp8_weight_update = None
        if skip_fp8_weight_update is not None:
            is_first_microbatch = False

        # Preprocess input tensor
        if isinstance(inp, QuantizedTensorStorage):
            raise TypeError("GroupedLinear doesn't support input tensor in FP8.")
        inp = self.prepare_forward(inp, num_gemms=self.num_gemms)

        try:
            weight_tensors = self._get_weight_tensors()
            bias_tensors = self._get_bias_tensors()

            quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers()

            if debug:
                if self.no_debug_features_active(list(chain(*quantizers))):
                    debug = False
                    quantizers = self._get_quantizers()

            (
                input_quantizers,
                weight_quantizers,
                output_quantizers,
                grad_input_quantizers,
                grad_weight_quantizers,
                grad_output_quantizers,
            ) = quantizers
            if not debug and weight_quantizers[0] is not None:
                # Experts share shape and recipe settings: compute once and broadcast.
                optimize_for_gemm = self._enable_weight_preswizzle(
                    weight_quantizers[0], weight_tensors[0]
                )
                for q in weight_quantizers:
                    q.optimize_for_gemm = optimize_for_gemm

            if is_grad_enabled:
                linear_fn = _GroupedLinear.apply
                autograd_ctx = []
            else:
                linear_fn = _GroupedLinear.forward
                autograd_ctx = [None]

            cache_weight = is_first_microbatch is not None
            weight_workspaces = (
                [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)]
                if cache_weight
                else [None] * num_gemms
            )

            non_tensor_args = (
                self.apply_bias,
                is_first_microbatch,
                self.fp8,
                self.fp8_calibration,
                self.wgrad_store,
                input_quantizers,
                weight_quantizers,
                output_quantizers,
                grad_input_quantizers,
                grad_weight_quantizers,
                grad_output_quantizers,
                self.fuse_wgrad_accumulation,
                is_cpu_offload_enabled(),
                self.sequence_parallel,
                self.activation_dtype,
                is_grad_enabled,
                weight_workspaces,
                cache_weight,
                skip_fp8_weight_update,
                self.save_original_input,
                self._delayed_scaling_input_quantizer,
                self._unsafe_requantization_input_quantizer,
                debug,
            )
            out, new_workspaces = linear_fn(
                *autograd_ctx,
                inp,
                m_splits,
                non_tensor_args,
                out,
                dgrad_out,
                *weight_tensors,
                *bias_tensors,
            )

            if cache_weight:
                for i, ws in enumerate(new_workspaces):
                    if ws is not None:
                        if isinstance(ws, torch.Tensor):
                            ws = ws.detach()
                        self._fp8_workspaces[f"weight{i}"] = ws

        finally:
            self.end_forward()

        if self.return_bias:
            return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors]
        return out

    def backward_dw(self):
        """
        Execute the delayed weight gradient computation.
        This method is called after the main backward pass to compute weight gradients.
        """
        if not self.need_backward_dw():
            return
        if self.wgrad_store.context is None or self.wgrad_store.context.empty():
            return
        with get_nvtx_range_context("_GroupedLinear_wgrad"):
            (_, grad_biases_, _), tensor_list = self.wgrad_store.pop()
            wgrad_list = tensor_list[2]
            weight_params = self._get_weight_tensors()
            if not self.fuse_wgrad_accumulation:
                for i in range(self.num_gemms):
                    weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype)
            has_grad_biases = [
                grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_
            ]
            if self.use_bias and any(has_grad_biases):
                grouped_bias = getattr(self, "bias", None)
                if grouped_bias is not None:
                    if not all(has_grad_biases):
                        raise RuntimeError("Expected all grouped bias gradients to be present.")
                    gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype)
                    if grouped_bias.grad is None:
                        grouped_bias.grad = gstack
                    else:
                        grouped_bias.grad.add_(gstack)
                else:
                    bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)]
                    for i in range(self.num_gemms):
                        if has_grad_biases[i] and bias_params[i].grad is None:
                            bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype)
            del grad_biases_
            del wgrad_list
            del tensor_list
            self._trigger_wgrad_accumulation_and_reduce_hooks()

    def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None:
        """Customize quantizers based on current scaling recipe + linear."""

        if self.tp_size > 1:
            raise ValueError(
                "GroupedLinear doesn't support TP > 1 with Float8 current scaling. "
                "Because the TP communication is handled outside of this module."
            )

        if fwd:
            for i in range(self.num_gemms):
                # set configs about amax epsilon and power_2_scale
                self.quantizers["scaling_fwd"][
                    self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"]
                ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale
                self.quantizers["scaling_fwd"][
                    self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"]
                ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon
                # also set weight quantizer with same amax_epsilon & power_2_scale
                self.quantizers["scaling_fwd"][
                    self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"]
                ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale
                self.quantizers["scaling_fwd"][
                    self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"]
                ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon
        else:
            for i in range(self.num_gemms):
                # set grad_output_quantizer with amax epsilon and power_2_scale
                self.quantizers["scaling_bwd"][
                    self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"]
                ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale
                self.quantizers["scaling_bwd"][
                    self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"]
                ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon

    def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]:
        """Get the weight tensors of the module."""
        grouped_weight = getattr(self, "weight", None)
        if grouped_weight is not None:
            weight_tensors = grouped_weight.quantized_tensors
            if weight_tensors is None:
                # TODO(ksivaman): Remove this after GEMM integration.
                weight_tensors = grouped_weight.split_into_quantized_tensors()
        else:
            weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)]
        if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors):
            warnings.warn(
                "You are using quantized weights without quantized compute. "
                "Please make sure this is intentional."
            )
            weight_tensors = [
                w.dequantize() if isinstance(w, QuantizedTensorStorage) else w
                for w in weight_tensors
            ]
        return weight_tensors

    def _get_bias_tensors(self) -> List[torch.Tensor]:
        """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``)."""
        grouped_bias = getattr(self, "bias", None)
        if grouped_bias is not None:
            parts = grouped_bias.quantized_tensors
            if parts is None:
                parts = grouped_bias.split_into_quantized_tensors()
            return [p.reshape(-1) for p in parts]
        return [getattr(self, f"bias{i}") for i in range(self.num_gemms)]

    def _get_weight_quantizers(self) -> List[Quantizer]:
        """Get the weight quantizers of the module."""
        if not self.fp8 and not self.fp8_calibration and not self.primary_weights_in_fp8:
            return [None] * self.num_gemms
        weight_quantizers = [
            self.quantizers["scaling_fwd"][
                self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"]
            ]
            for i in range(self.num_gemms)
        ]
        for i in range(self.num_gemms):
            weight_quantizers[i].internal = not self.primary_weights_in_fp8
        return weight_quantizers

    def _get_quantizers(self):
        if self.fp8:
            # Normally validated while installing recipe metadata. Keep this
            # O(1) generation guard so failed transitions cannot reuse stale
            # validation state if base metadata takes an early return on retry.
            self._validate_quantizer_generation(True)
            if torch.is_grad_enabled():
                self._validate_quantizer_generation(False)

        weight_quantizers = self._get_weight_quantizers()
        input_quantizers, output_quantizers = (
            [None] * self.num_gemms,
            [None] * self.num_gemms,
        )
        grad_input_quantizers, grad_weight_quantizers, grad_output_quantizers = (
            [None] * self.num_gemms,
            [None] * self.num_gemms,
            [None] * self.num_gemms,
        )
        if self.fp8:
            input_quantizers = [
                self.quantizers["scaling_fwd"][
                    self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"]
                ]
                for i in range(self.num_gemms)
            ]
            for i in range(self.num_gemms):
                input_quantizers[i].internal = True
                input_quantizers[i].optimize_for_gemm = True
            if torch.is_grad_enabled():
                grad_output_quantizers = [
                    self.quantizers["scaling_bwd"][
                        self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"]
                    ]
                    for i in range(self.num_gemms)
                ]
                for i in range(self.num_gemms):
                    grad_output_quantizers[i].internal = True
                    grad_output_quantizers[i].optimize_for_gemm = True
        return (
            input_quantizers,
            weight_quantizers,
            output_quantizers,
            grad_input_quantizers,
            grad_weight_quantizers,
            grad_output_quantizers,
        )

    def _get_debug_quantizers(self):
        original_quantizers = self._get_quantizers()
        if not TEDebugState.debug_enabled:
            raise RuntimeError("TEDebugState.debug_enabled must be True to get debug quantizers")

        names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"]
        return tuple(
            [
                DebugQuantizer(self.name + f".gemm_{q_id}", name, q, self.tp_group, self.tp_size)
                for q_id, q in enumerate(qs)
            ]
            for name, qs in zip(names, original_quantizers)
        )
