# SPDX-License-Identifier: Apache-2.0

# Adapted from https://github.com/sgl-project/sglang/pull/3730
import itertools
import unittest

import torch
import triton
import triton.language as tl
import pytest
import os
import numpy as np

from aiter.ops.triton.group_quant_int8 import per_token_group_quant_int8

device = "cuda"

def native_per_token_group_quant_int8(x, group_size, eps=1e-10, dtype=torch.int8):
    """per-token-group quantization on an input tensor `x` using native torch.

    It converts the tensor values into float8 values and returns the
    quantized tensor along with the scaling factor used for quantization.
    Note that only `torch.float8_e4m3fn` is supported for now.
    """
    assert (
        x.shape[-1] % group_size == 0
    ), "the last dimension of `x` cannot be divisible by `group_size`"
    assert x.is_contiguous(), "`x` is not contiguous"

    iinfo = torch.iinfo(dtype)
    int8_min = iinfo.min
    int8_max = iinfo.max

    x_ = x.reshape(x.numel() // group_size, group_size)
    amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
    x_s = amax / int8_max
    x_q = (x_ / x_s).clamp(min=int8_min, max=int8_max).to(dtype)
    x_q = x_q.reshape(x.shape)
    x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))

    return x_q, x_s

# Test cases for per_token_group_quant_int8
QUANT_TEST_CASES = [
    # rows, cols, group_size  # frequency
    (128, 7168, 128),  # 107326
    (128, 128, 128),  # 25812
    (1024, 128, 128),  # 25697
    (128, 1152, 128),  # 1399
    (4096, 7168, 128),  # 241
    (16384, 7168, 128),  # 241
    (512, 512, 128),  # 61
    (1024, 1024, 128),  # 61
    (1536, 1536, 128),  # 61
    (4096, 512, 128),  # 61
    (4096, 1024, 128),  # 61
    (4096, 1536, 128),  # 61
    (16384, 512, 128),  # 61
    (16384, 1024, 128),  # 61
    (16384, 1536, 128),  # 61
    (4096, 128, 128),  # 58
    (16384, 128, 128),  # 58
    (32768, 128, 128),  # 58
    (131072, 128, 128),  # 58
    (4096, 1152, 128),  # 3
    (16384, 1152, 128),  # 3
    (128, 7168, 1),  # 1
]

@pytest.mark.parametrize("rows,cols,group_size", QUANT_TEST_CASES)
def test_per_token_group_quant_int8(rows, cols, group_size):
    """Test per_token_group_quant_int8 implementation."""
    device = "cuda"
    x = torch.randn((rows, cols), dtype=torch.float16, device=device)
    
    # Run Triton implementation
    x_q_triton, x_s_triton = per_token_group_quant_int8(x, group_size)
    
    # Run native PyTorch implementation
    x_q_torch, x_s_torch = native_per_token_group_quant_int8(x, group_size)
    
    # Compare results
    torch.testing.assert_close(x_q_triton, x_q_torch, rtol=0, atol=1)
    torch.testing.assert_close(x_s_triton, x_s_torch, rtol=1e-3, atol=1e-3)