Feat:add kl_softmax #83

S1/wut0n_#83/kl_softmax_cudacode.py

@@ -0,0 +1,153 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+kl_softmax_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+// --- 修复后的融合内核：严格模拟PyTorch的计算路径 ---
+__global__ void kl_softmax_kernel_stable(
+    const float* __restrict__ logits_p,
+    const float* __restrict__ logits_q,
+    float* __restrict__ total_loss,
+    int batch_size,
+    int num_classes,
+    float epsilon
+) {
+    int sample_idx = blockIdx.x;
+    if (sample_idx >= batch_size) return;
+
+    int tid = threadIdx.x;
+    
+    // --- 步骤1: 计算softmax(Q)的log值，即log_softmax(Q) ---
+    // 这直接对应PyTorch的 torch.log(q.clamp(min=1e-8))
+    float max_logit_q = -FLT_MAX;
+    for (int i = tid; i < num_classes; i += blockDim.x) {
+        max_logit_q = fmaxf(max_logit_q, logits_q[sample_idx * num_classes + i]);
+    }
+    
+    extern __shared__ float shared_mem[];
+    shared_mem[tid] = max_logit_q;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (tid < stride) shared_mem[tid] = fmaxf(shared_mem[tid], shared_mem[tid + stride]);
+    }
+    __syncthreads();
+    max_logit_q = shared_mem[0];
+
+    float sum_exp_q = 0.0f;
+    for (int i = tid; i < num_classes; i += blockDim.x) {
+        sum_exp_q += expf(logits_q[sample_idx * num_classes + i] - max_logit_q);
+    }
+    shared_mem[tid] = sum_exp_q;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (tid < stride) shared_mem[tid] += shared_mem[tid + stride];
+    }
+    __syncthreads();
+    sum_exp_q = shared_mem[0];
+    float log_sum_exp_q = logf(sum_exp_q);
+
+    // --- 步骤2: 计算softmax(P) ---
+    float max_logit_p = -FLT_MAX;
+    for (int i = tid; i < num_classes; i += blockDim.x) {
+        max_logit_p = fmaxf(max_logit_p, logits_p[sample_idx * num_classes + i]);
+    }
+    shared_mem[tid] = max_logit_p;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (tid < stride) shared_mem[tid] = fmaxf(shared_mem[tid], shared_mem[tid + stride]);
+    }
+    __syncthreads();
+    max_logit_p = shared_mem[0];
+    
+    float sum_exp_p = 0.0f;
+    for (int i = tid; i < num_classes; i += blockDim.x) {
+        sum_exp_p += expf(logits_p[sample_idx * num_classes + i] - max_logit_p);
+    }
+    shared_mem[tid] = sum_exp_p;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (tid < stride) shared_mem[tid] += shared_mem[tid + stride];
+    }
+    __syncthreads();
+    sum_exp_p = shared_mem[0];
+
+    // --- 步骤3: 按照kl_div(input, target)的公式计算 ---
+    // kl_div(input, target) = sum(target * (log(target) - input))
+    // input = log_softmax(Q) = logits_q - max_q - log_sum_exp_q
+    // target = softmax(P) = exp(logits_p - max_p) / sum_exp_p
+    float sample_kl = 0.0f;
+    for (int i = tid; i < num_classes; i += blockDim.x) {
+        int idx = sample_idx * num_classes + i;
+        
+        float pi = expf(logits_p[idx] - max_logit_p) / sum_exp_p;
+        
+        if (pi > epsilon) {
+            float log_pi = (logits_p[idx] - max_logit_p) - logf(sum_exp_p);
+            float log_qi = (logits_q[idx] - max_logit_q) - log_sum_exp_q;
+            sample_kl += pi * (log_pi - log_qi);
+        }
+    }
+
+    // --- 步骤4: 块内归约求和 ---
+    shared_mem[tid] = sample_kl;
+    __syncthreads();
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+        if (tid < stride) shared_mem[tid] += shared_mem[tid + stride];
+    }
+    __syncthreads();
+
+    // --- 步骤5: 原子操作累加到全局损失 ---
+    if (tid == 0) {
+        atomicAdd(total_loss, shared_mem[0]);
+    }
+}
+
+torch::Tensor kl_softmax_cuda(torch::Tensor logits_p, torch::Tensor logits_q) {
+    TORCH_CHECK(logits_p.scalar_type() == torch::kFloat32, "Logits must be float32");
+    TORCH_CHECK(logits_p.sizes() == logits_q.sizes(), "Logits must have same shape");
+    
+    int batch_size = logits_p.size(0);
+    int num_classes = logits_p.size(1);
+
+    auto total_loss_tensor = torch::zeros(1, logits_p.options());
+
+    const int block_size = 256;
+    int num_blocks = batch_size;
+    
+    size_t shared_mem = block_size * sizeof(float);
+    
+    kl_softmax_kernel_stable<<<num_blocks, block_size, shared_mem>>>(
+        logits_p.data_ptr<float>(),
+        logits_q.data_ptr<float>(),
+        total_loss_tensor.data_ptr<float>(),
+        batch_size,
+        num_classes,
+        1e-8f
+    );
+    
+    return total_loss_tensor;
+}
+"""
+
+kl_softmax_cpp_source = """
+torch::Tensor kl_softmax_cuda(torch::Tensor logits_p, torch::Tensor logits_q);
+"""
+
+kl_softmax = load_inline(
+    name="kl_softmax",
+    cpp_sources=kl_softmax_cpp_source,
+    cuda_sources=kl_softmax_source,
+    functions=["kl_softmax_cuda"],
+    extra_cuda_cflags=["-O3"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self):
+        super(ModelNew, self).__init__()
+        self.kl_softmax = kl_softmax
+
+    def forward(self, logits_p, logits_q):
+        return self.kl_softmax.kl_softmax_cuda(logits_p, logits_q)

S1/wut0n_#83/kl_softmax_torchcode.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    基准实现：先计算Softmax，再计算KL散度。
+    """
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, logits_p, logits_q):
+        """
+        计算两个logits张量经过softmax后的KL散度。
+
+        Args:
+            logits_p (torch.Tensor): 真实分布的logits [batch_size, num_classes]
+            logits_q (torch.Tensor): 预测分布的logits [batch_size, num_classes]
+
+        Returns:
+            torch.Tensor: KL散度标量值
+        """
+        # 步骤1: 计算softmax概率分布
+        p = torch.softmax(logits_p, dim=1)
+        q = torch.softmax(logits_q, dim=1)
+        
+        # 步骤2: 计算KL散度
+        kl_loss = torch.nn.functional.kl_div(
+            torch.log(q.clamp(min=1e-8)),
+            p.clamp(min=1e-8),
+            reduction='sum'
+        )
+        
+        return kl_loss
+
+batch_size = 256
+num_classes = 1000
+
+def get_inputs():
+    """
+    生成合理的测试数据
+    """
+    logits_p = torch.randn(batch_size, num_classes)
+    logits_q = torch.randn(batch_size, num_classes)
+    return [logits_p, logits_q]
+
+def get_init_inputs():
+    return []  # 没有特殊的初始化输入需求

S1/wut0n_#83/prompt.txt

@@ -0,0 +1,141 @@
+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, which are derived from raw logits. This baseline implementation uses PyTorch's highly optimized built-in functions for correctness and performance.
+
+python
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+“”"
+基准实现：先计算Softmax，再计算KL散度。
+“”"
+def init(self):
+super(Model, self).init()
+
+def forward(self, logits_p, logits_q):
+    """
+    计算两个logits张量经过softmax后的KL散度。
+
+    Args:
+        logits_p (torch.Tensor): 真实分布的logits [batch_size, num_classes]
+        logits_q (torch.Tensor): 预测分布的logits [batch_size, num_classes]
+
+    Returns:
+        torch.Tensor: KL散度标量值
+    """
+    # 步骤1: 计算softmax概率分布
+    p = torch.softmax(logits_p, dim=1)
+    q = torch.softmax(logits_q, dim=1)
+    
+    # 步骤2: 计算KL散度
+    kl_loss = torch.nn.functional.kl_div(
+        torch.log(q.clamp(min=1e-8)),
+        p.clamp(min=1e-8),
+        reduction='sum'
+    )
+    
+    return kl_loss
+batch_size = 256
+num_classes = 1000
+
+def get_inputs():
+“”"
+生成合理的测试数据
+“”"
+logits_p = torch.randn(batch_size, num_classes)
+logits_q = torch.randn(batch_size, num_classes)
+return [logits_p, logits_q]
+
+def get_init_inputs():
+return [] # 没有特殊的初始化输入需求
+
+Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that computes the KL Divergence from logits in a fused manner. The implementation must be highly optimized and numerically stable.
+
+**CRITICAL REQUIREMENTS:**
+
+1.  **Operator Fusion:** The core task is to fuse the computation of `softmax` and `kl_div` into a **single CUDA kernel launch**. This means the intermediate probability tensors `p` and `q` should never be materialized in global memory.
+2.  **Algorithmic Optimization:** The kernel should compute `KL(softmax(logits_p) || softmax(logits_q))` directly from the input logits. To maximize numerical stability, use the standard `log-softmax` trick: `log_softmax(x)_i = x_i - max(x) - log(sum(exp(x_j - max(x))))`.
+3.  **Numerical Stability:** The KL part of the computation must be numerically stable. The final formula to be implemented for each element is `P_i * (log(P_i) - log(Q_i))`. If `P_i` is less than a small epsilon (e.g., `1e-8`), the term should be treated as 0.
+4.  **Performance Optimization:**
+    *   **Shared Memory Reduction:** Each thread block must compute its partial sum in shared memory before performing a single atomic add to the global output tensor.
+    *   **Per-Sample Parallelism:** A good strategy is to have each thread block compute the KL divergence for one sample in the batch.
+5.  **Kernel Logic:**
+    *   The kernel should launch a 1D grid of 1D blocks.
+    *   The final output is a single scalar tensor representing the total KL divergence sum.
+6.  **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.
+7.  **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_#83/run_code.py

@@ -0,0 +1,74 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from kl_softmax_torchcode import Model, get_inputs, get_init_inputs
+from kl_softmax_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_softmax 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA kl_softmax 平均执行时间: {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()