From 51cb90b9ab08bfce3c438d7fa141b2c3b44159fb Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 14 Aug 2026 10:19:51 +0000 Subject: [PATCH] fix: adapt xllm norm.cu for corex CUB (CUDA 10.2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Key change: replace CCCL 3.6 types with corex CUB equivalents - cuda::std::plus<> → cub::Sum - cuda::maximum<> → cub::Max - Remove #include Test results from real machine (3/4 passed): ✓ __shfl_down_sync works on ivcore10 ✓ manual SMEM+shuffle block reduce works ✓ corex CUB cub::BlockReduce compiles and runs correctly (32640) ✗ CCCL 3.6 variadic function issue — corex clang rejects device variadic Confirmed: use /usr/local/corex/include/cub/ for all kernel code cccl_upstream is reference only, NOT compilable on corex Build script: bash qwen3_6_scripts/build_xllm_kernels.sh --- ex_engine/xllm_kernels/cuda/norm.cu | 6 +- ex_engine/xllm_kernels/cuda/norm.cu.orig | 600 +++++++++++++++++++++++ qwen3_6_scripts/build_xllm_kernels.sh | 61 +++ 3 files changed, 664 insertions(+), 3 deletions(-) create mode 100644 ex_engine/xllm_kernels/cuda/norm.cu.orig create mode 100644 qwen3_6_scripts/build_xllm_kernels.sh diff --git a/ex_engine/xllm_kernels/cuda/norm.cu b/ex_engine/xllm_kernels/cuda/norm.cu index 30e70084..06c4a8c7 100644 --- a/ex_engine/xllm_kernels/cuda/norm.cu +++ b/ex_engine/xllm_kernels/cuda/norm.cu @@ -28,9 +28,9 @@ limitations under the License. // https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu #if CUB_VERSION >= 200800 -#include -using CubAddOp = ::cuda::std::plus<>; -using CubMaxOp = ::cuda::maximum<>; +// corex CUB (CUDA 10.2) — no cuda::std::functional +using CubAddOp = cub::Sum; +using CubMaxOp = cub::Max; #else // if CUB_VERSION < 200800 using CubAddOp = cub::Sum; using CubMaxOp = cub::Max; diff --git a/ex_engine/xllm_kernels/cuda/norm.cu.orig b/ex_engine/xllm_kernels/cuda/norm.cu.orig new file mode 100644 index 00000000..30e70084 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/norm.cu.orig @@ -0,0 +1,600 @@ +/* Copyright 2025 The vLLM Authors and 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 +#include + +#include +#include + +#include "cuda_ops_api.h" +#include "device_utils.cuh" +#include "fp8_quant_utils.cuh" +#include "type_convert.cuh" + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu + +#if CUB_VERSION >= 200800 +#include +using CubAddOp = ::cuda::std::plus<>; +using CubMaxOp = ::cuda::maximum<>; +#else // if CUB_VERSION < 200800 +using CubAddOp = cub::Sum; +using CubMaxOp = cub::Max; +#endif // CUB_VERSION + +namespace { + +using namespace xllm::kernel::cuda; + +template +__global__ void XLLM_KERNEL_ATTR(1024) + rms_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size] + const scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + const float x = static_cast(input[blockIdx.x * input_stride + idx]); + variance += x * x; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(input[blockIdx.x * input_stride + idx]); + out[blockIdx.x * hidden_size + idx] = + (static_cast(x * s_variance)) * weight[idx]; + } +} + +/* Function specialization in the case of FP16/BF16 tensors. + Additional optimizations we can make in this case are + packed and vectorized operations, which help with the + memory latency bottleneck. */ +template +__global__ std::enable_if_t<(width > 0) && _typeConvert::exists> +XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel( + scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [..., hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + // Sanity checks on our vector struct and type-punned pointer arithmetic + static_assert(std::is_pod_v<_f16Vec>); + static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); + + const int vec_hidden_size = hidden_size / width; + const int64_t vec_input_stride = input_stride / width; + __shared__ float s_variance; + float variance = 0.0f; + /* These and the argument pointers are all declared `restrict` as they are + not aliased in practice. Argument pointers should not be dereferenced + in this kernel as that would be undefined behavior */ + auto* __restrict__ input_v = + reinterpret_cast<_f16Vec*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(weight); + + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = residual_v[id]; + temp *= s_variance; + temp *= weight_v[idx]; + input_v[strided_id] = temp; + } +} + +/* Generic fused_add_rms_norm_kernel + The width field is not used here but necessary for other specializations. + */ +template +__global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> +XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel( + scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [..., hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + scalar_t z = input[blockIdx.x * input_stride + idx]; + z += residual[blockIdx.x * hidden_size + idx]; + float x = static_cast(z); + variance += x * x; + residual[blockIdx.x * hidden_size + idx] = z; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(residual[blockIdx.x * hidden_size + idx]); + input[blockIdx.x * input_stride + idx] = + (static_cast(x * s_variance)) * weight[idx]; + } +} + +#define LAUNCH_FUSED_ADD_RMS_NORM(width) \ + DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ + fused_add_rms_norm_kernel \ + <<>>(input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + epsilon, \ + num_tokens, \ + hidden_size); \ + }); + +// ============================================================================ +// Fused RMSNorm + Static FP8 Quantization Kernels +// ============================================================================ +// These kernels combine RMSNorm and FP8 quantization to reduce memory +// bandwidth by avoiding the intermediate write-back to global memory. + +// Dispatch macro for FP8 types +#define DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + [&] { \ + const auto& the_type = TYPE; \ + switch (the_type) { \ + case at::ScalarType::Float8_e4m3fn: { \ + using fp8_t = c10::Float8_e4m3fn; \ + return __VA_ARGS__(); \ + } \ + default: \ + AT_ERROR(#NAME, \ + " not implemented for FP8 type '", \ + toString(the_type), \ + "'"); \ + } \ + }() + +/** + * Fused RMSNorm + Static FP8 Quantization kernel (without residual) + * Combines RMSNorm and FP8 quantization in a single kernel to reduce + * memory bandwidth by avoiding intermediate write-back. + * + * @tparam scalar_t Input data type (float, half, bfloat16) + * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn) + * @param out Output FP8 tensor [num_tokens, hidden_size] + * @param input Input tensor [num_tokens, hidden_size] + * @param input_stride Stride of input tensor in the token dimension + * @param weight RMSNorm weight tensor [hidden_size] + * @param scale FP8 quantization scale (scalar) + * @param epsilon RMSNorm epsilon + * @param num_tokens Number of tokens + * @param hidden_size Hidden dimension size + */ +template +__global__ void rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + const scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + const scalar_t* input_row = input + blockIdx.x * input_stride; + + // Step 1: Compute variance for RMSNorm + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + const float x = static_cast(input_row[idx]); + variance += x * x; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse to avoid division + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(input_row[idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(out_norm, + scale_inv); + } +} + +/** + * Fused Add + RMSNorm + Static FP8 Quantization kernel (with residual) + * Optimized version with packed + vectorized operations for FP16/BF16. + * + * @tparam scalar_t Input data type (float, half, bfloat16) + * @tparam width Vector width for optimization (0, 8) + * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn) + */ +template +__global__ std::enable_if_t<(width > 0) && _typeConvert::exists> +fused_add_rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [num_tokens, hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + static_assert(std::is_pod_v<_f16Vec>); + static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); + + const int vec_hidden_size = hidden_size / width; + const int64_t vec_input_stride = input_stride / width; + __shared__ float s_variance; + float variance = 0.0f; + + auto* __restrict__ input_v = + reinterpret_cast<_f16Vec*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(weight); + + // Step 1: Fused add and compute variance + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + _f16Vec temp = residual_v[id]; + temp *= s_variance; + temp *= weight_v[idx]; + + // Convert each element to FP8 +#pragma unroll + for (int i = 0; i < width; ++i) { + float val = _typeConvert::convert(temp.data[i]); + out[id * width + i] = + xllm::kernel::cuda::scaled_fp8_conversion(val, + scale_inv); + } + } +} + +/** + * Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data) + */ +template +__global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> +fused_add_rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [num_tokens, hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + // Step 1: Fused add and compute variance + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + scalar_t z = input[blockIdx.x * input_stride + idx]; + z += residual[blockIdx.x * hidden_size + idx]; + float x = static_cast(z); + variance += x * x; + residual[blockIdx.x * hidden_size + idx] = z; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(residual[blockIdx.x * hidden_size + idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(out_norm, + scale_inv); + } +} + +#define LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(width) \ + DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \ + DISPATCH_FP8_TYPES( \ + out.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \ + fused_add_rms_norm_static_fp8_quant_kernel \ + <<>>(out.data_ptr(), \ + input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + scale.data_ptr(), \ + epsilon, \ + num_tokens, \ + hidden_size); \ + }); \ + }); + +} // namespace + +namespace xllm::kernel::cuda { + +// flashinfer rmsnorm ops +// void rmsnorm(torch::Tensor output, +// torch::Tensor input, +// torch::Tensor weight, +// double eps) { +// FunctionFactory::get_instance().rmsnorm_func("norm").call( +// output, input, weight, eps, support_pdl()); +// } + +void rms_norm(torch::Tensor output, // [..., hidden_size] + torch::Tensor input, // [..., hidden_size] + torch::Tensor weight, // [hidden_size] + double eps) { + CHECK(output.is_contiguous()); + CHECK(weight.is_contiguous()); + + // The kernel addresses tokens as `blockIdx.x * input_stride + idx`, which + // can only represent contiguous inputs or simple 2D strided rows. Flux q/k + // tensors reach this path as high-dimensional transposed views, so make that + // layout explicit before flattening tokens for the kernel. + if (input.dim() > 2 && !input.is_contiguous()) { + input = input.contiguous(); + } + CHECK(input.stride(-1) == 1); + + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + int64_t input_stride = input.stride(-2); + + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] { + rms_norm_kernel + <<>>(output.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + eps, + num_tokens, + hidden_size); + }); +} + +void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size] + torch::Tensor& residual, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + double epsilon) { + CHECK(weight.scalar_type() == input.scalar_type()); + CHECK(input.scalar_type() == residual.scalar_type()); + CHECK(residual.is_contiguous()); + CHECK(weight.is_contiguous()); + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + dim3 grid(num_tokens); + /* This kernel is memory-latency bound in many scenarios. + When num_tokens is large, a smaller block size allows + for increased block occupancy on CUs and better latency + hiding on global mem ops. */ + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 block(std::min(hidden_size, max_block_size)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + /*If the tensor types are FP16/BF16, try to use the optimized kernel + with packed + vectorized ops. + Max optimization is achieved with a width-8 vector of FP16/BF16s + since we can load at most 128 bits at once in a global memory op. + However, this requires each tensor's data to be aligned to 16 + bytes. + */ + auto inp_ptr = reinterpret_cast(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(weight.data_ptr()); + constexpr int kVectorWidth = 8; + constexpr int kReqAlignmentBytes = + kVectorWidth * 2; // kVectorWidth * sizeof(bfloat16 or float16) (float32 + // falls back to non-vectorized version anyway) + bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 && + res_ptr % kReqAlignmentBytes == 0 && + wt_ptr % kReqAlignmentBytes == 0; + bool offsets_are_multiple_of_vector_width = + hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) { + LAUNCH_FUSED_ADD_RMS_NORM(8); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0); + } +} + +// ============================================================================ +// Fused RMSNorm + Static FP8 Quantization Host Functions +// ============================================================================ + +void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size], FP8 + torch::Tensor& input, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + torch::Tensor& scale, // [1] + double epsilon) { + CHECK(out.is_contiguous()); + CHECK(input.stride(-1) == 1); + CHECK(weight.is_contiguous()); + CHECK(scale.is_contiguous()); + + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + // For large num_tokens, use smaller blocks to increase SM concurrency + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, max_block_size)); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_static_fp8_quant", [&] { + DISPATCH_FP8_TYPES(out.scalar_type(), "rms_norm_static_fp8_quant", [&] { + rms_norm_static_fp8_quant_kernel + <<>>(out.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + scale.data_ptr(), + epsilon, + num_tokens, + hidden_size); + }); + }); +} + +void fused_add_rms_norm_static_fp8_quant( + torch::Tensor& out, // [..., hidden_size], FP8 + torch::Tensor& input, // [..., hidden_size] + torch::Tensor& residual, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + torch::Tensor& scale, // [1] + double epsilon) { + CHECK(out.is_contiguous()); + CHECK(residual.is_contiguous()); + CHECK(weight.is_contiguous()); + CHECK(scale.is_contiguous()); + CHECK(residual.scalar_type() == input.scalar_type()); + CHECK(weight.scalar_type() == input.scalar_type()); + + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + dim3 grid(num_tokens); + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 block(std::min(hidden_size, max_block_size)); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // Check alignment for vectorized kernel + auto inp_ptr = reinterpret_cast(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(weight.data_ptr()); + constexpr int kVectorWidth = 8; + constexpr int kReqAlignmentBytes = kVectorWidth * 2; + + bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 && + res_ptr % kReqAlignmentBytes == 0 && + wt_ptr % kReqAlignmentBytes == 0; + bool offsets_are_multiple_of_vector_width = + hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0; + + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) { + LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(8); + } else { + LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(0); + } +} + +} // namespace xllm::kernel::cuda diff --git a/qwen3_6_scripts/build_xllm_kernels.sh b/qwen3_6_scripts/build_xllm_kernels.sh new file mode 100644 index 00000000..6b5f8108 --- /dev/null +++ b/qwen3_6_scripts/build_xllm_kernels.sh @@ -0,0 +1,61 @@ +#!/bin/bash +# build_xllm_kernels.sh — Compile xllm CUDA kernels into .so on BI-V100 +# +# Uses corex's own CUB (/usr/local/corex/include/cub/) NOT cccl_upstream +# Source: ex_engine/xllm_kernels/cuda/ +# Output: qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ +# +# Run: bash qwen3_6_scripts/build_xllm_kernels.sh + +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_DIR="$(dirname "$SCRIPT_DIR")" +CUDA_DIR="${PROJECT_DIR}/ex_engine/xllm_kernels/cuda" +HEADER_DIR="${CUDA_DIR}/headers" +PREBUILT_DIR="${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10" +mkdir -p "$PREBUILT_DIR" + +build_kernel() { + local name="$1" + local cu_file="$2" + echo "=== Building ${name}.so ===" + python3 -c " +import os, glob, shutil +from torch.utils.cpp_extension import load + +mod = load( + name='${name}', + sources=['${cu_file}'], + extra_cflags=['-std=c++17'], + extra_include_paths=['${HEADER_DIR}', '/usr/local/corex/include'], + verbose=True, +) + +import torch.utils.cpp_extension as ext +build_dir = ext._get_build_directory('${name}', verbose=False) +for f in glob.glob(os.path.join(build_dir, '*.so')): + dst = '${PREBUILT_DIR}/${name}.so' + shutil.copy2(f, dst) + sz = os.path.getsize(dst) + print(f'✓ ${name}.so ({sz} bytes) → {dst}') + break + +fns = [x for x in dir(mod) if not x.startswith('_')] +print(f'Functions: {fns}') +" +} + +echo "Building xllm CUDA kernels for BI-V100 (ivcore10)" +echo "Using corex CUB: /usr/local/corex/include/cub/" +echo "" + +# Build each kernel +build_kernel "xllm_norm" "${CUDA_DIR}/norm.cu" +build_kernel "xllm_activation" "${CUDA_DIR}/activation.cu" +build_kernel "xllm_rope" "${CUDA_DIR}/rope.cu" +build_kernel "xllm_block_copy" "${CUDA_DIR}/block_copy.cu" +build_kernel "xllm_cache" "${CUDA_DIR}/reshape_paged_cache.cu" + +echo "" +echo "=== All kernels built ===" +ls -la "${PREBUILT_DIR}"/xllm_*.so 2>/dev/null