        input_smem_cache.axis(1).parallelize(ParallelType.tma)
        input_smem_cache.axis(2).parallelize(ParallelType.tma)
        input_smem_cache.axis(3).parallelize(ParallelType.tma)
        input_smem_cache.axis(4).parallelize(ParallelType.tma)
        # [BIDx, Bulk, Bulk, Bulk, Bulk]

        if verbose_:
            print(fd.fusion.print_math())
            print(fd.fusion.print_kernel())

    index32bit = CompileParams(
        index_type=DataType.Int32, maxrregcount=255, enable_magic_zero=False
    )
    t0 = torch.randn(10000, 10000, dtype=torch.float, device="cuda:0")
    ke = KernelExecutor()
    ke.compile(fd.fusion, [t0], compile_params=index32bit)
    outputs = ke.run([t0])
    assert outputs[0].equal(t0.t())

def test_tutorial_compute_heuristics_and_schedule():
    """
    Demonstrate explicit scheduling: compute_heuristics, modify, then schedule.
    This shows how to customize automatically computed heuristics.
    """
    inputs = [
        torch.randn(4, 4, device="cuda"),
        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.exp(t2)
        fd.add_output(t3)

        # Step 1: Compute heuristics for pointwise scheduler
        heuristic_params = schedule.compute_heuristics(
            fd.fusion, SchedulerType.pointwise, inputs
        )

        before_modification = """
===== Pointwise Parameters ========
Tag: Pointwise heuristics Pointwise Characteristics:
 Gridx: 1 BlckY: 1 BlckX: 128
vectorization_factor: 1
unroll_factor_outer: 1
unroll_factor_inner: 1
====================================
"""
        assert str(heuristic_params) == before_modification

        # Step 2: Modify the computed heuristics
        # Example: Adjust vectorization and unroll factors
        heuristic_params.vectorization_factor = 1
        heuristic_params.unroll_factor_inner = 2

        after_modification = """
===== Pointwise Parameters ========
Tag: Pointwise heuristics Pointwise Characteristics:
 Gridx: 1 BlckY: 1 BlckX: 128
vectorization_factor: 1
unroll_factor_outer: 1
unroll_factor_inner: 2
====================================
"""
        assert str(heuristic_params) == after_modification

        # Step 3: Apply the schedule using modified heuristics
        schedule.schedule(fd.fusion, SchedulerType.pointwise, heuristic_params)

    schedule_fusion_math = """Inputs:
  T0_g_float[iS61{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iS62{1}, iS60{2}, iS58{128}]
  T1_g_float[iS47{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iS48{1}, iS46{2}, iS44{128}]
Outputs:
  T3_g_float[iblockIdx.x19{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS20{1}, iS18{2}, ithreadIdx.x16{128}] ca_pos( 2 ) produce_pos( 4 )

%kernel_math {
T4_l_float[iblockIdx.x54{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS55{1}, iS53{2}, ithreadIdx.x51{128}] ca_pos( 2 )
   = Set( T0_g_float[iS61{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iS62{1}, iS60{2}, iS58{128}], cache_op=Streaming )
T5_l_float[iblockIdx.x40{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iUS41{1}, iS39{2}, ithreadIdx.x37{128}] ca_pos( 2 )
   = Set( T1_g_float[iS47{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iS48{1}, iS46{2}, iS44{128}], cache_op=Streaming )
T2_l_float[iblockIdx.x33{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS34{1}, iS32{2}, ithreadIdx.x30{128}] ca_pos( 4 ) produce_pos( 2 )
   = T4_l_float[iblockIdx.x54{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS55{1}, iS53{2}, ithreadIdx.x51{128}] ca_pos( 2 )
   + T5_l_float[iblockIdx.x40{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iUS41{1}, iS39{2}, ithreadIdx.x37{128}] ca_pos( 2 );
T6_l_float[iblockIdx.x26{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS27{1}, iS25{2}, ithreadIdx.x23{128}] ca_pos( 4 ) produce_pos( 4 )
   = expf(T2_l_float[iblockIdx.x33{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS34{1}, iS32{2}, ithreadIdx.x30{128}] ca_pos( 4 ) produce_pos( 2 ));
T3_g_float[iblockIdx.x19{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS20{1}, iS18{2}, ithreadIdx.x16{128}] ca_pos( 2 ) produce_pos( 4 )
   = Set( T6_l_float[iblockIdx.x26{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS27{1}, iS25{2}, ithreadIdx.x23{128}] ca_pos( 4 ) produce_pos( 4 ), cache_op=Streaming )
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == schedule_fusion_math

    # Execute with the modified heuristic params
    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0] + inputs[1])
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_pointwise_auto_scheduler():
    """
    Implement a simple pointwise kernel with automatic scheduling.
    Uses nvfuser's PointwiseScheduler.
    """
    inputs = [
        torch.randn(4, 4, device="cuda"),
        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.exp(t2)
        fd.add_output(t3)

        # Apply selected scheduler


T4_l_float[iblockIdx.x54{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS55{1}, iS53{2}, ithreadIdx.x51{128}] ca_pos( 2 )
   = Set( T0_g_float[iS61{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iS62{1}, iS60{2}, iS58{128}], cache_op=Streaming )
T5_l_float[iblockIdx.x40{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iUS41{1}, iS39{2}, ithreadIdx.x37{128}] ca_pos( 2 )
   = Set( T1_g_float[iS47{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iS48{1}, iS46{2}, iS44{128}], cache_op=Streaming )
T2_l_float[iblockIdx.x33{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS34{1}, iS32{2}, ithreadIdx.x30{128}] ca_pos( 4 ) produce_pos( 2 )
   = T4_l_float[iblockIdx.x54{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS55{1}, iS53{2}, ithreadIdx.x51{128}] ca_pos( 2 )
   + T5_l_float[iblockIdx.x40{( ceilDiv(( ceilDiv(( i3 * i4 ), 128) ), 2) )}, iUS41{1}, iS39{2}, ithreadIdx.x37{128}] ca_pos( 2 );
T6_l_float[iblockIdx.x26{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS27{1}, iS25{2}, ithreadIdx.x23{128}] ca_pos( 4 ) produce_pos( 4 )
   = expf(T2_l_float[iblockIdx.x33{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS34{1}, iS32{2}, ithreadIdx.x30{128}] ca_pos( 4 ) produce_pos( 2 ));
T3_g_float[iblockIdx.x19{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS20{1}, iS18{2}, ithreadIdx.x16{128}] ca_pos( 2 ) produce_pos( 4 )
   = Set( T6_l_float[iblockIdx.x26{( ceilDiv(( ceilDiv(( i0 * i1 ), 128) ), 2) )}, iUS27{1}, iS25{2}, ithreadIdx.x23{128}] ca_pos( 4 ) produce_pos( 4 ), cache_op=Streaming )
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == schedule_fusion_math

    # Execute with the modified heuristic params
    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0] + inputs[1])
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_pointwise_auto_scheduler():
    """
    Implement a simple pointwise kernel with automatic scheduling.
    Uses nvfuser's PointwiseScheduler.
    """
    inputs = [
        torch.randn(4, 4, device="cuda"),
        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.exp(t2)
        fd.add_output(t3)

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.pointwise, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0] + inputs[1])
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_reduction_auto_scheduler():
    """
    Implement a simple reduction kernel with automatic scheduling.
    - Expects failure with PointwiseScheduler
    - Uses nvfuser's ReductionScheduler
    """
    inputs = [
        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.ops.sum(t0, dims=[1])
        t2 = fd.ops.exp(t1)
        fd.add_output(t2)

        # Test error msg for can_schedule
        pointwise_status, error_msg = schedule.can_schedule(


        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.exp(t2)
        fd.add_output(t3)

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.pointwise, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0] + inputs[1])
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_reduction_auto_scheduler():
    """
    Implement a simple reduction kernel with automatic scheduling.
    - Expects failure with PointwiseScheduler
    - Uses nvfuser's ReductionScheduler
    """
    inputs = [
        torch.randn(4, 4, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.ops.sum(t0, dims=[1])
        t2 = fd.ops.exp(t1)
        fd.add_output(t2)

        # Test error msg for can_schedule
        pointwise_status, error_msg = schedule.can_schedule(
            fd.fusion, SchedulerType.pointwise, inputs
        )
        assert not pointwise_status
        assert (
            error_msg.strip()
            == "Scheduler _pointwise_ ***rejected*** because : cannot find reference tensor"
        )

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.reduction, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0].sum(1))
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_inner_persistent_auto_scheduler():
    """
    Implement a simple normalization kernel with automatic scheduling.
    Uses nvfuser's InnerPersistentScheduler.
    """
    tensor_size = 4
    inputs = [torch.randn(tensor_size, tensor_size, device="cuda")]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(1e-6, dtype=DataType.Double)
        norm_const = fd.define_scalar(tensor_size, dtype=DataType.Int)

        bcast_sum0 = fd.ops.sum(t0, dims=[-1], keepdim=True)
        mean = fd.ops.div(bcast_sum0, norm_const)

        diff = fd.ops.sub(t0, mean)
        diff_sq = fd.ops.mul(diff, diff)


        # Test error msg for can_schedule
        pointwise_status, error_msg = schedule.can_schedule(
            fd.fusion, SchedulerType.pointwise, inputs
        )
        assert not pointwise_status
        assert (
            error_msg.strip()
            == "Scheduler _pointwise_ ***rejected*** because : cannot find reference tensor"
        )

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.reduction, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    eager_out = torch.exp(inputs[0].sum(1))
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_inner_persistent_auto_scheduler():
    """
    Implement a simple normalization kernel with automatic scheduling.
    Uses nvfuser's InnerPersistentScheduler.
    """
    tensor_size = 4
    inputs = [torch.randn(tensor_size, tensor_size, device="cuda")]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(1e-6, dtype=DataType.Double)
        norm_const = fd.define_scalar(tensor_size, dtype=DataType.Int)

        bcast_sum0 = fd.ops.sum(t0, dims=[-1], keepdim=True)
        mean = fd.ops.div(bcast_sum0, norm_const)

        diff = fd.ops.sub(t0, mean)
        diff_sq = fd.ops.mul(diff, diff)
        bcast_sum1 = fd.ops.sum(diff_sq, dims=[-1], keepdim=True)
        var = fd.ops.div(bcast_sum1, norm_const)

        t0_diff = fd.ops.sub(t0, mean)
        var_eps = fd.ops.sqrt(fd.ops.add(var, s0))
        t0_norm = fd.ops.div(t0_diff, var_eps)
        fd.add_output(t0_norm)

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.inner_persistent, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    var, mean = torch.var_mean(inputs[0], dim=-1, correction=0, keepdim=True)
    eager_out = (inputs[0] - mean) / torch.sqrt(var + 1e-6)
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_scheduling_layer_norm_with_profiling():
    """Test and profile layer norm fusion with manual scheduling."""

    def _definition_func(fd: FusionDefinition, inputs, tensor_size):
        """Define the layer norm fusion operations."""
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(1e-6, dtype=DataType.Double)
        norm_const = fd.define_scalar(tensor_size, dtype=DataType.Int)

        mean_cast = fd.ops.cast(t0, dtype=DataType.Float)
        bcast_sum0 = fd.ops.sum(mean_cast, dims=[-1], keepdim=True)
        mean = fd.ops.div(bcast_sum0, norm_const)

        var_cast = fd.ops.cast(t0, dtype=DataType.Float)
        diff = fd.ops.sub(var_cast, mean)
        diff_sq = fd.ops.mul(diff, diff)
        bcast_sum1 = fd.ops.sum(diff_sq, dims=[-1], keepdim=True)
        var = fd.ops.div(bcast_sum1, norm_const)


        diff_sq = fd.ops.mul(diff, diff)
        bcast_sum1 = fd.ops.sum(diff_sq, dims=[-1], keepdim=True)
        var = fd.ops.div(bcast_sum1, norm_const)

        t0_diff = fd.ops.sub(t0, mean)
        var_eps = fd.ops.sqrt(fd.ops.add(var, s0))
        t0_norm = fd.ops.div(t0_diff, var_eps)
        fd.add_output(t0_norm)

        # Apply selected scheduler
        heuristic_params = check_auto_schedule(schedule.schedule)(
            fd.fusion, SchedulerType.inner_persistent, inputs
        )

    nvf_out = fd.manual_execute(inputs, heuristic_params)
    var, mean = torch.var_mean(inputs[0], dim=-1, correction=0, keepdim=True)
    eager_out = (inputs[0] - mean) / torch.sqrt(var + 1e-6)
    torch.testing.assert_close(eager_out, nvf_out[0])

def test_tutorial_scheduling_layer_norm_with_profiling():
    """Test and profile layer norm fusion with manual scheduling."""

    def _definition_func(fd: FusionDefinition, inputs, tensor_size):
        """Define the layer norm fusion operations."""
        t0 = fd.from_pytorch(inputs[0])
        s0 = fd.define_scalar(1e-6, dtype=DataType.Double)
        norm_const = fd.define_scalar(tensor_size, dtype=DataType.Int)

        mean_cast = fd.ops.cast(t0, dtype=DataType.Float)
        bcast_sum0 = fd.ops.sum(mean_cast, dims=[-1], keepdim=True)
        mean = fd.ops.div(bcast_sum0, norm_const)

        var_cast = fd.ops.cast(t0, dtype=DataType.Float)
        diff = fd.ops.sub(var_cast, mean)
        diff_sq = fd.ops.mul(diff, diff)
        bcast_sum1 = fd.ops.sum(diff_sq, dims=[-1], keepdim=True)
        var = fd.ops.div(bcast_sum1, norm_const)

        t0_cast = fd.ops.cast(t0, dtype=DataType.Float)
        t0_diff = fd.ops.sub(t0_cast, mean)
        var_eps = fd.ops.sqrt(fd.ops.add(var, s0))
        t0_norm = fd.ops.div(t0_diff, var_eps)

        t0_norm_cast = fd.ops.cast(t0_norm, dtype=DataType.BFloat16)
        fd.add_output(t0_norm_cast)

    def _schedule_func(fd: FusionDefinition):
        """Schedule the layer norm fusion."""
        tv_inputs = list(filter(lambda v: v.is_tensor(), fd.fusion.inputs()))
        assert len(tv_inputs) == 1

        tv_outputs = list(filter(lambda v: v.is_tensor(), fd.fusion.outputs()))
        assert len(tv_outputs) == 1

        # create cache tensors
        cache_after_t0 = tv_inputs[0].cache_after()
        cache_after_t0.set_memory_type(MemoryType.shared)

        cache_before_t0_norm = tv_outputs[0].cache_before()
        cache_tvs = [cache_after_t0, cache_before_t0_norm]
        if verbose_:
            print("cache input:\t", cache_after_t0)
            print("cache output:\t", cache_before_t0_norm)

        # Schedule Reference Tensor
        reference_tv = tv_outputs[0]
        reference_tv.split(-1, 256 * 4)
        reference_tv.split(-1, 4)
        fd.sched.transform_like(reference_tv)
        if verbose_:
            print("scheduled reference tensor:\n", reference_tv)

        # Add rfactor TensorViews
        all_tvs = filter(lambda v: v.is_tensor(), fd.fusion.vals())
        reduction_tvs = list(
            filter(
                lambda tv: any(id.is_reduction() for id in tv.get_logical_domain()),
                all_tvs,
            )
        )
        assert len(reduction_tvs) == 2
        rfactor_tvs = [tv.rfactor([-1]) for tv in reduction_tvs]

        # Add common parallelization
        reference_tv.axis(0).parallelize(ParallelType.grid_x)
        reference_tv.axis(-2).parallelize(ParallelType.block_x)
        fd.sched.parallelize_like(reference_tv)
        if verbose_:
            print("parallelized reference tensor:\n", reference_tv)

        # Vectorize input load and output store
        cache_after_t0.axis(-1).parallelize(ParallelType.vectorize)
        tv_outputs[0].axis(-1).parallelize(ParallelType.vectorize)
        if verbose_:
            print("vectorized input load:\n", cache_after_t0)
            print("vectorized output store:\n", tv_outputs[0])

        # Add computeAt; inline_most automatically skips vectorized iterDomains
        fd.sched.inline_most()

    batch_size = 1024
    tensor_size = 4096
    inputs = [
        torch.randn(batch_size, tensor_size, dtype=torch.bfloat16, device="cuda"),
    ]

    print("\n\n===================== Schedule Layer Norm =========================")

    with FusionDefinition() as fd:
        _definition_func(fd, inputs, tensor_size)
        _schedule_func(fd)

    with PythonProfiler(auto_scheduled=False) as prof:
        nvf_out = fd.manual_execute(inputs)

    torch_out = torch.nn.functional.layer_norm(
        inputs[0], normalized_shape=inputs[0].shape[1:]
    )
    print("==================================================================")

    # Change rtol and atol for fp16 dtype
    results_match = torch.allclose(nvf_out[0], torch_out, rtol=1e-2, atol=1e-2)
    assert results_match, "Nvfuser and PyTorch results do not match!"

    print("====================== Profile Kernel =============================")
    assert len(prof.profile.kernel_profiles) == 1
    kp = prof.profile.kernel_profiles[0]

    basic_information = f"Name: {kp.name}, Schedule: {kp.scheduler}, \
    Segment id: {kp.segment_id}, Device: {kp.device}, Stream: {kp.stream}"
    print(basic_information)

    kernel_information = f"Compile time: {kp.compile_time_ms:.2f} ms, \
    Grid: {kp.grid_str}, Block: {kp.block_str}, Registers: {kp.registers}"
    print(kernel_information)

    runtime_information = f"Input size: {kp.input_bytes} bytes, \
    Output size: {kp.output_bytes} bytes, Time: {kp.time_ms:2f} ms"
    print(runtime_information)

    bandwidth_information = f"Effective Bandwidth: {kp.effective_bandwidth_gbs:.2f} GB/s, \
    Peak Bandwidth: {kp.percentage_peak_bandwidth:2f}%"
    print(bandwidth_information)
    print("===================================================================")

    # Validate profiler captured kernel info
    assert prof.profile.fusion_id >= 0
    assert len(prof.profile.kernel_profiles) > 0
    assert prof.profile.kernel_profiles[0].scheduler == "user"

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_warp_specialized_circular_buffering_pointwise():
    # NOTE: Only a functional test, so high performance is not expected.
    def _definition_func(fd: FusionDefinition, inputs):
        tv0 = fd.from_pytorch(inputs[0])
        tv1 = fd.from_pytorch(inputs[1])
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    def _schedule_func(fd: FusionDefinition):
        # Parameters
        number_of_stages = 4
        prefetch_distance = 1
        bulk_inner_dim = 128

        tv_inputs = list(filter(lambda v: v.is_tensor(), fd.fusion.inputs()))


    print(kernel_information)

    runtime_information = f"Input size: {kp.input_bytes} bytes, \
    Output size: {kp.output_bytes} bytes, Time: {kp.time_ms:2f} ms"
    print(runtime_information)

    bandwidth_information = f"Effective Bandwidth: {kp.effective_bandwidth_gbs:.2f} GB/s, \
    Peak Bandwidth: {kp.percentage_peak_bandwidth:2f}%"
    print(bandwidth_information)
    print("===================================================================")

    # Validate profiler captured kernel info
    assert prof.profile.fusion_id >= 0
    assert len(prof.profile.kernel_profiles) > 0
    assert prof.profile.kernel_profiles[0].scheduler == "user"

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_warp_specialized_circular_buffering_pointwise():
    # NOTE: Only a functional test, so high performance is not expected.
    def _definition_func(fd: FusionDefinition, inputs):
        tv0 = fd.from_pytorch(inputs[0])
        tv1 = fd.from_pytorch(inputs[1])
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    def _schedule_func(fd: FusionDefinition):
        # Parameters
        number_of_stages = 4
        prefetch_distance = 1
        bulk_inner_dim = 128

        tv_inputs = list(filter(lambda v: v.is_tensor(), fd.fusion.inputs()))
        assert len(tv_inputs) == 2
        tv0, tv1 = tv_inputs

        # Use TMA to load TV0 into shared memory
        tv3 = tv0.cache_after(LoadStoreOpType.tma)
        tv3.set_memory_type(MemoryType.shared)

        tv4 = tv1.cache_after(LoadStoreOpType.tma)
        tv4.set_memory_type(MemoryType.shared)

        tv_outputs = list(filter(lambda v: v.is_tensor(), fd.fusion.outputs()))
        assert len(tv_outputs) == 1
        reference = tv_outputs[0]

        # [M, N] -> [M, N/bid, bid]
        reference.split(-1, bulk_inner_dim)
        fd.sched.transform_like(reference)

        tv3.axis(0).parallelize(ParallelType.grid_x)
        tv4.axis(0).parallelize(ParallelType.grid_x)

        # Set computeAt position for circular buffering pipelining
        # The circular buffer axis is the first serial iterDomain to the left
        # of computeAt position. The domain is [M [BIDx], N/bid, bid], so the
        # circular buffer iterDomain is N/bid for computeAt position 2.
        fd.sched.inline_at(reference, pos=2)

        # Circular Buffer with TMA loads
        tv3.axis(2).parallelize(ParallelType.tma)
        tv4.axis(2).parallelize(ParallelType.tma)
        fd.sched.warp_specialize(
            tv3, number_of_stages, prefetch_distance, ParallelType.block_y
        )
        fd.sched.warp_specialize(
            tv4, number_of_stages, prefetch_distance, ParallelType.block_y
        )

        # Split reference to parallelize TMA tile
        reference.split(-1, bulk_inner_dim)
        reference.axis(0).parallelize(ParallelType.grid_x)
        reference.axis(-1).parallelize(ParallelType.block_x)

    # Inputs
    tensor_outer_dim = 128
    tensor_inner_dim = 1024
    t0 = torch.randn(
        tensor_outer_dim,
        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    t1 = torch.randn(
        tensor_outer_dim,
        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    inputs = [t0, t1]

    with FusionDefinition() as fd:
        _definition_func(fd, inputs)
        _schedule_func(fd)

    outputs = fd.manual_execute(inputs)
    warp_specialization_if_stmt = "if ((((nvfuser_index_t)threadIdx.y) >= 1LL)) {"
    assert warp_specialization_if_stmt in fd.fusion.print_kernel()
    assert torch.allclose(outputs[0], t0 + t1)

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_pdl():
    def nvf_fused_add_rmsnorm(fd: FusionDefinition, input_shape, weight_shape, eps):
        """
        Fused add root mean square normalization.
            * a [num_tokens, hidden_size]
            * residual [num_tokens, hidden_size]
            * weight [hidden_size]
            return rmsnorm(a + residual)
        """
        normalization_axis = -1
        keepdim = True
        a = fd.define_tensor(
            shape=input_shape,
            contiguity=[None, True],
            dtype=DataType.BFloat16,


        tensor_outer_dim,
        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    inputs = [t0, t1]

    with FusionDefinition() as fd:
        _definition_func(fd, inputs)
        _schedule_func(fd)

    outputs = fd.manual_execute(inputs)
    warp_specialization_if_stmt = "if ((((nvfuser_index_t)threadIdx.y) >= 1LL)) {"
    assert warp_specialization_if_stmt in fd.fusion.print_kernel()
    assert torch.allclose(outputs[0], t0 + t1)

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_pdl():
    def nvf_fused_add_rmsnorm(fd: FusionDefinition, input_shape, weight_shape, eps):
        """
        Fused add root mean square normalization.
            * a [num_tokens, hidden_size]
            * residual [num_tokens, hidden_size]
            * weight [hidden_size]
            return rmsnorm(a + residual)
        """
        normalization_axis = -1
        keepdim = True
        a = fd.define_tensor(
            shape=input_shape,
            contiguity=[None, True],
            dtype=DataType.BFloat16,
        )
        residual = fd.define_tensor(
            shape=input_shape,
            contiguity=[None, True],
            dtype=DataType.BFloat16,
        )
        weights = fd.define_tensor(
            shape=weight_shape,
            contiguity=[True],
            dtype=DataType.BFloat16,
        )
        norm_const = fd.define_scalar(input_shape[normalization_axis])
        eps_const = fd.define_scalar(eps)

        rmsnorm_input = fd.ops.add(a, residual)
        rmsnorm_input_sq = fd.ops.mul(rmsnorm_input, rmsnorm_input)
        sum0 = fd.ops.sum(rmsnorm_input_sq, dims=[normalization_axis], keepdim=keepdim)
        var = fd.ops.div(sum0, norm_const)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        pre_scale = fd.ops.mul(rmsnorm_input, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[1]
        )
        out = fd.ops.mul(pre_scale, weights_bcast)
        out_bf16 = fd.ops.cast(out, dtype=DataType.BFloat16)
        # NOTE alias_input does not check if output and alias_input have the same type
        fd.add_output(out_bf16, alias_input=a)

    def fused_add_rmsnorm(hidden_states, residual, weight, eps):
        inputs = [hidden_states, residual, weight]
        with FusionDefinition() as fd:
            nvf_fused_add_rmsnorm(
                fd,
                input_shape=hidden_states.shape,
                weight_shape=weight.shape,
                eps=eps,
            )
            # Apply selected scheduler
            heuristic_params = check_auto_schedule(schedule.schedule)(
                fd.fusion, SchedulerType.inner_persistent, inputs
            )
        fd.manual_execute(inputs, heuristic_params)
        assert fd.ke.is_programmatic_dependent_launch_enabled()

    def nvfuser(hidden_states, residual, layer, eps):
        fused_add_rmsnorm(hidden_states, residual, layer["post_attn_norm"], eps)
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    def baseline(hidden_states, residual, layer, eps, use_pdl):
        hidden_states += residual
        hidden_states = F.rms_norm(
            hidden_states, [hidden_states.shape[-1]], layer["post_attn_norm"], eps
        )
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    device = "cuda"
    dtype = torch.bfloat16
    num_tokens = 1  # decode phase: 1 token per sequence
    hidden_size = 2880
    num_experts = 32
    eps = 1e-6  # rms_norm_eps

    layer = {
        "post_attn_norm": torch.randn(
            hidden_size, dtype=dtype, device=device
        ).contiguous(),
        "router_weight": torch.randn(
            num_experts, hidden_size, dtype=dtype, device=device
        ).contiguous(),
        "router_bias": torch.randn(
            num_experts, dtype=dtype, device=device
        ).contiguous(),
    }

    # ========== Static tensors for CUDA graph capture ==========
    static_hidden_states = torch.randn(
        num_tokens, hidden_size, dtype=dtype, device=device
    ).contiguous()
    static_residual = torch.randn(
        num_tokens, hidden_size, dtype=dtype, device=device
    ).contiguous()

    baseline_result = baseline(
        static_hidden_states.clone().detach(),
        static_residual.clone().detach(),
        layer,
        eps,
        use_pdl=True,
    )
    nvf_result = nvfuser(static_hidden_states, static_residual, layer, eps)
    assert torch.allclose(nvf_result, baseline_result, rtol=1e-2, atol=1e-2)


verbose_ = True

# A helper function to test heuristic schedulers with automatic scheduling
def check_auto_schedule(schedule_fn):
    """
    A decorator to validate a schedule_fn before applying it to a fusion.

    Args:
        schedule_fn: The function to apply the scheduler
    """
    # List of all scheduler heuristics for testing
    # NOTE We cannot iterate pybind11 enum directly, so we extract the entries here.
    all_scheduler_heuristics = [
        heuristic
        for heuristic, _ in SchedulerType.__entries.values()
        if not SchedulerType.none
    ]

    def inner_fn(fusion, selected_heuristic, inputs):
        """
        Helper function to validate a schedule_fn.

        Args:
            fusion: The Fusion object to schedule
            selected_heuristic: The SchedulerType expected to work
            inputs: Input tensors for the fusion
        """
        available_heuristics = schedule.find_compatible_schedulers(fusion, inputs)

        # Assume that only a single heuristic is available for fusion
        assert len(available_heuristics) == 1

        # Check that only selected heuristic is available as a scheduler
        assert set(available_heuristics) == set([selected_heuristic])

        # Double-check with can_schedule
        status, _ = schedule.can_schedule(fusion, selected_heuristic, inputs)
        assert status

        # Check that the other schedulers are not compatible with this fusion
        assert all(
            [
                not schedule.can_schedule(fusion, h, inputs)[0]
                for h in all_scheduler_heuristics
                if h is not selected_heuristic
            ]
        )
        return schedule_fn(fusion, selected_heuristic, inputs)

    return inner_fn

def test_tutorial_memcpy():
    # First, we define a fusion. A common pattern is:
    # - Declare a Fusion, which works as a container of expressions using
    #   with context manager.
    # - Setup inputs. fd.define_tensor can be used to manually create tensors.
    #   fd.from_pytorch will create a TensorView given a pytorch tensor. Fusion
    #   registration is automatic.
    # - Define operations with the registered inputs.
    #   For supported operations, run:
    #   >>> import nvfuser_direct
    #   >>> fd = nvfuser_direct.FusionDefinition()
    #   >>> help(fd.ops)
    # - Most of operations that take tensors as inputs produce tensors
    #   as outputs, which can then be used as inputs to another
    #   operations.
    # - Final outputs should be set as fusion outputs with fd.add_output


    with FusionDefinition() as fd3:
        ref_tv = fusion_func(fd3)

        # We can of course mix BIDx and TIDx.
        ref_tv.axis(0).parallelize(ParallelType.grid_x)
        ref_tv.axis(1).parallelize(ParallelType.block_x)

        if verbose_:
            print(fd3.fusion.print_math())
            print(fd3.fusion.print_kernel())

    # The original input should not fail in this case. The kernel will be
    # launched with 10 thread blocks, each of which has 1024 threads. Try
    # running this test with NVFUSER_DUMP=launch_param to see the launch
    # configuration of each kernel lauch
    fd3.manual_validate([t0], [ref])

def test_tutorial_reduction_rfactor():
    # Just a very simple reduction of 1D tensor
    def fusion_func(fd: FusionDefinition) -> TensorView:
        tv0 = fd.define_tensor(shape=[-1])
        tv1 = fd.ops.sum(tv0, dims=[0])
        fd.add_output(tv1)
        return tv1

    # Create separate fusions because of multiple schedules
    with FusionDefinition() as fd0:
        tv1 = fusion_func(fd0)

        # A common pattern of reductions in CUDA involves multiple steps of
        # reductions, where the first step is a per-thread local reduction,
        # followed by a block reduction of the per-thread partial results,
        # and also potentially followed by a grid reduction of the
        # per-block partial results. Here's an example with a two-step
        # reduction:
        #
        # // Step 1: Per-thread reduction
        # float partial_result = 0;
        # for (int i = threadIdx.x; i += blockDim.x; i < N) {
        #   partial_result += input[i];
        # }
        #
        # // Step 2: Accumulation within each thread block
        # __shared__ float shared_buf[blockDim.x];


    runtime_information = f"Input size: {kp.input_bytes} bytes, \
    Output size: {kp.output_bytes} bytes, Time: {kp.time_ms:2f} ms"
    print(runtime_information)

    bandwidth_information = f"Effective Bandwidth: {kp.effective_bandwidth_gbs:.2f} GB/s, \
    Peak Bandwidth: {kp.percentage_peak_bandwidth:2f}%"
    print(bandwidth_information)
    print("===================================================================")

    # Validate profiler captured kernel info
    assert prof.profile.fusion_id >= 0
    assert len(prof.profile.kernel_profiles) > 0
    assert prof.profile.kernel_profiles[0].scheduler == "user"

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_warp_specialized_circular_buffering_pointwise():
    # NOTE: Only a functional test, so high performance is not expected.
    def _definition_func(fd: FusionDefinition, inputs):
        tv0 = fd.from_pytorch(inputs[0])
        tv1 = fd.from_pytorch(inputs[1])
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    def _schedule_func(fd: FusionDefinition):
        # Parameters
        number_of_stages = 4
        prefetch_distance = 1
        bulk_inner_dim = 128

        tv_inputs = list(filter(lambda v: v.is_tensor(), fd.fusion.inputs()))
        assert len(tv_inputs) == 2
        tv0, tv1 = tv_inputs

        # Use TMA to load TV0 into shared memory
        tv3 = tv0.cache_after(LoadStoreOpType.tma)
        tv3.set_memory_type(MemoryType.shared)

        tv4 = tv1.cache_after(LoadStoreOpType.tma)
        tv4.set_memory_type(MemoryType.shared)

        tv_outputs = list(filter(lambda v: v.is_tensor(), fd.fusion.outputs()))
        assert len(tv_outputs) == 1


    print(bandwidth_information)
    print("===================================================================")

    # Validate profiler captured kernel info
    assert prof.profile.fusion_id >= 0
    assert len(prof.profile.kernel_profiles) > 0
    assert prof.profile.kernel_profiles[0].scheduler == "user"

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_warp_specialized_circular_buffering_pointwise():
    # NOTE: Only a functional test, so high performance is not expected.
    def _definition_func(fd: FusionDefinition, inputs):
        tv0 = fd.from_pytorch(inputs[0])
        tv1 = fd.from_pytorch(inputs[1])
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    def _schedule_func(fd: FusionDefinition):
        # Parameters
        number_of_stages = 4
        prefetch_distance = 1
        bulk_inner_dim = 128

        tv_inputs = list(filter(lambda v: v.is_tensor(), fd.fusion.inputs()))
        assert len(tv_inputs) == 2
        tv0, tv1 = tv_inputs

        # Use TMA to load TV0 into shared memory
        tv3 = tv0.cache_after(LoadStoreOpType.tma)
        tv3.set_memory_type(MemoryType.shared)

        tv4 = tv1.cache_after(LoadStoreOpType.tma)
        tv4.set_memory_type(MemoryType.shared)

        tv_outputs = list(filter(lambda v: v.is_tensor(), fd.fusion.outputs()))
        assert len(tv_outputs) == 1
        reference = tv_outputs[0]

        # [M, N] -> [M, N/bid, bid]
        reference.split(-1, bulk_inner_dim)
        fd.sched.transform_like(reference)

        tv3.axis(0).parallelize(ParallelType.grid_x)
        tv4.axis(0).parallelize(ParallelType.grid_x)

        # Set computeAt position for circular buffering pipelining
        # The circular buffer axis is the first serial iterDomain to the left
        # of computeAt position. The domain is [M [BIDx], N/bid, bid], so the
        # circular buffer iterDomain is N/bid for computeAt position 2.
        fd.sched.inline_at(reference, pos=2)

        # Circular Buffer with TMA loads
        tv3.axis(2).parallelize(ParallelType.tma)
        tv4.axis(2).parallelize(ParallelType.tma)
        fd.sched.warp_specialize(
            tv3, number_of_stages, prefetch_distance, ParallelType.block_y
        )
        fd.sched.warp_specialize(
            tv4, number_of_stages, prefetch_distance, ParallelType.block_y
        )

        # Split reference to parallelize TMA tile
        reference.split(-1, bulk_inner_dim)
        reference.axis(0).parallelize(ParallelType.grid_x)
        reference.axis(-1).parallelize(ParallelType.block_x)

    # Inputs
    tensor_outer_dim = 128
    tensor_inner_dim = 1024
    t0 = torch.randn(
        tensor_outer_dim,
        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    t1 = torch.randn(
        tensor_outer_dim,
        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    inputs = [t0, t1]

    with FusionDefinition() as fd:
        _definition_func(fd, inputs)


        tensor_inner_dim,
        dtype=torch.float,
        device=torch.device("cuda:0"),
    )
    inputs = [t0, t1]

    with FusionDefinition() as fd:
        _definition_func(fd, inputs)
        _schedule_func(fd)

    outputs = fd.manual_execute(inputs)
    warp_specialization_if_stmt = "if ((((nvfuser_index_t)threadIdx.y) >= 1LL)) {"
    assert warp_specialization_if_stmt in fd.fusion.print_kernel()
    assert torch.allclose(outputs[0], t0 + t1)

@pytest.mark.skipif(
    is_pre_hopper(), reason="Only supported on Hopper and newer devices."
)
def test_pdl():
    def nvf_fused_add_rmsnorm(fd: FusionDefinition, input_shape, weight_shape, eps):
        """
        Fused add root mean square normalization.
            * a [num_tokens, hidden_size]
            * residual [num_tokens, hidden_size]
            * weight [hidden_size]
            return rmsnorm(a + residual)
        """
        normalization_axis = -1
        keepdim = True
        a = fd.define_tensor(
            shape=input_shape,
            contiguity=[None, True],
            dtype=DataType.BFloat16,
        )
        residual = fd.define_tensor(
            shape=input_shape,
            contiguity=[None, True],
            dtype=DataType.BFloat16,
        )
        weights = fd.define_tensor(
            shape=weight_shape,
            contiguity=[True],
            dtype=DataType.BFloat16,
        )
        norm_const = fd.define_scalar(input_shape[normalization_axis])
        eps_const = fd.define_scalar(eps)

        rmsnorm_input = fd.ops.add(a, residual)
        rmsnorm_input_sq = fd.ops.mul(rmsnorm_input, rmsnorm_input)
        sum0 = fd.ops.sum(rmsnorm_input_sq, dims=[normalization_axis], keepdim=keepdim)
        var = fd.ops.div(sum0, norm_const)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        pre_scale = fd.ops.mul(rmsnorm_input, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[1]
        )
        out = fd.ops.mul(pre_scale, weights_bcast)
        out_bf16 = fd.ops.cast(out, dtype=DataType.BFloat16)
        # NOTE alias_input does not check if output and alias_input have the same type
        fd.add_output(out_bf16, alias_input=a)

    def fused_add_rmsnorm(hidden_states, residual, weight, eps):
        inputs = [hidden_states, residual, weight]
        with FusionDefinition() as fd:
            nvf_fused_add_rmsnorm(
                fd,
                input_shape=hidden_states.shape,
                weight_shape=weight.shape,
                eps=eps,
            )
            # Apply selected scheduler
            heuristic_params = check_auto_schedule(schedule.schedule)(
                fd.fusion, SchedulerType.inner_persistent, inputs
            )
        fd.manual_execute(inputs, heuristic_params)
        assert fd.ke.is_programmatic_dependent_launch_enabled()

    def nvfuser(hidden_states, residual, layer, eps):
        fused_add_rmsnorm(hidden_states, residual, layer["post_attn_norm"], eps)
        router_logits = F.linear(


            dtype=DataType.BFloat16,
        )
        norm_const = fd.define_scalar(input_shape[normalization_axis])
        eps_const = fd.define_scalar(eps)

        rmsnorm_input = fd.ops.add(a, residual)
        rmsnorm_input_sq = fd.ops.mul(rmsnorm_input, rmsnorm_input)
        sum0 = fd.ops.sum(rmsnorm_input_sq, dims=[normalization_axis], keepdim=keepdim)
        var = fd.ops.div(sum0, norm_const)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        pre_scale = fd.ops.mul(rmsnorm_input, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[1]
        )
        out = fd.ops.mul(pre_scale, weights_bcast)
        out_bf16 = fd.ops.cast(out, dtype=DataType.BFloat16)
        # NOTE alias_input does not check if output and alias_input have the same type
        fd.add_output(out_bf16, alias_input=a)

    def fused_add_rmsnorm(hidden_states, residual, weight, eps):
        inputs = [hidden_states, residual, weight]
        with FusionDefinition() as fd:
            nvf_fused_add_rmsnorm(
                fd,
                input_shape=hidden_states.shape,
                weight_shape=weight.shape,
                eps=eps,
            )
            # Apply selected scheduler
            heuristic_params = check_auto_schedule(schedule.schedule)(
                fd.fusion, SchedulerType.inner_persistent, inputs
            )
        fd.manual_execute(inputs, heuristic_params)
        assert fd.ke.is_programmatic_dependent_launch_enabled()

    def nvfuser(hidden_states, residual, layer, eps):
        fused_add_rmsnorm(hidden_states, residual, layer["post_attn_norm"], eps)
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    def baseline(hidden_states, residual, layer, eps, use_pdl):
        hidden_states += residual
        hidden_states = F.rms_norm(
            hidden_states, [hidden_states.shape[-1]], layer["post_attn_norm"], eps
        )
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    device = "cuda"
    dtype = torch.bfloat16


nvfuser

        out_bf16 = fd.ops.cast(out, dtype=DataType.BFloat16)
        # NOTE alias_input does not check if output and alias_input have the same type
        fd.add_output(out_bf16, alias_input=a)

    def fused_add_rmsnorm(hidden_states, residual, weight, eps):
        inputs = [hidden_states, residual, weight]
        with FusionDefinition() as fd:
            nvf_fused_add_rmsnorm(
                fd,
                input_shape=hidden_states.shape,
                weight_shape=weight.shape,
                eps=eps,
            )
            # Apply selected scheduler
            heuristic_params = check_auto_schedule(schedule.schedule)(
                fd.fusion, SchedulerType.inner_persistent, inputs
            )
        fd.manual_execute(inputs, heuristic_params)
        assert fd.ke.is_programmatic_dependent_launch_enabled()

    def nvfuser(hidden_states, residual, layer, eps):
        fused_add_rmsnorm(hidden_states, residual, layer["post_attn_norm"], eps)
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    def baseline(hidden_states, residual, layer, eps, use_pdl):
        hidden_states += residual
        hidden_states = F.rms_norm(
            hidden_states, [hidden_states.shape[-1]], layer["post_attn_norm"], eps
        )
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    device = "cuda"
    dtype = torch.bfloat16
    num_tokens = 1  # decode phase: 1 token per sequence
    hidden_size = 2880
    num_experts = 32
    eps = 1e-6  # rms_norm_eps

    layer = {
        "post_attn_norm": torch.randn(


baseline

            nvf_fused_add_rmsnorm(
                fd,
                input_shape=hidden_states.shape,
                weight_shape=weight.shape,
                eps=eps,
            )
            # Apply selected scheduler
            heuristic_params = check_auto_schedule(schedule.schedule)(
                fd.fusion, SchedulerType.inner_persistent, inputs
            )
        fd.manual_execute(inputs, heuristic_params)
        assert fd.ke.is_programmatic_dependent_launch_enabled()

    def nvfuser(hidden_states, residual, layer, eps):
        fused_add_rmsnorm(hidden_states, residual, layer["post_attn_norm"], eps)
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    def baseline(hidden_states, residual, layer, eps, use_pdl):
        hidden_states += residual
        hidden_states = F.rms_norm(
            hidden_states, [hidden_states.shape[-1]], layer["post_attn_norm"], eps
        )
        router_logits = F.linear(
            hidden_states, layer["router_weight"], layer["router_bias"]
        )
        return hidden_states

    device = "cuda"
    dtype = torch.bfloat16
    num_tokens = 1  # decode phase: 1 token per sequence
    hidden_size = 2880
    num_experts = 32
    eps = 1e-6  # rms_norm_eps

    layer = {
        "post_attn_norm": torch.randn(
            hidden_size, dtype=dtype, device=device
        ).contiguous(),
        "router_weight": torch.randn(
            num_experts, hidden_size, dtype=dtype, device=device
        ).contiguous(),
        "router_bias": torch.randn(
            num_experts, dtype=dtype, device=device
        ).contiguous(),
    }


# 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,
    pytorch_nvfp4_quantize,
    round_up,
    activation_scale_to_nvfp4,
    fp4_to_fp32,
    unpack_fp4,
)

def get_ref_results(
    a_fp4,
    b_fp4,
    a_sf,
    b_sf,
    a_global_scale,
    b_global_scale,
    m,
    n,
    dtype,
    block_size,
    device,
):
    _, m_k = a_fp4.shape
    _, n_k = b_fp4.shape
    assert m_k == n_k
    a_in_dtype = dequantize_to_dtype(
        a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    b_in_dtype = dequantize_to_dtype(
        b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    return torch.matmul(a_in_dtype, b_in_dtype.t())

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    b_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    alpha = 1.0 / (a_global_scale * b_global_scale)

    a_fp4, a_scale_linear = pytorch_nvfp4_quantize(a_dtype, a_global_scale)
    b_fp4, b_scale_linear = pytorch_nvfp4_quantize(b_dtype, b_global_scale)
    a_scale_interleaved = linear_to_swizzled_128_4(a_scale_linear)
    b_scale_interleaved = linear_to_swizzled_128_4(b_scale_linear)

    expected_out = get_ref_results(
        a_fp4,
        b_fp4,
        a_scale_interleaved,
        b_scale_interleaved,
        a_global_scale,
        b_global_scale,
        m,
        n,
        dtype,
        block_size,
        "cuda",
    )
    out = nvf_cutlass.nvfp4_scaled_mm(
        a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
    )

    torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm_epilogue(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    b_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    alpha = 1.0 / (a_global_scale * b_global_scale)
    global_normconst = torch.tensor(2, dtype=torch.float, device="cuda")

    a_fp4, a_scale_linear = pytorch_nvfp4_quantize(a_dtype, a_global_scale)
    b_fp4, b_scale_linear = pytorch_nvfp4_quantize(b_dtype, b_global_scale)
    a_scale_interleaved = linear_to_swizzled_128_4(a_scale_linear)
    b_scale_interleaved = linear_to_swizzled_128_4(b_scale_linear)

    expected_out = get_ref_results(
        a_fp4,
        b_fp4,
        a_scale_interleaved,
        b_scale_interleaved,
        a_global_scale,
        b_global_scale,
        m,
        n,
        dtype,
        block_size,
        "cuda",
    )

    expected_out_fp4, expected_out_scale_linear = pytorch_nvfp4_quantize(
        expected_out, global_normconst
    )
    expected_out_scale_interleaved = linear_to_swizzled_128_4(expected_out_scale_linear)

    out_fp4, out_scale_interleaved = nvf_cutlass.nvfp4_scaled_mm_blockscale(
        a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, global_normconst
    )

    # Convert to unpacked fp32 to check nvfp4 tensors.
    expected_out_fp32 = fp4_to_fp32(unpack_fp4(expected_out_fp4.view(torch.uint8)))
    out_fp32 = fp4_to_fp32(unpack_fp4(out_fp4.view(torch.uint8)))

    # The absolute max difference is 2.0.
    abs_diff = torch.abs(expected_out_fp32 - out_fp32)
    assert torch.max(abs_diff) <= 2.0

    # The percentage of mismatched values is 1%.
    nonzero = torch.count_nonzero(torch.ne(abs_diff, 0.0))
    assert (nonzero / abs_diff.numel()) < 0.1

    # Compare scale factors
    # rtol = epsilon = 2**(-3) for fp8_m4e3
    # atol = 2**(num_exponent_bits + smallest_exponent_value)
    torch.testing.assert_close(
        out_scale_interleaved.to(torch.float),
        expected_out_scale_interleaved.to(torch.float),
        atol=2e-3,
        rtol=0.125,
    )

@pytest.mark.parametrize("config", [[1024, 128, 256]])
@pytest.mark.parametrize("tokens_per_expert_neg_one", [[115, 144, 8]])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
def test_nvfp4_grouped_mm(
    config,
    tokens_per_expert_neg_one,
    out_dtype,
):
    nvfp4_block_size = 16

    # k dimension is multiple of 128 to avoid padding
    m, n, k = config
    tokens_per_expert = list(tokens_per_expert_neg_one)
    tokens_per_expert.append(m - sum(tokens_per_expert))
    g = len(tokens_per_expert)

    mat1 = torch.testing.make_tensor((m, k), dtype=torch.float32, device="cuda:0")
    mat2 = torch.testing.make_tensor((g, n, k), dtype=torch.float32, device="cuda:0")
    ab_strides = torch.full((g,), k, dtype=torch.int64, device="cuda:0")
    c_strides = torch.full((g,), n, dtype=torch.int64, device="cuda:0")

    offsets = torch.empty((g,), dtype=torch.int32, device="cuda:0")
    blockscale_offsets = torch.empty((g,), dtype=torch.int32, device="cuda:0")
    problem_sizes = torch.empty((g, 3), dtype=torch.int32, device="cuda:0")
    mat2_gs = torch.empty((g,), dtype=torch.float32, device="cuda:0")
    scale2 = torch.empty(
        (g, n, k // nvfp4_block_size), dtype=torch.float8_e4m3fn, device="cuda:0"
    )

    accumulated_tokens = 0
    rounded_accumulated_tokens = 0
    mat2_fp4 = torch.empty(
        (g, n, k // 2), dtype=torch.float4_e2m1fn_x2, device="cuda:0"
    )

    # Use tokens_per_expert to calculate offsets into m dimension of input tensor.
    # The blockscale offset per expert is rounded to 128 to match the alignment
    # requirements for blockscale factor.
    for i in range(g):
        mat2_gs[i] = FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX / mat2[i].max()
        offsets[i] = accumulated_tokens
        blockscale_offsets[i] = rounded_accumulated_tokens
        accumulated_tokens += tokens_per_expert[i]
        # Note: we technically don't need to round up, since k is perfectly sized.
        rounded_accumulated_tokens += round_up(tokens_per_expert[i], 128)

        problem_sizes[i][0] = tokens_per_expert[i]
        problem_sizes[i][1] = n
        problem_sizes[i][2] = k

        mat2_fp4_i, mat2_sf_i = pytorch_nvfp4_quantize(mat2[i], mat2_gs[i])
        mat2_fp4[i] = mat2_fp4_i
        scale2[i] = linear_to_swizzled_128_4(mat2_sf_i)

    # prepare quantization for mat1
    # note: following sglang implementation, not computing global scaling factor for mat1
    #       similarly, we don't need to apply mat1_gs to alpha
    mat1_gs = torch.ones((g,), dtype=torch.float32, device="cuda:0")
    mat1_fp4, scale1 = activation_scale_to_nvfp4(
        mat1, mat1_gs, offsets, blockscale_offsets, nvfp4_block_size
    )

    out = nvf_cutlass.nvfp4_scaled_grouped_mm(
        mat1_fp4,
        mat2_fp4,
        scale1,
        scale2,
        mat2_gs,
        ab_strides,
        c_strides,
        problem_sizes,
        offsets,
        blockscale_offsets,
        out_dtype,
    )

    # Create pytorch expected output reference
    # For each expert, apply nvfp4 gemm. Slice the input matrix given the tokens_per_expert.
    # C[start:stop] = A[start:stop] @ B[expert].
    out_decomposed_ref = torch.empty(m, n, dtype=out_dtype, device="cuda:0")
    for i in range(g):
        l = offsets[i]
        l_sf = blockscale_offsets[i]
        if i == g - 1:
            r = m
        else:
            r = offsets[i + 1]
        r_sf = round_up(tokens_per_expert[i], 128) + l_sf
        # A cublas invalid value error occurs when passing in mat2_gs[i] as
        # alpha in the torch kernel.
        out_decomposed_ref[l:r] = (
            torch._scaled_mm(
                mat1_fp4[l:r],
                mat2_fp4[i].transpose(-1, -2),
                scale1[l_sf:r_sf],
                scale2[i],
                None,
                None,
                out_dtype,
            )
            * mat2_gs[i]
        )

    torch.testing.assert_close(out_decomposed_ref, out, atol=1e-2, rtol=1e-2)


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,
    pytorch_nvfp4_quantize,
    round_up,
    activation_scale_to_nvfp4,
    fp4_to_fp32,
    unpack_fp4,
)

def get_ref_results(
    a_fp4,
    b_fp4,
    a_sf,
    b_sf,
    a_global_scale,
    b_global_scale,
    m,
    n,
    dtype,
    block_size,
    device,
):
    _, m_k = a_fp4.shape
    _, n_k = b_fp4.shape
    assert m_k == n_k
    a_in_dtype = dequantize_to_dtype(
        a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    b_in_dtype = dequantize_to_dtype(
        b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    return torch.matmul(a_in_dtype, b_in_dtype.t())

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)


    block_size,
    device,
):
    _, m_k = a_fp4.shape
    _, n_k = b_fp4.shape
    assert m_k == n_k
    a_in_dtype = dequantize_to_dtype(
        a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    b_in_dtype = dequantize_to_dtype(
        b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
    )
    return torch.matmul(a_in_dtype, b_in_dtype.t())

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    b_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    alpha = 1.0 / (a_global_scale * b_global_scale)

    a_fp4, a_scale_linear = pytorch_nvfp4_quantize(a_dtype, a_global_scale)
    b_fp4, b_scale_linear = pytorch_nvfp4_quantize(b_dtype, b_global_scale)
    a_scale_interleaved = linear_to_swizzled_128_4(a_scale_linear)
    b_scale_interleaved = linear_to_swizzled_128_4(b_scale_linear)

    expected_out = get_ref_results(
        a_fp4,
        b_fp4,
        a_scale_interleaved,
        b_scale_interleaved,
        a_global_scale,
        b_global_scale,
        m,
        n,
        dtype,
        block_size,
        "cuda",
    )
    out = nvf_cutlass.nvfp4_scaled_mm(
        a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
    )

    torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm_epilogue(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)


        a_global_scale,
        b_global_scale,
        m,
        n,
        dtype,
        block_size,
        "cuda",
    )
    out = nvf_cutlass.nvfp4_scaled_mm(
        a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
    )

    torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "shape", [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
)
@torch.inference_mode()
def test_nvfp4_gemm_epilogue(
    dtype: torch.dtype,
    shape: tuple[int, int, int],
) -> None:
    m, n, packed_k = shape
    k = packed_k * 2
    block_size = 16
    a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
    b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")

    a_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    b_global_scale = (
        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
    ).to(torch.float32)
    alpha = 1.0 / (a_global_scale * b_global_scale)
    global_normconst = torch.tensor(2, dtype=torch.float, device="cuda")

    a_fp4, a_scale_linear = pytorch_nvfp4_quantize(a_dtype, a_global_scale)
    b_fp4, b_scale_linear = pytorch_nvfp4_quantize(b_dtype, b_global_scale)
    a_scale_interleaved = linear_to_swizzled_128_4(a_scale_linear)
    b_scale_interleaved = linear_to_swizzled_128_4(b_scale_linear)

    expected_out = get_ref_results(
        a_fp4,
        b_fp4,
        a_scale_interleaved,
        b_scale_interleaved,
        a_global_scale,
        b_global_scale,
        m,
        n,
        dtype,
        block_size,
        "cuda",
    )

    expected_out_fp4, expected_out_scale_linear = pytorch_nvfp4_quantize(
        expected_out, global_normconst
    )
    expected_out_scale_interleaved = linear_to_swizzled_128_4(expected_out_scale_linear)

    out_fp4, out_scale_interleaved = nvf_cutlass.nvfp4_scaled_mm_blockscale(
        a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, global_normconst
    )

    # Convert to unpacked fp32 to check nvfp4 tensors.
    expected_out_fp32 = fp4_to_fp32(unpack_fp4(expected_out_fp4.view(torch.uint8)))
    out_fp32 = fp4_to_fp32(unpack_fp4(out_fp4.view(torch.uint8)))

    # The absolute max difference is 2.0.
    abs_diff = torch.abs(expected_out_fp32 - out_fp32)
    assert torch.max(abs_diff) <= 2.0

    # The percentage of mismatched values is 1%.
    nonzero = torch.count_nonzero(torch.ne(abs_diff, 0.0))
    assert (nonzero / abs_diff.numel()) < 0.1

    # Compare scale factors
    # rtol = epsilon = 2**(-3) for fp8_m4e3
    # atol = 2**(num_exponent_bits + smallest_exponent_value)
    torch.testing.assert_close(
        out_scale_interleaved.to(torch.float),
        expected_out_scale_interleaved.to(torch.float),
        atol=2e-3,
        rtol=0.125,
    )

@pytest.mark.parametrize("config", [[1024, 128, 256]])
@pytest.mark.parametrize("tokens_per_expert_neg_one", [[115, 144, 8]])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
def test_nvfp4_grouped_mm(
    config,
    tokens_per_expert_neg_one,
    out_dtype,
):
    nvfp4_block_size = 16

    # k dimension is multiple of 128 to avoid padding
    m, n, k = config
    tokens_per_expert = list(tokens_per_expert_neg_one)
    tokens_per_expert.append(m - sum(tokens_per_expert))
    g = len(tokens_per_expert)

    mat1 = torch.testing.make_tensor((m, k), dtype=torch.float32, device="cuda:0")
    mat2 = torch.testing.make_tensor((g, n, k), dtype=torch.float32, device="cuda:0")


    assert torch.max(abs_diff) <= 2.0

    # The percentage of mismatched values is 1%.
    nonzero = torch.count_nonzero(torch.ne(abs_diff, 0.0))
    assert (nonzero / abs_diff.numel()) < 0.1

    # Compare scale factors
    # rtol = epsilon = 2**(-3) for fp8_m4e3
    # atol = 2**(num_exponent_bits + smallest_exponent_value)
    torch.testing.assert_close(
        out_scale_interleaved.to(torch.float),
        expected_out_scale_interleaved.to(torch.float),
        atol=2e-3,
        rtol=0.125,
    )

@pytest.mark.parametrize("config", [[1024, 128, 256]])
@pytest.mark.parametrize("tokens_per_expert_neg_one", [[115, 144, 8]])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
def test_nvfp4_grouped_mm(
    config,
    tokens_per_expert_neg_one,
    out_dtype,
):
    nvfp4_block_size = 16

    # k dimension is multiple of 128 to avoid padding
    m, n, k = config
    tokens_per_expert = list(tokens_per_expert_neg_one)
    tokens_per_expert.append(m - sum(tokens_per_expert))
    g = len(tokens_per_expert)

    mat1 = torch.testing.make_tensor((m, k), dtype=torch.float32, device="cuda:0")
    mat2 = torch.testing.make_tensor((g, n, k), dtype=torch.float32, device="cuda:0")
    ab_strides = torch.full((g,), k, dtype=torch.int64, device="cuda:0")
    c_strides = torch.full((g,), n, dtype=torch.int64, device="cuda:0")

    offsets = torch.empty((g,), dtype=torch.int32, device="cuda:0")
    blockscale_offsets = torch.empty((g,), dtype=torch.int32, device="cuda:0")
    problem_sizes = torch.empty((g, 3), dtype=torch.int32, device="cuda:0")
    mat2_gs = torch.empty((g,), dtype=torch.float32, device="cuda:0")
    scale2 = torch.empty(
        (g, n, k // nvfp4_block_size), dtype=torch.float8_e4m3fn, device="cuda:0"
    )

    accumulated_tokens = 0
    rounded_accumulated_tokens = 0
    mat2_fp4 = torch.empty(
        (g, n, k // 2), dtype=torch.float4_e2m1fn_x2, device="cuda:0"
    )

    # Use tokens_per_expert to calculate offsets into m dimension of input tensor.
    # The blockscale offset per expert is rounded to 128 to match the alignment
    # requirements for blockscale factor.
    for i in range(g):
        mat2_gs[i] = FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX / mat2[i].max()
        offsets[i] = accumulated_tokens
        blockscale_offsets[i] = rounded_accumulated_tokens
        accumulated_tokens += tokens_per_expert[i]
        # Note: we technically don't need to round up, since k is perfectly sized.
        rounded_accumulated_tokens += round_up(tokens_per_expert[i], 128)

        problem_sizes[i][0] = tokens_per_expert[i]
        problem_sizes[i][1] = n
        problem_sizes[i][2] = k

        mat2_fp4_i, mat2_sf_i = pytorch_nvfp4_quantize(mat2[i], mat2_gs[i])
        mat2_fp4[i] = mat2_fp4_i
        scale2[i] = linear_to_swizzled_128_4(mat2_sf_i)

    # prepare quantization for mat1
    # note: following sglang implementation, not computing global scaling factor for mat1
    #       similarly, we don't need to apply mat1_gs to alpha
    mat1_gs = torch.ones((g,), dtype=torch.float32, device="cuda:0")
    mat1_fp4, scale1 = activation_scale_to_nvfp4(
        mat1, mat1_gs, offsets, blockscale_offsets, nvfp4_block_size
    )

    out = nvf_cutlass.nvfp4_scaled_grouped_mm(
        mat1_fp4,
        mat2_fp4,
        scale1,
        scale2,
        mat2_gs,
        ab_strides,
        c_strides,
        problem_sizes,
        offsets,
        blockscale_offsets,
        out_dtype,
    )

    # Create pytorch expected output reference
    # For each expert, apply nvfp4 gemm. Slice the input matrix given the tokens_per_expert.
    # C[start:stop] = A[start:stop] @ B[expert].
    out_decomposed_ref = torch.empty(m, n, dtype=out_dtype, device="cuda:0")
    for i in range(g):
        l = offsets[i]
        l_sf = blockscale_offsets[i]
        if i == g - 1:
            r = m
        else:
            r = offsets[i + 1]
        r_sf = round_up(tokens_per_expert[i], 128) + l_sf
        # A cublas invalid value error occurs when passing in mat2_gs[i] as
        # alpha in the torch kernel.
        out_decomposed_ref[l:r] = (
            torch._scaled_mm(
                mat1_fp4[l:r],
                mat2_fp4[i].transpose(-1, -2),
                scale1[l_sf:r_sf],
                scale2[i],
                None,
                None,
                out_dtype,
            )
            * mat2_gs[i]
        )

    torch.testing.assert_close(out_decomposed_ref, out, atol=1e-2, rtol=1e-2)


# 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),
    ]

    inputs_tn = [
        inputs_tt[0].clone(),
        inputs_tt[1].clone().as_strided(size=[k, n], stride=[1, k]),
    ]

    inputs_nt = [
        inputs_tt[0].clone().as_strided(size=[m, k], stride=[1, m]),
        inputs_tt[1].clone(),
    ]

    inputs_tn = [inputs_tt[0].clone(), inputs_tn[1].clone()]

    inputs_nn = [inputs_nt[0].clone(), inputs_tn[1].clone()]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    for inps in [inputs_tt, inputs_tn, inputs_nt, inputs_nn]:
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inps=inps), inps
        )
        eager_out = torch.matmul(inps[0], inps[1])
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(eager_out, fp16_nvf_out)

def test_linear_without_bias(nvfuser_direct_test):
    """
    Test linear without bias. Check that bias keyword argument does not appear in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.linear(t0, t1)
        fd.add_output(t2)

    # Check that bias is not included with linear
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.ops.linear(tv0, tv1)
    fd.add_output(tv2)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_linear_with_bias(nvfuser_direct_test):
    """
    Test linear with bias. Check that bias keyword argument appears in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.from_pytorch(inputs[2])
        t3 = fd.ops.linear(t0, t1, t2)
        fd.add_output(t3)

    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.BFloat16, is_cpu=False)
    tv3 = fd.ops.linear(tv0, tv1, bias=tv2)
    fd.add_output(tv3)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1], inputs[2])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

@pytest.mark.parametrize("bias", [False, True])
def test_linear(nvfuser_direct_test, bias):
    m = 24
    n = 16
    k = 8
    bias1d = torch.randn(n, device="cuda", dtype=torch.float16) if bias else None

    inputs_mk_nk = [
        torch.randn(m, k, device="cuda", dtype=torch.float16),
        torch.randn(n, k, device="cuda", dtype=torch.float16),
    ]

    inputs_mk_kn = [
        inputs_mk_nk[0].clone(),
        inputs_mk_nk[1].clone().as_strided(size=[n, k], stride=[1, n]),
    ]

    inputs_km_nk = [
        inputs_mk_nk[0].clone().as_strided(size=[m, k], stride=[1, m]),
        inputs_mk_nk[1].clone(),
    ]

    inputs_km_kn = [
        inputs_km_nk[0].clone(),
        inputs_mk_kn[1].clone(),
    ]

    def fusion_func(
        fd: FusionDefinition,
        inp: torch.Tensor,
        wt: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        t0 = fd.from_pytorch(inp)
        t1 = fd.from_pytorch(wt)
        if bias is not None:
            t2 = fd.from_pytorch(bias)
            t_out = fd.ops.linear(t0, t1, t2)
        else:
            t_out = fd.ops.linear(t0, t1)
        fd.add_output(t_out)

    in_tensors = [inputs_mk_nk, inputs_mk_kn, inputs_km_nk, inputs_km_kn]
    for inp, wt in in_tensors:
        input_tensors = (inp, wt, bias1d) if bias1d is not None else (inp, wt)
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inp=inp, wt=wt, bias=bias1d),
            input_tensors,
        )
        eager_out = F.linear(inp, wt, bias1d)
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(fp16_nvf_out, eager_out, atol=1e-3, rtol=0)

def test_linear_slice(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([1, 2, 3])
        b = fd.define_tensor([4, 3])
        c = fd.ops.linear(a, b)
        d = fd.ops.slice(c, start_indices=[0, 0, 0], end_indices=[1, 2, 3])
        fd.add_output(d)

    inputs = [
        torch.randn(1, 2, 3, device="cuda:0"),
        torch.randn(4, 3, device="cuda:0"),
    ]
    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_matmul_operandA_2d_operandB_3d(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([2, 3])
        b = fd.define_tensor([7, 3, 5])
        c = fd.ops.matmul(a, b)
        assert c.ndim == 3
        fd.add_output(c)

    inputs = [
        torch.testing.make_tensor(2, 3, dtype=torch.float32, device="cuda"),
        torch.testing.make_tensor(7, 3, 5, dtype=torch.float32, device="cuda"),
    ]
    outputs, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert outputs[0].ndim == 3

def test_matmul_stride(nvfuser_direct_test):
    n, h, l, s, e = 4, 8, 16, 16, 8
    inputs = [
        torch.randn(n, h, l, e, device="cuda", dtype=torch.float16, requires_grad=True),
        torch.randn(n, h, s, e, device="cuda", dtype=torch.float16, requires_grad=True),
    ]
    for stride_order in itertools.permutations(range(4)):

        def fusion_func(fd: FusionDefinition) -> None:
            q = fd.from_pytorch(inputs[0])
            k = fd.from_pytorch(inputs[1])
            k_t = fd.ops.permute(k, [0, 1, 3, 2])
            out = fd.ops.matmul(q, k_t)
            strided_out = fd.ops.stride_order(out, stride_order)
            fd.add_output(strided_out)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = torch.matmul(inputs[0], torch.transpose(inputs[1], -2, -1))
        verify_stride_order(nvf_out[0].stride(), stride_order)
        nvfuser_direct_test.assertEqual(nvf_out[0], eager_out)


# 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),
    ]

    inputs_tn = [
        inputs_tt[0].clone(),
        inputs_tt[1].clone().as_strided(size=[k, n], stride=[1, k]),
    ]

    inputs_nt = [
        inputs_tt[0].clone().as_strided(size=[m, k], stride=[1, m]),
        inputs_tt[1].clone(),
    ]

    inputs_tn = [inputs_tt[0].clone(), inputs_tn[1].clone()]

    inputs_nn = [inputs_nt[0].clone(), inputs_tn[1].clone()]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    for inps in [inputs_tt, inputs_tn, inputs_nt, inputs_nn]:
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inps=inps), inps
        )
        eager_out = torch.matmul(inps[0], inps[1])
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(eager_out, fp16_nvf_out)

def test_linear_without_bias(nvfuser_direct_test):
    """
    Test linear without bias. Check that bias keyword argument does not appear in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.linear(t0, t1)
        fd.add_output(t2)


    inputs_tn = [inputs_tt[0].clone(), inputs_tn[1].clone()]

    inputs_nn = [inputs_nt[0].clone(), inputs_tn[1].clone()]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    for inps in [inputs_tt, inputs_tn, inputs_nt, inputs_nn]:
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inps=inps), inps
        )
        eager_out = torch.matmul(inps[0], inps[1])
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(eager_out, fp16_nvf_out)

def test_linear_without_bias(nvfuser_direct_test):
    """
    Test linear without bias. Check that bias keyword argument does not appear in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.linear(t0, t1)
        fd.add_output(t2)

    # Check that bias is not included with linear
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.ops.linear(tv0, tv1)
    fd.add_output(tv2)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_linear_with_bias(nvfuser_direct_test):
    """
    Test linear with bias. Check that bias keyword argument appears in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.from_pytorch(inputs[2])
        t3 = fd.ops.linear(t0, t1, t2)


    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.linear(t0, t1)
        fd.add_output(t2)

    # Check that bias is not included with linear
    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.ops.linear(tv0, tv1)
    fd.add_output(tv2)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_linear_with_bias(nvfuser_direct_test):
    """
    Test linear with bias. Check that bias keyword argument appears in python representation.
    """
    m = 24
    n = 16
    k = 8
    inputs = [
        torch.randn(m, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, k, device="cuda", dtype=torch.bfloat16),
        torch.randn(n, device="cuda", dtype=torch.bfloat16),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.from_pytorch(inputs[2])
        t3 = fd.ops.linear(t0, t1, t2)
        fd.add_output(t3)

    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.BFloat16, is_cpu=False)
    tv3 = fd.ops.linear(tv0, tv1, bias=tv2)
    fd.add_output(tv3)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1], inputs[2])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

@pytest.mark.parametrize("bias", [False, True])
def test_linear(nvfuser_direct_test, bias):
    m = 24
    n = 16
    k = 8
    bias1d = torch.randn(n, device="cuda", dtype=torch.float16) if bias else None

    inputs_mk_nk = [
        torch.randn(m, k, device="cuda", dtype=torch.float16),
        torch.randn(n, k, device="cuda", dtype=torch.float16),
    ]

    inputs_mk_kn = [
        inputs_mk_nk[0].clone(),
        inputs_mk_nk[1].clone().as_strided(size=[n, k], stride=[1, n]),
    ]

    inputs_km_nk = [


        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.from_pytorch(inputs[2])
        t3 = fd.ops.linear(t0, t1, t2)
        fd.add_output(t3)

    fd_str = """def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.BFloat16, is_cpu=False)
    tv3 = fd.ops.linear(tv0, tv1, bias=tv2)
    fd.add_output(tv3)"""

    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        fusion_func, inputs, expected_fd_str=fd_str
    )
    eager_out = torch.nn.functional.linear(inputs[0], inputs[1], inputs[2])
    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

@pytest.mark.parametrize("bias", [False, True])
def test_linear(nvfuser_direct_test, bias):
    m = 24
    n = 16
    k = 8
    bias1d = torch.randn(n, device="cuda", dtype=torch.float16) if bias else None

    inputs_mk_nk = [
        torch.randn(m, k, device="cuda", dtype=torch.float16),
        torch.randn(n, k, device="cuda", dtype=torch.float16),
    ]

    inputs_mk_kn = [
        inputs_mk_nk[0].clone(),
        inputs_mk_nk[1].clone().as_strided(size=[n, k], stride=[1, n]),
    ]

    inputs_km_nk = [
        inputs_mk_nk[0].clone().as_strided(size=[m, k], stride=[1, m]),
        inputs_mk_nk[1].clone(),
    ]

    inputs_km_kn = [
        inputs_km_nk[0].clone(),
        inputs_mk_kn[1].clone(),
    ]

    def fusion_func(
        fd: FusionDefinition,
        inp: torch.Tensor,
        wt: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        t0 = fd.from_pytorch(inp)
        t1 = fd.from_pytorch(wt)
        if bias is not None:
            t2 = fd.from_pytorch(bias)
            t_out = fd.ops.linear(t0, t1, t2)
        else:
            t_out = fd.ops.linear(t0, t1)
        fd.add_output(t_out)

    in_tensors = [inputs_mk_nk, inputs_mk_kn, inputs_km_nk, inputs_km_kn]
    for inp, wt in in_tensors:
        input_tensors = (inp, wt, bias1d) if bias1d is not None else (inp, wt)
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inp=inp, wt=wt, bias=bias1d),
            input_tensors,
        )
        eager_out = F.linear(inp, wt, bias1d)
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(fp16_nvf_out, eager_out, atol=1e-3, rtol=0)

def test_linear_slice(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([1, 2, 3])
        b = fd.define_tensor([4, 3])
        c = fd.ops.linear(a, b)
        d = fd.ops.slice(c, start_indices=[0, 0, 0], end_indices=[1, 2, 3])
        fd.add_output(d)

    inputs = [
        torch.randn(1, 2, 3, device="cuda:0"),
        torch.randn(4, 3, device="cuda:0"),
    ]
    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_matmul_operandA_2d_operandB_3d(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([2, 3])


        t1 = fd.from_pytorch(wt)
        if bias is not None:
            t2 = fd.from_pytorch(bias)
            t_out = fd.ops.linear(t0, t1, t2)
        else:
            t_out = fd.ops.linear(t0, t1)
        fd.add_output(t_out)

    in_tensors = [inputs_mk_nk, inputs_mk_kn, inputs_km_nk, inputs_km_kn]
    for inp, wt in in_tensors:
        input_tensors = (inp, wt, bias1d) if bias1d is not None else (inp, wt)
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
            partial(fusion_func, inp=inp, wt=wt, bias=bias1d),
            input_tensors,
        )
        eager_out = F.linear(inp, wt, bias1d)
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(fp16_nvf_out, eager_out, atol=1e-3, rtol=0)

def test_linear_slice(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([1, 2, 3])
        b = fd.define_tensor([4, 3])
        c = fd.ops.linear(a, b)
        d = fd.ops.slice(c, start_indices=[0, 0, 0], end_indices=[1, 2, 3])
        fd.add_output(d)

    inputs = [
        torch.randn(1, 2, 3, device="cuda:0"),
        torch.randn(4, 3, device="cuda:0"),
    ]
    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_matmul_operandA_2d_operandB_3d(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([2, 3])
        b = fd.define_tensor([7, 3, 5])
        c = fd.ops.matmul(a, b)
        assert c.ndim == 3
        fd.add_output(c)

    inputs = [
        torch.testing.make_tensor(2, 3, dtype=torch.float32, device="cuda"),
        torch.testing.make_tensor(7, 3, 5, dtype=torch.float32, device="cuda"),
    ]
    outputs, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert outputs[0].ndim == 3

def test_matmul_stride(nvfuser_direct_test):
    n, h, l, s, e = 4, 8, 16, 16, 8


        eager_out = F.linear(inp, wt, bias1d)
        fp16_nvf_out = nvf_out[0]
        nvfuser_direct_test.assertEqual(fp16_nvf_out, eager_out, atol=1e-3, rtol=0)

def test_linear_slice(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([1, 2, 3])
        b = fd.define_tensor([4, 3])
        c = fd.ops.linear(a, b)
        d = fd.ops.slice(c, start_indices=[0, 0, 0], end_indices=[1, 2, 3])
        fd.add_output(d)

    inputs = [
        torch.randn(1, 2, 3, device="cuda:0"),
        torch.randn(4, 3, device="cuda:0"),
    ]
    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_matmul_operandA_2d_operandB_3d(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([2, 3])
        b = fd.define_tensor([7, 3, 5])
        c = fd.ops.matmul(a, b)
        assert c.ndim == 3
        fd.add_output(c)

    inputs = [
        torch.testing.make_tensor(2, 3, dtype=torch.float32, device="cuda"),
        torch.testing.make_tensor(7, 3, 5, dtype=torch.float32, device="cuda"),
    ]
    outputs, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert outputs[0].ndim == 3

def test_matmul_stride(nvfuser_direct_test):
    n, h, l, s, e = 4, 8, 16, 16, 8
    inputs = [
        torch.randn(n, h, l, e, device="cuda", dtype=torch.float16, requires_grad=True),
        torch.randn(n, h, s, e, device="cuda", dtype=torch.float16, requires_grad=True),
    ]
    for stride_order in itertools.permutations(range(4)):

        def fusion_func(fd: FusionDefinition) -> None:
            q = fd.from_pytorch(inputs[0])
            k = fd.from_pytorch(inputs[1])
            k_t = fd.ops.permute(k, [0, 1, 3, 2])
            out = fd.ops.matmul(q, k_t)
            strided_out = fd.ops.stride_order(out, stride_order)
            fd.add_output(strided_out)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = torch.matmul(inputs[0], torch.transpose(inputs[1], -2, -1))


    ]
    nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

def test_matmul_operandA_2d_operandB_3d(nvfuser_direct_test):
    def fusion_func(fd: FusionDefinition) -> None:
        a = fd.define_tensor([2, 3])
        b = fd.define_tensor([7, 3, 5])
        c = fd.ops.matmul(a, b)
        assert c.ndim == 3
        fd.add_output(c)

    inputs = [
        torch.testing.make_tensor(2, 3, dtype=torch.float32, device="cuda"),
        torch.testing.make_tensor(7, 3, 5, dtype=torch.float32, device="cuda"),
    ]
    outputs, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert outputs[0].ndim == 3

def test_matmul_stride(nvfuser_direct_test):
    n, h, l, s, e = 4, 8, 16, 16, 8
    inputs = [
        torch.randn(n, h, l, e, device="cuda", dtype=torch.float16, requires_grad=True),
        torch.randn(n, h, s, e, device="cuda", dtype=torch.float16, requires_grad=True),
    ]
    for stride_order in itertools.permutations(range(4)):

        def fusion_func(fd: FusionDefinition) -> None:
            q = fd.from_pytorch(inputs[0])
            k = fd.from_pytorch(inputs[1])
            k_t = fd.ops.permute(k, [0, 1, 3, 2])
            out = fd.ops.matmul(q, k_t)
            strided_out = fd.ops.stride_order(out, stride_order)
            fd.add_output(strided_out)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = torch.matmul(inputs[0], torch.transpose(inputs[1], -2, -1))
        verify_stride_order(nvf_out[0].stride(), stride_order)
        nvfuser_direct_test.assertEqual(nvf_out[0], eager_out)


        c = fd.ops.matmul(a, b)
        assert c.ndim == 3
        fd.add_output(c)

    inputs = [
        torch.testing.make_tensor(2, 3, dtype=torch.float32, device="cuda"),
        torch.testing.make_tensor(7, 3, 5, dtype=torch.float32, device="cuda"),
    ]
    outputs, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
    assert outputs[0].ndim == 3

def test_matmul_stride(nvfuser_direct_test):
    n, h, l, s, e = 4, 8, 16, 16, 8
    inputs = [
        torch.randn(n, h, l, e, device="cuda", dtype=torch.float16, requires_grad=True),
        torch.randn(n, h, s, e, device="cuda", dtype=torch.float16, requires_grad=True),
    ]
    for stride_order in itertools.permutations(range(4)):

        def fusion_func(fd: FusionDefinition) -> None:
            q = fd.from_pytorch(inputs[0])
            k = fd.from_pytorch(inputs[1])
            k_t = fd.ops.permute(k, [0, 1, 3, 2])
            out = fd.ops.matmul(q, k_t)
            strided_out = fd.ops.stride_order(out, stride_order)
            fd.add_output(strided_out)

        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)
        eager_out = torch.matmul(inputs[0], torch.transpose(inputs[1], -2, -1))
        verify_stride_order(nvf_out[0].stride(), stride_order)
        nvfuser_direct_test.assertEqual(nvf_out[0], eager_out)