[init] baseline7 from project_6

ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu

@@ -0,0 +1,168 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <torch/cuda.h>
+
+#include <cmath>
+
+#include "kernels/cuda/utils.h"
+#include "xattention_ops_api.h"
+
+namespace {
+
+// Fused log-sum-exp combine kernel.
+//
+// Layout and strategy (aligned with the TileLang version):
+//   - Each block is responsible for one (batch_idx, head_idx) pair, i.e. one
+//     row in the flattened [B * H, D] layout.
+//   - Threads within a block parallelize along the head_dim (D) dimension to
+//     ensure coalesced global memory access.
+//
+// Tensors:
+//   shared_o    : [B, H, D] - shared attention output
+//   shared_lse  : [B, H, 1] - shared log-sum-exp (FP32)
+//   unshared_o  : [B, H, D] - unshared attention output
+//   unshared_lse: [B, H, 1] - unshared log-sum-exp (FP32)
+//   output      : [B, H, D] - combined output
+template <typename scalar_t, typename out_scalar_t>
+__global__ void lse_combine_kernel(
+    out_scalar_t* __restrict__ output,        // [B, H, D]
+    const scalar_t* __restrict__ shared_o,    // [B, H, D]
+    const float* __restrict__ shared_lse,     // [B, H, 1], always FP32
+    const scalar_t* __restrict__ unshared_o,  // [B, H, D]
+    const float* __restrict__ unshared_lse,   // [B, H, 1], always FP32
+    const int64_t B,                          // batch_size * beam_size
+    const int64_t H,                          // num_heads
+    const int64_t D) {                        // head_dim
+  const int64_t total_elements = B * H;
+  const int64_t idx = static_cast<int64_t>(blockIdx.y);
+
+  if (idx >= total_elements) {
+    return;
+  }
+
+  // Load LSE scalars for this (batch, head) pair.
+  const float shared_lse_val = shared_lse[idx];
+  const float unshared_lse_val = unshared_lse[idx];
+
+  // 1. Compute element-wise max LSE.
+  const float lse_max = fmaxf(shared_lse_val, unshared_lse_val);
+
+  // 2. Compute base-2 exponentials relative to max.
+  const float exp_shared = exp2f(shared_lse_val - lse_max);
+  const float exp_unshared = exp2f(unshared_lse_val - lse_max);
+
+  // 3. Compute merged LSE.
+  const float lse_new = lse_max + log2f(exp_shared + exp_unshared);
+
+  // 4. Compute normalized weights.
+  const float w_shared = exp2f(shared_lse_val - lse_new);
+  const float w_unshared = exp2f(unshared_lse_val - lse_new);
+
+  // 5. Weighted combine along the head_dim.
+  const int64_t base_idx = idx * D;
+  // Threads in the block parallelize along D with stride blockDim.x for
+  // coalesced global memory access.
+  for (int64_t d = threadIdx.x; d < D; d += blockDim.x) {
+    const float shared_val = static_cast<float>(shared_o[base_idx + d]);
+    const float unshared_val = static_cast<float>(unshared_o[base_idx + d]);
+    const float combined = w_shared * shared_val + w_unshared * unshared_val;
+    output[base_idx + d] = static_cast<out_scalar_t>(combined);
+  }
+}
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+// Host wrapper for the fused LSE combine kernel.
+//
+// All inputs are expected to be on the same CUDA device:
+//   shared_o    : [B, H, D], floating type (including Half/BFloat16)
+//   shared_lse  : [B, H, 1], float32
+//   unshared_o  : [B, H, D], same type/shape as shared_o
+//   unshared_lse: [B, H, 1], float32
+//   output      : [B, H, D], will be resized/allocated as needed.
+void lse_combine(torch::Tensor output,
+                 torch::Tensor shared_o,
+                 torch::Tensor shared_lse,
+                 torch::Tensor unshared_o,
+                 torch::Tensor unshared_lse) {
+  CHECK_EQ(shared_o.dim(), 3) << "shared_o must be 3D [B, H, D]";
+  CHECK_EQ(unshared_o.dim(), 3) << "unshared_o must be 3D [B, H, D]";
+  CHECK_EQ(shared_lse.dim(), 3) << "shared_lse must be 3D [B, H, 1]";
+  CHECK_EQ(unshared_lse.dim(), 3) << "unshared_lse must be 3D [B, H, 1]";
+
+  const int64_t B = shared_o.size(0);
+  const int64_t H = shared_o.size(1);
+  const int64_t D = shared_o.size(2);
+
+  CHECK_EQ(shared_o.sizes(), unshared_o.sizes())
+      << "shared_o and unshared_o must have same shape";
+  CHECK_EQ(shared_lse.scalar_type(), torch::kFloat32)
+      << "shared_lse must be float32";
+  CHECK_EQ(unshared_lse.scalar_type(), torch::kFloat32)
+      << "unshared_lse must be float32";
+  CHECK_EQ(shared_lse.size(0), B)
+      << "shared_lse shape mismatch, expected [B, H, 1]";
+  CHECK_EQ(shared_lse.size(1), H)
+      << "shared_lse shape mismatch, expected [B, H, 1]";
+  CHECK_EQ(shared_lse.size(2), 1)
+      << "shared_lse shape mismatch, expected [B, H, 1]";
+  CHECK_EQ(unshared_lse.size(0), B)
+      << "unshared_lse shape mismatch, expected [B, H, 1]";
+  CHECK_EQ(unshared_lse.size(1), H)
+      << "unshared_lse shape mismatch, expected [B, H, 1]";
+  CHECK_EQ(unshared_lse.size(2), 1)
+      << "unshared_lse shape mismatch, expected [B, H, 1]";
+
+  // Ensure output has the correct shape and dtype.
+  if (!output.defined() || output.sizes() != shared_o.sizes()) {
+    output = torch::empty_like(shared_o);
+  }
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(shared_o));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  // Launch kernel: one block per (batch, head) pair, threads along D.
+  const int64_t total_elements = B * H;
+  const int threads_per_block = 128;
+  dim3 block_dim(threads_per_block, 1, 1);
+  dim3 grid_dim(1, static_cast<unsigned int>(total_elements), 1);
+
+  DISPATCH_FLOATING_TYPES(
+      shared_o.scalar_type(), "lse_combine_kernel_input", [&] {
+        using in_t = scalar_t;
+        DISPATCH_FLOATING_TYPES(
+            output.scalar_type(), "lse_combine_kernel_output", [&] {
+              using out_t = scalar_t;
+              lse_combine_kernel<in_t, out_t>
+                  <<<grid_dim, block_dim, 0, stream>>>(
+                      output.data_ptr<out_t>(),
+                      shared_o.data_ptr<in_t>(),
+                      shared_lse.data_ptr<float>(),
+                      unshared_o.data_ptr<in_t>(),
+                      unshared_lse.data_ptr<float>(),
+                      B,
+                      H,
+                      D);
+            });
+      });
+
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+}  // namespace xllm::kernel::cuda