dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: import CUDA kernels from xllm/CCCL/FLA upstream repos

Browse Source 
Sources cloned and tree'd (no --depth):
  - jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels
  - NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives)
  - fla-org/flash-linear-attention: Triton GDN kernels
  - NVIDIA/cutlass: grouped GEMM reference (read, not copied)
  - Dao-AILab/flash-attention: attention kernel reference (SM80+, read only)

New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100):
  ex_engine/xllm_kernels/cuda/activation.cu    (188 lines) — silu_and_mul, gelu
  ex_engine/xllm_kernels/cuda/norm.cu          (600 lines) — rms_norm, fused_add_rms_norm
  ex_engine/xllm_kernels/cuda/rope.cu          (258 lines) — rotary_embedding
  ex_engine/xllm_kernels/cuda/block_copy.cu    (209 lines) — copy_blocks, swap_blocks
  ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops
  ex_engine/xllm_kernels/cuda/headers/         (5 headers for compilation)

ILU bridge kernel sources (from xllm, verified SAME as upstream):
  ex_engine/xllm_kernels/ilu/    (10 files, 925 lines total)
  — activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp,
    matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h

FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA):
  ex_engine/fla_kernels/gated_delta_rule/  (7 files, 2370 lines)
  — chunk_fwd.py (428), chunk.py (487), wy_fast.py (409),
    fused_recurrent.py (392), naive.py (161), gate.py (380)

CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl):
  cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced
  cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced

Compilation targets for real machine (ivcore10):
  1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16
  2. ILU bridges: torch.utils.cpp_extension linking ixformer .so
  3. FLA kernels: Triton JIT (if Triton works on BI-V100)
This commit is contained in:
claude
2026-08-14 07:48:52 +00:00
parent 051b02d3cd
commit 8d75652949
66 changed files with 12834 additions and 5458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

209
ex_engine/xllm_kernels/cuda/block_copy.cu Normal file
View File

@@ -0,0 +1,209 @@
/* Copyright 2025-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 <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <torch/extension.h>
#include <cstdint>
#include <type_traits>
#include "cuda_ops_api.h"
#include "utils.h"
namespace xllm::kernel::cuda {
namespace {
template <typename scalar_t>
struct VecType;
template <>
struct VecType<c10::Half> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<c10::BFloat16> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<float> {
using type = float4;
static constexpr int32_t vec_width = 4;
};
DEVICE_INLINE int32_t find_group_idx(const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int32_t dst_idx) {
int32_t left = 0;
int32_t right = num_groups - 1;
while (left < right) {
const int32_t mid = left + ((right - left) >> 1);
const bool move_left = dst_idx < cum_sum[mid];
right = move_left ? mid : right;
left = move_left ? left : mid + 1;
}
return left;
}
template <typename scalar_t, bool kVectorized>
__global__ void block_copy_kernel(const int64_t* __restrict__ key_cache_ptrs,
const int64_t* __restrict__ value_cache_ptrs,
const int32_t* __restrict__ src_block_indices,
const int32_t* __restrict__ dst_block_indices,
const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int64_t numel_per_block) {
const int64_t layer_idx = static_cast<int64_t>(blockIdx.x);
const int32_t dst_linear_idx = static_cast<int32_t>(blockIdx.y);
const int64_t tile_idx = static_cast<int64_t>(blockIdx.z);
scalar_t* __restrict__ key_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(key_cache_ptrs[layer_idx]));
scalar_t* __restrict__ value_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(value_cache_ptrs[layer_idx]));
const int32_t group_idx = find_group_idx(cum_sum, num_groups, dst_linear_idx);
const int32_t src_block = src_block_indices[group_idx];
const int32_t dst_block = dst_block_indices[dst_linear_idx];
const int64_t src_offset = static_cast<int64_t>(src_block) * numel_per_block;
const int64_t dst_offset = static_cast<int64_t>(dst_block) * numel_per_block;
if constexpr (kVectorized) {
using VecTypeT = typename VecType<scalar_t>::type;
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
const int64_t num_vecs_per_block = numel_per_block / kVecWidth;
const int64_t vec_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (vec_idx >= num_vecs_per_block) {
return;
}
const int64_t elem_offset = vec_idx * kVecWidth;
const auto* key_src_vec =
reinterpret_cast<const VecTypeT*>(key_cache + src_offset + elem_offset);
const auto* value_src_vec = reinterpret_cast<const VecTypeT*>(
value_cache + src_offset + elem_offset);
auto* key_dst_vec =
reinterpret_cast<VecTypeT*>(key_cache + dst_offset + elem_offset);
auto* value_dst_vec =
reinterpret_cast<VecTypeT*>(value_cache + dst_offset + elem_offset);
*key_dst_vec = *key_src_vec;
*value_dst_vec = *value_src_vec;
} else {
const int64_t elem_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (elem_idx >= numel_per_block) {
return;
}
key_cache[dst_offset + elem_idx] = key_cache[src_offset + elem_idx];
value_cache[dst_offset + elem_idx] = value_cache[src_offset + elem_idx];
}
}
} // namespace
void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices,
torch::Tensor dst_block_indices,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype) {
if (src_block_indices.numel() == 0) {
return;
}
CHECK(key_cache_ptrs.is_cuda());
CHECK(value_cache_ptrs.is_cuda());
CHECK(src_block_indices.is_cuda());
CHECK(dst_block_indices.is_cuda());
CHECK(cum_sum.is_cuda());
CHECK_EQ(key_cache_ptrs.scalar_type(), torch::kInt64);
CHECK_EQ(value_cache_ptrs.scalar_type(), torch::kInt64);
CHECK_EQ(src_block_indices.scalar_type(), torch::kInt32);
CHECK_EQ(dst_block_indices.scalar_type(), torch::kInt32);
CHECK_EQ(cum_sum.scalar_type(), torch::kInt32);
CHECK_EQ(key_cache_ptrs.dim(), 1);
CHECK_EQ(value_cache_ptrs.dim(), 1);
CHECK_EQ(src_block_indices.dim(), 1);
CHECK_EQ(dst_block_indices.dim(), 1);
CHECK_EQ(cum_sum.dim(), 1);
CHECK(key_cache_ptrs.is_contiguous());
CHECK(value_cache_ptrs.is_contiguous());
CHECK(src_block_indices.is_contiguous());
CHECK(dst_block_indices.is_contiguous());
CHECK(cum_sum.is_contiguous());
CHECK_EQ(key_cache_ptrs.size(0), value_cache_ptrs.size(0));
CHECK_EQ(src_block_indices.size(0), cum_sum.size(0));
CHECK_GT(numel_per_block, 0);
const at::cuda::OptionalCUDAGuard device_guard(key_cache_ptrs.device());
constexpr int32_t kThreadsPerBlock = 256;
const int32_t num_layers = static_cast<int32_t>(key_cache_ptrs.size(0));
const int32_t num_groups = static_cast<int32_t>(src_block_indices.size(0));
const int32_t num_dst_blocks =
static_cast<int32_t>(dst_block_indices.size(0));
const cudaStream_t stream =
c10::cuda::getCurrentCUDAStream(key_cache_ptrs.get_device());
DISPATCH_FLOATING_TYPES(cache_dtype, "block_copy_kernel", [&] {
constexpr bool kHasVecType = std::is_same_v<scalar_t, float> ||
std::is_same_v<scalar_t, c10::Half> ||
std::is_same_v<scalar_t, c10::BFloat16>;
if constexpr (kHasVecType) {
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
if (numel_per_block % kVecWidth == 0) {
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block / kVecWidth, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, true>
<<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return;
}
}
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, false><<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
}
} // namespace xllm::kernel::cuda