import torch

from kernels import get_kernel

# Tải xuống các kernel được tối ưu hóa từ Hugging Face hub
activation = get_kernel("kernels-community/activation")

# Tensor ngẫu nhiên
x = torch.randn((10, 10), dtype=torch.float16, device="cuda")

# Chạy kernel
y = torch.empty_like(x)
activation.gelu_fast(y, x)

print(y)

from kernels import get_kernel

flash_attention = get_kernel("kernels-community/flash-attn")

# /// script
# dependencies = [
#  "numpy",
#  "torch",
#  "kernels",
# ]
# ///

import torch
import torch.nn.functional as F
from kernels import get_kernel

DEVICE = "cuda"

# Tạo khả năng tái tạo
torch.manual_seed(42)

# Tải xuống các kernel kích hoạt được tối ưu hóa từ Hub
activation_kernels = get_kernel("kernels-community/activation")

# Tạo một tensor ngẫu nhiên trên GPU
x = torch.randn((4, 4), dtype=torch.float16, device=DEVICE)

# Chuẩn bị một tensor đầu ra
y = torch.empty_like(x)

# Chạy kernel GELU nhanh
activation_kernels.gelu_fast(y, x)

# Lấy đầu ra dự kiến bằng cách sử dụng GELU tích hợp của PyTorch
expected = F.gelu(x)

# So sánh đầu ra kernel với kết quả của PyTorch
torch.testing.assert_close(y, expected, rtol=1e-2, atol=1e-2)

print("✅ Đầu ra kernel khớp với PyTorch GELU!")

# Tùy chọn: in cả hai tensor để kiểm tra
print("\nTensor đầu vào:")
print(x)
print("\nĐầu ra kernel GELU nhanh:")
print(y)
print("\nĐầu ra PyTorch GELU:")
print(expected)

# Liệt kê các hàm khả dụng trong mô-đun kernel đã tải
print("\nCác hàm khả dụng trong 'kernels-community/activation':")
print(dir(activation_kernels))

# /// script
# dependencies = [
#  "numpy",
#  "torch",
#  "kernels",
# ]
# ///
import torch
import torch.nn as nn

DEVICE = "cuda"

DTYPE = torch.float16  # Sử dụng float16 để có tiềm năng hiệu suất kernel tốt hơn


# Implementation PyTorch đơn giản của RMSNorm để so sánh cơ bản
class RMSNorm(nn.Module):
    def __init__(self, hidden_size, variance_epsilon=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = variance_epsilon
        self.hidden_size = hidden_size

    def forward(self, x):
        # Giả sử x là (batch_size, ..., hidden_size)
        input_dtype = x.dtype
        # Tính phương sai trong float32 để ổn định
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)

        # Áp dụng trọng số và chuyển đổi trở lại dtype ban đầu
        return (self.weight * x).to(input_dtype)


class BaselineModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, eps=1e-5):
        super().__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.norm = RMSNorm(hidden_size, variance_epsilon=eps)
        self.activation = nn.GELU()
        self.linear2 = nn.Linear(hidden_size, output_size)

        # đảm bảo tất cả các trọng số của các lớp tuyến tính là 1 để kiểm tra
        with torch.no_grad():
            self.linear1.weight.fill_(1)
            self.linear1.bias.fill_(0)
            self.linear2.weight.fill_(1)
            self.linear2.bias.fill_(0)
            self.norm.weight.fill_(1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.norm(x)  # Áp dụng RMSNorm
        x = self.activation(x)
        x = self.linear2(x)
        return x


# Ví dụ sử dụng
input_size = 128
hidden_size = 256
output_size = 10
eps_val = 1e-5

baseline_model = (
    BaselineModel(input_size, hidden_size, output_size, eps=eps_val)
    .to(DEVICE)
    .to(DTYPE)
)
dummy_input = torch.randn(32, input_size, device=DEVICE, dtype=DTYPE)  # Batch gồm 32
output = baseline_model(dummy_input)
print("Hình dạng đầu ra của mô hình Baseline RMSNorm:", output.shape)

# /// script
# dependencies = [
#  "numpy",
#  "torch",
#  "kernels",
# ]
# ///
import torch
import torch.nn as nn
from kernels import get_kernel, use_kernel_forward_from_hub

# tái sử dụng mô hình từ đoạn mã trước hoặc sao chép lớp
# định nghĩa ở đây để chạy script này một cách độc lập
from rmsnorm_baseline import BaselineModel

DEVICE = "cuda"
DTYPE = torch.float16  # Sử dụng float16 để có tiềm năng hiệu suất kernel tốt hơn


layer_norm_kernel_module = get_kernel("kernels-community/triton-layer-norm")

# Đơn giản chỉ cần thêm decorator vào lớp LlamaRMSNorm để tự động thay thế hàm forward
# với phiên bản kernel được tối ưu hóa
#
# Lưu ý: lưu ý tất cả các kernel đều được ship với các lớp đã được ánh xạ và sẽ yêu cầu gọi hàm trực tiếp
# Tuy nhiên, trong trường hợp này, lớp LlamaRMSNorm đã được ánh xạ tới hàm kernel. Nếu không, chúng ta cần
# gọi hàm trực tiếp như sau:
# ```python
# layer_norm_kernel_module.rms_norm_fn(
#     hidden_states,
#     self.weight,
#     bias=None,
#     residual=None,
#     eps=self.variance_epsilon,
#     dropout_p=0.0,
#     prenorm=False,
#     residual_in_fp32=False,
# )
# ```
@use_kernel_forward_from_hub("LlamaRMSNorm")
class OriginalRMSNorm(nn.Module):
    def __init__(self, hidden_size, variance_epsilon=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = variance_epsilon
        self.hidden_size = hidden_size

    def forward(self, x):
        # Giả sử x là (batch_size, ..., hidden_size)
        input_dtype = x.dtype
        # Tính phương sai trong float32 để ổn định
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)

        # Áp dụng trọng số và chuyển đổi trở lại dtype ban đầu
        return (self.weight * x).to(input_dtype)


class KernelModel(nn.Module):
    def __init__(
        self,
        input_size,
        hidden_size,
        output_size,
        device="cuda",
        dtype=torch.float16,
        eps=1e-5,
    ):
        super().__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        # OriginalRMSNorm sẽ được thay thế bằng lớp kernel được tối ưu hóa
        # khi mô hình được tải
        self.norm = OriginalRMSNorm(hidden_size, variance_epsilon=eps)
        self.activation = nn.GELU()
        self.linear2 = nn.Linear(hidden_size, output_size)

        # đảm bảo tất cả các trọng số của các lớp tuyến tính là 1 để kiểm tra
        with torch.no_grad():
            self.linear1.weight.fill_(1)
            self.linear1.bias.fill_(0)
            self.linear2.weight.fill_(1)
            self.linear2.bias.fill_(0)
            self.norm.weight.fill_(1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.norm(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x


# Ví dụ sử dụng
input_size = 128
hidden_size = 256
output_size = 10
eps_val = 1e-5

kernel_model = (
    KernelModel(
        input_size, hidden_size, output_size, device=DEVICE, dtype=DTYPE, eps=eps_val
    )
    .to(DEVICE)
    .to(DTYPE)
)

baseline_model = (
    BaselineModel(input_size, hidden_size, output_size, eps=eps_val)
    .to(DEVICE)
    .to(DTYPE)
)

dummy_input = torch.randn(32, input_size, device=DEVICE, dtype=DTYPE)  # Batch gồm 32

output = baseline_model(dummy_input)
output_kernel = kernel_model(dummy_input)
print("Hình dạng đầu ra của mô hình Kernel RMSNorm:", output_kernel.shape)

# Xác minh các đầu ra gần nhau (các implementations RMSNorm phải gần nhau về mặt số)
try:
    torch.testing.assert_close(output, output_kernel, rtol=1e-2, atol=1e-2)
    print("\nCác đầu ra của mô hình Baseline và Kernel RMSNorm khớp nhau!")
except AssertionError as e:
    print("\nCác đầu ra của mô hình Baseline và Kernel RMSNorm hơi khác nhau:")
    print(e)
except NameError:
    print("\nBỏ qua so sánh đầu ra vì đầu ra của mô hình kernel không được tạo.")

# /// script
# dependencies = [
#  "numpy",
#  "torch",
#  "kernels",
# ]
# ///
import torch

# tái sử dụng các mô hình từ các đoạn mã trước hoặc sao chép lớp
# các định nghĩa ở đây để chạy script này một cách độc lập
from rmsnorm_baseline import BaselineModel
from rmsnorm_kernel import KernelModel

DEVICE = "cuda"
DTYPE = torch.float16  # Sử dụng float16 để có tiềm năng hiệu suất kernel tốt hơn


# Sử dụng torch.cuda.Event để có thời gian GPU chính xác (đảm bảo hàm được xác định)
def benchmark_model(model, input_tensor, num_runs=100, warmup_runs=10):
    model.eval()  # Đặt mô hình ở chế độ đánh giá
    dtype = input_tensor.dtype
    model = model.to(input_tensor.device).to(dtype)

    # Các lần chạy khởi động
    for _ in range(warmup_runs):
        _ = model(input_tensor)
    torch.cuda.synchronize()

    # Các lần chạy được tính thời gian
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    for _ in range(num_runs):
        _ = model(input_tensor)
    end_event.record()
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    avg_time_ms = elapsed_time_ms / num_runs
    return avg_time_ms


input_size_bench = 4096
hidden_size_bench = 4096  # Hiệu suất RMSNorm nhạy cảm với chiều này
output_size_bench = 10
eps_val_bench = 1e-5

# Tạo các mô hình lớn hơn và đầu vào để benchmark
# Đảm bảo cả hai mô hình đều được chuyển đổi hoàn toàn sang DEVICE và DTYPE mục tiêu
baseline_model_bench = (
    BaselineModel(
        input_size_bench, hidden_size_bench, output_size_bench, eps=eps_val_bench
    )
    .to(DEVICE)
    .to(DTYPE)
)
kernel_model_bench = (
    KernelModel(
        input_size_bench,
        hidden_size_bench,
        output_size_bench,
        device=DEVICE,
        dtype=DTYPE,
        eps=eps_val_bench,
    )
    .to(DEVICE)
    .to(DTYPE)
)

# gọi cả hai với kích thước batch lớn hơn để khởi động GPU
# và đảm bảo các mô hình được tải
warmup_input = torch.randn(4096, input_size_bench, device=DEVICE, dtype=DTYPE)
_ = kernel_model_bench(warmup_input)
_ = baseline_model_bench(warmup_input)

batch_sizes = [
    256,
    512,
    1024,
    2048,
    4096,
    8192,
    16384,
    32768,
]

print(
    f"{'Batch Size':<12} | {'Baseline Time (ms)':<18} | {'Kernel Time (ms)':<18} | {'Speedup'}"
)
print("-" * 74)

for batch_size in batch_sizes:
    # Gọi cuda synchronize để đảm bảo tất cả các hoạt động GPU trước đó đã hoàn thành
    torch.cuda.synchronize()

    # Tạo tensor đầu vào ngẫu nhiên
    # Đảm bảo tensor đầu vào nằm trên thiết bị và dtype chính xác
    bench_input = torch.randn(batch_size, input_size_bench, device=DEVICE, dtype=DTYPE)

    # Chạy benchmark chỉ khi kernel được tải thành công
    baseline_time = benchmark_model(baseline_model_bench, bench_input)

    kernel_time = -1  # Giá trị sentinel

    kernel_time = benchmark_model(kernel_model_bench, bench_input)

    baseline_time = round(baseline_time, 4)
    kernel_time = round(kernel_time, 4)
    speedup = round(baseline_time / kernel_time, 2) if kernel_time > 0 else "N/A"
    if kernel_time < baseline_time:
        speedup = f"{speedup:.2f}x"
    elif kernel_time == baseline_time:
        speedup = "1.00x (giống hệt nhau)"
    else:
        speedup = f"{kernel_time / baseline_time:.2f}x chậm hơn"
    print(f"{batch_size:<12} | {baseline_time:<18} | {kernel_time:<18} | {speedup}")

pip install kernels torch numpy