feat:add kl #80

S1/wut0n_#80/kl_cudacode.py

@@ -0,0 +1,186 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+kl_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+// 修复版本：使用数值稳定的公式 P * log(P/Q)，并优化性能
+__global__ void kl_kernel_stable(
+    const float* __restrict__ p,
+    const float* __restrict__ q,
+    float* __restrict__ kl_loss,
+    int size,
+    float epsilon
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    
+    float local_kl = 0.0f;
+    
+    int stride = blockDim.x * gridDim.x;
+    for (int i = idx; i < size; i += stride) {
+        float pi = p[i];
+        float qi = q[i];
+        
+        // --- 核心修复：数值稳定的KL散度计算 ---
+        // 当 pi 接近 0 时，pi * log(pi/qi) 的极限是 0，我们可以直接跳过
+        if (pi > epsilon) {
+            // 使用公式 P * log(P/Q) = P * (log(P) - log(Q))
+            // 优化为 P * log(P/Q) 只需一次log调用
+            // clamp qi 以避免除以零
+            qi = max(qi, epsilon);
+            local_kl += pi * logf(pi / qi);
+        }
+    }
+    
+    // 使用共享内存进行块内归约
+    extern __shared__ float shared_mem[];
+    shared_mem[threadIdx.x] = local_kl;
+    
+    __syncthreads();
+    
+    // 块内归约
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (threadIdx.x < stride) {
+            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
+        }
+        __syncthreads();
+    }
+    
+    // 只有第一个线程做原子操作
+    if (threadIdx.x == 0) {
+        atomicAdd(kl_loss, shared_mem[0]);
+    }
+}
+
+// 向量化版本：8元素并行处理 (同样应用了稳定性修复)
+__global__ void kl_kernel_vectorized_stable(
+    const float* __restrict__ p,
+    const float* __restrict__ q,
+    float* __restrict__ kl_loss,
+    int size,
+    float epsilon
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int vec_idx = idx * 8;
+    
+    float local_kl = 0.0f;
+    
+    if (vec_idx + 7 < size) {
+        // 加载8个元素（2个float4）
+        float4 p_vec1 = *reinterpret_cast<const float4*>(&p[vec_idx]);
+        float4 p_vec2 = *reinterpret_cast<const float4*>(&p[vec_idx + 4]);
+        float4 q_vec1 = *reinterpret_cast<const float4*>(&q[vec_idx]);
+        float4 q_vec2 = *reinterpret_cast<const float4*>(&q[vec_idx + 4]);
+        
+        // --- 应用稳定性修复的向量化计算 ---
+        float pi_val, qi_val;
+        #define CALC_KL(pi, qi) \
+            pi_val = pi; \
+            if (pi_val > epsilon) { \
+                qi_val = max(qi, epsilon); \
+                local_kl += pi_val * logf(pi_val / qi_val); \
+            }
+
+        CALC_KL(p_vec1.x, q_vec1.x);
+        CALC_KL(p_vec1.y, q_vec1.y);
+        CALC_KL(p_vec1.z, q_vec1.z);
+        CALC_KL(p_vec1.w, q_vec1.w);
+        CALC_KL(p_vec2.x, q_vec2.x);
+        CALC_KL(p_vec2.y, q_vec2.y);
+        CALC_KL(p_vec2.z, q_vec2.z);
+        CALC_KL(p_vec2.w, q_vec2.w);
+        
+        #undef CALC_KL
+
+    } else {
+        // 边界处理
+        for (int i = vec_idx; i < size && i < vec_idx + 8; i++) {
+            float pi = p[i];
+            if (pi > epsilon) {
+                float qi = max(q[i], epsilon);
+                local_kl += pi * logf(pi / qi);
+            }
+        }
+    }
+    
+    // 使用共享内存进行块内归约
+    extern __shared__ float shared_mem[];
+    shared_mem[threadIdx.x] = local_kl;
+    
+    __syncthreads();
+    
+    // 块内归约
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (threadIdx.x < stride) {
+            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
+        }
+        __syncthreads();
+    }
+    
+    // 只有第一个线程做原子操作
+    if (threadIdx.x == 0) {
+        atomicAdd(kl_loss, shared_mem[0]);
+    }
+}
+
+torch::Tensor kl_cuda(torch::Tensor p, torch::Tensor q) {
+    auto size = p.numel();
+    
+    // 创建结果张量
+    auto kl_loss = torch::zeros(1, p.options());
+    
+    // 根据数据规模选择策略
+    if (size >= 50000) {
+        // 大数据使用向量化版本
+        const int block_size = 256;
+        int num_blocks = (size + block_size * 8 - 1) / (block_size * 8);
+        size_t shared_mem = block_size * sizeof(float);
+        kl_kernel_vectorized_stable<<<num_blocks, block_size, shared_mem>>>(
+            p.data_ptr<float>(), 
+            q.data_ptr<float>(),
+            kl_loss.data_ptr<float>(),
+            size,
+            1e-8f
+        );
+    } else {
+        // 小数据使用优化版本
+        const int block_size = 256;
+        int temp = (size + block_size - 1) / block_size;
+        int num_blocks = (temp < 32) ? temp : 32;
+        size_t shared_mem = block_size * sizeof(float);
+        kl_kernel_stable<<<num_blocks, block_size, shared_mem>>>(
+            p.data_ptr<float>(), 
+            q.data_ptr<float>(),
+            kl_loss.data_ptr<float>(),
+            size,
+            1e-8f
+        );
+    }
+    
+    return kl_loss;
+}
+"""
+
+kl_cpp_source = """
+torch::Tensor kl_cuda(torch::Tensor p, torch::Tensor q);
+"""
+
+# --- 修复编译选项 ---
+# 移除硬编码的gencode和--use_fast_math，提升可移植性和精度
+kl = load_inline(
+    name="kl",
+    cpp_sources=kl_cpp_source,
+    cuda_sources=kl_source,
+    functions=["kl_cuda"],
+    extra_cuda_cflags=["-O3"], # 保留-O3优化，但移除有风险的标志
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self):
+        super(ModelNew, self).__init__()
+        self.kl = kl
+
+    def forward(self, p, q):
+        return self.kl.kl_cuda(p, q)

S1/wut0n_#80/kl_torchcode.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    合理优化的PyTorch KL Divergence实现
+    使用PyTorch内置优化函数，避免不必要的中间张量创建
+    """
+    def __init__(self):
+        super(Model, self).__init__()
+    
+    def forward(self, p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+        """
+        使用PyTorch内置的kl_div函数，这是最优化和标准的实现
+
+        Args:
+            p (torch.Tensor): 真实概率分布 [batch_size, num_classes]
+            q (torch.Tensor): 预测概率分布 [batch_size, num_classes]
+
+        Returns:
+            torch.Tensor: KL散度标量值
+        """
+        # PyTorch的kl_div函数要求输入是log概率
+        # KL(P||Q) = sum(P * (log(P) - log(Q)))
+        # kl_div(input, target) 计算的是 sum(target * (log(target) - input))
+        # 所以我们需要传入 log(Q) 作为 input，P 作为 target
+        return torch.nn.functional.kl_div(
+            torch.log(q.clamp(min=1e-8)),  # input: log(Q)
+            p.clamp(min=1e-8),             # target: P
+            reduction='sum'
+        )
+
+batch_size = 512
+num_classes = 1000
+
+def get_inputs():
+    """
+    生成合理的测试数据
+    使用softmax确保是有效的概率分布
+    """
+    p = torch.softmax(torch.randn(batch_size, num_classes), dim=1)
+    q = torch.softmax(torch.randn(batch_size, num_classes), dim=1)
+    return [p, q]
+
+def get_init_inputs():
+    return []  # 没有特殊的初始化输入需求

S1/wut0n_#80/prompt.txt

@@ -0,0 +1,143 @@
+You are an expert CUDA programmer tasked with accelerating a PyTorch model by replacing its operators with a highly optimized, custom CUDA kernel. You should consider operator fusion and algorithmic optimizations to achieve maximum speedup.
+
+Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple addition:
+
+python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+def init(self) -> None:
+super().init()
+
+def forward(self, a, b):
+    return a + b
+def get_inputs():
+a = torch.randn(1, 128).cuda()
+b = torch.randn(1, 128).cuda()
+return [a, b]
+
+def get_init_inputs():
+return []
+
+
+
+The example new architecture with a custom CUDA kernel looks like this:
+
+python
+import torch
+from torch.utils.cpp_extension import load_inline
+
+add_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+global void add_kernel(const float* a, const float* b, float* out, int size) {
+int idx = blockIdx.x * blockDim.x + threadIdx.x;
+if (idx < size) {
+out[idx] = a[idx] + b[idx];
+}
+}
+
+torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
+auto out = torch::empty_like(a);
+int size = a.numel();
+const int block_size = 256;
+int num_blocks = (size + block_size - 1) / block_size;
+add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
+return out;
+}
+"""
+
+add_cpp_source = """
+torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
+"""
+
+Compile the inline CUDA code
+add = load_inline(
+name=“add”,
+cpp_sources=add_cpp_source,
+cuda_sources=add_source,
+functions=[“add_cuda”],
+verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+def init(self):
+super(ModelNew, self).init()
+self.add = add
+
+def forward(self, a, b):
+    return self.add.add_cuda(a, b)
+
+
+---
+
+Now, you are given the following PyTorch architecture to accelerate. The model computes the Kullback-Leibler (KL) Divergence between two probability distributions, P and Q. This baseline implementation uses PyTorch's highly optimized built-in function for correctness and performance.
+
+python
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+“”"
+合理优化的PyTorch KL Divergence实现
+使用PyTorch内置优化函数，避免不必要的中间张量创建
+“”"
+def init(self):
+super(Model, self).init()
+
+def forward(self, p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+    """
+    使用PyTorch内置的kl_div函数，这是最优化和标准的实现
+
+    Args:
+        p (torch.Tensor): 真实概率分布 [batch_size, num_classes]
+        q (torch.Tensor): 预测概率分布 [batch_size, num_classes]
+
+    Returns:
+        torch.Tensor: KL散度标量值
+    """
+    # PyTorch的kl_div函数要求输入是log概率
+    # KL(P||Q) = sum(P * (log(P) - log(Q)))
+    # kl_div(input, target) 计算的是 sum(target * (log(target) - input))
+    # 所以我们需要传入 log(Q) 作为 input，P 作为 target
+    return torch.nn.functional.kl_div(
+        torch.log(q.clamp(min=1e-8)),  # input: log(Q)
+        p.clamp(min=1e-8),             # target: P
+        reduction='sum'
+    )
+batch_size = 512
+num_classes = 1000
+
+def get_inputs():
+“”"
+生成合理的测试数据
+使用softmax确保是有效的概率分布
+“”"
+p = torch.softmax(torch.randn(batch_size, num_classes), dim=1)
+q = torch.softmax(torch.randn(batch_size, num_classes), dim=1)
+return [p, q]
+
+def get_init_inputs():
+return [] # 没有特殊的初始化输入需求
+
+
+
+Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that computes the KL Divergence. The implementation must be highly optimized and numerically stable.
+
+**CRITICAL REQUIREMENTS:**
+
+1.  **Numerical Stability:** The implementation must be numerically stable. The direct formula `P * (log(P) - log(Q))` can result in NaNs when `P` is very close to zero. You must use the numerically stable equivalent: `P * log(P/Q)`. If `P` is less than a small epsilon (e.g., `1e-8`), the term should be treated as 0 to avoid `0 * log(0)` which is undefined.
+2.  **Performance Optimization:**
+    *   **Vectorization:** You must implement a vectorized kernel using `float4` loads to maximize memory bandwidth utilization. This kernel should be used for large tensors.
+    *   **Shared Memory Reduction:** To avoid the bottleneck of atomic operations, each thread block should compute a partial sum in shared memory and then have a single thread perform one atomic add to the global output tensor.
+    *   **Adaptive Strategy:** The host-side function should choose between a standard kernel and a vectorized kernel based on the size of the input tensors.
+3.  **Kernel Logic:**
+    *   The kernel should launch a 1D grid of 1D blocks.
+    *   Each thread should process multiple elements from the input tensors.
+    *   The final output is a single scalar tensor representing the total KL divergence sum.
+4.  **Edge Case Handling:** The `q` tensor must be clamped to a minimum `epsilon` before division to prevent division by zero.
+5.  **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return an empty list `[]` to match the baseline.
+6.  **Compilation Flags:** Use `-O3` for optimization but **do not** use `--use_fast_math` to ensure numerical accuracy with the PyTorch baseline. Avoid hardcoding compute capabilities to ensure portability.

S1/wut0n_#80/run_code.py

@@ -0,0 +1,74 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from kl_torchcode import Model, get_inputs, get_init_inputs
+from kl_cudacode import ModelNew
+
+def run_benchmark():
+    # 检查 CUDA 是否可用
+    if not torch.cuda.is_available():
+        print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
+        return
+    else:
+        device = torch.device("cuda")
+
+    # 初始化模型
+    init_inputs = get_init_inputs()
+    init_inputs = [
+        x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
+    ]
+    inputs = get_inputs()
+    inputs = [
+        x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
+    ]
+
+    torch_model = Model(*init_inputs).cuda()
+    cuda_model = ModelNew(*init_inputs).cuda()
+
+    torch_model.eval()
+    cuda_model.eval()
+
+    print("-------------------- 精度对齐验证 --------------------")
+    with torch.no_grad():
+        output_torch = torch_model( *inputs)
+        output_cuda = cuda_model(*inputs)
+
+    precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
+    if precision_flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        
+    print("\n-------------------- 性能加速比测试 --------------------")
+    num_iterations = 100
+    
+    # PyTorch 模型计时
+    torch.cuda.synchronize()
+    start_time = time.time()
+    for _ in range(num_iterations):
+        _ = torch_model(*inputs)
+    torch.cuda.synchronize()
+    torch_time = (time.time() - start_time) / num_iterations
+    
+    # 自定义 CUDA 内核计时
+    torch.cuda.synchronize()
+    start_time = time.time()
+    for _ in range(num_iterations):
+        _ = cuda_model(*inputs)
+    torch.cuda.synchronize()
+    cuda_time = (time.time() - start_time) / num_iterations
+    
+    print(f"PyTorch kl散度 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA kl散度 平均执行时间: {cuda_time:.6f} 秒")
+    speedup = 0
+    if cuda_time > 0:
+        speedup = torch_time / cuda_time
+        print(f"加速比 (Speedup): {speedup:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return precision_flag,speedup
+if __name__ == "__main__":
+    precision_flag,speedup = run_benchmark()