unslothai--unsloth/unsloth/kernels/moe/tests/test_grouped_gemm.py

# SPDX-License-Identifier: GNU Affero General Public License v3.0
# Copyright 2023-present the Unsloth team. All rights reserved.

from dataclasses import asdict

import pytest
import torch

from grouped_gemm.interface import (
    grouped_gemm,
    grouped_gemm_dW,
    grouped_gemm_dX,
    grouped_gemm_forward,
)
from grouped_gemm.kernels.tuning import (
    KernelConfig,
    KernelConfigBackward_dW,
    KernelConfigBackward_dX,
    KernelConfigForward,
)
from grouped_gemm.reference.moe_ops import (
    calculate_topk,
    get_routing_indices,
    permute,
    torch_grouped_gemm,
    unpermute,
)

from .common import (
    DATA_CONFIGS,
    KERNEL_CONFIGS_FWD,
    LLAMA_MODEL_CONFIG,
    QWEN_MODEL_CONFIG,
    SMALL_MODEL_CONFIGS,
    TOLERANCE,
    DataConfig,
    KERNEL_CONFIGS_BWD_dW,
    KERNEL_CONFIGS_BWD_dX,
    ModelConfig,
    make_inputs,
)

SEED = 0


# Only certain (permute_x, permute_y, use_W1) combinations are valid; see the
# module string below for the full rationale.
def check_valid_config(
    permute_x,
    permute_y,
    use_W1,
    fuse_mul_post = False,
    is_backward = False,
    verbose = False,
):
    use_W2 = not use_W1

    if permute_x and permute_y:
        if verbose:
            print(f"Skipping test: {permute_x = } {permute_y = }")
        return False
    if use_W2 and permute_x:
        if verbose:
            print(f"Skipping test: {permute_x = } {use_W2 = }")
        return False
    if use_W1 and permute_y:
        if verbose:
            print(f"Skipping test: {permute_y = } {use_W1 = }")
        return False
    if fuse_mul_post and use_W1:
        if verbose:
            print(f"Skipping test: {fuse_mul_post = } {use_W1 = }")
        return False
    if is_backward and fuse_mul_post:
        if verbose:
            print(f"Skipping test: {fuse_mul_post = } {is_backward = }")
        return False

    return True


"""
grouped_gemm_forward

permute_x: typically in a fused MoE MLP, we can fuse the permutation of hidden states (X) from token order to expert grouped order needed for grouped GEMM by directly loading X in permuted order rather than launching a separate permutation kernel.
permute_y: We can also fuse the unpermutation of tokens after the second grouped GEMM to restore to original token order.  This is fused into the second grouped GEMM by directly storing the output in unpermuted order.
fuse_mul: We can also fuse the multiplication of topk weights in the epilogue of the second grouped GEMM.  Note that this is only supported for inference and not training, although this may change in the future.
use_W1 test the shapes for the first grouped GEMM in a fused MoE MLP
use_W2 = `not use_W1` tests the shapes for the second grouped GEMM in a fused MoE MLP

Given the above, only certain combinations are valid:
- use_W1 is always False when permute_y is True since we only permute the second grouped GEMM
- use_W2 is always False when permute_x is True since we only permute the first grouped GEMM
- only one of permute_x and permute_y can be True
- fuse_mul is only True if permute_y is also True

See `check_valid_config` for more details.
"""


def _test_grouped_gemm_forward(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,  # W1 -> first grouped GEMM in a fused MoE MLP, not W1 -> second grouped GEMM in a fused MoE MLP
    fuse_mul_post: bool = False,
    flatten: bool = True,
    # Manually tuned parameters
    use_tma_load_w: bool = False,
    use_tma_load_x: bool = False,
    use_tma_store: bool = False,
    BLOCK_SIZE_M: int = None,
    BLOCK_SIZE_N: int = None,
    BLOCK_SIZE_K: int = None,
    num_warps: int = None,
    num_stages: int = None,
    # Autotuning parameters
    autotune: bool = False,
    num_autotune_configs: int = None,
    # Flag to manually enable TMA store
    allow_tma_store: bool = False,
    use_autograd: bool = False,
):
    if not check_valid_config(permute_x, permute_y, use_W1 = use_W1, fuse_mul_post = fuse_mul_post):
        pytest.skip(
            f"Skipping test due to invalid config: {permute_x = } {permute_y = } {use_W1 = } {fuse_mul_post = }"
        )

    if use_tma_store and not allow_tma_store:
        pytest.skip("TMA store needs to be debugged due to non-deterministic behavior")

    X1, X2, W1, W2, gating_output = make_inputs(
        M = data_config.bs * data_config.seq_len,
        N = model_config.intermediate_size,
        K = model_config.hidden_size,
        E = model_config.num_experts,
        topk = model_config.topk,
        dtype = data_config.dtype,
    )
    topk = model_config.topk
    use_sigmoid = model_config.use_sigmoid
    renormalize = model_config.renormalize

    X = X1 if use_W1 else X2
    num_tokens = data_config.bs * data_config.seq_len
    E, K, N = W2.shape  # E = num_experts, K = hidden_size, N = intermediate_size
    assert W1.shape == (E, 2 * N, K)
    W = W1 if use_W1 else W2

    if use_W1:
        assert X.shape == (
            num_tokens,
            K,
        ), f"X.shape: {X.shape}, num_tokens: {num_tokens}, K: {K}"
    else:
        assert X.shape == (
            num_tokens * topk,
            N,
        ), f"X.shape: {X.shape}, num_tokens: {num_tokens}, topk: {topk}, N: {N}"

    total_tokens = num_tokens * topk
    output_shape = (total_tokens, 2 * N) if use_W1 else (total_tokens, K)

    topk_weights, topk_ids = calculate_topk(
        gating_output, topk, use_sigmoid = use_sigmoid, renormalize = renormalize
    )
    topk_weights = topk_weights.view(-1)  # num_tokens * topk
    topk_ids = topk_ids.view(-1)  # num_tokens * topk

    expert_token_counts, gather_indices = get_routing_indices(topk_ids, num_experts = E)
    assert len(gather_indices) == total_tokens
    assert len(expert_token_counts) == E

    atol, rtol = TOLERANCE[X.dtype]

    Xperm = permute(X, gather_indices, topk)

    Xref = Xperm

    assert (
        Xperm.shape == (total_tokens, K) if use_W1 else (total_tokens, N)
    ), f"Xperm.shape: {Xperm.shape}, total_tokens: {total_tokens}, K: {K}"

    ref_output = torch_grouped_gemm(X = Xref, W = W, m_sizes = expert_token_counts)

    if permute_x:
        X_test = X
    else:
        X_test = Xperm

    # No need to run all configs for tests, otherwise takes too long
    if autotune:
        from grouped_gemm.kernels.forward import _autotuned_grouped_gemm_forward_kernel
        if num_autotune_configs is not None:
            _autotuned_grouped_gemm_forward_kernel.configs = (
                _autotuned_grouped_gemm_forward_kernel.configs[:num_autotune_configs]
            )

    # Use autograd.Function interface
    if use_autograd:
        from grouped_gemm.interface import grouped_gemm
        kernel_config_fwd = KernelConfigForward(
            BLOCK_SIZE_M = BLOCK_SIZE_M,
            BLOCK_SIZE_N = BLOCK_SIZE_N,
            BLOCK_SIZE_K = BLOCK_SIZE_K,
            num_warps = num_warps,
            num_stages = num_stages,
            permute_x = permute_x,
            permute_y = permute_y,
            fuse_mul_post = fuse_mul_post,
            use_tma_load_w = use_tma_load_w,
            use_tma_load_x = use_tma_load_x,
            use_tma_store = use_tma_store,
        )

        test_output = grouped_gemm(
            X = X_test,
            W = W,
            topk = topk,
            m_sizes = expert_token_counts,
            gather_indices = gather_indices,
            topk_weights = topk_weights if fuse_mul_post else None,
            permute_x = permute_x,
            permute_y = permute_y,
            fuse_mul_post = fuse_mul_post,
            kernel_config_fwd = kernel_config_fwd,
            autotune = autotune,
            is_first_gemm = use_W1,
        )
    # Use manual interface
    else:
        test_output = grouped_gemm_forward(
            X = X_test,
            W = W,
            topk = topk,
            m_sizes = expert_token_counts,
            gather_indices = gather_indices,
            topk_weights = topk_weights if fuse_mul_post else None,
            permute_x = permute_x,
            permute_y = permute_y,
            fuse_mul_post = fuse_mul_post,
            use_tma_load_w = use_tma_load_w,
            use_tma_load_x = use_tma_load_x,
            use_tma_store = use_tma_store,
            autotune = autotune,
            BLOCK_SIZE_M = BLOCK_SIZE_M,
            BLOCK_SIZE_N = BLOCK_SIZE_N,
            BLOCK_SIZE_K = BLOCK_SIZE_K,
            num_warps = num_warps,
            num_stages = num_stages,
            flatten = flatten,
        )
    assert ref_output.shape == output_shape
    assert test_output.shape == output_shape

    if permute_y:
        ref_output = unpermute(ref_output, gather_indices)
    if fuse_mul_post:
        # if we don't permute_y, then test output is permuted with topk weights applied
        # the ref output needs to be unpermuted before multiplying by topk weights since topk weights are in token order
        if not permute_y:
            ref_output = unpermute(ref_output, gather_indices)
            test_output = unpermute(test_output, gather_indices)
        ref_output = ref_output * topk_weights[:, None]

    assert torch.allclose(
        ref_output, test_output, atol = atol, rtol = rtol
    ), f"Grouped gemm forward failed: {(ref_output - test_output).abs().max().item():.6f}"


# NOTE: Fuse multiplication of topk weights is only supported for inference and not training, although this may change in the future; not currently tested.
@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_FWD,
    ids = lambda x: x.to_string(include_tuning_params = True, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS + [QWEN_MODEL_CONFIG, LLAMA_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_forward_manual(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfigForward,
    use_W1: bool,
):
    _test_grouped_gemm_forward(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_FWD,
    ids = lambda x: x.to_string(include_tuning_params = True, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS + [QWEN_MODEL_CONFIG, LLAMA_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_forward_manual_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfigForward,
    use_W1: bool,
):
    _test_grouped_gemm_forward(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        use_autograd = True,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize("num_autotune_configs", [10], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_forward_autotune(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    _test_grouped_gemm_forward(
        data_config = data_config,
        model_config = model_config,
        permute_x = permute_x,
        permute_y = permute_y,
        use_W1 = use_W1,
        num_autotune_configs = num_autotune_configs,
        autotune = True,
        use_autograd = False,
    )


@pytest.mark.parametrize("num_autotune_configs", [10], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_forward_autotune_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    _test_grouped_gemm_forward(
        data_config = data_config,
        model_config = model_config,
        permute_x = permute_x,
        permute_y = permute_y,
        use_W1 = use_W1,
        num_autotune_configs = num_autotune_configs,
        autotune = True,
        use_autograd = True,
    )


"""
grouped_gemm_backward_dX

use_W1 test the shapes for the first grouped GEMM in a fused MoE MLP
use_W2 = `not use_W1` tests the shapes for the second grouped GEMM in a fused MoE MLP

Only certain combinations of permute_x, permute_y, and fuse_mul are supported.

Typically in a fused MoE MLP, we can fuse the permutation of hidden states (X) from token order to expert grouped order needed for grouped GEMM by directly loading X in permuted order rather than launching a separate permutation kernel.
We can also fuse the unpermutation of tokens after the second grouped GEMM to restore to original token order.  This is fused into the second grouped GEMM by directly storing the output in unpermuted order.

Hence the following conditions:
- If use_W1 there are two cases:
    - permute_x is False and topk > 1:
    - dX_test is still in permuted order and has shape (total_tokens, K)
    - it needs to be unpermuted and summed across topk before comparing to ref_grad
- permute_x is True:
    - dX_test is already unpermuted and summed across topk with shape (num_tokens, K)
    - no further processing is needed
- permute_x is False and topk == 1:
    - dX_test needs to be permuted, no need to sum since topk == 1

- If use_W2:
    - permute_x is always False
    - if permute_y:
        - grad_output needs to be unpermuted before passing to grouped_gemm_dX
        - dX_test is permuted and has shape (total_tokens, N)
        - it needs to be unpermuted before comparing to ref_grad or can be compared directly to Xperm.grad
    - if not permute_y:
        - dX_test is not permuted and has shape (total_tokens, N)
        - no further processing is needed
"""


def _test_grouped_gemm_backward_dX(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool = False,
    permute_y: bool = False,
    use_tma_load_dy: bool = False,
    use_tma_load_w: bool = False,
    use_tma_store: bool = False,
    use_W1: bool = True,
    autotune: bool = False,
    num_autotune_configs: int = None,
    BLOCK_SIZE_M: int = None,
    BLOCK_SIZE_N: int = None,
    BLOCK_SIZE_K: int = None,
    num_warps: int = None,
    num_stages: int = None,
    flatten: bool = True,
    allow_tma_store: bool = False,
    use_autograd: bool = False,
    fuse_mul_post: bool = False,
):
    if not check_valid_config(permute_x, permute_y, use_W1 = use_W1, is_backward = True):
        pytest.skip(
            f"Skipping test due to invalid config: {permute_x = } {permute_y = } {use_W1 = }"
        )

    if use_tma_store and not allow_tma_store:
        pytest.skip("TMA store needs to be debugged due to non-deterministic behavior")

    if autotune and model_config.intermediate_size <= 128 and model_config.hidden_size <= 128:
        pytest.skip("Skipping autotuning for small model configs")

    # Prevent OOM for large intermediate sizes
    if model_config.intermediate_size > 2048:
        model_config.intermediate_size = 1024
    if model_config.hidden_size > 2048:
        model_config.hidden_size = 1024

    use_W2 = not use_W1
    X1, X2, W1, W2, gating_output = make_inputs(
        M = data_config.bs * data_config.seq_len,
        N = model_config.intermediate_size,
        K = model_config.hidden_size,
        E = model_config.num_experts,
        topk = model_config.topk,
        dtype = data_config.dtype,
        requires_grad = True,
    )
    topk = model_config.topk
    num_experts = model_config.num_experts
    use_sigmoid = model_config.use_sigmoid
    renormalize = model_config.renormalize

    X = X1 if use_W1 else X2
    num_tokens = data_config.bs * data_config.seq_len
    total_tokens = num_tokens * topk

    E, K, N = W2.shape  # E = num_experts, K = hidden_size, N = intermediate_size
    assert W1.shape == (E, 2 * N, K)
    W = W1 if use_W1 else W2

    if use_W1:
        assert X.shape == (
            num_tokens,
            K,
        ), f"X.shape: {X.shape}, num_tokens: {num_tokens}, K: {K}"
    else:
        assert X.shape == (
            total_tokens,
            N,
        ), f"X.shape: {X.shape}, total_tokens: {total_tokens}, N: {N}"

    W_test = W.detach().clone().requires_grad_(True)

    topk_weights, topk_ids = calculate_topk(
        gating_output, topk, use_sigmoid = use_sigmoid, renormalize = renormalize
    )
    topk_weights = topk_weights.view(-1)  # num_tokens * topk
    topk_ids = topk_ids.view(-1)  # num_tokens * topk

    expert_token_counts, gather_indices = get_routing_indices(topk_ids, num_experts = E)
    assert len(gather_indices) == total_tokens
    assert len(expert_token_counts) == num_experts

    atol, rtol = TOLERANCE[X.dtype]
    Xperm = permute(X, gather_indices, topk)

    # Need to retain grad otherwise grad is not propagated
    X.retain_grad()
    W.retain_grad()
    Xperm.retain_grad()

    assert Xperm.shape == (total_tokens, K) if use_W1 else (total_tokens, N)

    output_shape = (total_tokens, 2 * N) if use_W1 else (total_tokens, K)
    ref_output = torch_grouped_gemm(X = Xperm, W = W, m_sizes = expert_token_counts)
    assert (
        ref_output.shape == output_shape
    ), f"ref_output.shape: {ref_output.shape}, output_shape: {output_shape}"

    if permute_y:
        ref_output = unpermute(ref_output, gather_indices)

    grad_output = torch.randn_like(ref_output)
    ref_output.backward(grad_output)

    assert X.grad is not None
    assert W.grad is not None

    ref_grad = Xperm.grad

    if autotune:
        # No need to run all configs for autotuning
        from grouped_gemm.kernels.backward import _autotuned_grouped_gemm_dX_kernel
        if num_autotune_configs is not None:
            _autotuned_grouped_gemm_dX_kernel.configs = _autotuned_grouped_gemm_dX_kernel.configs[
                :num_autotune_configs
            ]

    if use_autograd:
        from grouped_gemm.interface import grouped_gemm

        if not autotune:
            kernel_config_fwd = KernelConfigForward()
            kernel_config_bwd_dX = KernelConfigBackward_dX(
                use_tma_load_dy = use_tma_load_dy,
                use_tma_load_w = use_tma_load_w,
                use_tma_store = use_tma_store,
                BLOCK_SIZE_M = BLOCK_SIZE_M,
                BLOCK_SIZE_N = BLOCK_SIZE_N,
                BLOCK_SIZE_K = BLOCK_SIZE_K,
                num_warps = num_warps,
                num_stages = num_stages,
            )
            kernel_config_bwd_dW = KernelConfigBackward_dW()
        else:
            from grouped_gemm.kernels.backward import (
                _autotuned_grouped_gemm_dW_kernel,
                _autotuned_grouped_gemm_dX_kernel,
            )
            from grouped_gemm.kernels.forward import (
                _autotuned_grouped_gemm_forward_kernel,
            )

            if num_autotune_configs is not None:
                _autotuned_grouped_gemm_dX_kernel.configs = (
                    _autotuned_grouped_gemm_dX_kernel.configs[:num_autotune_configs]
                )
                _autotuned_grouped_gemm_forward_kernel.configs = (
                    _autotuned_grouped_gemm_forward_kernel.configs[:num_autotune_configs]
                )

            kernel_config_fwd = None
            kernel_config_bwd_dX = None
        X_ = (
            X.detach().clone().requires_grad_(True)
            if permute_x
            else Xperm.detach().clone().requires_grad_(True)
        )
        test_output = grouped_gemm(
            X = X_,
            W = W_test,
            m_sizes = expert_token_counts,
            gather_indices = gather_indices,
            topk = topk,
            permute_x = permute_x,
            permute_y = permute_y,
            autotune = autotune,
            kernel_config_fwd = kernel_config_fwd,
            kernel_config_bwd_dX = kernel_config_bwd_dX,
            is_first_gemm = use_W1,
            dX_only = True,
        )
        assert (
            test_output.shape == ref_output.shape
        ), f"test_output.shape: {test_output.shape}, ref_output.shape: {ref_output.shape}"
        assert torch.allclose(
            test_output, ref_output, atol = atol, rtol = rtol
        ), f"Grouped gemm backward_dX forward outputs mismatch: {(test_output - ref_output).abs().max().item():.6f}"
        test_output.backward(grad_output)
        assert X_.grad is not None

        # NOTE:need to handle grad differenlty in this case due to errors arising to do how torch autograd handles unpermute and sum reduction
        # the grad of Xperm unpermuted and reduced across topk should match X_.grad
        # However, both will have a numerical difference with that of ref_grad
        # This is due to the fact that torch autograd handles unpermute and sum reduction differently see: https://discuss.pytorch.org/t/permute-unpermute-gradient/219557    else:
        if permute_x and use_W1:
            X_grad_unperm = unpermute(Xperm.grad, gather_indices)
            manual_grad_check = X_grad_unperm.view(num_tokens, topk, K).sum(dim = 1)
            assert (
                manual_grad_check.shape == X_.grad.shape
            ), f"manual_grad_check.shape: {manual_grad_check.shape}, X_.grad.shape: {X_.grad.shape}"
            assert torch.allclose(
                manual_grad_check, X_.grad, atol = atol, rtol = rtol
            ), f"Grouped gemm backward_dX forward outputs mismatch: {(manual_grad_check - X_.grad).abs().max().item():.6f}"
            manual_diff = (X_.grad - manual_grad_check).abs().max().item()
            autograd_diff = (X_.grad - X.grad).abs().max().item()
            print(f"manual_diff: {manual_diff:.6f}, autograd_diff: {autograd_diff:.6f}")
        else:
            assert torch.allclose(
                X_.grad, ref_grad, atol = atol, rtol = rtol
            ), f"Grouped gemm backward_dX forward outputs mismatch: {(X_.grad - ref_grad).abs().max().item():.6f}"
        return
    else:
        dX_test = grouped_gemm_dX(
            dY = grad_output,
            W = W_test,
            gather_indices = gather_indices,
            m_sizes = expert_token_counts,
            topk = topk,
            permute_x = permute_x,
            permute_y = permute_y,
            use_tma_load_w = use_tma_load_w,
            use_tma_load_dy = use_tma_load_dy,
            use_tma_store = use_tma_store,
            autotune = autotune,
            BLOCK_SIZE_M = BLOCK_SIZE_M,
            BLOCK_SIZE_N = BLOCK_SIZE_N,
            BLOCK_SIZE_K = BLOCK_SIZE_K,
            num_warps = num_warps,
            num_stages = num_stages,
            flatten = flatten,
            # debug=True,
        )

    # if permute_x and use_W1 (first grouped GEMM) then the kernel should have unpermuted the dX
    # therefore we need to unpermute the ref_grad to compare to the output of the kernel
    if permute_x and use_W1:
        ref_grad = unpermute(ref_grad, gather_indices)

    assert (
        ref_grad.shape == dX_test.shape
    ), f"Grouped gemm manual backward_dX outputs mismatch: ref_grad: {ref_grad.shape}, dX_test: {dX_test.shape}"
    diff = (ref_grad - dX_test).abs().max().item()

    assert torch.allclose(
        ref_grad, dX_test, atol = atol, rtol = rtol
    ), f"Grouped gemm manual backward_dX outputs mismatch: {diff:.6f}"

    if permute_x and use_W1:
        # Show that reduction results in diffs
        # First calculate X.grad manually by backpropping through unpermuted ref_grad
        dX_ref_check = ref_grad.view(num_tokens, topk, K).sum(dim = 1)
        # Do the same for the actual output of the kernel
        dX_test_check = dX_test.view(num_tokens, topk, K).sum(dim = 1)
        # Show diffs for each combination
        diff_ref_check = (X.grad - dX_ref_check).abs().max().item()
        diff_test_check = (X.grad - dX_test_check).abs().max().item()
        diff_check_test = (dX_ref_check - dX_test_check).abs().max().item()
        print(
            f"diff_ref_check: {diff_ref_check:.6f}, diff_test_check: {diff_test_check:.6f}, diff_check_test: {diff_check_test:.6f}"
        )


# NOTE: We reduce the size of the Llama4 model configs to prevent OOM
# Important to note that for the full model size (5120, 8192), the tests do result in diffs on the order of 1e-2.
@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_BWD_dX,
    ids = lambda x: x.to_string(include_tuning_params = True, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS[:1] + [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dX_manual(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfigBackward_dX,
    use_W1: bool,
):
    _test_grouped_gemm_backward_dX(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        use_autograd = False,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_BWD_dX,
    ids = lambda x: x.to_string(include_tuning_params = True, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS[:1] + [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dX_manual_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfigBackward_dX,
    use_W1: bool,
):
    _test_grouped_gemm_backward_dX(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        use_autograd = True,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize("num_autotune_configs", [20], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dX_autotune(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    # TMA loads / stores will be autotuned
    _test_grouped_gemm_backward_dX(
        data_config = data_config,
        model_config = model_config,
        permute_x = permute_x,
        permute_y = permute_y,
        use_W1 = use_W1,
        autotune = True,
        use_autograd = False,
        num_autotune_configs = num_autotune_configs,
    )


@pytest.mark.parametrize("num_autotune_configs", [20], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dX_autotune_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    # TMA loads / stores will be autotuned
    _test_grouped_gemm_backward_dX(
        data_config = data_config,
        model_config = model_config,
        permute_x = permute_x,
        permute_y = permute_y,
        use_W1 = use_W1,
        autotune = True,
        use_autograd = True,
        num_autotune_configs = num_autotune_configs,
    )


def _test_grouped_gemm_backward_dW(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    use_tma_load_dy: bool = False,
    use_tma_load_x: bool = False,
    use_tma_store: bool = False,
    BLOCK_SIZE_M: int = None,
    BLOCK_SIZE_N: int = None,
    BLOCK_SIZE_K: int = None,
    num_warps: int = None,
    num_stages: int = None,
    flatten: bool = True,
    autotune: bool = False,
    num_autotune_configs: int = None,
    allow_tma_store: bool = False,
    debug: bool = False,
    fuse_mul_post: bool = False,  # Unused for backward_dW
    use_autograd: bool = False,
):
    if not check_valid_config(
        permute_x,
        permute_y,
        fuse_mul_post = fuse_mul_post,
        use_W1 = use_W1,
        is_backward = True,
    ):
        pytest.skip(
            f"Skipping test due to invalid config: {permute_x = } {permute_y = } {use_W1 = }"
        )

    if use_tma_store and not allow_tma_store:
        pytest.skip("TMA store needs to be debugged due to non-deterministic behavior")

    X1, X2, W1, W2, gating_output = make_inputs(
        M = data_config.bs * data_config.seq_len,
        N = model_config.intermediate_size,
        K = model_config.hidden_size,
        E = model_config.num_experts,
        topk = model_config.topk,
        dtype = data_config.dtype,
        requires_grad = True,
    )
    topk = model_config.topk
    num_experts = model_config.num_experts
    use_sigmoid = model_config.use_sigmoid
    renormalize = model_config.renormalize

    X = X1 if use_W1 else X2
    num_tokens = data_config.bs * data_config.seq_len
    E, K, N = W2.shape  # E = num_experts, K = hidden_size, N = intermediate_size
    assert W1.shape == (E, 2 * N, K)
    W = W1 if use_W1 else W2

    if use_W1:
        assert X.shape == (
            num_tokens,
            K,
        ), f"X.shape: {X.shape}, num_tokens: {num_tokens}, K: {K}"
    else:
        assert X.shape == (
            num_tokens * topk,
            N,
        ), f"X.shape: {X.shape}, num_tokens: {num_tokens}, topk: {topk}, N: {N}"

    total_tokens = num_tokens * topk
    output_shape = (total_tokens, 2 * N) if use_W1 else (total_tokens, K)

    X_test = X.detach().clone().requires_grad_(True)
    W_test = W.detach().clone().requires_grad_(True)

    topk_weights, topk_ids = calculate_topk(
        gating_output, topk, use_sigmoid = use_sigmoid, renormalize = renormalize
    )
    topk_weights = topk_weights.view(-1)  # num_tokens * topk
    topk_ids = topk_ids.view(-1)  # num_tokens * topk

    expert_token_counts, gather_indices = get_routing_indices(topk_ids, num_experts = E)
    assert len(gather_indices) == total_tokens
    assert len(expert_token_counts) == num_experts

    atol, rtol = TOLERANCE[X.dtype]
    Xperm = permute(X, gather_indices, topk)
    Xperm_test = Xperm.detach().clone().requires_grad_(True)

    # Need to retain grad otherwise grad is not propagated
    X.retain_grad()
    W.retain_grad()
    Xperm.retain_grad()
    assert Xperm.shape == (total_tokens, K) if use_W1 else (total_tokens, N)

    output_shape = (total_tokens, 2 * N) if use_W1 else (total_tokens, K)

    ref_output = torch_grouped_gemm(X = Xperm, W = W, m_sizes = expert_token_counts)
    assert ref_output.shape == output_shape

    # if permute_y then the assumption is that the output of grouped_gemm was unpermuted on store
    # Therefore we have to unpermute before backpropping to ensure proper alignment
    if permute_y:
        ref_output = unpermute(ref_output, gather_indices)

    grad_output = torch.randn_like(ref_output)
    ref_output.backward(grad_output)
    assert X.grad is not None
    assert W.grad is not None

    # Test backward kernel directly
    X_ = X_test if permute_x else Xperm_test

    if debug:
        torch.set_printoptions(precision = 4)
        for i in range(num_experts):
            print(f"Expert {i} weight grad:\n{W.grad[i, :5, :5]}")

    if autotune:
        from grouped_gemm.kernels.backward import _autotuned_grouped_gemm_dW_kernel
        if num_autotune_configs is not None:
            _autotuned_grouped_gemm_dW_kernel.configs = _autotuned_grouped_gemm_dW_kernel.configs[
                :num_autotune_configs
            ]

    if use_autograd:
        from grouped_gemm.interface import grouped_gemm

        if not autotune:
            kernel_config_fwd = KernelConfigForward(
                # Only care about backward_dW config
                use_tma_load_w = False,
                use_tma_load_x = False,
                use_tma_store = False,
                BLOCK_SIZE_M = BLOCK_SIZE_M,
                BLOCK_SIZE_N = BLOCK_SIZE_N,
                BLOCK_SIZE_K = BLOCK_SIZE_K,
                num_warps = num_warps,
                num_stages = num_stages,
            )
            kernel_config_bwd_dW = KernelConfigBackward_dW(
                use_tma_load_dy = use_tma_load_dy,
                use_tma_load_x = use_tma_load_x,
                use_tma_store = use_tma_store,
                BLOCK_SIZE_M = BLOCK_SIZE_M,
                BLOCK_SIZE_N = BLOCK_SIZE_N,
                BLOCK_SIZE_K = BLOCK_SIZE_K,
                num_warps = num_warps,
                num_stages = num_stages,
            )
        else:
            from grouped_gemm.kernels.backward import _autotuned_grouped_gemm_dW_kernel
            from grouped_gemm.kernels.forward import (
                _autotuned_grouped_gemm_forward_kernel,
            )

            if num_autotune_configs is not None:
                _autotuned_grouped_gemm_forward_kernel.configs = (
                    _autotuned_grouped_gemm_forward_kernel.configs[:num_autotune_configs]
                )
                _autotuned_grouped_gemm_dW_kernel.configs = (
                    _autotuned_grouped_gemm_dW_kernel.configs[:num_autotune_configs]
                )
            kernel_config_fwd = None
            kernel_config_bwd_dW = None

        test_output = grouped_gemm(
            X = X_,
            W = W_test,
            m_sizes = expert_token_counts,
            gather_indices = gather_indices,
            topk = topk,
            permute_x = permute_x,
            permute_y = permute_y,
            kernel_config_fwd = kernel_config_fwd,
            kernel_config_bwd_dW = kernel_config_bwd_dW,
            autotune = autotune,
            is_first_gemm = use_W1,
            dW_only = True,
        )
        assert (
            test_output.shape == ref_output.shape
        ), f"Grouped gemm autograd backward_dW outputs mismatch: {test_output.shape} != {ref_output.shape}"
        assert torch.allclose(
            test_output, ref_output, atol = atol, rtol = rtol
        ), f"Grouped gemm autograd backward_dW forward outputs mismatch: {test_output.shape} != {ref_output.shape}"
        test_output.backward(grad_output)
        assert W_test.grad is not None
        dW_test = W_test.grad
    else:
        dW_test = grouped_gemm_dW(
            dY = grad_output,
            X = X_,
            m_sizes = expert_token_counts,
            gather_indices = gather_indices,
            topk = topk,
            permute_x = permute_x,
            permute_y = permute_y,
            use_tma_load_dy = use_tma_load_dy,
            use_tma_load_x = use_tma_load_x,
            use_tma_store = use_tma_store,
            BLOCK_SIZE_M = BLOCK_SIZE_M,
            BLOCK_SIZE_N = BLOCK_SIZE_N,
            BLOCK_SIZE_K = BLOCK_SIZE_K,
            num_warps = num_warps,
            num_stages = num_stages,
            flatten = flatten,
            autotune = autotune,
            debug = debug,
        )
    assert (
        W.grad.shape == dW_test.shape
    ), f"Grouped gemm manual backward_dW outputs mismatch: W.grad: {W.grad.shape}, dW_test: {dW_test.shape}"

    if debug:
        with torch.no_grad():
            if not torch.allclose(W.grad, dW_test, atol = atol, rtol = rtol):
                print(f"Ref Wgrad sum: {W.grad.sum().item():.4f}")
            print(f"Test Wgrad sum: {dW_test.sum().item():.4f}")

            for i in range(num_experts):
                print(f"Expert {i} weight grad:\n{W.grad[i, :5, :5]}")
                print(f"Expert {i} dW_test:\n{dW_test[i, :5, :5]}")
                expert_diff = (W.grad[i, :, :] - dW_test[i, :, :]).abs().max().item()
                print(f"Expert {i} diff: {expert_diff:.6f}")

            diff = (W.grad - dW_test).abs().max().item()
            assert False, f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"
    else:
        diff = (W.grad - dW_test).abs().max().item()
        assert torch.allclose(
            W.grad, dW_test, atol = atol, rtol = rtol
        ), f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"


@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_BWD_dW,
    ids = lambda x: x.to_string(include_tuning_params = False, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS + [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dW_manual(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfig,
    use_W1: bool,
    debug: bool = False,
):
    _test_grouped_gemm_backward_dW(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        use_autograd = False,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize(
    "kernel_config",
    KERNEL_CONFIGS_BWD_dW,
    ids = lambda x: x.to_string(include_tuning_params = False, include_tma = True),
)
@pytest.mark.parametrize(
    "model_config",
    SMALL_MODEL_CONFIGS + [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dW_manual_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    kernel_config: KernelConfig,
    use_W1: bool,
    debug: bool = False,
):
    _test_grouped_gemm_backward_dW(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        use_autograd = True,
        **asdict(kernel_config),
    )


@pytest.mark.parametrize("num_autotune_configs", [20], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dW_autotune(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    _test_grouped_gemm_backward_dW(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        permute_x = permute_x,
        permute_y = permute_y,
        autotune = True,
        use_autograd = False,
        num_autotune_configs = num_autotune_configs,
    )


@pytest.mark.parametrize("num_autotune_configs", [20], ids = lambda x: f"num_autotune_configs={x}")
@pytest.mark.parametrize("permute_x", [True, False], ids = lambda x: "permute_x" if x else "")
@pytest.mark.parametrize("permute_y", [True, False], ids = lambda x: "permute_y" if x else "")
@pytest.mark.parametrize(
    "model_config",
    [LLAMA_MODEL_CONFIG, QWEN_MODEL_CONFIG],
    ids = lambda x: f"topk={x.topk} num_experts={x.num_experts} hidden_size={x.hidden_size} intermediate_size={x.intermediate_size}",
)
@pytest.mark.parametrize(
    "data_config", DATA_CONFIGS, ids = lambda x: f"seq_len={x.seq_len} dtype={x.dtype}"
)
@pytest.mark.parametrize("use_W1", [True, False], ids = lambda x: f"use_W1={x}")
def test_grouped_gemm_backward_dW_autotune_autograd(
    data_config: DataConfig,
    model_config: ModelConfig,
    permute_x: bool,
    permute_y: bool,
    use_W1: bool,
    num_autotune_configs: int,
):
    _test_grouped_gemm_backward_dW(
        data_config = data_config,
        model_config = model_config,
        use_W1 = use_W1,
        permute_x = permute_x,
        permute_y = permute_y,
        autotune = True,
        use_autograd = True,
        num_autotune_configs = num_autotune_configs,
    )