            return f"{status_part} {error_part}"
        elif self.constraint_status == "not_available":
            status_part = colorize(colors.YELLOW, "[nvFuser] ○") + " " + name_padded
            message_part = colorize(colors.YELLOW, "Torch built without CUDA")
            return f"{status_part} {message_part}"
        else:
            return ""

    def is_failure(self) -> bool:
        """Check for both version failure and CUDA constraint failure."""
        # Check base version requirement
        if super().is_failure():
            return True

        # Check CUDA constraint
        if self.constraint_status == "mismatch":
            return True

        return False

    def generate_help(self, platform_info):
        """
        Generate PyTorch installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        version_min = self.version_required

        print(f"PyTorch {version_min}+ Required")
        print()
        print(
            "Why: nvFuser is a PyTorch extension and requires PyTorch with CUDA support."
        )
        print()
        print(f"Install PyTorch {version_min} or higher with CUDA:")
        print()

        # Show common CUDA versions
        print("  # For CUDA 13.1:")
        print("  pip install torch --index-url https://download.pytorch.org/whl/cu131")
        print("  # For CUDA 13.0:")
        print("  pip install torch --index-url https://download.pytorch.org/whl/cu130")
        print("  # For CUDA 12.8:")
        print("  pip install torch --index-url https://download.pytorch.org/whl/cu128")
        print()

        # If CUDA constraint mismatch, add additional help
        if self.constraint_status == "mismatch":
            print()
            print("IMPORTANT: Torch CUDA Version Mismatch Detected")
            print()
            print(
                "Why: PyTorch was built with a different CUDA version than your system's"
            )
            print("     CUDA Toolkit. This will cause runtime errors.")
            print()
            print("Resolution:")
            print()
            print("  You have two options:")
            print()
            print("  Option 1: Install matching CUDA Toolkit (recommended)")
            print()
            print(
                "    Install the CUDA Toolkit version that matches your PyTorch build."
            )
            print(
                "    Check PyTorch CUDA version: python -c 'import torch; print(torch.version.cuda)'"
            )
            print()
            print("  Option 2: Reinstall PyTorch for your CUDA version")
            print()
            print("    Reinstall PyTorch built for your system's CUDA Toolkit version.")
            print("    Check system CUDA version: nvcc --version")
            print()


python.tools.prereqs.requirements.nvmmh

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""nvidia-matmul-heuristics dependency requirement."""

from typing import Dict
from .base import BooleanRequirement

class NVMMHRequirement(BooleanRequirement):
    """
    nvidia-matmul-heuristics check.

    CMake variables used:
    - NVMMH_FOUND: Whether nvidia-matmul-heuristics is available
    - NVFUSER_REQUIREMENT_NVMMH_STATUS: Validation status
    - NVFUSER_REQUIREMENT_NVMMH_OPTIONAL: Whether NVMMH is optional

    No version checking - just verifies nvidia-matmul-heuristics headers are available.
    """

    def __init__(self, cmake_vars: Dict):
        """
        Initialize nvidia-matmul-heuristics requirement.

        Args:
            cmake_vars: Dictionary of all CMake variables
        """
        # Define dependency name and CMake variable names for this requirement
        name = "NVMMH"
        found_var = "NVMMH_FOUND"
        status_var = "NVFUSER_REQUIREMENT_NVMMH_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_NVMMH_OPTIONAL"
        location_var = "NVMMH_INCLUDE_DIR"

        super().__init__(
            name, cmake_vars, found_var, status_var, optional_var, location_var
        )

    def generate_help(self, platform_info):
        """
        Generate nvidia-matmul-heuristics installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        print("nvidia-matmul-heuristics (NVMMH)")
        print()
        print(
            "Why: nvidia-matmul-heuristics provides optimized matrix multiplication heuristics for nvFuser."
        )
        print()
        print("Install nvidia-matmul-heuristics:")
        print()
        print("  Recommended: pip installation:")
        print()
        print("    pip install nvidia-matmul-heuristics")
        print()
        print("  Note: This is an optional dependency. nvFuser will build without it,")
        print(
            "        but matmul operations may not have access to optimized heuristics."
        )
        print()


# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""nvidia-matmul-heuristics dependency requirement."""

from typing import Dict
from .base import BooleanRequirement

class NVMMHRequirement(BooleanRequirement):
    """
    nvidia-matmul-heuristics check.

    CMake variables used:
    - NVMMH_FOUND: Whether nvidia-matmul-heuristics is available
    - NVFUSER_REQUIREMENT_NVMMH_STATUS: Validation status
    - NVFUSER_REQUIREMENT_NVMMH_OPTIONAL: Whether NVMMH is optional

    No version checking - just verifies nvidia-matmul-heuristics headers are available.
    """

    def __init__(self, cmake_vars: Dict):
        """
        Initialize nvidia-matmul-heuristics requirement.

        Args:
            cmake_vars: Dictionary of all CMake variables
        """
        # Define dependency name and CMake variable names for this requirement
        name = "NVMMH"
        found_var = "NVMMH_FOUND"
        status_var = "NVFUSER_REQUIREMENT_NVMMH_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_NVMMH_OPTIONAL"
        location_var = "NVMMH_INCLUDE_DIR"

        super().__init__(
            name, cmake_vars, found_var, status_var, optional_var, location_var
        )

    def generate_help(self, platform_info):
        """
        Generate nvidia-matmul-heuristics installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        print("nvidia-matmul-heuristics (NVMMH)")
        print()
        print(
            "Why: nvidia-matmul-heuristics provides optimized matrix multiplication heuristics for nvFuser."
        )
        print()
        print("Install nvidia-matmul-heuristics:")
        print()
        print("  Recommended: pip installation:")
        print()
        print("    pip install nvidia-matmul-heuristics")
        print()


    """

    def __init__(self, cmake_vars: Dict):
        """
        Initialize nvidia-matmul-heuristics requirement.

        Args:
            cmake_vars: Dictionary of all CMake variables
        """
        # Define dependency name and CMake variable names for this requirement
        name = "NVMMH"
        found_var = "NVMMH_FOUND"
        status_var = "NVFUSER_REQUIREMENT_NVMMH_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_NVMMH_OPTIONAL"
        location_var = "NVMMH_INCLUDE_DIR"

        super().__init__(
            name, cmake_vars, found_var, status_var, optional_var, location_var
        )

    def generate_help(self, platform_info):
        """
        Generate nvidia-matmul-heuristics installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        print("nvidia-matmul-heuristics (NVMMH)")
        print()
        print(
            "Why: nvidia-matmul-heuristics provides optimized matrix multiplication heuristics for nvFuser."
        )
        print()
        print("Install nvidia-matmul-heuristics:")
        print()
        print("  Recommended: pip installation:")
        print()
        print("    pip install nvidia-matmul-heuristics")
        print()
        print("  Note: This is an optional dependency. nvFuser will build without it,")
        print(
            "        but matmul operations may not have access to optimized heuristics."
        )
        print()


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Python dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class PythonRequirement(VersionRequirement):
    """
    Python interpreter requirement.

    CMake variables used:
    - Python_FOUND: Whether Python was found
    - Python_VERSION: Detected version (e.g., "3.12.3")
    - Python_EXECUTABLE: Path to python binary
    - NVFUSER_REQUIREMENT_Python_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Python_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_Python_OPTIONAL: Whether Python is optional
    """

    def __init__(self, cmake_vars: Dict):
        """
        Initialize Python requirement.

        Args:
            cmake_vars: Dictionary of all CMake variables
        """
        # Define dependency name and CMake variable names for this requirement
        name = "Python"
        found_var = "Python_FOUND"
        status_var = "NVFUSER_REQUIREMENT_Python_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_Python_OPTIONAL"
        version_found_var = "Python_VERSION"
        version_required_var = "NVFUSER_REQUIREMENT_Python_VERSION_MIN"
        location_var = "Python_EXECUTABLE"

        super().__init__(
            name,
            cmake_vars,
            found_var,
            status_var,
            optional_var,
            version_found_var,
            version_required_var,
            location_var,
        )

    def generate_help(self, platform_info):
        """
        Generate Python installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        version_min = self.version_required

        print(f"Python {version_min}+ Required")
        print()
        print(
            "Why: nvFuser requires modern Python with type hints and language features."
        )
        print()
        print(f"Install Python {version_min} or higher:")
        print()

        os_type = platform_info["os"]

        if os_type == "Linux":
            if platform_info.get("ubuntu_based"):
                print("  Option 1: Ubuntu/Debian system package:")
                print()
                print("    sudo apt update")
                print(f"    sudo apt install python{version_min}")
                print()
            else:
                print("  Option 1: System package manager:")
                print()
                print("    # Example for RHEL/CentOS:")
                print(f"    # sudo yum install python{version_min}")
                print()

        elif os_type == "Darwin":
            print("  Option 1: Homebrew:")
            print()
            print(f"    brew install python@{version_min}")
            print()

        print("  Option 2: Conda:")
        print()
        print(f"    conda create -n nvfuser python={version_min}")
        print("    conda activate nvfuser")
        print()


# SPDX-License-Identifier: BSD-3-Clause
"""Python dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class PythonRequirement(VersionRequirement):
    """
    Python interpreter requirement.

    CMake variables used:
    - Python_FOUND: Whether Python was found
    - Python_VERSION: Detected version (e.g., "3.12.3")
    - Python_EXECUTABLE: Path to python binary
    - NVFUSER_REQUIREMENT_Python_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Python_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_Python_OPTIONAL: Whether Python is optional
    """

    def __init__(self, cmake_vars: Dict):
        """
        Initialize Python requirement.

        Args:
            cmake_vars: Dictionary of all CMake variables
        """
        # Define dependency name and CMake variable names for this requirement
        name = "Python"
        found_var = "Python_FOUND"
        status_var = "NVFUSER_REQUIREMENT_Python_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_Python_OPTIONAL"
        version_found_var = "Python_VERSION"
        version_required_var = "NVFUSER_REQUIREMENT_Python_VERSION_MIN"
        location_var = "Python_EXECUTABLE"

        super().__init__(
            name,
            cmake_vars,
            found_var,
            status_var,
            optional_var,
            version_found_var,
            version_required_var,
            location_var,
        )

    def generate_help(self, platform_info):
        """
        Generate Python installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        version_min = self.version_required

        print(f"Python {version_min}+ Required")
        print()
        print(
            "Why: nvFuser requires modern Python with type hints and language features."
        )
        print()
        print(f"Install Python {version_min} or higher:")
        print()

        os_type = platform_info["os"]


        # Define dependency name and CMake variable names for this requirement
        name = "Python"
        found_var = "Python_FOUND"
        status_var = "NVFUSER_REQUIREMENT_Python_STATUS"
        optional_var = "NVFUSER_REQUIREMENT_Python_OPTIONAL"
        version_found_var = "Python_VERSION"
        version_required_var = "NVFUSER_REQUIREMENT_Python_VERSION_MIN"
        location_var = "Python_EXECUTABLE"

        super().__init__(
            name,
            cmake_vars,
            found_var,
            status_var,
            optional_var,
            version_found_var,
            version_required_var,
            location_var,
        )

    def generate_help(self, platform_info):
        """
        Generate Python installation help.

        Args:
            platform_info: Platform detection dict from detect_platform()
        """
        version_min = self.version_required

        print(f"Python {version_min}+ Required")
        print()
        print(
            "Why: nvFuser requires modern Python with type hints and language features."
        )
        print()
        print(f"Install Python {version_min} or higher:")
        print()

        os_type = platform_info["os"]

        if os_type == "Linux":
            if platform_info.get("ubuntu_based"):
                print("  Option 1: Ubuntu/Debian system package:")
                print()
                print("    sudo apt update")
                print(f"    sudo apt install python{version_min}")
                print()
            else:
                print("  Option 1: System package manager:")
                print()
                print("    # Example for RHEL/CentOS:")
                print(f"    # sudo yum install python{version_min}")
                print()

        elif os_type == "Darwin":
            print("  Option 1: Homebrew:")
            print()
            print(f"    brew install python@{version_min}")
            print()

        print("  Option 2: Conda:")
        print()
        print(f"    conda create -n nvfuser python={version_min}")
        print("    conda activate nvfuser")
        print()


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import os

__all__ = [
    "cmake_prefix_path",
]

cmake_prefix_path = os.path.join(
    os.path.dirname(os.path.dirname(__file__)),
    "nvfuser_common",
    "share",
    "cmake",
    "nvfuser",
)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os

def pytest_configure(config):
    """
    Hook called after command line options have been parsed and all plugins
    and initial conftest files have been loaded.
    """
    # Append to NVFUSER_ENABLE environment variable for all tests in this directory
    # Note: id_model is now enabled by default, so we only need extra validation
    existing = os.environ.get("NVFUSER_ENABLE", "")
    new_options = "id_model_extra_validation"

    if existing:
        os.environ["NVFUSER_ENABLE"] = f"{existing},{new_options}"
    else:
        os.environ["NVFUSER_ENABLE"] = new_options


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import math
from contextlib import contextmanager
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from thunder.dynamo import thunderfx

# Sizes used in Llama 4 Maverick
@dataclass
class Config:
    hidden_size: int = 5120
    intermediate_size: int = 8192
    num_routed_experts: int = 128


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

from opinfo_core import OpInfo
from opinfo_utils import ArgumentType, is_tensor

import sys

if "nvfuser" in sys.modules:
    from nvfuser import FusionDefinition
    from nvfuser.pytorch_utils import (
        python_scalar_to_nvfuser_dtype,
        torch_dtype_to_nvfuser_dtype,
    )
else:
    from nvfuser_direct import FusionDefinition
    from nvfuser_direct.pytorch_utils import (


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

from opinfo_utils import (
    all_dtypes_except_reduced,
    ArgumentType,
    torch_to_python_dtype_map,
)
from typing import Callable, Optional
import torch
from enum import Enum
from dataclasses import dataclass, field

class ReferenceType(Enum):
    Pytorch = 0
    Jax = 1
    Numpy = 2
    Python = 3


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import itertools
from functools import partial, wraps

import math
import torch
from torch.testing import make_tensor
import random
from numbers import Number

from opinfo_core import OpInfo, SampleInput, ErrorSample, Domain
from opinfo_utils import (
    make_number,
    find_nonmatching_dtype,
    is_floating_dtype,
    float_complex_dtypes,
    complex_dtypes,


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

# import nvfuser_direct first to conditionally avoid importing nvfuser
import nvfuser_direct  # noqa: F401,F403

import torch
import pytest
import numpy as np
from copy import deepcopy
from typing import Callable

from opinfo_fusion_definitions import default_fd_fn
from opinfo_framework import create_op_test
from opinfo_core import ReferenceType, OpInfo, SampleInput
from opinfos import opinfos
from opinfo_utils import (
    ArgumentType,
    is_tensor,


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import inspect
import torch
from typing import Callable
from opinfo_utils import map_dtype_to_str
import pytest

def _instantiate_opinfo_test_template(
    template: Callable, *, opinfo, dtype: torch.dtype
) -> Callable:
    """Instantiates a test template for an operator."""

    def test():
        # Ref: https://github.com/pytorch/pytorch/blob/aa8ea1d787a9d21b064b664c5344376265feea6c/torch/testing/_internal/common_utils.py#L2251-L2263
        # > CUDA device side error will cause subsequence test cases to fail.
        # > stop entire test suite if catches RuntimeError during torch.cuda.synchronize().


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import math
import torch
from looseversion import LooseVersion
from opinfo_core import OpInfo, ReferenceType, Domain
from opinfo_fusion_definitions import (
    api_test_fd_fn,
    tensor_input_fd_fn,
    tensor_api_test_fd_fn,
    vector_api_test_fd_fn,
)
from opinfo_input_generators import (
    argsort_generator,
    topk_generator,
    topk_error_generator,
    broadcast_error_generator,
    broadcast_in_dim_generator,


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
import torch
from torch.testing import make_tensor
from enum import Enum, auto
from typing import Optional
from functools import wraps

try:
    # flake8: noqa
    import jax

    JAX_AVAILABLE = True
except ImportError as e:
    JAX_AVAILABLE = False
    pass


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
from copy import deepcopy
import torch
from torch.testing._internal.common_utils import TestCase

from nvfuser_direct import FusionDefinition, LRUCache
from python.direct_utils import is_pre_volta, check_captured_python_definition

class NVFuserTest(TestCase):
    def __init__(self, cache=None):
        super().__init__()
        self.cache = cache

    # Helper function to verify the nvfuser output and make sure the string
    # definition based on the FusionDefinition is executable and matches the


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

from nvfuser_direct import (
    FusionDefinition,
    DataType,
)
from python.direct_utils import (
    FLOAT4_E2M1_MAX,
    FLOAT8_E4M3_EPS,
    FLOAT8_E4M3_MAX,
    pytorch_nvfp4_quantize,
    microarchitecture_is,
    linear_to_swizzled_128_4,
    round_up,
    activation_scale_to_nvfp4,
)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

def test_import_correct():
    try:
        import nvfuser_direct  # noqa: F401
    except Exception as e:
        raise RuntimeError("Failed to import nvfuser_direct.")


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
import torch
from nvfuser_direct import nvf_cutlass
from python.direct_utils import microarchitecture_is

if not microarchitecture_is(10, 0):
    pytest.skip(
        reason="MxFp8 Requires compute capability 10.",
        allow_module_level=True,
    )

from python.direct_utils import (
    linear_to_swizzled_128_4,
    swizzled_to_linear_128_4,
)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
import torch
from python.direct_utils import microarchitecture_is
from nvfuser_direct import nvf_cutlass

# GPU Compute Capability: https://developer.nvidia.com/cuda/gpus
# tested on blackwell compute 10.0 (B200 and GB200)
# doesn't support 12.0 (RTX PRO 6000 and RTX 50XX)
# Not tested on 10.3 (B300 and GB300)
# Not tested on 12.1 (DGX Spark)
@pytest.mark.skipif(
    not microarchitecture_is(10, 0),
    reason="Does not support blackwell compute 12.0, other arches are not tested.",
)
@pytest.mark.parametrize("config", [[1024, 128, 256], [267, 128, 256]])


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
from nvfuser_direct import FusionDefinition, ParallelType, DataType

def test_matmul():
    c = 3

    with FusionDefinition() as fd:
        inp = fd.define_tensor([-1, -1], contiguity=True, dtype=DataType.Float)
        w = fd.define_tensor([-1, -1], contiguity=True, dtype=DataType.Float)
        out = fd.ops.matmul(inp, w)
        fd.add_output(out)

        out.outer_split(1, c)
        out.axis(1).parallelize(ParallelType.stream)
        # With NVFUSER_DUMP=host_ir, you'll see the host IR container like the


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import itertools
import math
import pytest
import torch
import torch.nn.functional as F
from enum import Enum, auto
from functools import partial

from nvfuser_direct import (
    FusionDefinition,
    DataType,
)
from python.direct_utils import (
    is_pre_ampere,
    define_sdpa_rng_state,
    verify_stride_order,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import torch._refs as refs
import torch._prims as prims

from nvfuser_direct import (
    FusionDefinition,
    DataType,
    compute_tensor_descriptor,
)
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

import pytest
import itertools
from python.direct_utils import (
    is_pre_ampere,
    is_pre_hopper,


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

from nvfuser_direct import (
    FusionDefinition,
    DataType,
)
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from python.direct_utils import (
    FLOAT4_E2M1_MAX,
    FLOAT8_E4M3_MAX,
    pytorch_nvfp4_quantize,
    microarchitecture_is,
    is_pre_blackwell,
    microarchitecture_is_pre,
    linear_to_swizzled_128_4,
    round_up,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# This file contains certain building blocks of the AlphaFold3 model.

import pytest
import torch
from dataclasses import dataclass
from enum import Enum, auto

from nvfuser_direct import FusionDefinition, DataType, TensorView

@dataclass
class ModelConfig:
    c_z: int = 128
    c_hidden: int = 32
    n_heads: int = 4


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import pytest
from nvfuser_direct import FusionDefinition, DataType
from python.direct_utils import skip_if_global_memory_below_gb

# Use smaller range for torch.testing.make_tensor for nvfuser_direct.validate
LOW_VAL = -2
HIGH_VAL = 2

def test_issue1129(nvfuser_direct_test):
    """
    Test for issue 1129 - tests reshape and index_select operations with strided tensors.

    This test verifies that reshape and index_select operations work correctly
    with tensors that have non-standard strides, particularly when reshaping


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
import torch
import torch.nn.functional as F
from nvfuser_direct import (
    FusionDefinition,
    IdMappingMode,
    ParallelType,
    TensorView,
    Merge,
    Split,
    BroadcastOp,
    SqueezeOp,
    ReshapeOp,
    LoadStoreOpType,
    MemoryType,
    DataType,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import pytest
import torch
from nvfuser_direct import nvf_cutlass
from python.direct_utils import microarchitecture_is

if not microarchitecture_is(10, 0):
    pytest.skip(
        reason="Nvfp4 Requires compute capability 10.0, other arches are not tested.",
        allow_module_level=True,
    )

from python.direct_utils import (
    FLOAT4_E2M1_MAX,
    FLOAT8_E4M3_MAX,
    dequantize_to_dtype,
    linear_to_swizzled_128_4,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import pytest
from python.direct_utils import verify_stride_order
from nvfuser_direct import FusionDefinition
from functools import partial
import itertools
import torch.nn.functional as F

def test_matmul(nvfuser_direct_test):
    m = 24
    n = 16
    k = 8
    inputs_tt = [
        torch.randn(m, k, device="cuda", dtype=torch.float16),
        torch.randn(k, n, device="cuda", dtype=torch.float16),


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import torch._refs as refs
import torch._prims as prims

import io
from contextlib import redirect_stdout, redirect_stderr
import pytest

from nvfuser_direct import FusionDefinition, DataType
import random
import itertools
import math
from functools import partial
from typing import List


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch

from nvfuser_direct import FusionDefinition, DataType

def test_contiguous():
    with FusionDefinition() as fd:
        inp = fd.define_tensor([-1, -1], dtype=DataType.Float, contiguity=True)
        out = fd.ops.sum(inp, [-1])
        fd.add_output(out)

    # It might not be obvious but this tensor is indeed contiguous by
    # definition. inp.size(0) == 1 so inp.stride(0) shouldn't matter.
    #
    # I ran into such a tensor when I shard a reference tensor for a particular
    # device using slicing:
    # ```


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

from nvfuser_direct import (
    FusionDefinition,
    PythonProfiler,
    DataType,
    version,
    LruFusionCache,
    LRUCache,
)
import torch
import pytest
import io
import re
from contextlib import redirect_stdout, redirect_stderr
from typing import Callable


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

import transformer_engine.pytorch as te

from . import Parallelism
from .benchmark_utils import get_benchmark_fns
from enum import auto, Enum

compute_cap = torch.cuda.get_device_capability()

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition
from collections.abc import Iterable
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement, Shard
from typing import Callable, cast, TypeAlias

DTensorsKey: TypeAlias = tuple[tuple[str, str], ...]

def make_key_from_dtensors(dtensors: Iterable[DTensor]) -> DTensorsKey:
    key = tuple((repr(dt.device_mesh), repr(dt.placements)) for dt in dtensors)
    return key

class FusionDefinitionWrapper:


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch
import torch.distributed as dist

from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Shard, Replicate
from .fusion_definition_wrapper import FusionDefinitionWrapper
from .linear import LinearFunction
from nvfuser_direct import FusionDefinition, DataType

@pytest.mark.mpi
def test_plus_one(setup_default_process_group, multidevice_test):
    def define_fusion(fd: FusionDefinition):
        inp = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        one = fd.define_scalar(1.0, dtype=DataType.Float)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Run command:
# mpirun -np [num_devices] pytest tests/python/multidevice/[test_name].py --only-mpi -s

import os
import pytest
from typing import Iterable

import torch
import torch.distributed as dist

import nvfuser_direct as nvfuser

class MultideviceTest:
    def __init__(self):
        os.environ["NVIDIA_TF32_OVERRIDE"] = "0"


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest

import torch

import nvfuser_direct as nvfuser

# Unit tests for multidevice_test.shard_tensor.
class TestShardTensor:
    @pytest.mark.mpi
    def test_inner_split(self, multidevice_test):
        d = multidevice_test.size

        with nvfuser.FusionDefinition() as fd:
            tv = fd.define_tensor([-1])
            fd.add_output(tv)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import os

import torch
import torch.distributed as dist
from torch.distributed.tensor import distribute_tensor, Shard

import nvfuser_direct as nvfuser
from .benchmark_utils import get_benchmark_fns
from nvfuser_direct import DataType, FusionDefinition, CommunicatorBackend, TensorView
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

def row_parallel_linear_forward(
    h: int, num_devices: int, num_chunks: int
) -> FusionDefinition:
    with FusionDefinition() as fd:


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from dataclasses import dataclass
from enum import auto, Enum
from functools import lru_cache
from typing import Iterable

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement

from nvfuser_direct import DataType, FusionDefinition
from .fusion_definition_wrapper import FusionDefinitionWrapper

@dataclass(frozen=True)
class LinearConfig:
    in_features: int


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Run command:
# mpirun -np 1 pytest tests/python/multidevice/test_deepseek_v3.py --only-mpi -s

from contextlib import contextmanager
from enum import Enum, auto
from functools import wraps
from typing import Optional

import pytest

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ParallelStyle,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from enum import auto, Enum

class Parallelism(Enum):
    # https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/parallelisms.html#tensor-parallelism
    TENSOR_PARALLEL = auto()
    # https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/parallelisms.html#sequence-parallelism
    SEQUENCE_PARALLEL = auto()


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# This file contains certain building blocks of the AlphaFold3 model.

import pytest
import torch
from dataclasses import dataclass
from enum import Enum, auto

import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition, DataType, TensorView

@dataclass
class ModelConfig:
    c_z: int = 128
    c_hidden: int = 32
    n_heads: int = 4


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch

import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition

@pytest.mark.mpi
def test_allgather(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((d * 4,), contiguity=True, dtype=DataType.Float)
        out = fd.ops.set(inp_tv)
        fd.add_output(out)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch
import torch.distributed as dist

# Partitions n into k integers whose sum is n.
def random_partition(n: int, k: int) -> torch.Tensor:
    offsets = torch.randint(0, n + 1, [k - 1]).tolist()
    # The last partition ends at index n.
    offsets.append(n)
    offsets.sort()

    sizes = []
    prev = 0
    for offset in offsets:
        sizes.append(offset - prev)
        prev = offset


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest

import torch
import torch.nn.functional as F

import nvfuser_direct as nvfuser
from . import Parallelism
from .benchmark_utils import get_benchmark_fns
from nvfuser_direct import DataType, FusionDefinition
from python.direct_utils import (
    create_sdpa_rng_tensors,
    is_pre_ampere,
)

@pytest.mark.mpi
def test_grouped_mlp(multidevice_test):
    d = multidevice_test.size


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch

import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition, PythonProfiler

# Avoid doing this when possible. This test started to exist before nvFuser
# supports DID loop split. As a result of that, the weight in this test has to be
# 3D, different from a normal linear.
@pytest.mark.mpi
def test_linear_logical_split(multidevice_test):
    d = multidevice_test.size

    b, s, e = 2, 1024, 768

    with FusionDefinition() as fd:


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch

def get_benchmark_fn(func, /, profile: bool):
    def wrapper(*args, **kwargs):
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        result = func(*args, **kwargs)
        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStop()
        return result

    return wrapper

# Returns two functors, the first with profiler off and the second with profiler


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch
from enum import Enum, auto
from torch.nn.attention import SDPBackend

import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition
from python.direct_utils import is_pre_ampere

@pytest.mark.mpi
def test_sizes_and_ranks(multidevice_test):
    size, rank, local_size, local_rank = (
        multidevice_test.size,
        multidevice_test.rank,
        multidevice_test.local_size,
        multidevice_test.local_rank,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

# restore swizzled on block scaling factor:
# 1. restore swizzle
# 2. removes padding via slicing to [:mn, :k]
def swizzled_to_linear_128_4(a_sf_swizzled: torch.Tensor, mn, k):
    mn_padded, sf_k_padded = a_sf_swizzled.shape
    m_tiles = mn_padded // 128
    k_tiles = sf_k_padded // 4
    tmp = torch.reshape(a_sf_swizzled, (m_tiles, k_tiles, 32, 4, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, sf_k_padded)[:mn, :k]


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from . import utils, narrow_precision  # noqa: F401,F403
from .utils import *  # noqa: F401,F403
from .narrow_precision import *  # noqa: F401,F403


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import pytest
from nvfuser_direct import FusionDefinition, DataType, TensorView
from looseversion import LooseVersion

def microarchitecture_is(major, minor):
    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
    return prop.major == major and prop.minor == minor

def microarchitecture_is_pre(major):
    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
    return prop.major < major


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import os
import pathlib
import importlib.util
from setuptools import setup

try:
    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
except ImportError:
    raise RuntimeError("PyTorch cpp_extension is required for building this package.")

nvfuser_csrc_dir = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..", "csrc"
)
dynamic_type_dir = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..", "lib", "dynamic_type", "src"
)
flatbuffers_dir = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import torch
import nvfuser_extension  # noqa: F401

torch.manual_seed(0)
t = torch.randn((5, 5), device="cuda")
expected = torch.sinh(t)
output = torch.ops.myop.sinh_nvfuser(t)

print("Expected:", expected)
print("Output:", output)

assert torch.allclose(output, expected)
print("They match!")


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from . import nvfuser_extension  # noqa: F401
from nvfuser_extension import *  # noqa: F401,F403


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# TODO Update script to use nvfuser_direct module
from nvfuser import (
    FusionDefinition,
    SchedulerType,
    ClusterDims,
    MatMulTileOptions,
    GemmTile,
    MmaMacroEncode,
    MmaMacroArch,
    MatmulTileRasterizationOrder,
    MatmulCircularBufferingStrategy,
)
import torch
from torch.autograd import DeviceType
from torch.profiler import profile, record_function, ProfilerActivity
import math
import itertools
from enum import IntEnum


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import math
import itertools
from nvfuser import FusionCache, FusionDefinition
from dataclasses import dataclass, astuple

# ================================ Description ================================
# This file contains the utility function for autotuning scripts.
# =============================================================================

# Returns the result of a/b rounded to the nearest integer in the direction of
# positive infinity.
def ceil_div(a, b):
    return int(math.ceil(a / b))


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import itertools
import random
from nvfuser import FusionCache, FusionDefinition, SchedulerType, DataType
from nvfuser.pytorch_utils import torch_dtype_to_nvfuser_dtype
from copy import deepcopy

# ============================ Description ============================

# 1. Define a nvfuser fusion and its pytorch eager mode reference.
#
# 2. Profile the CUDA kernel performance by iterating over a set of input
# arguments and scheduler configurations.
#
# 3. Train a regression model to predict the desired performance metric given
# some input arguments and a scheduler configuration.


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import itertools
import math

# TODO Update script to use nvfuser_direct module
from nvfuser import FusionDefinition, SchedulerType, DataType
from dataclasses import dataclass
from enum import Enum

from autotune_utils import (
    ScriptConfiguration,
    collect_data,
    separate_data,
    test_model_rmse,
    test_model,
)


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import itertools

# TODO Update script to use nvfuser_direct module
from nvfuser import FusionDefinition, SchedulerType

# Description of the problem
M = 512
N = 512
K = 4096
dtype = torch.bfloat16
# TODO: layout

# These are the parameters we'll optimize
parameter_configurations = [
    splitk_factors := list(range(1, 8)),


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch
import itertools

# TODO Update script to use nvfuser_direct module
from nvfuser import FusionDefinition, SchedulerType, DataType, ParallelType
from enum import Enum
from dataclasses import dataclass

from autotune_utils import (
    ScriptConfiguration,
    collect_data,
    separate_data,
    test_model_rmse,
    test_model,
    ceil_div,
    floor_div,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import os
import sys
import datetime
import pathlib

sys.path.insert(0, os.path.abspath("../.."))

project = "nvFuser"
author = "NVIDIA CORPORATION & AFFILIATES"

# Copyright statement
release_year = 2023
current_year = datetime.date.today().year
if current_year == release_year:
    copyright_year = release_year
else:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# "python benchmark_thunder.py -h" for help.

import argparse
import os
import subprocess
from typing import Iterable

# Switches to the given branch, syncs to head, and returns the short hash of
# the head.
def switch_to(branch: str, sync: bool) -> str:
    subprocess.check_call(f"git fetch origin {branch}", shell=True)

    # `advice.detachedHead=false` silences the detached HEAD warning.
    subprocess.check_call(
        f"git -c advice.detachedHead=false checkout {branch}", shell=True
    )


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# "compare_benchmark.py -h" for help.

import argparse
from dataclasses import dataclass
import json
import matplotlib.pyplot as plt
import numpy as np
import os
import subprocess
from typing import Iterable

# If `arg` uses a forbidden option, returns that forbidden option; otherwise,
# returns an empty string.
def uses_forbidden_option(arg: str) -> str:
    for forbidden_option in (
        "--benchmark_out",


tools/memory.py

# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

def get_available_memory_gb():
    """Returns the available memory in GB."""
    try:
        import psutil

        return psutil.virtual_memory().available / 1024 / 1024 / 1024
    except:  # noqa: E722
        pass

    try:
        with open("/proc/meminfo", "r") as f:
            while True:
                line = f.readline()
                if line.startswith("MemAvailable:"):
                    mem = line.split()[1]
                    assert line.split()[2] == "kB"


tools/cpp-repro-gen.py

# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import ast
import typing
import sys
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--symbolic_sizes", nargs="+", type=int)
cmd_args = parser.parse_args()
symbolic_sizes_index = 0

input = sys.stdin.read()

lines = input.strip().splitlines()

is_aten = False

def parse(l: str):


import os

ext_map = {
    "cpp": ("c", "cpp", "cu", "cc", "cuh", "h"),
    "py": ("py",),
    "txt": ("txt",),
}

licence_style_0 = """\
// clang-format off
/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 */
// clang-format on
"""

licence_style_1 = """\
# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Generates a C++ header files embedding the original input as a string literal

import argparse
import pathlib
from datetime import datetime

arg_parser = argparse.ArgumentParser(
    description="Converts source files to C++ string literals", allow_abbrev=False
)

arg_parser.add_argument("-i", "--input", required=True, help="Input source file")

arg_parser.add_argument(
    "-o", "--output", required=True, help="Name of the generated header file"
)

args = arg_parser.parse_args()


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import subprocess
import os
import glob
import datetime
from pathlib import Path
import time
import argparse

def setup_logging_dir():
    """Create a timestamped directory for test logs and update latest symlink"""
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_dir = Path(f"test_log_{timestamp}")
    log_dir.mkdir(exist_ok=True)

    # Check and handle the 'latest' symlink
    latest_link = Path("test_log_latest")
    if latest_link.exists():


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import subprocess
import sys
from pathlib import Path

UNKNOWN = "Unknown"
nvfuser_root = Path(__file__).parent.parent

# note that this root currently is still part of pytorch.
def get_sha() -> str:
    try:
        return (
            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=nvfuser_root)
            .decode("ascii")
            .strip()
        )
    except Exception:
        import os


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

"""
This script reproduces the computation used for generating the validation tolerances in csrc/validator_utils.h:
https://github.com/NVIDIA/Fuser/blob/892b7ac5646b3b873f13f2b2ea1db15fa9b9effb/csrc/validator_utils.h#L25-L46

The computation (randn + mul + sum) is repeated a large number of times.
The base tolerance is the max error seen when comparing the desired precision type with the next higher precision.
The final tolerances are computed by adding a safety factor of 3/4 to this base tolerance.

The output of this file is a .npy file corresponding to each dtype, containing a {size: safety_factor * base_tolerances} dict.

Note: The script may take a long time for a high number of num_iters, so it may be useful to only run the script for one dtype at a time. Modify `dtype_to_ref_dtypes` accordingly.

To run: python tools/generate_validation_tolerances.py
To load the files:
    ```
    import numpy as np
    tol_dict = np.load(file_name.npy, allow_pickle=True).item()


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Find corresponding .cu files for matching tests, even when new tests are
introduced between two commits. Diffs are displayed and the return value is the
number of mismatched corresponding tests.

Tests are skipped if they produce different numbers of .cu files, or if they
exist in only one of the given runs.

Example usage:
    python tools/diff_codegen_nvfuser_tests.py \
            codegen_comparison/{$commit1,$commit2}/binary_tests
"""

from dataclasses import asdict, dataclass, field, InitVar
import difflib
from enum import Enum
import os
import re


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Find corresponding .cu files for matching tests, even when new tests are
introduced between two commits. Diffs are displayed and the return value is the
number of mismatched corresponding tests.

Tests are skipped if they produce different numbers of .cu files, or if they
exist in only one of the given runs.

Example usage:
    python tools/diff_codegen_nvfuser_tests.py \
            codegen_comparison/{$commit1,$commit2}/binary_tests
"""

import os

from codediff import TestRun, TestDifferences

if __name__ == "__main__":


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
T0 = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float)
T1 = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float)
T2 = fd.define_tensor(symbolic_sizes=[-1, -1], contiguous=[True, True], dtype=DataType.Half)
T3 = fd.ops.broadcast_in_dim(T0, output_shape=[1, 1024, 768], broadcast_dims=[2])
T4 = fd.ops.broadcast_in_dim(T1, output_shape=[1, 1024, 768], broadcast_dims=[2])
T5 = fd.ops.view(T2, original_shape=[1024, 768], new_shape=[1, 1024, 768])
T6 = fd.ops.cast(T5, dtype=DataType.Float)
S7 = fd.define_scalar(0.500000)
T8 = fd.ops.mul(T6, S7)
S9 = fd.define_scalar(0.707107)
T10 = fd.ops.mul(T6, S9)
T11 = fd.ops.erf(T10)
S12 = fd.define_scalar(1.00000)
T13 = fd.ops.add(T11, S12)
T14 = fd.ops.mul(T8, T13)
T15 = fd.ops.cast(T14, dtype=DataType.Half)
T16 = fd.ops.cast(T15, dtype=DataType.Float)
T17, T18 = fd.ops.var_mean(T16, dims=[2], correction=0, keepdim=False)


#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

"""
Load environment options from YAML configuration.

This module provides a loader for env_options.yaml that replaces the
hardcoded Python definitions in configure_env.py.
"""

from __future__ import annotations

import yaml
from pathlib import Path
from dataclasses import dataclass, field
from typing import Literal

@dataclass


#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

"""
Curses-based TUI for nvFuser environment configuration.
Provides a ccmake-like interface for navigating and configuring options.
"""

from __future__ import annotations

import curses
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from configure_env import EnvVarConfig

from configure_env import CATEGORY_NAMES


#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

"""
Interactive nvFuser Environment Configuration Tool

This tool provides an interactive interface (similar to ccmake) for configuring
nvFuser environment variables. It helps users set up debug flags, feature toggles,
and runtime options without needing to remember complex NVFUSER_* variable names.

Usage:
    python tools/configure_env.py                    # Interactive TUI mode
    python tools/configure_env.py --simple           # Simple prompt mode
    python tools/configure_env.py --generate-script  # Generate shell script
"""

from __future__ import annotations

import argparse


import argparse
import concurrent.futures
import json
import logging
import os
import re
import shutil
import subprocess
import sys
import sysconfig
import time
from enum import Enum
from pathlib import Path
from typing import Any, List, NamedTuple, Optional, Pattern

# Nvfuser directory root
result = subprocess.run(
    ["git", "rev-parse", "--show-toplevel"],
    stdout=subprocess.PIPE,
    check=True,
)


import argparse
import concurrent.futures
import json
import logging
import os
import subprocess
import sys
import time
from enum import Enum
from pathlib import Path
from typing import Any, List, NamedTuple, Optional

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"


import argparse
import json
import logging
import os
import re
import subprocess
import sys
import time
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, NamedTuple, Optional, Pattern

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):


import argparse
import concurrent.futures
import json
import logging
import os
import subprocess
import sys
import time
from enum import Enum
from typing import Any, List, NamedTuple, Optional, BinaryIO

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"


"""
Generic linter that greps for a pattern and optionally suggests replacements.
"""

import argparse
import json
import logging
import os
import subprocess
import sys
import time
from enum import Enum
from typing import Any, List, NamedTuple, Optional

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)


"""
EXEC: Ensure that source files are not executable.
"""
import argparse
import json
import logging
import os
import sys

from enum import Enum
from typing import NamedTuple, Optional

LINTER_CODE = "EXEC"

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"


"""
Initializer script that installs stuff to pip.
"""
import os
import argparse
import logging
import subprocess
import sys
import time

from typing import List

def run_command(args: List[str]) -> "subprocess.CompletedProcess[bytes]":
    logging.debug("$ %s", " ".join(args))
    start_time = time.monotonic()
    try:
        return subprocess.run(args, check=True)
    finally:
        end_time = time.monotonic()
        logging.debug("took %dms", (end_time - start_time) * 1000)


import argparse
import json
import logging
import os
import re
import subprocess
import sys
import time
from enum import Enum
from typing import Any, Dict, List, NamedTuple, Optional, Set, Pattern

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"


"""
NEWLINE: Checks files to make sure there are no trailing newlines.
"""
import argparse
import json
import logging
import os
import sys

from enum import Enum
from typing import NamedTuple, Optional

NEWLINE = 10  # ASCII "\n"
LINTER_CODE = "NEWLINE"

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
from .normalization import norm_bwd_nvf_benchmark, norm_bwd_baseline_benchmark
from .core import DEFAULT_EXECUTORS

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_batchnorm_bwd_nvf_benchmark(
    benchmark,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import silu_mul

def silu_mul_fwd_fusion(fd: FusionDefinition, dtype: DataType):
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import gelu

def gelu_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_layernorm

def dropout_layernorm_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os

def test_env_vars_are_set():
    for var in [
        "TORCHINDUCTOR_COORDINATE_DESCENT_TUNING",
        "TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS",
    ]:
        assert os.environ[var] == "1"


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    with_executor,
    DEFAULT_EXECUTORS,
    check_module_available,
)
import torch
from .global_params import FLOAT_DTYPES, BENCHMARK_CONFIG
from .torch_ops import rmsnorm_add
import itertools
from random import sample
from typing import List, Tuple

if check_module_available("flashinfer"):
    from flashinfer import fused_add_rmsnorm


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import BENCHMARK_CONFIG
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
import os

ORIGINAL_ENV_VARS = {}

def pytest_sessionstart(session):
    for var, value in [
        ("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING", "1"),
        ("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS", "1"),
    ]:
        ORIGINAL_ENV_VARS[var] = os.environ.get(var)
        os.environ[var] = value

def pytest_sessionfinish(session):


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest

import torch

from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
from .global_params import FLOAT_DTYPES
from .torch_ops import scatter_reduce

# (top_k, hidden) configurations seen in models.
TEST_CONFIGS = [
    (8, 7168),  # deepseek r1
]


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from nvfuser_direct import FusionDefinition, DataType
from .core import run_benchmark, with_executor

# These tests are sourced from automation added into thunder.
#
# To recreate:
#   * Use the tfogal/benchmarking-dumps branch in thunder.
#   * Pass nv_store_fusion_inputs=True to your thunderfx call
#   * Disable the torch.compile executor in the thunderfx call, which ensures
#     that all fusions go to nvFuser.
#   * After running the network of interest, invoke thunder.tests.dump.traces()
# This particular file is aimed at concatenation operations, so it is an
# agglomeration of cases that included `torch.cat` in the fusion.

def cat_qwen2_fwd_11(t17353, t17351, cos_2, sin_2, query_states_1, key_states_1):


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import torch

from torch import nn

from typing import Tuple
from functools import partial

from .model_configs import configs

SEQ_LENGTHS = (
    1024,
    2048,
    4096,
    8192,
    12288,
    16384,
    20480,
    24576,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def bcast_add_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    bcast_axis: int,
    contiguous: bool,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False, stride_order=[0]
    )
    stride_order = [0, 1] if not contiguous else [1, 0]


# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.
#
# NOTE: `down_size`, and `pack_uint4` are copied from PyTorch's test code.
#
# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# NOTE: `pytorch_nvfp4_quantize` and `linear_to_swizzled_128_4` are copied from NVIDIA's Fuser's test code.
#
# Pulled from the lightning-thunder repo. Reference:
# https://github.com/Lightning-AI/lightning-thunder/blob/4d3a3c3a7481efdc6a23cdeea99c3ffd31af5e78/thunder/benchmarks/layers_for_inference_benchmark.py

# fmt: off

from __future__ import annotations
from typing import TYPE_CHECKING


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import torch

from .model_configs import configs

class CrossEntropyLossBase:
    def __init__(self, model_name, dtype):
        self.config = configs[model_name]()
        self.dtype = dtype

    def model(self):
        raise NotImplementedError

    def inputs(self):
        hidden_states = torch.randn(
            self.config.batch_size,
            self.config.seq_len,
            self.config.hidden_size,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
from .normalization import norm_bwd_nvf_benchmark, norm_bwd_baseline_benchmark

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_instancenorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

"""Inference benchmark focusing on throughput and latency metrics of prefill and decode phases.

AutoModelForCausalLM from Hugging Face transformers is used for model implementation.

Key metrics:
- Throughput (tokens/second)
- Latency (ms/token)
- Time to First Token (TTFT)
- Time Between Output Tokens (TBOT)

Pulled from the lightning-thunder repo. Reference:
https://github.com/Lightning-AI/lightning-thunder/blob/4d3a3c3a7481efdc6a23cdeea99c3ffd31af5e78/thunder/benchmarks/benchmark_inference.py
"""

# fmt: off

from __future__ import annotations


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest

import torch

from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
from .global_params import FLOAT_DTYPES
from .torch_ops import embedding

# (vocab, hidden) configurations seen in models.
EMBEDDING_CONFIGS = [
    (152064, 3584),  # hf_qwen2


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, head_size: int, dropout_p: float
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation.
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    dropout_p: float,
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

# test the influence of epilogue on the performance of reduction.
# current reduction scheduler only allows epilogue to be fused with outer reduction without post reduction broadcast.
# So, in this test, only outer reduction is tested. [reduction_axis] is kept to allow the extension to inner reduction.

def reduction_epilogue_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    reduction_axis: int,
) -> None:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import rmsnorm

def rmsnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    eps: float = 1e-5,
):
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch
import torch.nn.functional as F

def dropout_layernorm(inputs: list):
    inp1, inp2, weights, bias, dropout_p = inputs
    return F.layer_norm(
        inp2 + torch.nn.functional.dropout(inp1, p=dropout_p),
        normalized_shape=inp1.shape[1:],
        weight=weights,
        bias=bias,
    )

def dropout_rmsnorm(inputs: list):
    inp1, inp2, weights, dropout_p = inputs
    x = inp2 + F.dropout(inp1, p=dropout_p)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import run_benchmark, with_executor, unary_bwd_torch, clear_dynamo_cache

from .rope_ops import rope_setup, SEQ_LENGTHS

@pytest.mark.parametrize(
    "variation",
    [
        "llama_2_7b_hf",
        "llama_3_8B",
        "hf_qwen2",
        "hf_phi3",
        "hf_mistral_nemo",
        "litgpt-gemma-2-9b",
        "litgpt-mistral-7b",
        "litgpt-meta-llama-3-8B",
        "litgpt-phi3.5-mini",


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import scale_bias_relu

def sbr_bwd_fusion(
    fd: FusionDefinition,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import silu_mul

def silu_mul_bwd_fusion(fd: FusionDefinition, dtype: DataType):
    T0 = fd.define_tensor(


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from collections.abc import Iterable
import pytest_benchmark
import torch
from typing import List, Callable, Union
import numpy as np
from nvfuser_direct import FusionDefinition, PythonProfiler
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
from nvfuser_direct.benchmark_utils import FusionProfileTimer, CuptiTimer
import warnings
import thunder
from thunder.executors.nvfuserex import nvfuserex
import importlib.util

# These variables can be overwritten through CLI commands
# --benchmark-rounds=rounds --benchmark-warmup-rounds=warmup_rounds
# --benchmark-num-inputs=num_inputs
BENCHMARK_CONFIG = {
    "rounds": 10,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
import numpy as np
from .torch_ops import softmax

def softmax_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest

import torch
from nvfuser_direct import FusionDefinition, DataType
from .core import (
    run_benchmark,
    with_executor,
    unary_bwd_torch,
    clear_dynamo_cache,
    check_module_available,
)
from .cross_entropy_loss import (
    cross_entropy_loss_setup,
    SyntheticMiniModel,
)
from .torch_ops import cross_entropy as torch_cross_entropy_fwd


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from nvfuser_direct import FusionDefinition, DataType
from .global_params import PROMOTE_DTYPES
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
import torch
from .core import run_benchmark, unary_bwd_torch, clear_dynamo_cache, with_executor
import numpy as np
from functools import reduce

def norm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    norm: str,
    num_dims: int,
    channels_last: bool,
    eps: float = 1e-5,
    momentum: float = 0.01,
) -> None:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def pointwise_mul_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
    T2 = fd.ops.mul(T0, T0)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
from .normalization import norm_fwd_nvf_benchmark, norm_fwd_baseline_benchmark
from .core import DEFAULT_EXECUTORS

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_batchnorm_fwd_nvf_benchmark(
    benchmark,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def reduction_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    reduction_axis: int,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def get_n_groups(C):
    # start num_groups from 1 and increase to 32 at max
    # 32 is a widely used value for num_groups
    # it doesn't make sense to use num_groups > C
    num_groups = 1
    while num_groups * 2 <= 32 and C % (num_groups * 2) == 0:
        num_groups *= 2
    return num_groups


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_rmsnorm

def dropout_rmsnorm_bwd_fusion(
    fd: FusionDefinition,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os
import pytest
import torch
import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition, CommunicatorBackend
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import BENCHMARK_CONFIG, clear_l2_cache

# Run command:
# export NVFUSER_BUILD_WITH_UCC=1
# pip install --no-build-isolation -e python -v
# mpirun -np 1 pytest benchmarks/python/benchmark_overlap.py --with-mpi

class CUDAEventTimer:
    """Custom CUDA event-based timer for accurate GPU timing.


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import gelu

def gelu_bwd_fusion(
    fd: FusionDefinition,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import softmax

def softmax_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1],
        contiguity=[True, True],
        dtype=dtype,
        is_cpu=False,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np

def gelu_bwd_reduction_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    grad = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from functools import partial

def llama_hf_cfg(config_str):
    class Config:
        def __init__(
            self, n_head, head_size, n_query_groups, rope_n_elem, batches, seq_length
        ):
            self.n_head = n_head
            self.head_size = head_size
            self.n_query_groups = n_query_groups
            self.rope_n_elem = rope_n_elem
            self.batches = batches
            self.seq_length = seq_length

    configs = {}
    configs["llama_2_7b_hf"] = Config(
        n_head=32,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import layernorm

def layernorm_bwd_fusion(
    fd: FusionDefinition,


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .benchmark_inference import (
    InferenceBenchmarkConfig,
    InferenceBenchmark,
    _register_nvfp4_ops,
)
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES

# NOTE: for some reason, changing the order of nvfp4 and cudagraph parameter breaks thunder benchmark. I suspect there's something with thunder's cache.
@pytest.mark.parametrize("input_length", [4096])
@pytest.mark.parametrize("output_length", [4])
@pytest.mark.parametrize("mode", ["thunder", "inductor"])
@pytest.mark.parametrize("enable_nvfp4", [True, False])
@pytest.mark.parametrize("enable_cudagraph", [False, True])
def test_llama4_inference_benchmark(
    benchmark,
    input_length: int,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import scale_bias_relu

def sbr_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
):
    T0 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(
        shape=[-1, -1],


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition
from .core import run_benchmark
import torch

import csv
import functools
import os

def matmul_fusion(fd: FusionDefinition, inputs: list[torch.Tensor]) -> None:
    a = fd.from_pytorch(inputs[0])
    b = fd.from_pytorch(inputs[1])
    out = fd.ops.matmul(a, b)
    fd.add_output(out)

def load_matmul_problems():


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import layernorm

def layernorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    eps: float = 1e-5,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import rmsnorm

def rmsnorm_bwd_fusion(
    fd: FusionDefinition,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
from .normalization import norm_fwd_nvf_benchmark, norm_fwd_baseline_benchmark

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_instancenorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_layernorm

def dropout_layernorm_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float, eps: float = 1e-5
) -> None:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_bwd_fusion(


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_rmsnorm

def dropout_rmsnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_bwd_fusion(


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def transpose_fusion_input_smem(
    fd: FusionDefinition,
    dtype: DataType,
    is_copy_transpose: bool,
    axes: list,
    rank: int,
):
    """Single input: the transposed input is read through shared memory."""
    shape = [-1] * rank
    contiguity = [True] * rank


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import torch
from typing import Union, List, Tuple
from nvfuser_direct import DataType
from .core import BENCHMARK_CONFIG
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
import itertools
import os
from random import sample

# BENCHMARK_MODE = weekly/nightly.
BENCHMARK_MODE = os.getenv("BENCHMARK_MODE")
if not BENCHMARK_MODE:
    BENCHMARK_MODE = "nightly"

# Datatypes to benchmark
FLOAT_DTYPES = [torch.float32]
# Run only one of float16 / bfloat16.
if DEVICE_PROPERTIES["gpu_compute_capability_major"] >= 8:


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType, PythonProfiler
from ..core import run_benchmark
import torch

def many_matmul_fusion(fd: FusionDefinition) -> None:
    x = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
    )
    y = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False
    )
    a = fd.ops.add(x, y)
    for _ in range(5):
        a_transpose = fd.ops.permute(a, [1, 0])
        matmul_out = fd.ops.matmul(a_transpose, y)
        add_out = fd.ops.add(a_transpose, y)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from ..core import run_benchmark
import torch
from ..global_params import PROMOTE_DTYPES
from functools import partial

def pointwise_ops_fusion(fd: FusionDefinition, dtype: DataType, num_iters: int):
    x = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    y = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    if dtype in PROMOTE_DTYPES:
        x = fd.ops.cast(x, dtype=DataType.Float)
        y = fd.ops.cast(y, dtype=DataType.Float)

    a = fd.ops.add(x, y)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from ..core import run_benchmark
import torch

def adaptive_layernorm_fwd_fusion(fd: FusionDefinition, eps: float = 1e-6) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1, -1],
        contiguity=[True, True, True],
        dtype=DataType.Half,
        is_cpu=False,
        stride_order=[2, 1, 0],
    )
    T1 = fd.define_tensor(
        shape=[-1, -1],
        contiguity=[True, True],
        dtype=DataType.Half,


python/setup.py

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Usage:
# pip install --no-build-isolation -e python -v
#   This build command is equivalent to: python setup.py develop
#   Options:
#     -v: verbose output
#     --no-build-isolation: don't build in a temporary directory
#     -e: install in development mode
#
# Environment variables used during build:
#   MAX_JOBS
#     maximum number of compile jobs we should use to compile your code
#
#   NVFUSER_BUILD_CMAKE_ONLY
#     Only generate ./build directory with cmake setup
#
#   NVFUSER_BUILD_NO_PYTHON
#     Skips python API target `libnvfuser.so`, i.e. `_C.cpython-xxx.so`