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

import abc
import contextlib

import pytest
import torch
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode

try:
    from torch._opaque_base import OpaqueBaseMeta
    from torch._library.opaque_object import (
        get_opaque_type_name,
        register_opaque_type,
        MemberType,
    )

    _opaque_available = True
except ImportError:
    _opaque_available = False

import transformer_engine.pytorch as te
import transformer_engine_torch as tex
from transformer_engine.common import recipe
from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx
from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
from transformer_engine.pytorch.quantization import QuantizerRole
from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer
from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
from transformer_engine.pytorch import (
    is_fp8_available,
    is_mxfp8_available,
    is_fp8_block_scaling_available,
    is_nvfp4_available,
    Float8BlockQuantizer,
    MXFP8Quantizer,
    NVFP4Quantizer,
)
from utils import recipe_id
from transformer_engine.pytorch.attention.dot_product_attention.backends import (
    UnfusedDotProductAttention,
)

fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True)
mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True)
fp8_block_scaling_available, reason_for_no_fp8_block_scaling = is_fp8_block_scaling_available(
    return_reason=True
)
nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True)


def nvfp4_row_scaled():
    nvfp4_recipe = recipe.NVFP4BlockScaling(
        disable_rht=True,
        disable_stochastic_rounding=True,
        disable_2d_quantization=True,
        row_scaled_activation=True,
        backward_override="dequantized",
    )
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams()
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
    return nvfp4_recipe


def nvfp4_4over6():
    nvfp4_recipe = recipe.NVFP4BlockScaling(
        disable_rht=True,
        disable_stochastic_rounding=True,
        nvfp4_4over6="all",
    )
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True)
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
    return nvfp4_recipe


_all_recipes: list = []
if fp8_available:
    _all_recipes.append(recipe.Float8CurrentScaling())
if fp8_block_scaling_available:
    _all_recipes.append(recipe.Float8BlockScaling())
if mxfp8_available:
    _all_recipes.append(recipe.MXFP8BlockScaling())
if nvfp4_available:
    _all_recipes.append(recipe.NVFP4BlockScaling())
    _all_recipes.append(nvfp4_4over6())
    _all_recipes.append(nvfp4_row_scaled())


# ---------------------------------------------------------------------------
# ToyQuantizer – opaque value-type quantizer for torch.compile
# (requires torch opaque object support, not available in older PyTorch)
# ---------------------------------------------------------------------------

if _opaque_available:

    class _ToyQuantizerMeta(OpaqueBaseMeta, abc.ABCMeta):
        pass

    class ToyQuantizer(Float8CurrentScalingQuantizer, metaclass=_ToyQuantizerMeta):
        """Quantizer with a string tag, registered as an
        opaque value type so torch.compile can treat it as a baked-in constant."""

        def __init__(self, tag: str):
            super().__init__(fp8_dtype=te.DType.kFloat8E4M3, device=torch.device("cuda"))
            self.tag = tag

        def __eq__(self, other):
            if not isinstance(other, ToyQuantizer):
                return NotImplemented
            return self.tag == other.tag and self.dtype == other.dtype

        def __hash__(self):
            return hash((type(self), self.tag, self.dtype))

        def __fx_repr__(self):
            return (
                f"ToyQuantizer(tag={self.tag!r})",
                {"ToyQuantizer": ToyQuantizer},
            )

    register_opaque_type(
        ToyQuantizer,
        typ="value",
        members={
            "__setattr__": MemberType.USE_REAL,
            "set_usage": MemberType.USE_REAL,
        },
    )

    _Q = get_opaque_type_name(ToyQuantizer)

    def _make_qfactory(tag: str):
        """Return a qfactory that produces ToyQuantizer instances tagged with *tag*.

        The factory dispatches on ``QuantizerRole.tensor_type``; the roles are
        supplied by :meth:`ToyLinear.get_quantizer_roles`.
        """

        quantizers = {
            tensor_type: ToyQuantizer(tag=f"{tag}:{tensor_type}")
            for tensor_type in (
                "input",
                "weight",
                "output",
                "grad_output",
                "grad_input",
            )
        }

        def qfactory(role: QuantizerRole):
            return quantizers[role.tensor_type]

        return qfactory

    # ---------------------------------------------------------------------------
    # ToyLinear – minimal TE module backed by BasicLinear functional ops
    # ---------------------------------------------------------------------------

    class ToyLinear(TransformerEngineBaseModule):
        """Minimal TE-compatible linear module used for torch.compile tests."""

        def __init__(
            self,
            in_features: int,
            out_features: int,
            device: str = "cuda",
            dtype: torch.dtype = torch.bfloat16,
        ) -> None:
            super().__init__()
            self.in_features = in_features
            self.out_features = out_features
            self.weight = torch.nn.Parameter(
                torch.empty(out_features, in_features, dtype=dtype, device=device)
            )
            torch.nn.init.normal_(self.weight)

        def get_quantizer_roles(self, *, fwd: bool, num_quantizers: int):
            # Supplying explicit roles keeps CustomRecipeState from emitting a
            # warning (which would graph-break under fullgraph=True) and lets the
            # qfactory dispatch per tensor slot. Order must match the module's
            # quantizer array (FP8FwdTensorIdx / FP8BwdTensorIdx).
            if fwd:
                return [
                    QuantizerRole(module_type="linear", tensor_type="input"),
                    QuantizerRole(module_type="linear", tensor_type="weight"),
                    QuantizerRole(module_type="linear", tensor_type="output"),
                ]
            return [
                QuantizerRole(module_type="linear", tensor_type="grad_output"),
                QuantizerRole(module_type="linear", tensor_type="grad_input"),
            ]

        def _get_weight_tensors(self):
            return [self.weight]

        def _get_weight_quantizers(self):
            if not self.fp8 and not self.fp8_calibration:
                return [None]
            weight_q = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
            weight_q.internal = True
            return [weight_q]

        def forward(self, inp: torch.Tensor) -> torch.Tensor:
            inp = self.prepare_forward(inp, num_gemms=1)
            try:
                input_q = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT]
                input_q.internal = True
                input_q.optimize_for_gemm = True
                (weight_q,) = self._get_weight_quantizers()
                grad_output_q = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1]
                grad_output_q.internal = True
                grad_output_q.optimize_for_gemm = True

                return torch.ops.test_te.toy_linear(
                    inp,
                    self.weight,
                    input_q,
                    weight_q,
                    grad_output_q,
                )
            finally:
                self.end_forward()

    # ---------------------------------------------------------------------------
    # Opaque custom ops (torch.library)
    # ---------------------------------------------------------------------------

    _lib = torch.library.Library("test_te", "DEF")

    _lib.define(
        f"toy_linear(Tensor inp, Tensor weight, {_Q} input_q, {_Q} weight_q, {_Q} grad_output_q)"
        " -> Tensor"
    )

    _lib.define(
        "toy_linear_backward(Tensor grad_output, Tensor inp, Tensor weight,"
        f" {_Q} grad_output_q) -> (Tensor, Tensor)"
    )

    last_fwd_quantizers: list[dict[str, "ToyQuantizer"]] = []
    last_bwd_quantizers: list[dict[str, "ToyQuantizer"]] = []

    @torch.library.impl("test_te::toy_linear", "CompositeExplicitAutograd", lib=_lib)
    def _toy_linear_fwd_impl(inp, weight, input_q, weight_q, grad_output_q):
        last_fwd_quantizers.append(
            {
                "input_q": input_q,
                "weight_q": weight_q,
                "grad_output_q": grad_output_q,
            }
        )
        out, _, _ = BasicLinear._functional_forward(
            input=inp,
            weight=weight,
            dtype=inp.dtype,
            input_quantizer=input_q,
            weight_quantizer=weight_q,
        )
        return out

    @torch.library.register_fake("test_te::toy_linear", lib=_lib)
    def _toy_linear_fwd_fake(inp, weight, input_q, weight_q, grad_output_q):
        return inp @ weight.T

    def _toy_linear_setup_context(ctx, inputs, output):
        inp, weight, _input_q, _weight_q, grad_output_q = inputs
        ctx.save_for_backward(inp, weight)
        ctx.grad_output_q = grad_output_q

    @torch.library.impl("test_te::toy_linear_backward", "CompositeExplicitAutograd", lib=_lib)
    def _toy_linear_bwd_impl(grad_output, inp, weight, grad_output_q):
        last_bwd_quantizers.append({"grad_output_q": grad_output_q})
        dx, dw = BasicLinear._functional_backward(
            grad_output=grad_output,
            input=inp,
            weight=weight,
            grad_output_quantizer=grad_output_q,
            grad_input_quantizer=None,
        )
        return dx, dw

    @torch.library.register_fake("test_te::toy_linear_backward", lib=_lib)
    def _toy_linear_bwd_fake(grad_output, inp, weight, grad_output_q):
        return torch.empty_like(inp), torch.empty_like(weight)

    def _toy_linear_backward(ctx, grad_output):
        inp, weight = ctx.saved_tensors
        dx, dw = torch.ops.test_te.toy_linear_backward(
            grad_output,
            inp,
            weight,
            ctx.grad_output_q,
        )
        return dx, dw, None, None, None

    torch.library.register_autograd(
        "test_te::toy_linear",
        _toy_linear_backward,
        setup_context=_toy_linear_setup_context,
        lib=_lib,
    )


# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------


@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
def test_autocast_nested_custom():
    """One ToyLinear model used under nested te.autocast with 3 distinct
    CustomRecipe instances (each producing differently-tagged ToyQuantizers).

    Layout:
        with autocast(recipe0):           # outer
            out = model(inp)
            with autocast(recipe1):       # nested inside outer
                out = model(out)
        with autocast(recipe2):           # separate, after the nested pair
            out = model(out)

    fullgraph=True makes torch.compile raise if any graph break occurs.
    """
    dtype = torch.bfloat16
    device = "cuda"

    model = ToyLinear(32, 32, device=device, dtype=dtype)

    recipe0 = recipe.CustomRecipe(qfactory=_make_qfactory("R0"))
    recipe1 = recipe.CustomRecipe(qfactory=_make_qfactory("R1"))
    recipe2 = recipe.CustomRecipe(qfactory=_make_qfactory("R2"))

    inp = torch.randn(8, 32, dtype=dtype, device=device, requires_grad=True)

    def fn(inp):
        with te.autocast(recipe=recipe0):
            out = model(inp)
            with te.autocast(recipe=recipe1):
                out = model(out)
        with te.autocast(recipe=recipe2):
            out = model(out)
        return out

    torch._dynamo.reset()

    compiled = torch.compile(fn, fullgraph=True)
    last_fwd_quantizers.clear()
    last_bwd_quantizers.clear()

    out = compiled(inp)
    out.sum().backward()

    # Forward: 3 calls — R0, R1, R2
    assert len(last_fwd_quantizers) == 3, f"Expected 3 fwd calls, got {len(last_fwd_quantizers)}"
    for i, tag in enumerate(["R0", "R1", "R2"]):
        fq = last_fwd_quantizers[i]
        assert fq["input_q"].tag.startswith(f"{tag}:"), f"fwd[{i}] input_q: {fq['input_q'].tag}"
        assert fq["weight_q"].tag.startswith(f"{tag}:"), f"fwd[{i}] weight_q: {fq['weight_q'].tag}"
        assert fq["grad_output_q"].tag.startswith(
            f"{tag}:"
        ), f"fwd[{i}] grad_output_q: {fq['grad_output_q'].tag}"

    # Backward: 3 calls — reverse order R2, R1, R0
    assert len(last_bwd_quantizers) == 3, f"Expected 3 bwd calls, got {len(last_bwd_quantizers)}"
    for i, tag in enumerate(["R2", "R1", "R0"]):
        bq = last_bwd_quantizers[i]
        assert bq["grad_output_q"].tag.startswith(
            f"{tag}:"
        ), f"bwd[{i}] grad_output_q: {bq['grad_output_q'].tag}"


@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
@pytest.mark.parametrize("fp8_recipe", _all_recipes, ids=recipe_id)
def test_autocast_sanity(fp8_recipe):
    """Smoke test: torch.nn.Linear inside a single te.autocast with each
    built-in recipe. Forward + backward under torch.compile(fullgraph=True)."""
    dtype = torch.bfloat16
    device = "cuda"

    model = torch.nn.Linear(32, 32, dtype=dtype, device=device)
    inp = torch.randn(8, 32, dtype=dtype, device=device, requires_grad=True)

    def fn(inp):
        with te.autocast(recipe=fp8_recipe):
            return model(inp)

    torch._dynamo.reset()
    compiled = torch.compile(fn, fullgraph=True)

    out = compiled(inp)
    out.sum().backward()


_UNFUSED_DPA_CONFIG = dict(
    batch_size=2,
    num_heads=4,
    head_dim=64,
    max_seqlen_q=128,
    max_seqlen_kv=128,
)


def _make_unfused_attention(dtype: torch.dtype) -> UnfusedDotProductAttention:
    cfg = _UNFUSED_DPA_CONFIG
    softmax_scale = cfg["head_dim"] ** -0.5
    module = UnfusedDotProductAttention(
        softmax_scale=softmax_scale,
        attention_type="self",
        attention_dropout=0.0,
        layer_number=1,
        softmax_type="vanilla",
        return_max_logit=False,
    )
    return module.to(dtype=dtype, device="cuda")


_EMPTY_ALIBI_CACHE = {
    "_num_heads": None,
    "_alibi_slopes": None,
    "_max_seqlen_q": None,
    "_max_seqlen_kv": None,
    "_bottom_right_alignment": True,
    "_alibi_bias": None,
    "_alibi_slopes_require_update": False,
    "_alibi_bias_require_update": False,
}


def _make_unfused_qkv(qkv_layout: str, dtype: torch.dtype, requires_grad: bool = True):
    """Build (q, k, v) tensors matching `qkv_layout`. Returns also the
    extra kwargs (`cu_seqlens_*`, `max_seqlen_*`) that the unfused module
    needs for `thd` layouts (empty dict otherwise)."""
    cfg = _UNFUSED_DPA_CONFIG
    b, s_q, s_kv = cfg["batch_size"], cfg["max_seqlen_q"], cfg["max_seqlen_kv"]
    h, d = cfg["num_heads"], cfg["head_dim"]
    qkv_format = "".join(c for c in qkv_layout.split("_")[0] if c.isalpha())

    extra: dict = {}

    def _separate(shape):
        return tuple(
            torch.randn(shape, dtype=dtype, device="cuda", requires_grad=requires_grad)
            for _ in range(3)
        )

    if qkv_layout == "bshd_bshd_bshd":
        q, k, v = _separate((b, s_q, h, d))
    elif qkv_layout == "sbhd_sbhd_sbhd":
        q, k, v = _separate((s_q, b, h, d))
    elif qkv_layout == "thd_thd_thd":
        # All sequences in the batch have the maximum length; no padding.
        cu = torch.arange(0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda")
        q, k, v = _separate((b * s_q, h, d))
        extra = dict(
            cu_seqlens_q=cu,
            cu_seqlens_kv=cu,
            max_seqlen_q=s_q,
            max_seqlen_kv=s_kv,
        )
    elif qkv_layout == "bs3hd":
        # Packed: shape (b, s, 3, h, d), q/k/v are views along dim=-3.
        qkv = torch.randn(
            (b, s_q, 3, h, d),
            dtype=dtype,
            device="cuda",
            requires_grad=requires_grad,
        )
        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
        # q/k/v are non-leaf views; retain their grads so the assertions in
        # the test (`q.grad is not None` etc.) work for packed layouts.
        if requires_grad:
            for t in (q, k, v):
                t.retain_grad()
    elif qkv_layout == "sbh3d":
        # Packed: shape (s, b, h, 3, d), q/k/v are views along dim=-2.
        qkv = torch.randn(
            (s_q, b, h, 3, d),
            dtype=dtype,
            device="cuda",
            requires_grad=requires_grad,
        )
        q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2]
        if requires_grad:
            for t in (q, k, v):
                t.retain_grad()
    else:
        raise ValueError(f"Unsupported qkv_layout in test: {qkv_layout}")

    return q, k, v, extra, qkv_format


def _call_unfused(
    module: UnfusedDotProductAttention,
    qkv_layout: str,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    extra: dict,
) -> torch.Tensor:
    return module(
        _EMPTY_ALIBI_CACHE,
        q,
        k,
        v,
        qkv_layout=qkv_layout,
        attn_mask_type="causal",
        **extra,
    )


@pytest.mark.parametrize(
    "qkv_layout",
    [
        "bshd_bshd_bshd",
        "sbhd_sbhd_sbhd",
        "thd_thd_thd",
        "bs3hd",
        "sbh3d",
    ],
)
def test_unfused_dpa_torch_compile(qkv_layout):
    """Compile UnfusedDotProductAttention.forward with
    `torch.compile(fullgraph=True, mode="reduce-overhead")` for several
    qkv layouts.

    - `fullgraph=True` makes the test fail on any graph break inside the
      unfused attention path.
    - `mode="reduce-overhead"` uses the inductor cudagraphs backend, so
      forward+backward are captured into CUDA graphs and replayed on
      subsequent iterations."""
    dtype = torch.bfloat16

    module = _make_unfused_attention(dtype)

    def fn(q, k, v, extra):
        return _call_unfused(module, qkv_layout, q, k, v, extra)

    torch._dynamo.reset()
    compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead")

    for _ in range(3):
        q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True)
        out = compiled(q, k, v, extra)
        out.sum().backward()
        torch.cuda.synchronize()
        assert torch.isfinite(out).all()
        assert q.grad is not None
        assert k.grad is not None
        assert v.grad is not None


# ---------------------------------------------------------------------------
# get_attention_backend under torch.compile
# ---------------------------------------------------------------------------


# Scalars in AttentionParams must stay concrete: assume_constant_result cannot
# convert symbolic scalars (dynamo's automatic dynamic would make changed ints
# symbolic on recompilation), so pin them static explicitly.
@torch._dynamo.config.patch(specialize_int=True, specialize_float=True, recompile_limit=32)
def test_get_attention_backend_traceable(monkeypatch):
    """get_attention_backend must trace under torch.compile(fullgraph=True)
    without graph breaks. The compiled selection must stay consistent with
    eager when NVTE_* env vars flip (dynamo guards on os.environ) and when
    attention params change, and the baked tex.get_fused_attn_backend result
    must drive the selection."""
    from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils

    def fn(x, params):
        (
            use_flash_attention,
            _,
            use_fused_attention,
            fused_attention_backend,
            use_unfused_attention,
            _,
        ) = dpa_utils.get_attention_backend(params)
        # Encode the full selection (enabled backends + fused sub-backend) in
        # the tensor value: without a tensor op dynamo skips the frame entirely
        # (nothing gets compiled or guarded), and the output makes compiled vs
        # eager selection directly comparable.
        return (
            x
            + (1 if use_flash_attention else 0)
            + (2 if use_fused_attention else 0)
            + (4 if use_unfused_attention else 0)
            + (8 * int(fused_attention_backend) if fused_attention_backend is not None else 0)
        )

    # Dynamo only guards os.environ entries that exist at trace time (reads of
    # absent keys are not guarded yet), so set the vars explicitly.
    for env_var, value in (
        ("NVTE_FLASH_ATTN", "1"),
        ("NVTE_FUSED_ATTN", "1"),
        ("NVTE_UNFUSED_ATTN", "1"),
        ("NVTE_FP8_DPA_BWD", "1"),
        ("NVTE_DPA_FP8CS_O_in_F16", "1"),
        ("NVTE_DPA_FP8_RECIPE", ""),
        ("NVTE_UnfusedDPA_Emulate_FP8", "0"),
    ):
        monkeypatch.setenv(env_var, value)

    torch._dynamo.reset()
    compiled = torch.compile(fn, fullgraph=True)
    x = torch.zeros(8, device="cuda")
    params = dpa_utils.AttentionParams()

    torch.testing.assert_close(compiled(x, params), fn(x, params))

    # Flip env vars one by one: the compiled function must recompile (guards
    # on os.environ) and keep matching eager.
    for env_var in ("NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN", "NVTE_FLASH_ATTN"):
        monkeypatch.setenv(env_var, "0")
        torch.testing.assert_close(compiled(x, params), fn(x, params))
        monkeypatch.setenv(env_var, "1")

    # FP8 attention (fp8_dpa recipe): covers the FP8-only branch (run_config
    # env reads, recipe filters, get_fp8_te_dtype). Flipping an FP8-only env
    # var (emulation enables UnfusedDotProductAttention) must recompile too.
    fp8_params = dpa_utils.AttentionParams(
        fp8=True, fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}
    )
    torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params))
    monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1")
    torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params))
    monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "0")

    # Changing attention params (ints, layout string, dtype) must recompile
    # and keep matching eager, still with no graph break.
    for changed_params in (
        dpa_utils.AttentionParams(head_dim_qk=128, head_dim_v=128),
        dpa_utils.AttentionParams(max_seqlen_q=512, max_seqlen_kv=512),
        dpa_utils.AttentionParams(qkv_layout="bshd_bshd_bshd"),
        dpa_utils.AttentionParams(qkv_dtype=torch.float16),
    ):
        torch.testing.assert_close(compiled(x, changed_params), fn(x, changed_params))

    # The baked probe result must drive the selection: report no fused
    # sub-backend and expect UnfusedDotProductAttention (flash disabled, so the
    # outcome is deterministic). Use a fresh frame: already-compiled frames
    # keep the previously baked constant (assume_constant_result installs no
    # guard on the wrapped function).
    monkeypatch.setenv("NVTE_FLASH_ATTN", "0")
    monkeypatch.setattr(
        dpa_utils.tex,
        "get_fused_attn_backend",
        lambda *args: dpa_utils.FusedAttnBackend["No_Backend"],
    )

    def fn_no_backend(x, params):
        return fn(x, params)

    compiled_no_backend = torch.compile(fn_no_backend, fullgraph=True)
    torch.testing.assert_close(compiled_no_backend(x, params), x + 4.0)


# ---------------------------------------------------------------------------
# Value-opaque quantizers
# ---------------------------------------------------------------------------


def _mxfp8(dtype=tex.DType.kFloat8E4M3):
    return MXFP8Quantizer(fp8_dtype=dtype)


def _blockwise(force_pow_2_scales=True):
    return Float8BlockQuantizer(
        fp8_dtype=tex.DType.kFloat8E4M3,
        rowwise=True,
        columnwise=True,
        force_pow_2_scales=force_pow_2_scales,
    )


def _current_scaling(amax_epsilon=0.0):
    return Float8CurrentScalingQuantizer(
        fp8_dtype=tex.DType.kFloat8E4M3,
        device=torch.device("cpu"),
        amax_epsilon=amax_epsilon,
    )


def _nvfp4(with_rht=True):
    # Default with_rht=True so the quantize round-trip below exercises the
    # derived ``rht_matrix`` tensor (the field most likely to be dropped on
    # value-key reconstruction). Post-RHT amax is required by the kernel
    # whenever RHT is on (pre-RHT amax is unsupported).
    return NVFP4Quantizer(
        fp4_dtype=tex.DType.kFloat4E2M1,
        rowwise=True,
        columnwise=True,
        with_rht=with_rht,
        with_post_rht_amax=with_rht,
    )


def _hw_available(quantizer):
    """Whether this HW can actually run the quantize kernel for *quantizer*."""
    if isinstance(quantizer, MXFP8Quantizer):
        return mxfp8_available
    if isinstance(quantizer, NVFP4Quantizer):
        return nvfp4_available
    if isinstance(quantizer, Float8BlockQuantizer):
        return fp8_block_scaling_available
    return fp8_available  # Float8CurrentScalingQuantizer


# (factory, kwargs producing a different-but-valid config)
_VALUE_QUANTIZERS = [
    pytest.param(_mxfp8, id="mxfp8"),
    pytest.param(_blockwise, id="float8_blockwise"),
    pytest.param(_current_scaling, id="float8_current_scaling"),
    pytest.param(_nvfp4, id="nvfp4"),
]


@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS)
def test_quantizer_value_object(factory):
    """Value semantics + ``__fx_repr__`` round-trip via the production FX path."""
    a = factory()

    # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object.
    repr_str, globals_ = a.__fx_repr__()
    rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
    assert rebuilt == a and rebuilt is not a
    assert hash(rebuilt) == hash(a)
    # The deprecated amax-reduction group is never part of the value.
    assert getattr(rebuilt, "amax_reduction_group", None) is None

    # The rebuilt quantizer must also *behave* identically, not just compare
    # equal: equality only looks at the value key, so a field the kernel needs
    # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would
    # slip through the checks above and only blow up at quantize time. Run the
    # real quantize kernel on both and require bit-exact results.
    if _hw_available(a):
        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
        torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0)


def test_value_quantizer_rejects_process_group():
    """A value quantizer holding a live ProcessGroup must refuse to be turned
    into a value key / FX constant (raise), not silently drop the group."""
    import torch.distributed as dist  # pylint: disable=import-outside-toplevel

    created = not dist.is_initialized()
    if created:
        dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1)
    try:
        q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)
        q.amax_reduction_group = dist.group.WORLD
        # Every value-materialization path must reject it (hash, eq, __fx_repr__).
        with pytest.raises(TypeError):
            hash(q)
        with pytest.raises(TypeError):
            q.__fx_repr__()
    finally:
        if created:
            dist.destroy_process_group()


if _opaque_available:
    # A minimal custom op taking a tensor and a value-opaque quantizer that
    # quantizes + dequantizes inside it, one per production quantizer class.
    # ``test_quantizer_value_object_fullgraph`` drives this under
    # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the
    # graph -- proving the opaque-type registration took effect (a graph break
    # would make ``fullgraph=True`` raise).
    _qdq_lib = torch.library.Library("test_te_qdq", "DEF")
    _QDQ_OPS = {}
    for _qcls in (
        MXFP8Quantizer,
        Float8BlockQuantizer,
        Float8CurrentScalingQuantizer,
        NVFP4Quantizer,
    ):
        _op = f"qdq_{_qcls.__name__}"
        _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor")

        @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib)
        def _qdq_impl(x, q):
            return q(x).dequantize()

        @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib)
        def _qdq_fake(x, q):
            return torch.empty_like(x)

        _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op)


@pytest.mark.skipif(
    not _opaque_available,
    reason="torch.compile opaque-object support requires PyTorch >= 2.11",
)
@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS)
def test_quantizer_value_object_fullgraph(factory):
    """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph.

    A custom op quantizes+dequantizes with the (opaque value) quantizer; the
    compiled result must match eager. ``fullgraph=True`` raises on any graph
    break, so this proves the opaque-type registration actually took effect --
    unlike merely passing the quantizer through.
    """
    q = factory()
    if not _hw_available(q):
        pytest.skip("format not supported on this HW")

    op = _QDQ_OPS[type(q)]
    x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")

    def fn(inp):
        return op(inp, q)

    ref = fn(x)
    torch._dynamo.reset()
    out = torch.compile(fn, fullgraph=True)(x)
    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)


# ---------------------------------------------------------------------------
# torch.compile-traceable allocation primitives + TensorSpec
# ---------------------------------------------------------------------------


# (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128)
# / NVFP4 (mult. of 16) constraints.
# Format support is gated at runtime, in the tests that run a kernel; the rest is
# pure Python and works on any HW.
_SPEC_QUANTIZERS = [
    pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"),
    pytest.param(_mxfp8, (64, 128), id="mxfp8"),
    pytest.param(_blockwise, (128, 256), id="fp8_blockwise"),
    pytest.param(_nvfp4, (64, 128), id="nvfp4"),
]


def _build_from_primitives(quantizer, shape, dtype, device="cpu"):
    """Assemble a quantized tensor straight from the quantizer primitives:
    ``alloc_tensors`` (inner tensors) + ``create_metadata`` (ctx) + the storage's
    ``__tensor_unflatten__`` -- i.e. exactly what ``TensorSpec.create_tensor``
    does, but without going through :class:`TensorSpec`.
    """
    names = tuple(quantizer.inner_tensor_specs(shape))
    ctx = quantizer.create_metadata(shape, dtype=dtype)
    allocated = quantizer.alloc_tensors(shape, device=device)
    inner_tensors = {name: allocated[name] for name in names}
    storage_cls = ctx["cls"]
    # Row-major (contiguous) outer stride for ``__tensor_unflatten__``; ``meta``
    # device computes it without allocating storage.
    outer_stride = torch.empty(tuple(shape), device="meta").stride()
    return storage_cls.__tensor_unflatten__(inner_tensors, ctx, tuple(shape), outer_stride)


def _signature(tensor, names):
    """Comparable shape/dtype fingerprint of a tensor and its inner tensors."""
    sig = {"__shape__": tuple(tensor.shape), "__dtype__": tensor.dtype}
    for name in names:
        inner = getattr(tensor, name)
        sig[name] = (tuple(inner.shape), inner.dtype)
    return sig


def _skip_if_dequantize_unsupported(q):
    """Skip when this HW can't run ``dequantize()`` for the quantizer's format.

    ``dequantize()`` runs the real kernel on CUDA, so each format has its own
    availability gate (mirrors the ``is_*_available`` checks in test_numerics).
    """
    if isinstance(q, MXFP8Quantizer):
        if not mxfp8_available:
            pytest.skip(reason_for_no_mxfp8)
    elif isinstance(q, NVFP4Quantizer):
        if not nvfp4_available:
            pytest.skip(reason_for_no_nvfp4)
    elif isinstance(q, Float8BlockQuantizer):
        if not fp8_block_scaling_available:
            pytest.skip(reason_for_no_fp8_block_scaling)
    elif not fp8_available:  # Float8 current scaling
        pytest.skip(reason_for_no_fp8)


# ----- Quantizer primitives -----


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
def test_alloc_tensors_fake(factory, shape):
    """``alloc_tensors`` produces FakeTensors with the described shapes/dtypes."""
    q = factory()
    specs = q.inner_tensor_specs(shape)
    with FakeTensorMode():
        alloc = q.alloc_tensors(shape, device="cpu")
    assert set(alloc) == set(specs)
    for name, (spec_shape, spec_dtype) in specs.items():
        assert isinstance(alloc[name], FakeTensor)
        assert tuple(alloc[name].shape) == tuple(spec_shape)
        assert alloc[name].dtype == spec_dtype


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
def test_storage_flatten_unflatten_roundtrip(factory, shape):
    """Storage ``__tensor_flatten__`` / ``__tensor_unflatten__`` round-trips.

    Build a tensor from ``alloc_tensors`` + ``create_metadata``, flatten it, then
    unflatten and verify shape/dtype and every inner buffer match before vs after.
    """
    q = factory()
    _skip_if_dequantize_unsupported(q)

    tensor = _build_from_primitives(q, shape, torch.bfloat16)
    names = tuple(q.inner_tensor_specs(shape))
    # Fill inner tensors with deterministic data (empty() may contain NaNs) so
    # the round-trip can be checked by value via dequantize().
    for name in names:
        inner = getattr(tensor, name)
        inner.copy_(torch.arange(inner.numel(), device=inner.device).reshape(inner.shape))
    before = _signature(tensor, names)
    expected = tensor.dequantize()

    flat_names, flat_ctx = tensor.__tensor_flatten__()
    assert set(flat_names) == set(names)
    inner = {name: getattr(tensor, name) for name in flat_names}
    rebuilt = type(tensor).__tensor_unflatten__(
        inner, flat_ctx, tuple(tensor.shape), tensor.stride()
    )

    assert isinstance(rebuilt, QuantizedTensor)
    assert _signature(rebuilt, flat_names) == before
    # The reconstructed tensor dequantizes to the same values.
    torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True)


_USAGE_COMBOS = [
    pytest.param(True, True, id="rowwise_columnwise"),
    pytest.param(True, False, id="rowwise_only"),
    pytest.param(False, True, id="columnwise_only"),
]


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
@pytest.mark.parametrize("rowwise, columnwise", _USAGE_COMBOS)
@pytest.mark.parametrize("internal", [False, True], ids=["wrapper", "internal"])
def test_python_alloc_matches_cpp_make_empty(factory, shape, rowwise, columnwise, internal):
    """Pure-Python allocation is interchangeable with the C++ allocation.

    Builds the same quantized tensor twice: via ``Quantizer.make_empty``
    (``tex.create_empty_quantized_tensor``, the C++ path) and via the Python
    primitives ``inner_tensor_specs`` + ``create_metadata`` + ``alloc_tensors``
    + ``__tensor_unflatten__`` (exactly what ``TensorSpec.create_tensor``
    does). Checks:

    * structural parity -- same concrete class, buffer set, per-buffer
      shape/dtype/device, logical shape/dtype and flatten context;
    * functional parity -- the real C++ quantize kernel writes bit-identical
      results into the Python-allocated inner tensors as into the C++-allocated
      ones, proving the Python buffer description matches the layout
      (padding/alignment) the kernels expect.
    """

    # Two independent, identically-configured quantizers so no state can leak
    # between the two allocation paths.
    def make_quantizer():
        q = factory()
        q.set_usage(rowwise=rowwise, columnwise=columnwise)
        q.internal = internal
        return q

    q_ref = make_quantizer()
    if not _hw_available(q_ref):
        pytest.skip("format not supported on this HW")
    q_py = make_quantizer()

    ref = q_ref.make_empty(shape, dtype=torch.bfloat16, device="cuda")
    py = _build_from_primitives(q_py, shape, torch.bfloat16, device="cuda")

    # --- Structural parity ---
    assert type(py) is type(ref)
    ref_names, ref_ctx = ref.__tensor_flatten__()
    py_names, py_ctx = py.__tensor_flatten__()
    assert set(py_names) == set(ref_names)
    for name in ref_names:
        ref_inner, py_inner = getattr(ref, name), getattr(py, name)
        assert tuple(py_inner.shape) == tuple(ref_inner.shape), name
        assert py_inner.dtype == ref_inner.dtype, name
        assert py_inner.device == ref_inner.device, name

    # Logical shape / dtype (bare storages are not torch.Tensors: they expose
    # size() and _dtype instead of .shape / .dtype).
    if isinstance(ref, QuantizedTensor):
        assert tuple(py.shape) == tuple(ref.shape) == tuple(shape)
        assert py.dtype == ref.dtype == torch.bfloat16
    else:
        assert tuple(py.size()) == tuple(ref.size()) == tuple(shape)
        # pylint: disable=protected-access
        assert py._dtype == ref._dtype == torch.bfloat16

    # Flatten context. The quantizer entry needs special handling: production
    # quantizers get a value-based __eq__ from register_value_opaque_quantizer,
    # but fall back to field-wise comparison for classes that don't define one
    # (plain object.__eq__ is identity, which would spuriously fail).
    assert set(py_ctx) == set(ref_ctx)
    for key in ("cls", "is_tensor", "requires_grad"):
        assert py_ctx[key] == ref_ctx[key], key
    ref_kwargs, py_kwargs = ref_ctx["nontensor_kwargs"], py_ctx["nontensor_kwargs"]
    assert set(py_kwargs) == set(ref_kwargs)
    for key in ref_kwargs:
        rv, pv = ref_kwargs[key], py_kwargs[key]
        if isinstance(rv, Quantizer) or isinstance(pv, Quantizer):
            assert type(pv) is type(rv), key
            assert (pv.rowwise_usage, pv.columnwise_usage, pv.internal) == (
                rv.rowwise_usage,
                rv.columnwise_usage,
                rv.internal,
            ), key
            if type(rv).__eq__ is not object.__eq__:
                assert pv == rv, key
        else:
            assert pv == rv, key

    # --- Functional parity: run the real C++ quantize kernel into both ---
    x = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")

    def _quantize_into(quantizer, dst):
        if internal:
            # update_quantized() only accepts the wrapper classes; internal
            # (bare storage) tensors are filled through the same underlying
            # kernel binding directly.
            tex.quantize(x, quantizer, dst, None)
        else:
            quantizer.update_quantized(x, dst)

    # Scale-inv padding is never written by the kernel and both paths allocate it
    # uninitialized; zero it so the comparison below covers only kernel output.
    for name in ref_names:
        getattr(ref, name).zero_()
        getattr(py, name).zero_()

    # Some combos are rejected by the quantize kernel itself regardless of who
    # allocated the tensor (e.g. FP8 current-scaling columnwise-only on
    # TN-capable archs: there is no rowwise data buffer and
    # nvte_compute_scale_from_amax asserts on it). Parity then means the
    # Python-allocated tensor is rejected the same way -- not a silent skip.
    try:
        _quantize_into(q_ref, ref)
    except RuntimeError:
        with pytest.raises(RuntimeError):
            _quantize_into(q_py, py)
        return
    _quantize_into(q_py, py)
    for name in ref_names:
        torch.testing.assert_close(
            getattr(py, name), getattr(ref, name), rtol=0.0, atol=0.0, equal_nan=True
        )

    # Value check through dequantize(). Some layouts cannot dequantize at all
    # (e.g. FP8 columnwise-only raises NotImplementedError) -- the C++-allocated
    # reference defines what is supported, and when it raises, the bitwise
    # buffer equality above already proves value parity.
    try:
        expected = ref.dequantize()
    except NotImplementedError:
        expected = None
    if expected is not None:
        torch.testing.assert_close(py.dequantize(), expected, rtol=0.0, atol=0.0)


# ----- TensorSpec -----


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
def test_tensor_spec_matches_primitives(factory, shape):
    """TensorSpec is a thin wrapper: its ``create_metadata`` /
    ``create_inner_tensors`` / ``create_tensor`` match building everything
    directly from the quantizer primitives."""
    q = factory()
    spec = TensorSpec(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
    assert spec.is_quantized

    # Metadata matches the quantizer's.
    assert spec.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16)

    # inner_names + create_inner_tensors match inner_tensor_specs.
    specs = q.inner_tensor_specs(shape)
    names = tuple(specs)
    assert spec.inner_names() == names
    inner_tensors = spec.create_inner_tensors()
    assert len(inner_tensors) == len(names)
    for name, inner in zip(names, inner_tensors):
        exp_shape, exp_dtype = specs[name]
        assert tuple(inner.shape) == tuple(exp_shape)
        assert inner.dtype == exp_dtype

    # The assembled tensor matches one built directly from the primitives.
    direct = _build_from_primitives(q, shape, torch.bfloat16)
    assert _signature(spec.create_tensor(), names) == _signature(direct, names)


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
@pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"])
def test_tensor_spec_create_tensor(factory, shape, fake):
    """``create_tensor`` yields a quantized tensor with the right shape/dtype;
    its inner tensors are fake exactly under ``FakeTensorMode``."""
    q = factory()
    spec = TensorSpec(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
    with FakeTensorMode() if fake else contextlib.nullcontext():
        out = spec.create_tensor()
    assert isinstance(out, QuantizedTensor)
    assert tuple(out.shape) == tuple(shape)
    assert out.dtype == torch.bfloat16
    for name in spec.inner_names():
        assert isinstance(getattr(out, name), FakeTensor) == fake


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
def test_tensor_spec_create_tensor_compiles(factory, shape):
    """``TensorSpec.create_tensor`` traces under ``fullgraph=True`` (CPU)."""
    q = factory()

    def fn(x):
        spec = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, quantizer=q, device=x.device)
        t = spec.create_tensor()
        acc = x.new_zeros(())
        for name in spec.inner_names():
            acc = acc + getattr(t, name).float().sum()
        return acc

    x = torch.zeros(*shape, dtype=torch.bfloat16)
    torch._dynamo.reset()
    out = torch.compile(fn, fullgraph=True)(x)
    assert out.shape == ()


def test_to_tensor_spec_plain():
    """``to_tensor_spec`` describes a plain tensor."""
    t = torch.empty(2, 3, dtype=torch.float32)
    spec = to_tensor_spec(t)
    assert not spec.is_quantized
    assert spec.shape == (2, 3)
    assert spec.dtype == torch.float32
    assert spec.inner_names() == ("data",)


@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS)
def test_to_tensor_spec_quantized(factory, shape):
    """``to_tensor_spec`` round-trips a quantized tensor back into a spec."""
    q = factory()
    tensor = TensorSpec(
        shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")
    ).create_tensor()

    spec = to_tensor_spec(tensor)
    assert spec.is_quantized
    assert spec.shape == tuple(shape)
    assert spec.dtype == torch.bfloat16
    # Same buffer layout as the original tensor.
    assert spec.inner_names() == tuple(q.inner_tensor_specs(shape))
    # Rebuilding from the derived spec matches the original tensor's structure.
    assert _signature(spec.create_tensor(), spec.inner_names()) == _signature(
        tensor, spec.inner_names()
    )
