Automated PR - 2026-07-07

2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
@@ -0,0 +1,111 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <torch/python.h>
+
+#include <vector>
+
+// Forward declaration of CUDA kernel template
+template<typename out_t>
+void rms_norm_split_rope_cuda(
+    void* x,
+    void* sin_freqs,
+    void* cos_freqs,
+    void* weights,
+    int b,
+    int s,
+    int n,
+    int h,
+    long cos_sb, long cos_sn, long cos_ss,
+    long sin_sb, long sin_sn, long sin_ss,
+    void* out,
+    cudaStream_t stream
+);
+
+at::Tensor rms_norm_split_rope(
+    at::Tensor &x,
+    at::Tensor &sin_freqs,
+    at::Tensor &cos_freqs,
+    at::Tensor &weights,
+    bool out_fp8
+) {
+    TORCH_CHECK(x.scalar_type() == at::ScalarType::BFloat16, "Input must be BFloat16");
+    TORCH_CHECK(sin_freqs.scalar_type() == at::ScalarType::BFloat16, "sin_freqs must be BFloat16");
+    TORCH_CHECK(cos_freqs.scalar_type() == at::ScalarType::BFloat16, "cos_freqs must be BFloat16");
+    TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
+    TORCH_CHECK(sin_freqs.is_cuda(), "sin_freqs must be on CUDA");
+    TORCH_CHECK(cos_freqs.is_cuda(), "cos_freqs must be on CUDA");
+
+    // Get dimensions
+    // x: [b, s, h]
+    // cos, sin: [b, n, s, d] where n*d = h/2
+    int b = x.size(0);
+    int s = x.size(1);
+    int h = x.size(2);
+
+    TORCH_CHECK(cos_freqs.dim() == 4, "cos_freqs must be 4D");
+    TORCH_CHECK(sin_freqs.dim() == 4, "sin_freqs must be 4D");
+
+    int n = cos_freqs.size(1);
+    int d = h / n;
+
+    
+    // Require a contiguous innermost (d/2) dim for the vectorized int4 freq load,
+    // but keep the outer (b, n, s) strides: apply_split_rotary_emb hands us a
+    // swapaxes view (logical [b, n, s, d/2], physical [b, s, n, d/2]) whose inner
+    // stride is already 1, so this never copies it. The strides are forwarded to
+    // the kernel so the read is correct regardless of the physical layout.
+    if (x.stride(-1) != 1) { x = x.contiguous(); }
+    if (cos_freqs.stride(-1) != 1) { cos_freqs = cos_freqs.contiguous(); }
+    if (sin_freqs.stride(-1) != 1) { sin_freqs = sin_freqs.contiguous(); }
+
+    long cos_sb = cos_freqs.stride(0), cos_sn = cos_freqs.stride(1), cos_ss = cos_freqs.stride(2);
+    long sin_sb = sin_freqs.stride(0), sin_sn = sin_freqs.stride(1), sin_ss = sin_freqs.stride(2);
+
+    // Create output tensor
+    at::Tensor out;
+    if (out_fp8) {
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
+    } else {
+        out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
+    }
+    
+    // Setup CUDA
+    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // Launch kernel
+    if (out_fp8) {
+        rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
+            x.data_ptr(),
+            sin_freqs.data_ptr(),
+            cos_freqs.data_ptr(),
+            weights.data_ptr(),  // weights (optional, not used yet)
+            b,
+            s,
+            n,
+            h,
+            cos_sb, cos_sn, cos_ss,
+            sin_sb, sin_sn, sin_ss,
+            (void*)out.data_ptr(),
+            stream
+        );
+    } else {
+        rms_norm_split_rope_cuda<at::BFloat16>(
+            x.data_ptr(),
+            sin_freqs.data_ptr(),
+            cos_freqs.data_ptr(),
+            weights.data_ptr(),  // weights (optional, not used yet)
+            b,
+            s,
+            n,
+            h,
+            cos_sb, cos_sn, cos_ss,
+            sin_sb, sin_sn, sin_ss,
+            (void*)out.data_ptr(),
+            stream
+        );
+    }
+
+    return out;
+}