data: complete SGEMM upstream from 3 repos (siboehm+wangzyon+edtallison) + xllm fused_qknorm_rope + xattention kernels
SGEMM repos (upstream_ref/sgemm_cuda/, 41 files): siboehm/SGEMM_CUDA: kernel 1-12, runner, CMake, cuBLAS benchmark wangzyon/NVIDIA_SGEMM_PRACTICE: kernel 1-7 (Chinese comments), utils edtallison/sgemm-cuda: kernel 01-09 (learning notes), Makefile xllm kernels (ex_engine/xllm_kernels/cuda/): fused_qknorm_rope.cu + bind — saves 128 kernel launches/fwd xattention/ — 6 files from upstream xllm headers: corex_compat_utils.h, topk_last_dim.cuh ilu/CMakeLists.txt SO_BUILD_MANIFEST.md — complete .so inventory and call chain analysis
This commit is contained in:
Claude
107
ex_engine/build_xllm_kernels.sh
Executable file
107
ex_engine/build_xllm_kernels.sh
Executable file
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env bash
|
||||
# build_xllm_kernels.sh — Compile xllm CUDA kernels into .so for BI-V100
|
||||
#
|
||||
# Architecture (CCCL compile pattern):
|
||||
# CCCL: CMakePresets.json → cmake --preset cub-cpp20 → ninja → .so
|
||||
# EX: torch.utils.cpp_extension → clang --cuda-gpu-arch=ivcore10 → .so
|
||||
#
|
||||
# Usage:
|
||||
# bash ex_engine/build_xllm_kernels.sh [--output-dir /path/to/output]
|
||||
#
|
||||
# Prerequisites:
|
||||
# - BI-V100 machine with corex SDK
|
||||
# - PyTorch with CUDA support
|
||||
# - corex clang/16 compiler
|
||||
#
|
||||
# Outputs:
|
||||
# xllm_fused_qknorm_rope.so — Fused QK-Norm + RoPE (saves 128 kernel launches/fwd)
|
||||
|
||||
set -eo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
KERNELS_DIR="${SCRIPT_DIR}/xllm_kernels/cuda"
|
||||
HEADERS_DIR="${KERNELS_DIR}/headers"
|
||||
BINDINGS_DIR="${KERNELS_DIR}/bindings"
|
||||
OUTPUT_DIR="${1:-${SCRIPT_DIR}/../qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10}"
|
||||
|
||||
mkdir -p "${OUTPUT_DIR}"
|
||||
|
||||
echo "[build] KERNELS_DIR=${KERNELS_DIR}"
|
||||
echo "[build] HEADERS_DIR=${HEADERS_DIR}"
|
||||
echo "[build] OUTPUT_DIR=${OUTPUT_DIR}"
|
||||
|
||||
# Common compile flags for BI-V100 (ivcore10 = SM70-class)
|
||||
CUDA_FLAGS="-O2 --cuda-gpu-arch=ivcore10"
|
||||
CXX_FLAGS="-O2 -std=c++17"
|
||||
INCLUDE_FLAGS="-I${HEADERS_DIR}"
|
||||
|
||||
# Use torch's cpp_extension for JIT compile
|
||||
build_so() {
|
||||
local name=$1
|
||||
local sources=$2
|
||||
local extra_flags="${3:-}"
|
||||
|
||||
echo "[build] Building ${name}.so from: ${sources}"
|
||||
|
||||
python3 -c "
|
||||
import os, sys
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
sources = '${sources}'.split()
|
||||
abs_sources = [os.path.join('${SCRIPT_DIR}', '..', s) if not os.path.isabs(s) else s for s in sources]
|
||||
abs_sources = [os.path.abspath(s) for s in abs_sources]
|
||||
|
||||
for s in abs_sources:
|
||||
if not os.path.exists(s):
|
||||
print(f'ERROR: source not found: {s}', file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
mod = load(
|
||||
name='${name}',
|
||||
sources=abs_sources,
|
||||
extra_cuda_cflags=['-O2'],
|
||||
extra_cflags=['-O2', '-std=c++17'],
|
||||
extra_include_paths=['${HEADERS_DIR}'],
|
||||
build_directory='/tmp/build_${name}',
|
||||
verbose=True,
|
||||
)
|
||||
# Find the compiled .so
|
||||
import glob
|
||||
sos = glob.glob('/tmp/build_${name}/${name}*.so')
|
||||
if sos:
|
||||
import shutil
|
||||
dst = os.path.join('${OUTPUT_DIR}', '${name}.so')
|
||||
shutil.copy2(sos[0], dst)
|
||||
print(f'[build] SUCCESS: {dst}')
|
||||
else:
|
||||
print('[build] WARN: .so not found after build', file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f'[build] FAIL ${name}: {e}', file=sys.stderr)
|
||||
sys.exit(1)
|
||||
" || echo "[build] FAILED: ${name}"
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Build targets
|
||||
# ============================================================================
|
||||
|
||||
# 1. xllm_fused_qknorm_rope — Fused QK-Norm + RoPE
|
||||
# Source: upstream xllm fused_qknorm_rope.cu
|
||||
# Note: Requires corex_compat_utils.h instead of glog-dependent utils.h
|
||||
# The .cu includes "cuda_ops_api.h" and "utils.h" — we need to make sure
|
||||
# the include path resolves to our corex-compat headers first.
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " 1. xllm_fused_qknorm_rope.so"
|
||||
echo "============================================================"
|
||||
build_so "xllm_fused_qknorm_rope" \
|
||||
"ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Build complete. Output:"
|
||||
echo "============================================================"
|
||||
ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30
|
||||
echo ""
|
||||
echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"
|
||||
@@ -0,0 +1,38 @@
|
||||
// xllm_fused_qknorm_rope_bind.cpp — pybind11 for fused QK-Norm + RoPE kernel
|
||||
// Source: upstream_ref/xllm/xllm/core/kernels/cuda/fused_qknorm_rope.cu
|
||||
// Saves 4 kernel launches per layer (separate q_norm, k_norm, q_rope, k_rope)
|
||||
// Qwen3.5 has 32 full-attention layers → saves 128 kernel launches per forward
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
void fused_qk_norm_rope(
|
||||
torch::Tensor& qkv,
|
||||
int64_t num_heads_q,
|
||||
int64_t num_heads_k,
|
||||
int64_t num_heads_v,
|
||||
int64_t head_dim,
|
||||
double eps,
|
||||
const torch::Tensor& q_weight,
|
||||
const torch::Tensor& k_weight,
|
||||
const torch::Tensor& cos_sin_cache,
|
||||
bool interleaved,
|
||||
const torch::Tensor& position_ids);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("fused_qk_norm_rope",
|
||||
&xllm::kernel::cuda::fused_qk_norm_rope,
|
||||
"Fused QK-Norm + RoPE (xllm CUDA kernel)",
|
||||
py::arg("qkv"),
|
||||
py::arg("num_heads_q"),
|
||||
py::arg("num_heads_k"),
|
||||
py::arg("num_heads_v"),
|
||||
py::arg("head_dim"),
|
||||
py::arg("eps") = 1e-6,
|
||||
py::arg("q_weight"),
|
||||
py::arg("k_weight"),
|
||||
py::arg("cos_sin_cache"),
|
||||
py::arg("interleaved") = false,
|
||||
py::arg("position_ids"));
|
||||
}
|
||||
463
ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu
Normal file
463
ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu
Normal file
@@ -0,0 +1,463 @@
|
||||
/* 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/cuda.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <type_traits>
|
||||
|
||||
#include "cuda_ops_api.h"
|
||||
#include "type_convert.cuh"
|
||||
#include "utils.h"
|
||||
|
||||
using at::device_of;
|
||||
|
||||
// Borrowed from:
|
||||
// https://github.com/vllm-project/vllm/blob/022f3cea5327cc720a325c50931e1edcfdf2d32b/csrc/fused_qknorm_rope_kernel.cu
|
||||
|
||||
constexpr uint32_t kFinalMask = 0xffffffffu;
|
||||
|
||||
namespace {
|
||||
|
||||
using namespace xllm::kernel::cuda;
|
||||
|
||||
template <typename T, int num>
|
||||
struct packed_as;
|
||||
// Specialization for packed_as used in this kernel.
|
||||
template <>
|
||||
struct packed_as<uint, 1> {
|
||||
using type = uint;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct packed_as<uint, 2> {
|
||||
using type = uint2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct packed_as<uint, 4> {
|
||||
using type = uint4;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__inline__ __device__ T warp_reduce_sum(T val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val += __shfl_xor_sync(kFinalMask, val, mask, 32);
|
||||
return val;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline __device__ __host__ T div_up(T m, T n) {
|
||||
return (m + n - 1) / n;
|
||||
}
|
||||
|
||||
// Perform per-head QK Norm and RoPE in a single kernel.
|
||||
// scalar_t_in: data type of QKV and RMSNorm weights
|
||||
// scalar_t_cache: data type of cos/sin cache
|
||||
// head_dim: the dimension of each head
|
||||
// interleave: interleave=!is_neox.
|
||||
template <typename scalar_t_in,
|
||||
typename scalar_t_cache,
|
||||
int head_dim,
|
||||
bool interleave>
|
||||
__global__ void fused_qknorm_rope_kernel(
|
||||
void* qkv_void, // Combined QKV tensor
|
||||
int const num_heads_q, // Number of query heads
|
||||
int const num_heads_k, // Number of key heads
|
||||
int const num_heads_v, // Number of value heads
|
||||
float const eps, // Epsilon for RMS normalization
|
||||
void const* q_weight_void, // RMSNorm weights for query
|
||||
void const* k_weight_void, // RMSNorm weights for key
|
||||
void const* cos_sin_cache_void, // Pre-computed cos/sin cache
|
||||
int64_t const* position_ids, // Position IDs for RoPE
|
||||
int const num_tokens, // Number of tokens
|
||||
int const rotary_dim // Dimension for RoPE
|
||||
) {
|
||||
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
|
||||
if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
|
||||
std::is_same_v<scalar_t_cache, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
|
||||
using Converter = _typeConvert<scalar_t_in>;
|
||||
static_assert(Converter::exists,
|
||||
"Input QKV data type is not supported for this CUDA "
|
||||
"architecture or toolkit version.");
|
||||
using T_in = typename Converter::hip_type;
|
||||
using T2_in = typename Converter::packed_hip_type;
|
||||
|
||||
using CacheConverter = _typeConvert<scalar_t_cache>;
|
||||
static_assert(CacheConverter::exists,
|
||||
"Cache data type is not supported for this CUDA architecture "
|
||||
"or toolkit version.");
|
||||
using T_cache = typename CacheConverter::hip_type;
|
||||
|
||||
T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
|
||||
T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
|
||||
T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
|
||||
T_cache const* cos_sin_cache =
|
||||
reinterpret_cast<T_cache const*>(cos_sin_cache_void);
|
||||
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const warpId = threadIdx.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
|
||||
// Calculate global warp index to determine which head/token this warp
|
||||
// processes
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
|
||||
|
||||
// Total number of attention heads (Q and K)
|
||||
int const total_qk_heads = num_heads_q + num_heads_k;
|
||||
|
||||
// Determine which token and head type (Q or K) this warp processes
|
||||
int const tokenIdx = globalWarpIdx / total_qk_heads;
|
||||
int const localHeadIdx = globalWarpIdx % total_qk_heads;
|
||||
|
||||
// Skip if this warp is assigned beyond the number of tokens
|
||||
if (tokenIdx >= num_tokens) return;
|
||||
|
||||
bool const isQ = localHeadIdx < num_heads_q;
|
||||
int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
|
||||
|
||||
int const num_heads = num_heads_q + num_heads_k + num_heads_v;
|
||||
|
||||
static_assert(head_dim % (32 * 2) == 0,
|
||||
"head_dim must be divisible by 64 (each warp processes one "
|
||||
"head, and each thread gets even number of "
|
||||
"elements)");
|
||||
constexpr int numElemsPerThread = head_dim / 32;
|
||||
float elements[numElemsPerThread];
|
||||
constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
|
||||
static_assert(elemSizeBytes % 4 == 0,
|
||||
"numSizeBytes must be a multiple of 4");
|
||||
constexpr int vecSize =
|
||||
elemSizeBytes /
|
||||
4; // Use packed_as<uint, vecSize> to perform loading/saving.
|
||||
using vec_T = typename packed_as<uint, vecSize>::type;
|
||||
|
||||
int offsetWarp; // Offset for the warp
|
||||
if (isQ) {
|
||||
// Q segment: token offset + head offset within Q segment
|
||||
offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
|
||||
} else {
|
||||
// K segment: token offset + entire Q segment + head offset within K
|
||||
// segment
|
||||
offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
|
||||
headIdx * head_dim;
|
||||
}
|
||||
int offsetThread = offsetWarp + laneId * numElemsPerThread;
|
||||
|
||||
// Sum of squares for RMSNorm
|
||||
float sumOfSquares = 0.0f;
|
||||
|
||||
// Load.
|
||||
{
|
||||
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
|
||||
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < num_packed_elems; i++) {
|
||||
// Interpret the generic vector chunk as the specific packed type
|
||||
T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
|
||||
// Convert to float2 for computation
|
||||
float2 vals = Converter::convert(packed_val);
|
||||
sumOfSquares += vals.x * vals.x;
|
||||
sumOfSquares += vals.y * vals.y;
|
||||
|
||||
elements[2 * i] = vals.x;
|
||||
elements[2 * i + 1] = vals.y;
|
||||
}
|
||||
}
|
||||
|
||||
// Reduce sum across warp using the utility function
|
||||
sumOfSquares = warp_reduce_sum(sumOfSquares);
|
||||
|
||||
// Compute RMS normalization factor
|
||||
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
|
||||
|
||||
// Normalize elements
|
||||
#pragma unroll
|
||||
for (int i = 0; i < numElemsPerThread; i++) {
|
||||
int dim = laneId * numElemsPerThread + i;
|
||||
float weight = isQ ? Converter::convert(q_weight[dim])
|
||||
: Converter::convert(k_weight[dim]);
|
||||
elements[i] *= rms_rcp * weight;
|
||||
}
|
||||
|
||||
// Apply RoPE to normalized elements
|
||||
float elements2[numElemsPerThread]; // Additional buffer required for RoPE.
|
||||
|
||||
int64_t pos_id = position_ids[tokenIdx];
|
||||
|
||||
// Calculate cache pointer for this position - similar to
|
||||
// pos_encoding_kernels.cu
|
||||
T_cache const* cache_ptr = cos_sin_cache + pos_id * rotary_dim;
|
||||
int const embed_dim = rotary_dim / 2;
|
||||
T_cache const* cos_ptr = cache_ptr;
|
||||
T_cache const* sin_ptr = cache_ptr + embed_dim;
|
||||
int const rotary_lanes = rotary_dim / numElemsPerThread; // rotary range
|
||||
if (laneId < rotary_lanes) {
|
||||
if constexpr (interleave) {
|
||||
// Perform interleaving. Use pre-computed cos/sin values.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < numElemsPerThread / 2; ++i) {
|
||||
int const idx0 = 2 * i;
|
||||
int const idx1 = 2 * i + 1;
|
||||
// Global dimension index in the head
|
||||
int const dim_idx = laneId * numElemsPerThread + idx0;
|
||||
|
||||
float const val0 = elements[idx0];
|
||||
float const val1 = elements[idx1];
|
||||
|
||||
int const half_dim = dim_idx / 2;
|
||||
float const cos_val =
|
||||
CacheConverter::convert(__ldg(cos_ptr + half_dim));
|
||||
float const sin_val =
|
||||
CacheConverter::convert(__ldg(sin_ptr + half_dim));
|
||||
|
||||
elements[idx0] = val0 * cos_val - val1 * sin_val;
|
||||
elements[idx1] = val0 * sin_val + val1 * cos_val;
|
||||
}
|
||||
} else {
|
||||
// Before data exchange with in warp, we need to sync.
|
||||
__syncwarp();
|
||||
int pairOffset = (rotary_dim / 2) / numElemsPerThread;
|
||||
// Get the data from the other half of the warp. Use pre-computed
|
||||
// cos/sin values.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < numElemsPerThread; i++) {
|
||||
elements2[i] = __shfl_xor_sync(kFinalMask, elements[i], pairOffset);
|
||||
|
||||
if (laneId < pairOffset) {
|
||||
elements2[i] = -elements2[i];
|
||||
}
|
||||
int dim_idx = laneId * numElemsPerThread + i;
|
||||
|
||||
dim_idx = (dim_idx * 2) % rotary_dim;
|
||||
int half_dim = dim_idx / 2;
|
||||
float cos_val = CacheConverter::convert(__ldg(cos_ptr + half_dim));
|
||||
float sin_val = CacheConverter::convert(__ldg(sin_ptr + half_dim));
|
||||
|
||||
elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
|
||||
}
|
||||
// __shfl_xor_sync does not provide memfence. Need to sync again.
|
||||
__syncwarp();
|
||||
}
|
||||
}
|
||||
// Store.
|
||||
{
|
||||
vec_T vec;
|
||||
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < num_packed_elems; i++) {
|
||||
// Convert from float2 back to the specific packed type
|
||||
float2 vals = {elements[2 * i], elements[2 * i + 1]};
|
||||
T2_in packed_val = Converter::convert(vals);
|
||||
// Place it into the generic vector
|
||||
*(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
|
||||
}
|
||||
*reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
|
||||
}
|
||||
|
||||
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Borrowed from
|
||||
// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
|
||||
#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
|
||||
if (interleave) { \
|
||||
const bool INTERLEAVE = true; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
const bool INTERLEAVE = false; \
|
||||
__VA_ARGS__ \
|
||||
}
|
||||
|
||||
template <typename scalar_t_in, typename scalar_t_cache>
|
||||
void launch_fused_qknorm_rope(void* qkv,
|
||||
int const num_tokens,
|
||||
int const num_heads_q,
|
||||
int const num_heads_k,
|
||||
int const num_heads_v,
|
||||
int const head_dim,
|
||||
int const rotary_dim,
|
||||
float const eps,
|
||||
void const* q_weight,
|
||||
void const* k_weight,
|
||||
void const* cos_sin_cache,
|
||||
bool const interleave,
|
||||
int64_t const* position_ids,
|
||||
cudaStream_t stream) {
|
||||
constexpr int blockSize = 256;
|
||||
|
||||
int const warpsPerBlock = blockSize / 32;
|
||||
int const totalQKHeads = num_heads_q + num_heads_k;
|
||||
int const totalWarps = num_tokens * totalQKHeads;
|
||||
|
||||
int const gridSize = div_up(totalWarps, warpsPerBlock);
|
||||
dim3 gridDim(gridSize);
|
||||
dim3 blockDim(blockSize);
|
||||
|
||||
switch (head_dim) {
|
||||
case 64:
|
||||
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
|
||||
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
|
||||
<<<gridDim, blockDim, 0, stream>>>(qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
cos_sin_cache,
|
||||
position_ids,
|
||||
num_tokens,
|
||||
rotary_dim);
|
||||
});
|
||||
break;
|
||||
case 128:
|
||||
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
|
||||
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
|
||||
<<<gridDim, blockDim, 0, stream>>>(qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
cos_sin_cache,
|
||||
position_ids,
|
||||
num_tokens,
|
||||
rotary_dim);
|
||||
});
|
||||
break;
|
||||
case 256:
|
||||
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
|
||||
fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
|
||||
<<<gridDim, blockDim, 0, stream>>>(qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
cos_sin_cache,
|
||||
position_ids,
|
||||
num_tokens,
|
||||
rotary_dim);
|
||||
});
|
||||
break;
|
||||
default:
|
||||
CHECK(false) << "Unsupported head dimension for fusedQKNormRope: "
|
||||
<< head_dim;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
void fused_qk_norm_rope(
|
||||
torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
|
||||
// (num_heads_q+num_heads_k+num_heads_v)*head_dim]
|
||||
int64_t num_heads_q, // Number of query heads
|
||||
int64_t num_heads_k, // Number of key heads
|
||||
int64_t num_heads_v, // Number of value heads
|
||||
int64_t head_dim, // Dimension per head
|
||||
double eps, // Epsilon for RMS normalization
|
||||
const torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
|
||||
const torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
|
||||
const torch::Tensor&
|
||||
cos_sin_cache, // Cos/sin cache [max_position, rotary_dim]
|
||||
bool interleaved, // Whether RoPE is applied in interleaved style
|
||||
const torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
|
||||
) {
|
||||
// Input validation
|
||||
CHECK(qkv.is_cuda()) << "qkv must be a CUDA tensor";
|
||||
CHECK(qkv.is_contiguous()) << "qkv must be contiguous";
|
||||
CHECK(position_ids.is_cuda()) << "position_ids must be a CUDA tensor";
|
||||
CHECK(position_ids.is_contiguous()) << "position_ids must be contiguous";
|
||||
CHECK(q_weight.is_cuda()) << "q_weight must be a CUDA tensor";
|
||||
CHECK(q_weight.is_contiguous()) << "q_weight must be contiguous";
|
||||
CHECK(k_weight.is_cuda()) << "k_weight must be a CUDA tensor";
|
||||
CHECK(k_weight.is_contiguous()) << "k_weight must be contiguous";
|
||||
CHECK(cos_sin_cache.is_cuda()) << "cos_sin_cache must be a CUDA tensor";
|
||||
CHECK(cos_sin_cache.is_contiguous()) << "cos_sin_cache must be contiguous";
|
||||
CHECK(position_ids.scalar_type() == torch::kInt64)
|
||||
<< "position_ids dtype is " << position_ids.scalar_type()
|
||||
<< ", while Int64 is expected";
|
||||
|
||||
CHECK(qkv.dim() == 2) << "QKV tensor must be 2D: [num_tokens, "
|
||||
<< "(num_heads_q+num_heads_k+num_heads_v)*head_dim]";
|
||||
CHECK(position_ids.dim() == 1) << "Position IDs must be 1D: [num_tokens]";
|
||||
CHECK(q_weight.dim() == 1) << "Query weights must be 1D: [head_dim]";
|
||||
CHECK(k_weight.dim() == 1) << "Key weights must be 1D: [head_dim]";
|
||||
CHECK(cos_sin_cache.dim() == 2)
|
||||
<< "Cos/sin cache must be 2D: [max_position, rotary_dim]";
|
||||
CHECK(q_weight.size(0) == head_dim)
|
||||
<< "Query weights size must match head dimension";
|
||||
CHECK(k_weight.size(0) == head_dim)
|
||||
<< "Key weights size must match head dimension";
|
||||
|
||||
CHECK(cos_sin_cache.size(1) % 2 == 0) << "rotary_dim must be even";
|
||||
CHECK(cos_sin_cache.size(1) <= head_dim)
|
||||
<< "rotary_dim must be less than or equal to head_dim";
|
||||
|
||||
CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
|
||||
qkv.scalar_type() == k_weight.scalar_type())
|
||||
<< "qkv, q_weight and k_weight must have the same dtype";
|
||||
|
||||
int64_t num_tokens = qkv.size(0);
|
||||
CHECK(position_ids.size(0) == num_tokens)
|
||||
<< "Number of tokens in position_ids must match QKV";
|
||||
|
||||
int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
|
||||
CHECK(qkv.size(1) == total_heads * head_dim)
|
||||
<< "QKV tensor size must match total number of heads and head dimension";
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(qkv));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
|
||||
using qkv_scalar_t = scalar_t;
|
||||
DISPATCH_FLOATING_TYPES(
|
||||
cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
|
||||
using cache_scalar_t = scalar_t;
|
||||
launch_fused_qknorm_rope<qkv_scalar_t, cache_scalar_t>(
|
||||
qkv.data_ptr(),
|
||||
static_cast<int>(num_tokens),
|
||||
static_cast<int>(num_heads_q),
|
||||
static_cast<int>(num_heads_k),
|
||||
static_cast<int>(num_heads_v),
|
||||
static_cast<int>(head_dim),
|
||||
static_cast<int>(cos_sin_cache.size(1)),
|
||||
static_cast<float>(eps),
|
||||
q_weight.data_ptr(),
|
||||
k_weight.data_ptr(),
|
||||
cos_sin_cache.data_ptr(),
|
||||
interleaved,
|
||||
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
|
||||
stream);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
37
ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h
Normal file
37
ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h
Normal file
@@ -0,0 +1,37 @@
|
||||
// corex_compat_utils.h — Lightweight replacement for xllm's utils.h
|
||||
// Removes glog/tvm dependencies for BI-V100 corex compilation
|
||||
// Provides CHECK macro via TORCH_CHECK and DISPATCH macros from device_utils.cuh
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
// Replace glog CHECK with TORCH_CHECK
|
||||
#ifndef CHECK
|
||||
#define CHECK(cond) TORCH_CHECK(cond)
|
||||
#endif
|
||||
|
||||
#ifndef CHECK_EQ
|
||||
#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b))
|
||||
#endif
|
||||
|
||||
#ifndef CHECK_GE
|
||||
#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b))
|
||||
#endif
|
||||
|
||||
// Include device_utils for DISPATCH_HALF_TYPES etc
|
||||
#include "device_utils.cuh"
|
||||
|
||||
// ffi namespace stub (some headers reference it)
|
||||
namespace ffi {
|
||||
template <typename T>
|
||||
using Array = std::vector<T>;
|
||||
}
|
||||
|
||||
// HOST_DEVICE_INLINE
|
||||
#if defined(__CUDACC__) || defined(_NVHPC_CUDA)
|
||||
#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
|
||||
#else
|
||||
#define HOST_DEVICE_INLINE inline
|
||||
#endif
|
||||
2114
ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh
Normal file
2114
ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh
Normal file
File diff suppressed because it is too large
Load Diff
129
ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp
Normal file
129
ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp
Normal file
@@ -0,0 +1,129 @@
|
||||
/* 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/script.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "cuda.h"
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
void beam_search(torch::Tensor acc_logprob,
|
||||
torch::Tensor in_sequence_group,
|
||||
torch::Tensor top_tokens,
|
||||
torch::Tensor top_logprobs,
|
||||
torch::Tensor out_acc_logprob,
|
||||
torch::Tensor out_token_ids,
|
||||
torch::Tensor out_token_index,
|
||||
torch::Tensor out_beam_count_prefix_sums,
|
||||
torch::Tensor out_sequence_group,
|
||||
uint32_t batch_size,
|
||||
uint32_t current_step) {
|
||||
torch::Device device = acc_logprob.device();
|
||||
|
||||
uint32_t beam_size = in_sequence_group.size(1);
|
||||
|
||||
uint32_t top_k = top_tokens.size(1);
|
||||
uint32_t total_rounds = in_sequence_group.size(2);
|
||||
|
||||
CHECK_EQ(beam_size, top_k) << "beam_size must be equal with top_k.";
|
||||
|
||||
if (current_step == 0) {
|
||||
auto tokens_view =
|
||||
top_tokens.view({batch_size, top_k}).slice(1, 0, beam_size);
|
||||
auto init_probs_view =
|
||||
top_logprobs.view({batch_size, top_k}).slice(1, 0, beam_size);
|
||||
|
||||
out_token_ids.view({batch_size, beam_size}).copy_(tokens_view);
|
||||
out_acc_logprob.view({batch_size, beam_size}).copy_(init_probs_view);
|
||||
|
||||
auto indices =
|
||||
torch::arange(
|
||||
beam_size,
|
||||
torch::TensorOptions().dtype(torch::kInt32).device(device))
|
||||
.unsqueeze(0)
|
||||
.expand({batch_size, -1})
|
||||
.reshape({-1, 1});
|
||||
out_token_index.copy_(indices);
|
||||
|
||||
auto sequence_view =
|
||||
out_sequence_group.view({batch_size, beam_size, total_rounds});
|
||||
sequence_view.slice(2, 0, 1).squeeze(2).copy_(tokens_view);
|
||||
|
||||
} else {
|
||||
auto combined_probs =
|
||||
(acc_logprob + top_logprobs).view({batch_size, beam_size * top_k});
|
||||
|
||||
auto topk_result = torch::topk(combined_probs, beam_size, -1);
|
||||
auto new_probs = std::get<0>(topk_result); // [batch_size, beam_size]
|
||||
auto new_indices = std::get<1>(topk_result); // [batch_size, beam_size]
|
||||
|
||||
auto ordered_indices = new_indices.argsort(static_cast<int64_t>(1), false);
|
||||
// Reorder new_probs (and corresponding new_indices) by ordered_indices to
|
||||
// keep alignment.
|
||||
if (current_step < total_rounds - 1) {
|
||||
new_probs = new_probs.gather(1, ordered_indices);
|
||||
new_indices = new_indices.gather(1, ordered_indices);
|
||||
}
|
||||
|
||||
auto parent_beam = (new_indices / top_k).to(torch::kLong);
|
||||
auto token_in_beam = (new_indices % top_k).to(torch::kLong);
|
||||
|
||||
auto top_tokens_reshaped = top_tokens.view({batch_size, beam_size, top_k});
|
||||
|
||||
auto batch_idx =
|
||||
torch::arange(batch_size,
|
||||
torch::TensorOptions().dtype(torch::kLong).device(device))
|
||||
.unsqueeze(1)
|
||||
.expand_as(parent_beam);
|
||||
|
||||
using torch::indexing::TensorIndex;
|
||||
auto new_tokens = top_tokens_reshaped.index({TensorIndex(batch_idx),
|
||||
TensorIndex(parent_beam),
|
||||
TensorIndex(token_in_beam)});
|
||||
|
||||
out_acc_logprob.view({batch_size, beam_size}).copy_(new_probs);
|
||||
out_token_index.view({batch_size, beam_size})
|
||||
.copy_(new_indices.to(torch::kInt32));
|
||||
out_token_ids.view({batch_size, beam_size}).copy_(new_tokens);
|
||||
|
||||
auto batch_range =
|
||||
torch::arange(
|
||||
batch_size,
|
||||
torch::TensorOptions().dtype(torch::kInt32).device(device))
|
||||
.unsqueeze(1)
|
||||
.expand({-1, beam_size});
|
||||
auto beam_range =
|
||||
torch::arange(
|
||||
beam_size,
|
||||
torch::TensorOptions().dtype(torch::kInt32).device(device))
|
||||
.unsqueeze(0)
|
||||
.expand({batch_size, -1});
|
||||
|
||||
using torch::indexing::Slice;
|
||||
using torch::indexing::TensorIndex;
|
||||
out_sequence_group.slice(2, 0, current_step) =
|
||||
in_sequence_group.index({TensorIndex(batch_range),
|
||||
TensorIndex(parent_beam.to(torch::kInt32)),
|
||||
Slice(0, current_step)});
|
||||
|
||||
out_sequence_group.slice(2, current_step, current_step + 1) =
|
||||
new_tokens.unsqueeze(2);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
312
ex_engine/xllm_kernels/cuda/xattention/cache_select.cu
Normal file
312
ex_engine/xllm_kernels/cuda/xattention/cache_select.cu
Normal file
@@ -0,0 +1,312 @@
|
||||
/* 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <glog/logging.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
#include "xattention_ops_api.h"
|
||||
|
||||
namespace {
|
||||
|
||||
// In-place cache selection kernel for Xattention.
|
||||
// Reorders KV cache entries based on beam search results. After beam search,
|
||||
// the beam indices may have changed, and this kernel copies KV cache data from
|
||||
// old beam positions to new beam positions to maintain consistency.
|
||||
// Inputs:
|
||||
// k_ptrs_i64 : [Layer] - pointers to K cache tensors for each layer
|
||||
// v_ptrs_i64 : [Layer] - pointers to V cache tensors for each layer
|
||||
// beam_index : [B*Beam] - mapping from new beam index to old beam index
|
||||
// block_table : [B] - request ID per batch item (extracted from [B*Beam,
|
||||
// 1]) B : batch size (actual batch size, not batch_size *
|
||||
// beam_size) Beam : beam width Kv : number of KV
|
||||
// heads MaxStep : maximum decode steps D : head
|
||||
// dimension MaxReq : maximum number of requests Layer :
|
||||
// number of transformer layers decode_step : current decode step
|
||||
// (0-indexed)
|
||||
// Cache layout: [MaxReq, Beam, MaxStep, Kv, D]
|
||||
// The kernel performs two passes to avoid overwriting data:
|
||||
// pass-1: copy from old_beam > new_beam (increasing new_beam)
|
||||
// pass-2: copy from old_beam < new_beam (decreasing new_beam)
|
||||
template <typename scalar_t>
|
||||
__global__ void cache_select_inplace_ptrs_kernel(
|
||||
const int64_t* __restrict__ k_ptrs_i64, // [Layer]
|
||||
const int64_t* __restrict__ v_ptrs_i64, // [Layer]
|
||||
const int32_t* __restrict__ beam_index, // [B*Beam]
|
||||
const int32_t* __restrict__ block_table, // [B]
|
||||
int32_t B,
|
||||
int32_t Beam,
|
||||
int32_t Kv,
|
||||
int32_t MaxStep,
|
||||
int32_t D,
|
||||
int32_t MaxReq,
|
||||
int32_t Layer,
|
||||
int32_t decode_step) {
|
||||
const int32_t b = static_cast<int32_t>(blockIdx.x);
|
||||
const int32_t kv = static_cast<int32_t>(blockIdx.y);
|
||||
const int32_t layer = static_cast<int32_t>(blockIdx.z);
|
||||
|
||||
if (b >= B || kv >= Kv || layer >= Layer) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t step_end =
|
||||
decode_step < (MaxStep - 1) ? decode_step : (MaxStep - 1);
|
||||
|
||||
const int32_t req = block_table[b];
|
||||
if (req < 0 || req >= MaxReq) {
|
||||
return;
|
||||
}
|
||||
|
||||
scalar_t* __restrict__ k_cache =
|
||||
reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(k_ptrs_i64[layer]));
|
||||
scalar_t* __restrict__ v_cache =
|
||||
reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(v_ptrs_i64[layer]));
|
||||
|
||||
// base(req, beam, s, kv, d) = ((((req*Beam + beam)*MaxStep + s)*Kv + kv) * D
|
||||
// + d)
|
||||
const int64_t req_base = static_cast<int64_t>(req) * Beam;
|
||||
const int64_t step_kv_stride = static_cast<int64_t>(Kv) * D;
|
||||
const int64_t kv_d_base = static_cast<int64_t>(kv) * D;
|
||||
|
||||
// grid_step is typically small; loop over s in-kernel to reduce launch
|
||||
// blocks.
|
||||
for (int32_t s = 0; s <= step_end; ++s) {
|
||||
// pass-1: new_beam increasing, copy if old_beam > new_beam
|
||||
for (int32_t new_beam = 0; new_beam < Beam; ++new_beam) {
|
||||
const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
|
||||
if (old_beam >= 0 && old_beam < Beam && old_beam > new_beam) {
|
||||
const int64_t dst_base =
|
||||
((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
|
||||
const int64_t src_base =
|
||||
((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
|
||||
for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
|
||||
d += static_cast<int32_t>(blockDim.x)) {
|
||||
k_cache[dst_base + d] = k_cache[src_base + d];
|
||||
v_cache[dst_base + d] = v_cache[src_base + d];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pass-2: new_beam decreasing, copy if old_beam < new_beam
|
||||
for (int32_t new_beam = Beam - 1; new_beam >= 0; --new_beam) {
|
||||
const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
|
||||
if (old_beam >= 0 && old_beam < Beam && old_beam < new_beam) {
|
||||
const int64_t dst_base =
|
||||
((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
|
||||
const int64_t src_base =
|
||||
((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
|
||||
for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
|
||||
d += static_cast<int32_t>(blockDim.x)) {
|
||||
k_cache[dst_base + d] = k_cache[src_base + d];
|
||||
v_cache[dst_base + d] = v_cache[src_base + d];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void cache_select_cuda_launch_ptrs(
|
||||
torch::Tensor k0,
|
||||
torch::Tensor v0,
|
||||
torch::Tensor k_ptrs_i64, // [Layer] int64 (CUDA)
|
||||
torch::Tensor v_ptrs_i64, // [Layer] int64 (CUDA)
|
||||
torch::Tensor beam_index_i32, // [B*Beam, 1] int32
|
||||
torch::Tensor block_table_i32, // [B] int32
|
||||
int64_t decode_step,
|
||||
int64_t layer_num) {
|
||||
CHECK(k_ptrs_i64.is_cuda() && v_ptrs_i64.is_cuda())
|
||||
<< "k_ptrs_i64/v_ptrs_i64 must be CUDA";
|
||||
CHECK_EQ(k_ptrs_i64.scalar_type(), torch::kInt64)
|
||||
<< "k_ptrs_i64/v_ptrs_i64 must be int64";
|
||||
CHECK_EQ(v_ptrs_i64.scalar_type(), torch::kInt64)
|
||||
<< "k_ptrs_i64/v_ptrs_i64 must be int64";
|
||||
CHECK(k_ptrs_i64.is_contiguous() && v_ptrs_i64.is_contiguous())
|
||||
<< "k_ptrs_i64/v_ptrs_i64 must be contiguous";
|
||||
|
||||
const int64_t B64 = block_table_i32.size(0);
|
||||
const int64_t Beam64 = k0.size(1);
|
||||
const int64_t MaxStep64 = k0.size(2);
|
||||
const int64_t Kv64 = k0.size(3);
|
||||
const int64_t D64 = k0.size(4);
|
||||
const int64_t MaxReq64 = k0.size(0);
|
||||
const int64_t Layer64 = layer_num;
|
||||
|
||||
const int32_t B = static_cast<int32_t>(B64);
|
||||
const int32_t Beam = static_cast<int32_t>(Beam64);
|
||||
const int32_t Kv = static_cast<int32_t>(Kv64);
|
||||
const int32_t MaxStep = static_cast<int32_t>(MaxStep64);
|
||||
const int32_t D = static_cast<int32_t>(D64);
|
||||
const int32_t MaxReq = static_cast<int32_t>(MaxReq64);
|
||||
const int32_t Layer = static_cast<int32_t>(Layer64);
|
||||
const int32_t decode_step_i32 = static_cast<int32_t>(decode_step);
|
||||
|
||||
// Warp-aligned threads, capped to keep occupancy reasonable.
|
||||
int threads_per_block = ((D + 31) / 32) * 32;
|
||||
if (threads_per_block < 32) {
|
||||
threads_per_block = 32;
|
||||
}
|
||||
if (threads_per_block > 256) {
|
||||
threads_per_block = 256;
|
||||
}
|
||||
dim3 block_dim(static_cast<unsigned int>(threads_per_block), 1, 1);
|
||||
|
||||
CHECK_LE(Kv64, static_cast<int64_t>(UINT32_MAX)) << "Kv too large for grid.y";
|
||||
CHECK_LE(Layer64, 65535) << "layer_num too large for grid.z";
|
||||
dim3 grid_dim(static_cast<unsigned int>(B),
|
||||
static_cast<unsigned int>(Kv),
|
||||
static_cast<unsigned int>(Layer));
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND2(torch::ScalarType::Half,
|
||||
torch::ScalarType::BFloat16,
|
||||
k0.scalar_type(),
|
||||
"cache_select_inplace_ptrs_kernel",
|
||||
[&] {
|
||||
cache_select_inplace_ptrs_kernel<scalar_t>
|
||||
<<<grid_dim, block_dim, 0, stream>>>(
|
||||
k_ptrs_i64.data_ptr<int64_t>(),
|
||||
v_ptrs_i64.data_ptr<int64_t>(),
|
||||
beam_index_i32.data_ptr<int32_t>(),
|
||||
block_table_i32.data_ptr<int32_t>(),
|
||||
B,
|
||||
Beam,
|
||||
Kv,
|
||||
MaxStep,
|
||||
D,
|
||||
MaxReq,
|
||||
Layer,
|
||||
decode_step_i32);
|
||||
});
|
||||
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
void cache_select(const torch::Tensor& beam_index, // [B*Beam, 1]
|
||||
std::vector<torch::Tensor>& unshared_k_cache,
|
||||
std::vector<torch::Tensor>& unshared_v_cache,
|
||||
const torch::Tensor& block_table, // [B*Beam, 1]
|
||||
int64_t decode_step,
|
||||
int64_t beam_size,
|
||||
int64_t layer_num) {
|
||||
CHECK_GE(layer_num, 0) << "layer_num must be >= 0";
|
||||
if (layer_num == 0) {
|
||||
return;
|
||||
}
|
||||
CHECK_EQ(static_cast<int64_t>(unshared_k_cache.size()), layer_num)
|
||||
<< "unshared_k_cache length mismatch";
|
||||
CHECK_EQ(static_cast<int64_t>(unshared_v_cache.size()), layer_num)
|
||||
<< "unshared_v_cache length mismatch";
|
||||
|
||||
CHECK(beam_index.is_cuda()) << "beam_index must be CUDA";
|
||||
CHECK(block_table.is_cuda()) << "block_table must be CUDA";
|
||||
CHECK_EQ(block_table.dim(), 2) << "block_table must be [B*Beam, 1]";
|
||||
CHECK_EQ(block_table.size(1), 1) << "block_table must be [B*Beam, 1]";
|
||||
CHECK_EQ(beam_index.dim(), 2) << "beam_index must be [B*Beam, 1]";
|
||||
CHECK_EQ(beam_index.size(1), 1) << "beam_index must be [B*Beam, 1]";
|
||||
CHECK_GE(decode_step, 0) << "decode_step must be >= 0";
|
||||
CHECK_GT(beam_size, 0) << "beam_size must be > 0";
|
||||
|
||||
// block_table is [B*Beam, 1] with sequential values [0,1,2,3,...]
|
||||
// Infer actual batch_size
|
||||
CHECK_EQ(block_table.size(0) % beam_size, 0)
|
||||
<< "block_table.size(0) must be divisible by beam_size";
|
||||
const int64_t B = block_table.size(0) / beam_size;
|
||||
CHECK_EQ(beam_index.size(0), B * beam_size)
|
||||
<< "beam_index size mismatch with B*beam_size";
|
||||
|
||||
// Prepare indices (int32, contiguous).
|
||||
auto beam_index_i32 = beam_index.to(torch::kInt32).contiguous();
|
||||
auto block_table_i32 = torch::arange(
|
||||
0,
|
||||
B,
|
||||
torch::TensorOptions().dtype(torch::kInt32).device(block_table.device()));
|
||||
// Validate shapes/dtypes against layer 0.
|
||||
const auto& k0 = unshared_k_cache[0];
|
||||
const auto& v0 = unshared_v_cache[0];
|
||||
CHECK(k0.is_cuda() && v0.is_cuda()) << "cache must be CUDA";
|
||||
CHECK(k0.is_contiguous() && v0.is_contiguous()) << "cache must be contiguous";
|
||||
CHECK_EQ(k0.dim(), 5) << "cache must be 5D [MaxReq, Beam, MaxStep, Kv, D]";
|
||||
CHECK_EQ(v0.sizes(), k0.sizes()) << "k/v cache shapes must match";
|
||||
CHECK_EQ(k0.size(1), beam_size) << "beam_size mismatch with cache";
|
||||
CHECK_LT(decode_step, k0.size(2)) << "decode_step must be < max_decode_step";
|
||||
|
||||
// Pack layer pointers into CUDA int64 tensors so we can launch once.
|
||||
// Note: pointer values are produced on host (data_ptr()), then copied to GPU.
|
||||
c10::cuda::CUDAGuard device_guard(k0.device());
|
||||
auto ptr_cuda_opts =
|
||||
torch::TensorOptions().dtype(torch::kInt64).device(k0.device());
|
||||
auto k_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
|
||||
auto v_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
|
||||
std::vector<int64_t> k_ptrs_host(static_cast<size_t>(layer_num));
|
||||
std::vector<int64_t> v_ptrs_host(static_cast<size_t>(layer_num));
|
||||
|
||||
for (int64_t layer = 0; layer < layer_num; ++layer) {
|
||||
auto k = unshared_k_cache[static_cast<size_t>(layer)];
|
||||
auto v = unshared_v_cache[static_cast<size_t>(layer)];
|
||||
CHECK(k.is_cuda() && v.is_cuda()) << "cache must be CUDA";
|
||||
CHECK(k.is_contiguous() && v.is_contiguous()) << "cache must be contiguous";
|
||||
CHECK_EQ(k.sizes(), k0.sizes()) << "all layers must have same cache shape";
|
||||
CHECK_EQ(v.sizes(), k0.sizes()) << "all layers must have same cache shape";
|
||||
CHECK_EQ(k.scalar_type(), k0.scalar_type())
|
||||
<< "all layers must have same dtype";
|
||||
CHECK_EQ(v.scalar_type(), k0.scalar_type())
|
||||
<< "all layers must have same dtype";
|
||||
CHECK_EQ(k.get_device(), k0.get_device())
|
||||
<< "all layers must be on the same CUDA device";
|
||||
CHECK_EQ(v.get_device(), k0.get_device())
|
||||
<< "all layers must be on the same CUDA device";
|
||||
|
||||
k_ptrs_host[static_cast<size_t>(layer)] =
|
||||
static_cast<int64_t>(reinterpret_cast<uintptr_t>(k.data_ptr()));
|
||||
v_ptrs_host[static_cast<size_t>(layer)] =
|
||||
static_cast<int64_t>(reinterpret_cast<uintptr_t>(v.data_ptr()));
|
||||
}
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
C10_CUDA_CHECK(
|
||||
cudaMemcpyAsync(k_ptrs_i64.data_ptr<int64_t>(),
|
||||
k_ptrs_host.data(),
|
||||
static_cast<size_t>(layer_num) * sizeof(int64_t),
|
||||
cudaMemcpyHostToDevice,
|
||||
stream));
|
||||
C10_CUDA_CHECK(
|
||||
cudaMemcpyAsync(v_ptrs_i64.data_ptr<int64_t>(),
|
||||
v_ptrs_host.data(),
|
||||
static_cast<size_t>(layer_num) * sizeof(int64_t),
|
||||
cudaMemcpyHostToDevice,
|
||||
stream));
|
||||
|
||||
cache_select_cuda_launch_ptrs(k0,
|
||||
v0,
|
||||
k_ptrs_i64,
|
||||
v_ptrs_i64,
|
||||
beam_index_i32,
|
||||
block_table_i32,
|
||||
decode_step,
|
||||
layer_num);
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
@@ -0,0 +1,298 @@
|
||||
/* 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 <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
#include "kernels/cuda/utils.h"
|
||||
#include "xattention_ops_api.h"
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t>
|
||||
struct VecType;
|
||||
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using type = uint4; // 8 elements * 2 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using type = uint4; // 8 elements * 2 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using type = float4; // 4 elements * 4 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 4;
|
||||
};
|
||||
|
||||
// decoder reshape and cache kernel.
|
||||
// Copies proj_k and proj_v into unshared_k_cache / unshared_v_cache.
|
||||
// Inputs:
|
||||
// proj_k : [batch_size, beam_size, kv_heads, head_dim]
|
||||
// proj_v : [batch_size, beam_size, kv_heads, head_dim]
|
||||
// step : [1] - current decode step
|
||||
// batch_size : batch size
|
||||
// beam_size : beam size
|
||||
// kv_heads : number of kv heads
|
||||
// head_dim : head dimension
|
||||
// k_stride0 : proj_k.stride(0)
|
||||
// k_stride1 : proj_k.stride(1)
|
||||
// v_stride0 : proj_v.stride(0)
|
||||
// v_stride1 : proj_v.stride(1)
|
||||
// cache_stride0 : unshared_k_cache.stride(0)
|
||||
// cache_stride1 : unshared_k_cache.stride(1)
|
||||
// cache_stride2 : unshared_k_cache.stride(2)
|
||||
// cache_stride3 : unshared_k_cache.stride(3)
|
||||
// Outputs:
|
||||
// unshared_k_cache : [max_batch_size, beam_size, max_step, kv_heads,
|
||||
// head_dim]
|
||||
// unshared_v_cache : [max_batch_size, beam_size, max_step, kv_heads,
|
||||
// head_dim]
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void decoder_reshape_and_cache_kernel(
|
||||
const scalar_t* __restrict__ proj_k,
|
||||
const scalar_t* __restrict__ proj_v,
|
||||
scalar_t* __restrict__ unshared_k_cache,
|
||||
scalar_t* __restrict__ unshared_v_cache,
|
||||
const int32_t* __restrict__ step,
|
||||
const int64_t batch_size,
|
||||
const int64_t beam_size,
|
||||
const int64_t kv_heads,
|
||||
const int64_t head_dim,
|
||||
const int64_t k_stride0,
|
||||
const int64_t k_stride1,
|
||||
const int64_t v_stride0,
|
||||
const int64_t v_stride1,
|
||||
const int64_t cache_stride0,
|
||||
const int64_t cache_stride1,
|
||||
const int64_t cache_stride2,
|
||||
const int64_t cache_stride3) {
|
||||
using VecTypeT = typename VecType<scalar_t>::type;
|
||||
constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
|
||||
|
||||
const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
|
||||
const int64_t total_tokens = batch_size * beam_size;
|
||||
if (token_idx >= total_tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t batch_idx = token_idx / beam_size;
|
||||
const int64_t beam_idx = token_idx - batch_idx * beam_size;
|
||||
|
||||
__shared__ int32_t current_step_s;
|
||||
if (threadIdx.x == 0) {
|
||||
current_step_s = __ldg(step);
|
||||
}
|
||||
__syncthreads();
|
||||
const int64_t current_step = static_cast<int64_t>(current_step_s);
|
||||
|
||||
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
|
||||
const int64_t total_vecs = kv_heads * vecs_per_head;
|
||||
|
||||
const int64_t k_token_base = batch_idx * k_stride0 + beam_idx * k_stride1;
|
||||
const int64_t v_token_base = batch_idx * v_stride0 + beam_idx * v_stride1;
|
||||
const int64_t dst_token_base = batch_idx * cache_stride0 +
|
||||
beam_idx * cache_stride1 +
|
||||
current_step * cache_stride2;
|
||||
|
||||
for (int64_t linear_idx = static_cast<int64_t>(threadIdx.x);
|
||||
linear_idx < total_vecs;
|
||||
linear_idx += static_cast<int64_t>(blockDim.x)) {
|
||||
const int64_t head_idx = linear_idx / vecs_per_head;
|
||||
const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
|
||||
const int64_t vec_offset = vec_idx * VEC_WIDTH;
|
||||
|
||||
const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
|
||||
proj_k + k_token_base + head_idx * head_dim + vec_offset);
|
||||
const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
|
||||
proj_v + v_token_base + head_idx * head_dim + vec_offset);
|
||||
auto* k_dst_vec =
|
||||
reinterpret_cast<VecTypeT*>(unshared_k_cache + dst_token_base +
|
||||
head_idx * cache_stride3 + vec_offset);
|
||||
auto* v_dst_vec =
|
||||
reinterpret_cast<VecTypeT*>(unshared_v_cache + dst_token_base +
|
||||
head_idx * cache_stride3 + vec_offset);
|
||||
|
||||
*k_dst_vec = *k_src_vec;
|
||||
*v_dst_vec = *v_src_vec;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
void decoder_reshape_and_cache(torch::Tensor proj_k,
|
||||
torch::Tensor proj_v,
|
||||
torch::Tensor unshared_k_cache,
|
||||
torch::Tensor unshared_v_cache,
|
||||
torch::Tensor step) {
|
||||
CHECK_EQ(proj_k.dim(), 4) << "proj_k must be 4-dimensional";
|
||||
CHECK_EQ(proj_v.dim(), 4) << "proj_v must be 4-dimensional";
|
||||
CHECK_EQ(unshared_k_cache.dim(), 5)
|
||||
<< "unshared_k_cache must be 5-dimensional";
|
||||
CHECK_EQ(unshared_v_cache.dim(), 5)
|
||||
<< "unshared_v_cache must be 5-dimensional";
|
||||
CHECK(proj_k.is_cuda() && proj_v.is_cuda() && unshared_k_cache.is_cuda() &&
|
||||
unshared_v_cache.is_cuda() && step.is_cuda())
|
||||
<< "all tensors must be CUDA tensors";
|
||||
CHECK_EQ(step.dim(), 1) << "step must be 1-dimensional";
|
||||
CHECK_EQ(step.size(0), 1) << "step must have shape [1]";
|
||||
CHECK_EQ(step.scalar_type(), at::ScalarType::Int)
|
||||
<< "step must be int32 (torch::kInt32)";
|
||||
|
||||
const int64_t batch_size = proj_k.size(0);
|
||||
const int64_t beam_size = proj_k.size(1);
|
||||
const int64_t kv_heads = proj_k.size(2);
|
||||
const int64_t head_dim = proj_k.size(3);
|
||||
|
||||
CHECK_EQ(proj_v.sizes(), proj_k.sizes())
|
||||
<< "proj_v and proj_k must have same shape";
|
||||
CHECK_EQ(unshared_k_cache.size(3), kv_heads)
|
||||
<< "unshared_k_cache kv_heads mismatch";
|
||||
CHECK_EQ(unshared_k_cache.size(4), head_dim)
|
||||
<< "unshared_k_cache head_dim mismatch";
|
||||
CHECK(unshared_v_cache.sizes() == unshared_k_cache.sizes())
|
||||
<< "unshared_v_cache and unshared_k_cache must have same shape";
|
||||
|
||||
// This kernel is specialized for qkv-slice layouts:
|
||||
// last dim contiguous and kv head stride tightly packed by head_dim.
|
||||
CHECK_EQ(proj_k.stride(3), 1) << "proj_k must satisfy stride(3)=1";
|
||||
CHECK_EQ(proj_v.stride(3), 1) << "proj_v must satisfy stride(3)=1";
|
||||
CHECK_EQ(proj_k.stride(2), head_dim)
|
||||
<< "proj_k must satisfy stride(2)=head_dim";
|
||||
CHECK_EQ(proj_v.stride(2), head_dim)
|
||||
<< "proj_v must satisfy stride(2)=head_dim";
|
||||
CHECK_EQ(unshared_k_cache.stride(4), 1)
|
||||
<< "unshared_k_cache must satisfy stride(4)=1";
|
||||
CHECK_EQ(unshared_v_cache.stride(4), 1)
|
||||
<< "unshared_v_cache must satisfy stride(4)=1";
|
||||
CHECK_EQ(unshared_k_cache.stride(3), head_dim)
|
||||
<< "unshared_k_cache must satisfy stride(3)=head_dim";
|
||||
CHECK_EQ(unshared_v_cache.stride(3), head_dim)
|
||||
<< "unshared_v_cache must satisfy stride(3)=head_dim";
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
const int64_t k_stride0 = proj_k.stride(0);
|
||||
const int64_t k_stride1 = proj_k.stride(1);
|
||||
const int64_t v_stride0 = proj_v.stride(0);
|
||||
const int64_t v_stride1 = proj_v.stride(1);
|
||||
const int64_t cache_stride0 = unshared_k_cache.stride(0);
|
||||
const int64_t cache_stride1 = unshared_k_cache.stride(1);
|
||||
const int64_t cache_stride2 = unshared_k_cache.stride(2);
|
||||
const int64_t cache_stride3 = unshared_k_cache.stride(3);
|
||||
|
||||
// Launch kernel: one block per (batch, beam), threads cover
|
||||
// kv_heads*head_dim.
|
||||
const int64_t total_tokens = batch_size * beam_size;
|
||||
dim3 grid_dim(1, static_cast<unsigned int>(total_tokens), 1);
|
||||
|
||||
DISPATCH_FLOATING_TYPES(
|
||||
proj_k.scalar_type(), "decoder_reshape_and_cache_kernel", [&] {
|
||||
constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
|
||||
std::is_same_v<scalar_t, c10::BFloat16>)
|
||||
? 8
|
||||
: 4; // FP16/BF16: 8, Float: 4
|
||||
constexpr int32_t kWarpSize = 32;
|
||||
constexpr int32_t kMaxThreadsPerBlock = 256;
|
||||
constexpr int32_t kAlignmentBytes = 16; // 128-bit alignment
|
||||
|
||||
CHECK(head_dim % VEC_WIDTH == 0)
|
||||
<< "head_dim must be divisible by vector width: " << VEC_WIDTH;
|
||||
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
|
||||
const int64_t total_vecs = kv_heads * vecs_per_head;
|
||||
CHECK(total_vecs > 0) << "total_vecs must be > 0";
|
||||
|
||||
int32_t threads_per_block = static_cast<int32_t>(
|
||||
total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
|
||||
: total_vecs);
|
||||
threads_per_block =
|
||||
((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
|
||||
if (threads_per_block < kWarpSize) {
|
||||
threads_per_block = kWarpSize;
|
||||
}
|
||||
dim3 block_dim(threads_per_block, 1, 1);
|
||||
|
||||
const auto proj_k_ptr =
|
||||
reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
|
||||
const auto proj_v_ptr =
|
||||
reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
|
||||
const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
|
||||
unshared_k_cache.data_ptr<scalar_t>());
|
||||
const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
|
||||
unshared_v_cache.data_ptr<scalar_t>());
|
||||
CHECK(proj_k_ptr % kAlignmentBytes == 0)
|
||||
<< "proj_k data_ptr must be 16-byte aligned";
|
||||
CHECK(proj_v_ptr % kAlignmentBytes == 0)
|
||||
<< "proj_v data_ptr must be 16-byte aligned";
|
||||
CHECK(k_cache_ptr % kAlignmentBytes == 0)
|
||||
<< "unshared_k_cache data_ptr must be 16-byte aligned";
|
||||
CHECK(v_cache_ptr % kAlignmentBytes == 0)
|
||||
<< "unshared_v_cache data_ptr must be 16-byte aligned";
|
||||
|
||||
const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
|
||||
CHECK((k_stride0 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "proj_k stride(0) bytes must be 16-byte aligned";
|
||||
CHECK((k_stride1 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "proj_k stride(1) bytes must be 16-byte aligned";
|
||||
CHECK((v_stride0 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "proj_v stride(0) bytes must be 16-byte aligned";
|
||||
CHECK((v_stride1 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "proj_v stride(1) bytes must be 16-byte aligned";
|
||||
CHECK((cache_stride0 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "cache stride(0) bytes must be 16-byte aligned";
|
||||
CHECK((cache_stride1 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "cache stride(1) bytes must be 16-byte aligned";
|
||||
CHECK((cache_stride2 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "cache stride(2) bytes must be 16-byte aligned";
|
||||
CHECK((cache_stride3 * scalar_bytes) % kAlignmentBytes == 0)
|
||||
<< "cache stride(3) bytes must be 16-byte aligned";
|
||||
|
||||
decoder_reshape_and_cache_kernel<scalar_t>
|
||||
<<<grid_dim, block_dim, 0, stream>>>(
|
||||
proj_k.data_ptr<scalar_t>(),
|
||||
proj_v.data_ptr<scalar_t>(),
|
||||
unshared_k_cache.data_ptr<scalar_t>(),
|
||||
unshared_v_cache.data_ptr<scalar_t>(),
|
||||
step.data_ptr<int32_t>(),
|
||||
batch_size,
|
||||
beam_size,
|
||||
kv_heads,
|
||||
head_dim,
|
||||
k_stride0,
|
||||
k_stride1,
|
||||
v_stride0,
|
||||
v_stride1,
|
||||
cache_stride0,
|
||||
cache_stride1,
|
||||
cache_stride2,
|
||||
cache_stride3);
|
||||
});
|
||||
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
168
ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu
Normal file
168
ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu
Normal file
@@ -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
|
||||
@@ -0,0 +1,220 @@
|
||||
/* 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 <glog/logging.h>
|
||||
#include <torch/cuda.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
#include "kernels/cuda/cuda_ops_api.h"
|
||||
#include "kernels/cuda/utils.h"
|
||||
using at::device_of;
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t>
|
||||
struct VecType;
|
||||
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using type = uint4; // 8 elements * 2 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using type = uint4; // 8 elements * 2 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using type = float4; // 4 elements * 4 bytes = 16 bytes
|
||||
static constexpr int32_t vec_width = 4;
|
||||
};
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void prefill_reshape_and_cache_kernel(
|
||||
const scalar_t* __restrict__ proj_k, // [shared_len, kv_heads, head_dim]
|
||||
const scalar_t* __restrict__ proj_v, // [shared_len, kv_heads, head_dim]
|
||||
scalar_t* __restrict__ shared_k_cache, // [shared_len, kv_heads, head_dim]
|
||||
scalar_t* __restrict__ shared_v_cache, // [shared_len, kv_heads, head_dim]
|
||||
const int64_t shared_len,
|
||||
const int64_t kv_heads,
|
||||
const int64_t head_dim,
|
||||
const int64_t k_stride0, // proj_k.stride(0)
|
||||
const int64_t v_stride0, // proj_v.stride(0)
|
||||
const int64_t v_stride1) { // proj_v.stride(1), same as head_dim
|
||||
using VecTypeT = typename VecType<scalar_t>::type;
|
||||
constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
|
||||
const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
|
||||
if (token_idx >= shared_len) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
|
||||
const int64_t total_vecs = kv_heads * vecs_per_head;
|
||||
const int64_t k_token_base = token_idx * k_stride0;
|
||||
const int64_t v_token_base = token_idx * v_stride0;
|
||||
const int64_t dst_token_base = token_idx * kv_heads * head_dim;
|
||||
|
||||
for (int64_t linear_idx = threadIdx.x; linear_idx < total_vecs;
|
||||
linear_idx += blockDim.x) {
|
||||
const int64_t head_idx = linear_idx / vecs_per_head;
|
||||
const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
|
||||
const int64_t head_offset = head_idx * head_dim;
|
||||
const int64_t vec_offset = vec_idx * VEC_WIDTH;
|
||||
|
||||
const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
|
||||
proj_k + k_token_base + head_offset + vec_offset);
|
||||
const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
|
||||
proj_v + v_token_base + head_idx * v_stride1 + vec_offset);
|
||||
auto* k_dst_vec = reinterpret_cast<VecTypeT*>(
|
||||
shared_k_cache + dst_token_base + head_offset + vec_offset);
|
||||
auto* v_dst_vec = reinterpret_cast<VecTypeT*>(
|
||||
shared_v_cache + dst_token_base + head_offset + vec_offset);
|
||||
|
||||
*k_dst_vec = *k_src_vec;
|
||||
*v_dst_vec = *v_src_vec;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
void prefill_reshape_and_cache(
|
||||
torch::Tensor proj_k, // [shared_len, kv_heads, head_dim]
|
||||
torch::Tensor proj_v, // [shared_len, kv_heads, head_dim]
|
||||
torch::Tensor
|
||||
shared_k_cache, // [num_shared_kv_seq_len, kv_heads, head_dim]
|
||||
torch::Tensor shared_v_cache) {
|
||||
CHECK(proj_k.dim() == 3) << "proj_k must be 3-dimensional";
|
||||
CHECK(proj_v.dim() == 3) << "proj_v must be 3-dimensional";
|
||||
CHECK(shared_k_cache.dim() == 3) << "shared_k_cache must be 3-dimensional";
|
||||
CHECK(shared_v_cache.dim() == 3) << "shared_v_cache must be 3-dimensional";
|
||||
CHECK(proj_k.is_cuda() && proj_v.is_cuda() && shared_k_cache.is_cuda() &&
|
||||
shared_v_cache.is_cuda())
|
||||
<< "all tensors must be CUDA tensors";
|
||||
|
||||
const int64_t shared_len = proj_k.size(0);
|
||||
const int64_t kv_heads = proj_k.size(1);
|
||||
const int64_t head_dim = proj_k.size(2);
|
||||
CHECK(proj_v.sizes() == proj_k.sizes())
|
||||
<< "proj_v and proj_k must have same shape";
|
||||
CHECK(shared_k_cache.size(0) >= shared_len &&
|
||||
shared_k_cache.size(1) == kv_heads &&
|
||||
shared_k_cache.size(2) == head_dim)
|
||||
<< "shared_k_cache shape mismatch";
|
||||
CHECK(shared_v_cache.size(0) >= shared_len &&
|
||||
shared_v_cache.size(1) == kv_heads &&
|
||||
shared_v_cache.size(2) == head_dim)
|
||||
<< "shared_v_cache shape mismatch";
|
||||
|
||||
shared_k_cache = shared_k_cache.slice(0, 0, shared_len);
|
||||
shared_v_cache = shared_v_cache.slice(0, 0, shared_len);
|
||||
|
||||
// This kernel is specialized for qkv-slice layouts:
|
||||
// last dim contiguous and head stride tightly packed by head_dim.
|
||||
CHECK(proj_k.stride(2) == 1 && proj_v.stride(2) == 1)
|
||||
<< "proj_k/proj_v must be contiguous on head_dim (stride(2)=1)";
|
||||
CHECK(proj_k.stride(1) == head_dim && proj_v.stride(1) == head_dim)
|
||||
<< "proj_k/proj_v must satisfy stride(1)=head_dim for qkv-slice layout";
|
||||
CHECK(shared_k_cache.stride(2) == 1 && shared_v_cache.stride(2) == 1)
|
||||
<< "shared caches must be contiguous on head_dim (stride(2)=1)";
|
||||
CHECK(shared_k_cache.stride(1) == head_dim &&
|
||||
shared_v_cache.stride(1) == head_dim)
|
||||
<< "shared caches must satisfy stride(1)=head_dim";
|
||||
CHECK(shared_k_cache.stride(0) == kv_heads * head_dim &&
|
||||
shared_v_cache.stride(0) == kv_heads * head_dim)
|
||||
<< "shared caches must be contiguous on token stride";
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
const int64_t k_stride0 = proj_k.stride(0);
|
||||
const int64_t v_stride0 = proj_v.stride(0);
|
||||
const int64_t v_stride1 = proj_v.stride(1);
|
||||
dim3 grid_dim(1, static_cast<unsigned int>(shared_len), 1);
|
||||
|
||||
DISPATCH_FLOATING_TYPES(
|
||||
proj_k.scalar_type(), "prefill_reshape_and_cache_kernel", [&] {
|
||||
constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
|
||||
std::is_same_v<scalar_t, c10::BFloat16>)
|
||||
? 8
|
||||
: 4; // FP16/BF16: 8, Float: 4
|
||||
constexpr int32_t kWarpSize = 32;
|
||||
constexpr int32_t kMaxThreadsPerBlock = 256;
|
||||
|
||||
CHECK(head_dim % VEC_WIDTH == 0)
|
||||
<< "head_dim must be divisible by vector width: " << VEC_WIDTH;
|
||||
const int64_t vecs_per_head = head_dim / VEC_WIDTH;
|
||||
const int64_t total_vecs = kv_heads * vecs_per_head;
|
||||
CHECK(total_vecs > 0) << "total_vecs must be > 0";
|
||||
|
||||
int32_t threads_per_block = static_cast<int32_t>(
|
||||
total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
|
||||
: total_vecs);
|
||||
threads_per_block =
|
||||
((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
|
||||
if (threads_per_block < kWarpSize) {
|
||||
threads_per_block = kWarpSize;
|
||||
}
|
||||
dim3 block_dim(threads_per_block, 1, 1);
|
||||
|
||||
const auto proj_k_ptr =
|
||||
reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
|
||||
const auto proj_v_ptr =
|
||||
reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
|
||||
const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
|
||||
shared_k_cache.data_ptr<scalar_t>());
|
||||
const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
|
||||
shared_v_cache.data_ptr<scalar_t>());
|
||||
|
||||
constexpr int32_t alignment_bytes = 16; // 128-bit alignment
|
||||
CHECK(proj_k_ptr % alignment_bytes == 0)
|
||||
<< "proj_k data_ptr must be 16-byte aligned";
|
||||
CHECK(proj_v_ptr % alignment_bytes == 0)
|
||||
<< "proj_v data_ptr must be 16-byte aligned";
|
||||
CHECK(k_cache_ptr % alignment_bytes == 0)
|
||||
<< "shared_k_cache data_ptr must be 16-byte aligned";
|
||||
CHECK(v_cache_ptr % alignment_bytes == 0)
|
||||
<< "shared_v_cache data_ptr must be 16-byte aligned";
|
||||
|
||||
const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
|
||||
CHECK((k_stride0 * scalar_bytes) % alignment_bytes == 0)
|
||||
<< "proj_k stride(0) bytes must be 16-byte aligned";
|
||||
CHECK((v_stride0 * scalar_bytes) % alignment_bytes == 0)
|
||||
<< "proj_v stride(0) bytes must be 16-byte aligned";
|
||||
CHECK((v_stride1 * scalar_bytes) % alignment_bytes == 0)
|
||||
<< "proj_v stride(1) bytes must be 16-byte aligned";
|
||||
|
||||
prefill_reshape_and_cache_kernel<scalar_t>
|
||||
<<<grid_dim, block_dim, 0, stream>>>(
|
||||
proj_k.data_ptr<scalar_t>(),
|
||||
proj_v.data_ptr<scalar_t>(),
|
||||
shared_k_cache.data_ptr<scalar_t>(),
|
||||
shared_v_cache.data_ptr<scalar_t>(),
|
||||
shared_len,
|
||||
kv_heads,
|
||||
head_dim,
|
||||
k_stride0,
|
||||
v_stride0,
|
||||
v_stride1);
|
||||
});
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
} // namespace xllm::kernel::cuda
|
||||
63
ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h
Normal file
63
ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h
Normal file
@@ -0,0 +1,63 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
void decoder_reshape_and_cache(torch::Tensor proj_k,
|
||||
torch::Tensor proj_v,
|
||||
torch::Tensor unshared_k_cache,
|
||||
torch::Tensor unshared_v_cache,
|
||||
torch::Tensor step);
|
||||
|
||||
void cache_select(const torch::Tensor& beam_index,
|
||||
std::vector<torch::Tensor>& unshared_k_cache,
|
||||
std::vector<torch::Tensor>& unshared_v_cache,
|
||||
const torch::Tensor& block_table,
|
||||
int64_t decode_step,
|
||||
int64_t beam_size,
|
||||
int64_t layer_num);
|
||||
|
||||
void lse_combine(torch::Tensor output,
|
||||
torch::Tensor shared_o,
|
||||
torch::Tensor shared_lse,
|
||||
torch::Tensor unshared_o,
|
||||
torch::Tensor unshared_lse);
|
||||
|
||||
void prefill_reshape_and_cache(
|
||||
torch::Tensor proj_k, // [shared_len, kv_heads, head_dim]
|
||||
torch::Tensor proj_v, // [shared_len, kv_heads, head_dim]
|
||||
torch::Tensor
|
||||
shared_k_cache, // [num_shared_kv_seq_len, kv_heads, head_dim]
|
||||
torch::Tensor shared_v_cache);
|
||||
|
||||
void beam_search(torch::Tensor acc_logprob,
|
||||
torch::Tensor in_sequence_group,
|
||||
torch::Tensor top_tokens,
|
||||
torch::Tensor top_logprobs,
|
||||
torch::Tensor out_acc_logprob,
|
||||
torch::Tensor out_token_ids,
|
||||
torch::Tensor out_token_index,
|
||||
torch::Tensor out_beam_count_prefix_sums,
|
||||
torch::Tensor out_sequence_group,
|
||||
uint32_t batch_size,
|
||||
uint32_t current_step);
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
28
ex_engine/xllm_kernels/ilu/CMakeLists.txt
Normal file
28
ex_engine/xllm_kernels/ilu/CMakeLists.txt
Normal file
@@ -0,0 +1,28 @@
|
||||
include(cc_library)
|
||||
set(CMAKE_CUDA_ARCHITECTURES ivcore11)
|
||||
file(GLOB_RECURSE ILU_HEADER_FILES
|
||||
"${CMAKE_CURRENT_LIST_DIR}/*.h"
|
||||
)
|
||||
|
||||
file(GLOB_RECURSE ILU_SOURCE_FILES
|
||||
"${CMAKE_CURRENT_LIST_DIR}/*.cpp"
|
||||
"${CMAKE_CURRENT_LIST_DIR}/*.cu"
|
||||
)
|
||||
|
||||
find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
|
||||
|
||||
cc_library(
|
||||
NAME
|
||||
ilu_kernels
|
||||
HDRS
|
||||
${ILU_HEADER_FILES}
|
||||
SRCS
|
||||
${ILU_SOURCE_FILES}
|
||||
DEPS
|
||||
torch
|
||||
:util
|
||||
ixformer_kernels
|
||||
ixformer
|
||||
${Python3_LIBRARIES}
|
||||
cuinfer
|
||||
)
|
||||
Reference in New Issue
Block a user