diff --git a/SO_BUILD_MANIFEST.md b/SO_BUILD_MANIFEST.md new file mode 100644 index 00000000..246a57f8 --- /dev/null +++ b/SO_BUILD_MANIFEST.md @@ -0,0 +1,127 @@ +# 动态链接库完整清单与调用链 + +## 1. 已有预编译 .so（22 个）→ 调用链状态 + +### A. 已接入模型调用链（15 个） + +| .so | 来源 | 模型中的环境变量 | 状态 | +|-----|------|-----------------|------| +| corex_gdn_causal_conv | 自研 CUDA | `BI100_GDN_COREX_CAUSAL_CONV` (default=True) | ✅ 代码引用 4 处 | +| corex_gdn_gated_norm | 自研 CUDA | `BI100_GDN_COREX_GATED_NORM` (default=True) | ✅ 代码引用 4 处 | +| corex_gdn_beta_decay | 自研 CUDA | `BI100_GDN_COREX_BETA_DECAY` (default=True) | ✅ 代码引用 4 处 | +| corex_gdn_qk_map | 自研 CUDA | `BI100_GDN_COREX_QK_MAP` (default=True) | ✅ 代码引用 4 处 | +| corex_gdn_packed_decode | 自研 CUDA | `BI100_GDN_COREX_PACKED_DECODE` (default=False) | ✅ yaml 已开 | +| corex_gdn_chunk_recurrent | 自研 CUDA | 自动检测 | ✅ 代码引用 4 处 | +| corex_attn_head_rms_norm | 自研 CUDA | `BI100_ATTN_COREX_HEAD_RMS_NORM` (default=True) | ✅ 代码引用 5 处 | +| corex_moe_direct_routed | 自研 CUDA | `BI100_MOE_COREX_DIRECT_ROUTED` (default=False) | ✅ yaml 已开 | +| corex_moe_exact_reduce | 自研 CUDA | `BI100_MOE_COREX_EXACT_REDUCE` (default=True) | ✅ 代码引用 4 处 | +| corex_moe_weight_gather | 自研 CUDA | `BI100_MOE_COREX_WEIGHT_GATHER` (default=True) | ✅ 代码引用 4 处 | +| corex_moe_topk_softmax | 自研 CUDA | `BI100_MOE_COREX_TOPK_SOFTMAX` (default=True) | ✅ yaml 已开 | +| corex_moe_index_combine | 自研 CUDA | `BI100_MOE_COREX_INDEX_COMBINE` (default=True) | ✅ 代码引用 4 处 | +| xllm_moe | 搬自 xllm upstream | `BI100_MOE_XLLM` (default=True) | ✅ 代码引用 7 处 | +| xllm_activation | 搬自 xllm upstream | 无直接 env | ❌ 编了但没接入 | +| xllm_norm | 搬自 xllm upstream | 无直接 env | ❌ 编了但没接入 | + +### B. 已编译但未接入（7 个） — 需要修复 + +| .so | 来源 | 提供的函数 | 为什么没接入 | 接入方案 | +|-----|------|-----------|------------|---------| +| **ix_full_bridge** | ix_full_bridge.cpp → ixformer::infer | silu_and_mul, rms_norm, fused_add_rms_norm, ix_linear, ix_linear_ex | qwen3_5.py 没有 import | patch_vllm_ops.py 已写好（最新 commit），通过 ix_startup_patch.py 自动 hook | +| **xllm_activation** | xllm activation.cu | silu_and_mul, gelu_and_mul, act_and_mul | 与 _custom_ops→ixf_F 冗余 | 作为 backup，当 ixf_F 不可用时走 xllm kernel | +| **xllm_norm** | xllm norm.cu | rms_norm, fused_add_rms_norm | 与 _custom_ops→ixf_F 冗余 | 同上 | +| **xllm_rope** | xllm rope.cu | rotary_embedding | 与 _custom_ops→ixf_F 冗余 | 同上 | +| **xllm_cache** | xllm reshape_paged_cache.cu | reshape_paged_cache | paged_attn.py 没有调用 | 需要在 cache 写入路径接入 | +| **corex_fused_paged_prefill** | 自研 CUDA | fused prefill attention | paged_attn.py 有代码但 env 没开 | computility-run.yaml 加 `BI100_ATTN_COREX_FUSED_PAGED_PREFILL=1` | +| **corex_paged_kv_gather** | 自研 CUDA | paged KV gather | paged_attn.py 有代码但 env 没开 | 同上 | +| **corex_block_major_kv_transfer** | 自研 CUDA | block-major KV copy | 完全没有调用点 | 需要在 worker/cache_engine 接入 | + +## 2. 需要从 upstream 搬过来编译的代码 + +### 来源: upstream_ref/xllm/xllm/core/kernels/cuda/ + +| 文件 | 功能 | 对应 .so | 优先级 | +|------|------|---------|--------| +| xattention/decoder_reshape_and_cache.cu | fused KV cache write | xllm_xattn_cache | P0 | +| xattention/prefill_reshape_and_cache.cu | prefill cache write | xllm_xattn_cache | P0 | +| xattention/cache_select.cu | cache select | xllm_xattn_cache | P1 | +| xattention/lse_combine.cu | LSE combine | xllm_xattn_cache | P1 | +| fused_qknorm_rope.cu | fused QK norm + RoPE | xllm_fused_qknorm_rope | P0（每层省 4 kernel launch） | +| matmul.cpp | ixformer GEMM wrapper | 已在 ilu/matmul.cpp | ✅ 已搬 | +| fp8_quant.cu | FP8 quantization | xllm_fp8 | P2 | + +### 来源: upstream_ref/xllm/xllm/core/kernels/ilu/ + +**全部已搬到 ex_engine/xllm_kernels/ilu/**（对比确认只差 CMakeLists.txt） + +### 来源: upstream_ref/ds_vllm/csrc/libtorch_stable/ + +| 文件 | 功能 | 可用性 | +|------|------|--------| +| attention/paged_attention_v1.cu | paged attention v1 | SM70 兼容，但依赖 vllm C++ build | +| attention/paged_attention_v2.cu | paged attention v2 | 同上 | +| layernorm_kernels.cu | RMSNorm kernel | SM70 兼容 | +| activation_kernels.cu | SiLU kernel | SM70 兼容 | +| pos_encoding_kernels.cu | RoPE kernel | SM70 兼容 | +| moe/topk_softmax_kernels.cu | topk+softmax fused | SM70 兼容 | +| moe/moe_align_sum_kernels.cu | MoE align+sum | SM70 兼容 | + +## 3. ixformer::infer 可用 API（base 镜像已有） + +来自 `upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h`: + +``` +ixformer::infer::silu_and_mul(input, output) +ixformer::infer::rms_norm(input, weight, output, bias, eps) +ixformer::infer::residual_rms_norm(input, residual, weight, output, residual_out, bias, alpha, eps, is_post) +ixformer::infer::ixformer_linear(input, weight, act_type, bias, out, persistent) +ixformer::infer::ixformer_linear_ex(input, weight, bias, out) +ixformer::infer::xllm_rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox) +ixformer::infer::xllm_reshape_and_cache(key, value, key_cache, value_cache, slot_mapping, key_stride, value_stride) +ixformer::infer::xllm_paged_attention(out, query, key_cache, value_cache, ...) +ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(query, key_cache, value_cache, ...) +ixformer::infer::topk_softmax(weights, indices, token_expert_indices, gating_output, renormalize) +ixformer::infer::moe_compute_token_index_api(topk_ids, src_dst, dst_src, expert_sizes, ...) +ixformer::infer::moe_expand_input(output, input, dst_to_src, src_to_dst, dst_tokens, expand_factor) +ixformer::infer::moe_w16a16_group_gemm(output, input, weights, tokens_per_experts, ...) +ixformer::infer::moe_output_reduce_sum(output, input, weight, mask, extra_residual, scaling) +``` + +这些函数通过 `ix_full_bridge.so` pybind11 暴露给 Python 侧。 + +## 4. 调用链完整性检查 + +### 当前断裂点: + +1. **ix_full_bridge.so 的 group_gemm → MoE Python for-loop** + - `ixformer::infer::moe_w16a16_group_gemm` 在 ix_full_bridge.so 中可用 + - 但 qwen3_5.py MoE prefill 路径 (L1813-1825) 还是 `F.linear` per-expert loop + - 需要: ix_fused_moe.py 的 7 步 pipeline 走 group_gemm 而非 per-expert linear + +2. **corex_fused_paged_prefill → paged_attn.py env 没开** + - .so 已编译已部署 + - paged_attn.py 已有完整调用代码 (L2030) + - computility-run.yaml 缺少 `BI100_ATTN_COREX_FUSED_PAGED_PREFILL=1` + +3. **xllm_cache → reshape_and_cache 没接入** + - base 镜像 ixformer 已有 `xllm_reshape_and_cache` + - vllm 的 cache_ops 走的是另一条路径 + +## 5. 需要编出的新 .so + +| 目标 .so | 源文件 | 编译方式 | 依赖 | +|---------|--------|---------|------| +| xllm_fused_qknorm_rope.so | upstream fused_qknorm_rope.cu + bind | corex clang --cuda-gpu-arch=ivcore10 | libcudart, torch | +| xllm_xattn_cache.so | upstream xattention/*.cu + bind | 同上 | 同上 | + +## 6. computility-run.yaml 需要补全的 env + +```yaml +- name: BI100_ATTN_COREX_FUSED_PAGED_PREFILL + value: '1' +- name: BI100_ATTN_COREX_PAGED_KV_GATHER + value: '1' +- name: IX_OPS_AUTO_PATCH + value: '1' +- name: PYTORCH_CUDA_ALLOC_CONF + value: 'expandable_segments:True' +``` diff --git a/ex_engine/build_xllm_kernels.sh b/ex_engine/build_xllm_kernels.sh new file mode 100755 index 00000000..78b94825 --- /dev/null +++ b/ex_engine/build_xllm_kernels.sh @@ -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)" diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp new file mode 100644 index 00000000..53978f32 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp @@ -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 + +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")); +} diff --git a/ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu b/ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu new file mode 100644 index 00000000..37fad188 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu @@ -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 +#include +#include +#include + +#include +#include + +#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 +struct packed_as; +// Specialization for packed_as used in this kernel. +template <> +struct packed_as { + using type = uint; +}; + +template <> +struct packed_as { + using type = uint2; +}; + +template <> +struct packed_as { + using type = uint4; +}; + +template +__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 +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 +__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) || + std::is_same_v) { + return; + } else { +#endif + + using Converter = _typeConvert; + 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; + 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(qkv_void); + T_in const* q_weight = reinterpret_cast(q_weight_void); + T_in const* k_weight = reinterpret_cast(k_weight_void); + T_cache const* cos_sin_cache = + reinterpret_cast(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 to perform loading/saving. + using vec_T = typename packed_as::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(&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(&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(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(&vec) + i) = packed_val; + } + *reinterpret_cast(&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 +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 + <<>>(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 + <<>>(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 + <<>>(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.data_ptr(), + static_cast(num_tokens), + static_cast(num_heads_q), + static_cast(num_heads_k), + static_cast(num_heads_v), + static_cast(head_dim), + static_cast(cos_sin_cache.size(1)), + static_cast(eps), + q_weight.data_ptr(), + k_weight.data_ptr(), + cos_sin_cache.data_ptr(), + interleaved, + reinterpret_cast(position_ids.data_ptr()), + stream); + }); + }); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h b/ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h new file mode 100644 index 00000000..269540e0 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/corex_compat_utils.h @@ -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 +#include + +// 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 +using Array = std::vector; +} + +// HOST_DEVICE_INLINE +#if defined(__CUDACC__) || defined(_NVHPC_CUDA) +#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__ +#else +#define HOST_DEVICE_INLINE inline +#endif diff --git a/ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh b/ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh new file mode 100644 index 00000000..5bd3b96e --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/topk_last_dim.cuh @@ -0,0 +1,2114 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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. + */ + +// refers to +// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/topkLastDim.cu +// refers to +// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/topkLastDim.h + +/** + * This file contains a specialized implementation of AIR TopK + * introduced in https://dl.acm.org/doi/pdf/10.1145/3581784.3607062 . + * Another variant can be found in TopP sampling: + * cpp/tensorrt_llm/kernels/samplingAirTopPKernels.cu . + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "moe/moe_topk.cuh" +#include "platform/device.h" +// #include "topk_last_dim.h" + +using SizeType32 = int32_t; + +namespace xllm::kernel::cuda { + +namespace reduce_topk { + +/////////////// + +// AIR TopK Kernel + +#if 1 + +namespace air_topk_stable { +using WideT = float4; +constexpr int VECTORIZED_READ_SIZE = 16; +constexpr int WARP_SIZE = 32; + +// constexpr unsigned FULL_WARP_MASK = 0xffffffff; + +template +struct ComputeOffset { + __host__ __device__ explicit ComputeOffset(IdxT const& cols) : cols_(cols) {} + + __host__ __device__ IdxT operator()(IdxT const& x) const { return cols_ * x; } + + IdxT cols_; +}; + +template +__host__ __device__ constexpr int calc_num_buckets() { + return 1 << BitsPerPass; +} + +/** + * @brief Provide a ceiling division operation ie. ceil(a / b) + * @tparam IntType supposed to be only integers for now! + */ +template +constexpr __host__ __device__ IntType ceildiv(IntType a, IntType b) { + return (a + b - 1) / b; +} + +/** + * @brief Provide an alignment function ie. ceil(a / b) * b + * @tparam IntType supposed to be only integers for now! + */ +template +constexpr __host__ __device__ IntType alignTo(IntType a, IntType b) { + return ceildiv(a, b) * b; +} + +template +__host__ __device__ constexpr int calc_num_passes() { + return ceildiv(sizeof(T) * 8, BitsPerPass); +} + +__host__ __device__ __forceinline__ int round(int num, int round_value) { + return ((num - 1) / round_value + 1) * round_value; +} + +/** + * Bit 0 is the least significant (rightmost); + * this implementation processes input from the most to the least significant + * bit. This way, we can skip some passes in the end at the cost of having an + * unsorted output. + * + * NB: Use pass=-1 for calc_mask(). + */ +template +__device__ constexpr int calc_start_bit(int pass) { + int start_bit = static_cast(sizeof(T) * 8) - (pass + 1) * BitsPerPass; + if (start_bit < 0) { + start_bit = 0; + } + return start_bit; +} + +template +__device__ constexpr unsigned calc_mask(int pass) { + static_assert(BitsPerPass <= 31); + int num_bits = calc_start_bit(pass - 1) - + calc_start_bit(pass); + return (1 << num_bits) - 1; +} + +/** + * Use CUB to twiddle bits - so that we can correctly compare bits of + * floating-point values as well as of integers. + */ +template +__device__ typename cub::Traits::UnsignedBits twiddle_in(T key, + bool select_min) { + auto bits = reinterpret_cast::UnsignedBits&>(key); + bits = cub::Traits::TwiddleIn(bits); + if (!select_min) { + bits = ~bits; + } + return bits; +} + +template +__device__ T twiddle_out(typename cub::Traits::UnsignedBits bits, + bool select_min) { + if (!select_min) { + bits = ~bits; + } + bits = cub::Traits::TwiddleOut(bits); + return reinterpret_cast(bits); +} + +template +__device__ int calc_bucket(T x, int start_bit, unsigned mask, bool select_min) { + static_assert( + BitsPerPass <= sizeof(int) * 8 - 1, + "BitsPerPass is too large that the result type could not be int"); + return (twiddle_in(x, select_min) >> start_bit) & mask; +} + +template +constexpr inline std::enable_if_t::value, bool> +is_a_power_of_two(I val) noexcept { + return ((val - 1) & val) == 0; +} + +template +__host__ __device__ IdxT calc_buf_len(IdxT len) { + // When writing is skipped, only read `in`(type T). + // When writing is not skipped, read `in_buf`(T) and `in_idx_buf`(IdxT), and + // write `out_buf`(T) and `out_idx_buf`(IdxT). The ratio between these cases + // determines whether to skip writing and hence the buffer size. + constexpr RATIO_T ratio = 2 + sizeof(IdxT) * 2 / sizeof(T); + // Even such estimation is too conservative, so further decrease buf_len by + // 1/8 + IdxT buf_len = len / (ratio * 8); + + // one-block kernel splits one large buffer into smaller ones, so round buf + // size to 256 bytes to avoid alignment issues + static_assert(is_a_power_of_two(sizeof(T))); + static_assert(is_a_power_of_two(sizeof(IdxT))); + constexpr IdxT aligned = 256 / std::min(sizeof(T), sizeof(IdxT)); + buf_len = buf_len & (~(aligned - 1)); + return buf_len; +} + +/** + * Map a Func over the input data, using vectorized load instructions if + * possible. + * + * NB: in future, we should move this to + * cpp/include/raft/linalg/detail/unary_op.cuh, which currently does not support + * the second lambda argument (index of an element) + * + * @tparam T element type + * @tparam IdxT indexing type + * @tparam Func void (T x, IdxT idx) + * + * @param thread_rank rank of the calling thread among all participating threads + * @param num_threads number of the threads that participate in processing + * @param in the input data + * @param len the number of elements to read + * @param f the lambda taking two arguments (T x, IdxT idx) + */ +template +__device__ void vectorized_process(size_t thread_rank, + size_t num_threads, + T const* in, + IdxT len, + Func f) { + if constexpr (sizeof(T) >= sizeof(WideT)) { + for (IdxT i = thread_rank; i < len; i += num_threads) { + f(in[i], i); + } + } else { + static_assert(sizeof(WideT) % sizeof(T) == 0); + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); + + // TODO: it's UB + union { + WideT scalar; + T array[items_per_scalar]; + } wide; + + int skip_cnt = + (reinterpret_cast(in) % sizeof(WideT)) + ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / + sizeof(T)) + : 0; + if (skip_cnt > len) { + skip_cnt = len; + } + WideT const* in_cast = reinterpret_cast(in + skip_cnt); + const IdxT len_cast = (len - skip_cnt) / items_per_scalar; + + for (IdxT i = thread_rank; i < len_cast; i += num_threads) { + wide.scalar = in_cast[i]; + const IdxT real_i = skip_cnt + i * items_per_scalar; +#pragma unroll + for (int j = 0; j < items_per_scalar; ++j) { + f(wide.array[j], real_i + j); + } + } + + static_assert(WARP_SIZE >= items_per_scalar); + // and because items_per_scalar > skip_cnt, WARP_SIZE > skip_cnt + // no need to use loop + if (thread_rank < skip_cnt) { + f(in[thread_rank], thread_rank); + } + // because len_cast = (len - skip_cnt) / items_per_scalar, + // len_cast * items_per_scalar + items_per_scalar > len - skip_cnt; + // and so + // len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <= + // WARP_SIZE no need to use loop + const IdxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank; + if (remain_i < len) { + f(in[remain_i], remain_i); + } + } +} + +// sync_width should >= WARP_SIZE +template +__device__ void vectorized_process(T const* in, + IdxT len, + Func f, + int sync_width) { + const IdxT stride = blockDim.x * gridDim.x; + const IdxT tid = blockIdx.x * blockDim.x + threadIdx.x; + if constexpr (sizeof(T) >= sizeof(WideT)) { + for (IdxT i = tid; i < len; i += stride) { + f(in[i], i, true); + } + } else { + static_assert(sizeof(WideT) % sizeof(T) == 0); + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); + + union { + WideT scalar; + T array[items_per_scalar]; + } wide; + + int skip_cnt = + (reinterpret_cast(in) % sizeof(WideT)) + ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / + sizeof(T)) + : 0; + if (skip_cnt > len) { + skip_cnt = len; + } + WideT const* in_cast = reinterpret_cast(in + skip_cnt); + const IdxT len_cast = (len - skip_cnt) / items_per_scalar; + + const IdxT len_cast_for_sync = + ((len_cast - 1) / sync_width + 1) * sync_width; + for (IdxT i = tid; i < len_cast_for_sync; i += stride) { + bool valid = i < len_cast; + if (valid) { + wide.scalar = in_cast[i]; + } + const IdxT real_i = skip_cnt + i * items_per_scalar; +#pragma unroll + for (int j = 0; j < items_per_scalar; ++j) { + f(wide.array[j], real_i + j, valid); + } + } + + static_assert(WARP_SIZE >= items_per_scalar); + // need at most one warp for skipped and remained elements, + // and sync_width >= WARP_SIZE + if (tid < sync_width) { + bool valid = tid < skip_cnt; + T value = valid ? in[tid] : T(); + f(value, tid, valid); + + const IdxT remain_i = skip_cnt + len_cast * items_per_scalar + tid; + valid = remain_i < len; + value = valid ? in[remain_i] : T(); + f(value, remain_i, valid); + } + } +} + +template +struct alignas(128) Counter { + // We are processing the values in multiple passes, from most significant to + // least significant. In each pass, we keep the length of input (`len`) and + // the `k` of current pass, and update them at the end of the pass. + IdxT k; + IdxT len; + + // `previous_len` is the length of input in previous pass. Note that + // `previous_len` rather than `len` is used for the filtering step because + // filtering is indeed for previous pass (see comments before + // `radix_kernel`). + IdxT previous_len; + + // We determine the bits of the k_th value inside the mask processed by the + // pass. The already known bits are stored in `kth_value_bits`. It's used to + // discriminate a element is a result (written to `out`), a candidate for next + // pass (written to `out_buf`), or not useful (discarded). The bits that are + // not yet processed do not matter for this purpose. + typename cub::Traits::UnsignedBits kth_value_bits; + + // Record how many elements have passed filtering. It's used to determine the + // position in the `out_buf` where an element should be written. + alignas(128) IdxT filter_cnt; + + // For a row inside a batch, we may launch multiple thread blocks. This + // counter is used to determine if the current block is the last running + // block. If so, this block will execute scan() and choose_bucket(). + alignas(128) unsigned int finished_block_cnt; + + // Record how many elements have been written to the front of `out`. Elements + // less (if select_min==true) than the k-th value are written from front to + // back. + alignas(128) IdxT out_cnt; + + // Record how many elements have been written to the back of `out`. Elements + // equal to the k-th value are written from back to front. We need to keep + // count of them separately because the number of elements that <= the k-th + // value might exceed k. + alignas(128) IdxT out_back_cnt; +}; + +/** + * Fused filtering of the current pass and building histogram for the next pass + * (see steps 4 & 1 in `radix_kernel` description). + */ +template +__device__ void filter_and_histogram(T const* in_buf, + IdxT const* in_idx_buf, + T* out_buf, + IdxT* out_idx_buf, + T* out, + IdxT* out_idx, + IdxT previous_len, + Counter* counter, + IdxT* histogram, + bool select_min, + int pass, + bool early_stop) { + constexpr int num_buckets = calc_num_buckets(); + __shared__ IdxT histogram_smem[num_buckets]; + for (IdxT i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram_smem[i] = 0; + } + __syncthreads(); + + int const start_bit = calc_start_bit(pass); + unsigned const mask = calc_mask(pass); + + if (pass == 0) { + // Passed to vectorized_process, this function executes in all blocks in + // parallel, i.e. the work is split along the input (both, in batches and + // chunks of a single row). Later, the histograms are merged using + // atomicAdd. + auto f = [select_min, start_bit, mask](T value, IdxT) { + int bucket = + calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram_smem + bucket, static_cast(1)); + }; + vectorized_process( + static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, + in_buf, + previous_len, + f); + } else { + IdxT* p_filter_cnt = &counter->filter_cnt; + IdxT* p_out_cnt = &counter->out_cnt; + auto const kth_value_bits = counter->kth_value_bits; + int const previous_start_bit = calc_start_bit(pass - 1); + + // See the remark above on the distributed execution of `f` using + // vectorized_process. + auto f = [in_idx_buf, + out_buf, + out_idx_buf, + out, + out_idx, + select_min, + start_bit, + mask, + previous_start_bit, + kth_value_bits, + p_filter_cnt, + p_out_cnt, + early_stop](T value, IdxT i) { + const auto previous_bits = + (twiddle_in(value, select_min) >> previous_start_bit) + << previous_start_bit; + if (previous_bits == kth_value_bits) { + if (early_stop) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + out[pos] = value; + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else { + if (out_buf) { + IdxT pos = atomicAdd(p_filter_cnt, static_cast(1)); + out_buf[pos] = value; + out_idx_buf[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + + int bucket = + calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram_smem + bucket, static_cast(1)); + } + } + // the condition `(out_buf || early_stop)` is a little tricky: + // If we skip writing to `out_buf` (when `out_buf` is nullptr), we should + // skip writing to `out` too. So we won't write the same value to `out` + // multiple times in different passes. And if we keep skipping the + // writing, values will be written in `last_filter_kernel()` at last. But + // when `early_stop` is true, we need to write to `out` since it's the + // last chance. + else if ((out_buf || early_stop) && previous_bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + out[pos] = value; + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + }; + vectorized_process( + static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, + in_buf, + previous_len, + f); + } + if (early_stop) { + return; + } + __syncthreads(); + + // merge histograms produced by individual blocks + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + if (histogram_smem[i] != 0) { + atomicAdd(histogram + i, histogram_smem[i]); + } + } +} + +/** + * Replace histogram with its own prefix sum + * (step 2 in `radix_kernel` description) + */ +template +__device__ void scan(IdxT volatile* histogram) { + constexpr int num_buckets = calc_num_buckets(); + if constexpr (num_buckets >= BlockSize) { + static_assert(num_buckets % BlockSize == 0); + constexpr int items_per_thread = num_buckets / BlockSize; + typedef cub:: + BlockLoad + BlockLoad; + typedef cub::BlockStore + BlockStore; + typedef cub::BlockScan BlockScan; + + __shared__ union { + typename BlockLoad::TempStorage load; + typename BlockScan::TempStorage scan; + typename BlockStore::TempStorage store; + } temp_storage; + + IdxT thread_data[items_per_thread]; + + BlockLoad(temp_storage.load).Load(histogram, thread_data); + __syncthreads(); + + BlockScan(temp_storage.scan).InclusiveSum(thread_data, thread_data); + __syncthreads(); + + BlockStore(temp_storage.store).Store(histogram, thread_data); + } else { + typedef cub::BlockScan BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + IdxT thread_data = 0; + if (threadIdx.x < num_buckets) { + thread_data = histogram[threadIdx.x]; + } + + BlockScan(temp_storage).InclusiveSum(thread_data, thread_data); + __syncthreads(); + + if (threadIdx.x < num_buckets) { + histogram[threadIdx.x] = thread_data; + } + } +} + +/** + * Calculate in which bucket the k-th value will fall + * (steps 3 in `radix_kernel` description) + */ +template +__device__ void choose_bucket(Counter* counter, + IdxT const* histogram, + const IdxT k, + int const pass) { + constexpr int num_buckets = calc_num_buckets(); + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + IdxT prev = (i == 0) ? 0 : histogram[i - 1]; + IdxT cur = histogram[i]; + + // one and only one thread will satisfy this condition, so counter is + // written by only one thread + if (prev < k && cur >= k) { + counter->k = k - prev; // how many values still are there to find + counter->len = cur - prev; // number of values in next pass + typename cub::Traits::UnsignedBits bucket = i; + int start_bit = calc_start_bit(pass); + counter->kth_value_bits |= bucket << start_bit; + } + } +} + +// For one-block version, last_filter() could be called when pass < num_passes +// - 1. So `pass` could not be constexpr +template +__device__ void last_filter(T const* in_buf, + IdxT const* in_idx_buf, + T* out, + IdxT* out_idx, + IdxT current_len, + IdxT k, + Counter* counter, + bool const select_min, + int const pass) { + auto const kth_value_bits = counter->kth_value_bits; + int const start_bit = calc_start_bit(pass); + + // changed in choose_bucket(); need to reload + const IdxT num_of_kth_needed = counter->k; + IdxT* p_out_cnt = &counter->out_cnt; + IdxT* p_out_back_cnt = &counter->out_back_cnt; + IdxT* p_equal = out_idx + k - num_of_kth_needed; + ::cuda::atomic_ref ref_last( + p_equal[num_of_kth_needed - 1]); + for (IdxT i = threadIdx.x; i < current_len; i += blockDim.x) { + const T value = in_buf[i]; + auto const bits = (twiddle_in(value, select_min) >> start_bit) << start_bit; + if (bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + out[pos] = value; + // For one-block version, `in_idx_buf` could be nullptr at pass 0. + // For non one-block version, if writing has been skipped, `in_idx_buf` + // could be nullptr if `in_buf` is `in` + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else if (bits == kth_value_bits) { + IdxT new_idx = in_idx_buf ? in_idx_buf[i] : i; + IdxT back_pos = atomicAdd(p_out_back_cnt, static_cast(1)); + if (back_pos < num_of_kth_needed) { + IdxT pos = k - 1 - back_pos; + out[pos] = value; + if constexpr (!prioritize_smaller_indice) { + out_idx[pos] = new_idx; + } + } + if constexpr (prioritize_smaller_indice) { + if (new_idx < ref_last.load(::cuda::memory_order_relaxed)) { + for (int j = 0; j < num_of_kth_needed; j++) { + IdxT pre_idx = atomicMin(&p_equal[j], new_idx); + if (pre_idx > new_idx) { + new_idx = pre_idx; + } + } + } + } + } + } +} + +template +__global__ void last_filter_kernel(T const* in, + IdxT const* in_idx, + T const* in_buf, + IdxT const* in_idx_buf, + T* out, + IdxT* out_idx, + IdxT len, + IdxT k, + Counter* counters, + bool const select_min) { + const size_t batch_id = + blockIdx.y; // size_t to avoid multiplication overflow + + Counter* counter = counters + batch_id; + IdxT previous_len = counter->previous_len; + if (previous_len == 0) { + return; + } + const IdxT buf_len = calc_buf_len(len); + if (previous_len > buf_len || in_buf == in) { + in_buf = in + batch_id * len; + in_idx_buf = in_idx ? (in_idx + batch_id * len) : nullptr; + previous_len = len; + } else { + in_buf += batch_id * buf_len; + in_idx_buf += batch_id * buf_len; + } + out += batch_id * k; + out_idx += batch_id * k; + + constexpr int pass = calc_num_passes() - 1; + constexpr int start_bit = calc_start_bit(pass); + + auto const kth_value_bits = counter->kth_value_bits; + const IdxT num_of_kth_needed = counter->k; + IdxT* p_out_cnt = &counter->out_cnt; + IdxT* p_out_back_cnt = &counter->out_back_cnt; + IdxT* p_equal = out_idx + k - num_of_kth_needed; + ::cuda::atomic_ref ref_last(p_equal[num_of_kth_needed - 1]); + auto f = [k, + select_min, + kth_value_bits, + num_of_kth_needed, + p_out_cnt, + p_out_back_cnt, + in_idx_buf, + out, + out_idx, + p_equal, + ref_last](T value, IdxT i) { + const auto bits = (twiddle_in(value, select_min) >> start_bit) << start_bit; + if (bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + out[pos] = value; + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else if (bits == kth_value_bits) { + IdxT new_idx = in_idx_buf ? in_idx_buf[i] : i; + IdxT back_pos = atomicAdd(p_out_back_cnt, static_cast(1)); + if (back_pos < num_of_kth_needed) { + IdxT pos = k - 1 - back_pos; + out[pos] = value; + if constexpr (!prioritize_smaller_indice) { + out_idx[pos] = new_idx; + } + } + if constexpr (prioritize_smaller_indice) { + if (new_idx < ref_last.load(::cuda::memory_order_relaxed)) { + for (int j = 0; j < num_of_kth_needed; j++) { + IdxT pre_idx = atomicMin(&p_equal[j], new_idx); + if (pre_idx > new_idx) { + new_idx = pre_idx; + } + } + } + } + } + }; + + vectorized_process(static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, + in_buf, + previous_len, + f); +} + +/** + * + * It is expected to call this kernel multiple times (passes), in each pass we + * process a radix, going from the most significant towards the least + * significant bits (MSD). + * + * Conceptually, each pass consists of 4 steps: + * + * 1. Calculate histogram + * First, transform bits into a digit, the value of which is in the range + * [0, 2^{BITS_PER_PASS}-1]. Then count the frequency of each digit value + * and the result is a histogram. That is, histogram[i] contains the count of + * inputs having value i. + * + * 2. Scan the histogram + * Inclusive prefix sum is computed for the histogram. After this step, + * histogram[i] contains the count of inputs having value <= i. + * + * 3. Find the bucket j of the histogram that the k-th value falls into + * + * 4. Filtering + * Input elements whose digit value +__global__ void radix_kernel(T const* in, + IdxT const* in_idx, + T const* in_buf, + IdxT const* in_idx_buf, + T* out_buf, + IdxT* out_idx_buf, + T* out, + IdxT* out_idx, + Counter* counters, + IdxT* histograms, + const IdxT len, + const IdxT k, + bool const select_min, + int const pass) { + const size_t batch_id = blockIdx.y; + auto counter = counters + batch_id; + IdxT current_k; + IdxT previous_len; + IdxT current_len; + if (pass == 0) { + current_k = k; + previous_len = len; + // Need to do this so setting counter->previous_len for the next pass is + // correct. This value is meaningless for pass 0, but it's fine because pass + // 0 won't be the last pass in this implementation so pass 0 won't hit the + // "if (pass == num_passes - 1)" branch. Maybe it's better to reload + // counter->previous_len and use it rather than current_len in last_filter() + current_len = len; + } else { + current_k = counter->k; + current_len = counter->len; + previous_len = counter->previous_len; + } + if (current_len == 0) { + return; + } + + // When k=len, early_stop will be true at pass 0. It means + // filter_and_histogram() should handle correctly the case that pass=0 and + // early_stop=true. However, this special case of k=len is handled in other + // way in select_k() so such case is not possible here. + bool const early_stop = (current_len == current_k); + const IdxT buf_len = calc_buf_len(len); + + // "previous_len > buf_len" means previous pass skips writing buffer + if (pass == 0 || pass == 1 || previous_len > buf_len) { + in_buf = in + batch_id * len; + in_idx_buf = in_idx ? (in_idx + batch_id * len) : nullptr; + previous_len = len; + } else { + in_buf += batch_id * buf_len; + in_idx_buf += batch_id * buf_len; + } + // "current_len > buf_len" means current pass will skip writing buffer + if (pass == 0 || current_len > buf_len) { + out_buf = nullptr; + out_idx_buf = nullptr; + } else { + out_buf += batch_id * buf_len; + out_idx_buf += batch_id * buf_len; + } + out += batch_id * k; + out_idx += batch_id * k; + + constexpr int num_buckets = calc_num_buckets(); + auto histogram = histograms + batch_id * num_buckets; + + filter_and_histogram(in_buf, + in_idx_buf, + out_buf, + out_idx_buf, + out, + out_idx, + previous_len, + counter, + histogram, + select_min, + pass, + early_stop); + __threadfence(); + + bool isLastBlock = false; + if (threadIdx.x == 0) { + unsigned int finished = + atomicInc(&counter->finished_block_cnt, gridDim.x - 1); + isLastBlock = (finished == (gridDim.x - 1)); + } + + if (__syncthreads_or(isLastBlock)) { + if (early_stop) { + if (threadIdx.x == 0) { + // `last_filter_kernel()` requires setting previous_len + counter->previous_len = 0; + counter->len = 0; + } + return; + } + + scan(histogram); + __syncthreads(); + choose_bucket(counter, histogram, current_k, pass); + __syncthreads(); + + constexpr int num_passes = calc_num_passes(); + // reset for next pass + if (pass != num_passes - 1) { + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram[i] = 0; + } + } + if (threadIdx.x == 0) { + // `last_filter_kernel()` requires setting previous_len even in the last + // pass + counter->previous_len = current_len; + // not necessary for the last pass, but put it here anyway + counter->filter_cnt = 0; + } + + // if constexpr (fused_last_filter) { + // if (pass == num_passes - 1) { + // last_filter(out_buf ? out_buf : in_buf, + // out_idx_buf ? out_idx_buf : + // in_idx_buf, out, out_idx, + // out_buf ? current_len : len, k, + // counter, select_min, pass); + // } + // } + if (pass == num_passes - 1) { + const volatile IdxT num_of_kth_needed = counter->k; + for (IdxT i = threadIdx.x; i < num_of_kth_needed; i += blockDim.x) { + out_idx[k - num_of_kth_needed + i] = + ::cuda::std::numeric_limits::max(); + } + __syncthreads(); + if constexpr (fused_last_filter) { + last_filter( + out_buf ? out_buf : in_buf, + out_idx_buf ? out_idx_buf : in_idx_buf, + out, + out_idx, + out_buf ? current_len : len, + k, + counter, + select_min, + pass); + } + } + } +} + +template +unsigned calc_grid_dim(int batch_size, IdxT len, int sm_cnt) { + static_assert(VECTORIZED_READ_SIZE / sizeof(T) >= 1); + + int active_blocks; + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &active_blocks, + radix_kernel, + BlockSize, + 0); + active_blocks *= sm_cnt; + + IdxT best_num_blocks = 0; + float best_tail_wave_penalty = 1.0f; + const IdxT max_num_blocks = + ceildiv(len, VECTORIZED_READ_SIZE / sizeof(T) * BlockSize); + for (int num_waves = 1;; ++num_waves) { + IdxT num_blocks = std::min( + max_num_blocks, + static_cast(std::max(num_waves * active_blocks / batch_size, 1))); + IdxT items_per_thread = ceildiv(len, num_blocks * BlockSize); + items_per_thread = + alignTo(items_per_thread, VECTORIZED_READ_SIZE / sizeof(T)); + num_blocks = ceildiv(len, items_per_thread * BlockSize); + float actual_num_waves = + static_cast(num_blocks) * batch_size / active_blocks; + float tail_wave_penalty = + (ceilf(actual_num_waves) - actual_num_waves) / ceilf(actual_num_waves); + + // 0.15 is determined experimentally. It also ensures breaking the loop + // early, e.g. when num_waves > 7, tail_wave_penalty will always <0.15 + if (tail_wave_penalty < 0.15) { + best_num_blocks = num_blocks; + break; + } else if (tail_wave_penalty < best_tail_wave_penalty) { + best_num_blocks = num_blocks; + best_tail_wave_penalty = tail_wave_penalty; + } + + if (num_blocks == max_num_blocks) { + break; + } + } + return best_num_blocks; +} + +template +__host__ __device__ void set_buf_pointers(T const* in, + IdxT const* in_idx, + T* buf1, + IdxT* idx_buf1, + T* buf2, + IdxT* idx_buf2, + int pass, + T const*& in_buf, + IdxT const*& in_idx_buf, + T*& out_buf, + IdxT*& out_idx_buf) { + if (pass == 0) { + in_buf = in; + in_idx_buf = nullptr; + out_buf = nullptr; + out_idx_buf = nullptr; + } else if (pass == 1) { + in_buf = in; + in_idx_buf = in_idx; + out_buf = buf1; + out_idx_buf = idx_buf1; + } else if (pass % 2 == 0) { + in_buf = buf1; + in_idx_buf = idx_buf1; + out_buf = buf2; + out_idx_buf = idx_buf2; + } else { + in_buf = buf2; + in_idx_buf = idx_buf2; + out_buf = buf1; + out_idx_buf = idx_buf1; + } +} + +template +__device__ void set_buf_pointers(T const* in, + IdxT const* in_idx, + char* bufs, + IdxT buf_len, + int pass, + T const*& in_buf, + IdxT const*& in_idx_buf, + T*& out_buf, + IdxT*& out_idx_buf) { + // bufs consists of 4 pieces in order: buf1, buf2, idx_buf1, idx_buf2 + if (pass == 0) { + in_buf = in; + in_idx_buf = nullptr; + out_buf = nullptr; + out_idx_buf = nullptr; + } else if (pass == 1) { + in_buf = in; + in_idx_buf = in_idx; + out_buf = reinterpret_cast(bufs); + out_idx_buf = reinterpret_cast(bufs + sizeof(T) * 2 * buf_len); + } else if (pass % 2 == 0) { + in_buf = reinterpret_cast(bufs); + in_idx_buf = reinterpret_cast(bufs + sizeof(T) * 2 * buf_len); + out_buf = const_cast(in_buf + buf_len); + out_idx_buf = const_cast(in_idx_buf + buf_len); + } else { + out_buf = reinterpret_cast(bufs); + out_idx_buf = reinterpret_cast(bufs + sizeof(T) * 2 * buf_len); + in_buf = out_buf + buf_len; + in_idx_buf = out_idx_buf + buf_len; + } +} + +// The following a few functions are for the one-block version, which uses +// single thread block for each row of a batch. +template +__device__ void filter_and_histogram_for_one_block(T const* in_buf, + IdxT const* in_idx_buf, + T* out_buf, + IdxT* out_idx_buf, + T* out, + IdxT* out_idx, + const IdxT previous_len, + Counter* counter, + IdxT* histogram, + bool select_min, + int pass) { + constexpr int num_buckets = calc_num_buckets(); + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram[i] = 0; + } + IdxT* p_filter_cnt = &counter->filter_cnt; + if (threadIdx.x == 0) { + *p_filter_cnt = 0; + } + __syncthreads(); + + int const start_bit = calc_start_bit(pass); + unsigned const mask = calc_mask(pass); + + if (pass == 0) { + auto f = [histogram, select_min, start_bit, mask](T value, IdxT) { + int bucket = + calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + }; + vectorized_process(threadIdx.x, blockDim.x, in_buf, previous_len, f); + } else if (!out_buf) { + // not use vectorized_process here because it increases #registers a lot + auto const kth_value_bits = counter->kth_value_bits; + int const previous_start_bit = calc_start_bit(pass - 1); + + for (IdxT i = threadIdx.x; i < previous_len; i += blockDim.x) { + const T value = in_buf[i]; + auto const previous_bits = + (twiddle_in(value, select_min) >> previous_start_bit) + << previous_start_bit; + if (previous_bits == kth_value_bits) { + int bucket = + calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + } + } + } else { + // not use vectorized_process here because it increases #registers a lot + IdxT* p_out_cnt = &counter->out_cnt; + auto const kth_value_bits = counter->kth_value_bits; + int const previous_start_bit = calc_start_bit(pass - 1); + + for (IdxT i = threadIdx.x; i < previous_len; i += blockDim.x) { + const T value = in_buf[i]; + auto const previous_bits = + (twiddle_in(value, select_min) >> previous_start_bit) + << previous_start_bit; + if (previous_bits == kth_value_bits) { +#if CUDART_VERSION < 12000 + // Avoiding potential compiler bug in CUDA 11 + volatile +#endif + IdxT pos = atomicAdd(p_filter_cnt, static_cast(1)); + out_buf[pos] = value; + out_idx_buf[pos] = in_idx_buf ? in_idx_buf[i] : i; + + int bucket = + calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + } else if (previous_bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + out[pos] = value; + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + } + } +} + +template +__global__ void radix_topk_one_block_kernel(T const* in, + IdxT const* in_idx, + const IdxT len, + const IdxT k, + T* out, + IdxT* out_idx, + bool const select_min, + char* bufs) { + constexpr int num_buckets = calc_num_buckets(); + __shared__ Counter counter; + __shared__ IdxT histogram[num_buckets]; + + if (threadIdx.x == 0) { + counter.k = k; + counter.len = len; + counter.previous_len = len; + counter.kth_value_bits = 0; + counter.out_cnt = 0; + counter.out_back_cnt = 0; + } + __syncthreads(); + + const size_t batch_id = + blockIdx.x; // size_t to avoid multiplication overflow + in += batch_id * len; + if (in_idx) { + in_idx += batch_id * len; + } + + out += batch_id * k; + out_idx += batch_id * k; + const IdxT buf_len = calc_buf_len(len); + bufs += batch_id * buf_len * 2 * (sizeof(T) + sizeof(IdxT)); + + constexpr int num_passes = calc_num_passes(); + for (int pass = 0; pass < num_passes; ++pass) { + T const* in_buf = nullptr; + IdxT const* in_idx_buf = nullptr; + T* out_buf = nullptr; + IdxT* out_idx_buf = nullptr; + set_buf_pointers(in, + in_idx, + bufs, + buf_len, + pass, + in_buf, + in_idx_buf, + out_buf, + out_idx_buf); + + const IdxT current_len = counter.len; + const IdxT current_k = counter.k; + IdxT previous_len = counter.previous_len; + if (previous_len > buf_len) { + in_buf = in; + in_idx_buf = in_idx; + previous_len = len; + } + if (current_len > buf_len) { + // so "out_buf==nullptr" denotes skipping writing buffer in current pass + out_buf = nullptr; + out_idx_buf = nullptr; + } + + filter_and_histogram_for_one_block( + in_buf, + in_idx_buf, + out_buf, + out_idx_buf, + out, + out_idx, + previous_len, + &counter, + histogram, + select_min, + pass); //@TODO CHECK UPDATE CODE + __syncthreads(); + + scan(histogram); + __syncthreads(); + + choose_bucket(&counter, histogram, current_k, pass); + if (threadIdx.x == 0) { + counter.previous_len = current_len; + } + __syncthreads(); + + if ((pass == num_passes - 1)) { + if constexpr (prioritize_smaller_indice) { + const IdxT num_of_kth_needed = counter.k; + for (IdxT i = threadIdx.x; i < num_of_kth_needed; i += blockDim.x) { + out_idx[k - num_of_kth_needed + i] = + ::cuda::std::numeric_limits::max(); + } + __syncthreads(); + } + last_filter( + out_buf ? out_buf : in, + out_buf ? out_idx_buf : in_idx, + out, + out_idx, + out_buf ? current_len : len, + k, + &counter, + select_min, + pass); + break; + } else if (counter.len == counter.k) { + last_filter(out_buf ? out_buf : in, + out_buf ? out_idx_buf : in_idx, + out, + out_idx, + out_buf ? current_len : len, + k, + &counter, + select_min, + pass); + break; + } + } +} +} // namespace air_topk_stable + +//} +namespace moe_topk { +namespace cg = cooperative_groups; +static constexpr int kBLOCK_SIZE = 1024; +static constexpr int kWARP_SIZE = 32; +static constexpr int kWARPS_PER_BLOCK = kBLOCK_SIZE / kWARP_SIZE; + +template +__device__ __forceinline__ T negativeInfinity() { + return -INFINITY; +} + +template <> +__device__ __forceinline__ half negativeInfinity() { + return -CUDART_INF_FP16; +} + +template <> +__device__ __forceinline__ __nv_bfloat16 negativeInfinity<__nv_bfloat16>() { + return -CUDART_INF_BF16; +} + +/****************TopK kernel for candidate number<= 128 and K <= 8 + * **************** */ +template +__global__ void moe_topk_kernel(InputT const* in, + OutputT* out, + IdxT* outIdx, + int32_t const batchSize, + int32_t const len, + int32_t const topK) { + uint32_t const blockRank = blockIdx.x; + uint32_t const tIdx = kBLOCK_SIZE * blockRank + threadIdx.x; + uint32_t const warpIdx = tIdx / kWARP_SIZE; + uint32_t const laneIdx = tIdx % kWARP_SIZE; + uint32_t const warpNum = gridDim.x * kWARPS_PER_BLOCK; + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + InputT minScore = negativeInfinity(); + + for (uint32_t tokenId = warpIdx; tokenId < batchSize; tokenId += warpNum) { + auto scoreOffset = tokenId * len; + auto outputOffset = tokenId * topK; + InputT inputScore[MaxLen / kWARP_SIZE]; + IdxT inputIndex[MaxLen / kWARP_SIZE]; + + InputT warpTopKScore[MaxTopK]; + IdxT warpTopKExpertIdx[MaxTopK]; + + // Load scores and indices for this warp + for (uint32_t i = 0; i < MaxLen / kWARP_SIZE; ++i) { + auto expertIdx = i * kWARP_SIZE + laneIdx; + inputScore[i] = expertIdx < len + ? static_cast(in[scoreOffset + expertIdx]) + : minScore; + inputIndex[i] = expertIdx; + } + + // Reduce topK scores and indices for this warp + reduce_topk::reduceTopK(warp, + warpTopKScore, + warpTopKExpertIdx, + inputScore, + inputIndex, + minScore); + + if (laneIdx < topK) { + out[outputOffset + laneIdx] = + static_cast(warpTopKScore[laneIdx]); + outIdx[outputOffset + laneIdx] = warpTopKExpertIdx[laneIdx]; + } + } // end for tokenId +} +} // namespace moe_topk + +/***************Runtime API****************/ + +inline size_t calc_aligned_size(std::vector const& sizes) { + const size_t ALIGN_BYTES = 256; + const size_t ALIGN_MASK = ~(ALIGN_BYTES - 1); + size_t total = 0; + for (auto sz : sizes) { + total += (sz + ALIGN_BYTES - 1) & ALIGN_MASK; + } + return total + ALIGN_BYTES - 1; +} + +inline std::vector calc_aligned_pointers( + void const* p, + std::vector const& sizes) { + const size_t ALIGN_BYTES = 256; + const size_t ALIGN_MASK = ~(ALIGN_BYTES - 1); + + char* ptr = reinterpret_cast( + (reinterpret_cast(p) + ALIGN_BYTES - 1) & ALIGN_MASK); + + std::vector aligned_pointers; + aligned_pointers.reserve(sizes.size()); + for (auto sz : sizes) { + aligned_pointers.push_back(ptr); + ptr += (sz + ALIGN_BYTES - 1) & ALIGN_MASK; + } + + return aligned_pointers; +} + +template +void standalone_stable_radix_topk_(void* buf, + size_t& buf_size, + T const* in, + IdxT const* in_idx, + int batch_size, + IdxT len, + IdxT k, + T* out, + IdxT* out_idx, + bool select_min, + bool fused_last_filter, + unsigned grid_dim, + cudaStream_t stream, + bool sorted = false) { + static_assert(air_topk_stable::calc_num_passes() > 1); + constexpr int num_buckets = air_topk_stable::calc_num_buckets(); + + air_topk_stable::Counter* counters = nullptr; + IdxT* histograms = nullptr; + T* buf1 = nullptr; + IdxT* idx_buf1 = nullptr; + T* buf2 = nullptr; + IdxT* idx_buf2 = nullptr; + + void* sort_temp_storage = nullptr; + size_t temp_storage_bytes = 0; + size_t temp_storage_bytes_sort = 0; + T* topk_out = nullptr; + IdxT* topk_out_idx = nullptr; + T* sort_in = nullptr; + IdxT* sort_in_idx = nullptr; + + air_topk_stable::ComputeOffset computeoffset(k); + + thrust::counting_iterator counting_iter(0); + thrust::transform_iterator, + thrust::counting_iterator> + transform_iter(counting_iter, computeoffset); + + cub::DeviceSegmentedSort::SortPairs(NULL, + temp_storage_bytes, + out_idx, + out_idx, + out, + out, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + if (sorted) { + if (select_min) { + cub::DeviceSegmentedSort::StableSortPairs(NULL, + temp_storage_bytes_sort, + out, + out, + out_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } else { + cub::DeviceSegmentedSort::StableSortPairsDescending( + NULL, + temp_storage_bytes_sort, + out, + out, + out_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } + } + temp_storage_bytes = max(temp_storage_bytes, temp_storage_bytes_sort); + + { + IdxT len_candidates = air_topk_stable::calc_buf_len(len); + size_t sort_buffer_size = 0; + if (sorted) { + sort_buffer_size = k * batch_size; + } + std::vector sizes = { + sizeof(*counters) * batch_size, + sizeof(*histograms) * num_buckets * batch_size, + sizeof(*buf1) * len_candidates * batch_size, + sizeof(*idx_buf1) * len_candidates * batch_size, + sizeof(*buf2) * len_candidates * batch_size, + sizeof(*idx_buf2) * len_candidates * batch_size, + temp_storage_bytes, + sizeof(*topk_out) * k * batch_size, + sizeof(*topk_out_idx) * k * batch_size, + sizeof(*sort_in) * sort_buffer_size, + sizeof(*sort_in_idx) * sort_buffer_size}; + size_t total_size = calc_aligned_size(sizes); + if (!buf) { + buf_size = total_size; + return; + } + + std::vector aligned_pointers = calc_aligned_pointers(buf, sizes); + counters = static_cast(aligned_pointers[0]); + histograms = static_cast(aligned_pointers[1]); + buf1 = static_cast(aligned_pointers[2]); + idx_buf1 = static_cast(aligned_pointers[3]); + buf2 = static_cast(aligned_pointers[4]); + idx_buf2 = static_cast(aligned_pointers[5]); + sort_temp_storage = aligned_pointers[6]; + topk_out = static_cast(aligned_pointers[7]); + topk_out_idx = static_cast(aligned_pointers[8]); + if (sorted) { + sort_in = static_cast(aligned_pointers[9]); + sort_in_idx = static_cast(aligned_pointers[10]); + } + cudaMemsetAsync(aligned_pointers[0], + 0, + static_cast(aligned_pointers[2]) - + static_cast(aligned_pointers[0]), + stream); + } + + T const* in_buf = nullptr; + IdxT const* in_idx_buf = nullptr; + T* out_buf = nullptr; + IdxT* out_idx_buf = nullptr; + + dim3 blocks(grid_dim, batch_size); + + constexpr int num_passes = air_topk_stable::calc_num_passes(); + + auto kernel = air_topk_stable:: + radix_kernel; + + for (int pass = 0; pass < num_passes; ++pass) { + air_topk_stable::set_buf_pointers(in, + in_idx, + buf1, + idx_buf1, + buf2, + idx_buf2, + pass, + in_buf, + in_idx_buf, + out_buf, + out_idx_buf); + + if (fused_last_filter && pass == num_passes - 1) { + kernel = air_topk_stable:: + radix_kernel; + } + + kernel<<>>(in, + in_idx, + in_buf, + in_idx_buf, + out_buf, + out_idx_buf, + topk_out, + topk_out_idx, + counters, + histograms, + len, + k, + select_min, + pass); + } + + if (!fused_last_filter) { + air_topk_stable::last_filter_kernel + <<>>(in, + in_idx, + out_buf, + out_idx_buf, + topk_out, + topk_out_idx, + len, + k, + counters, + select_min); + } + + T* idx_sort_out = sorted ? sort_in : out; + IdxT* idx_sort_out_idx = sorted ? sort_in_idx : out_idx; + + cub::DeviceSegmentedSort::SortPairs(sort_temp_storage, + temp_storage_bytes, + topk_out_idx, + idx_sort_out_idx, + topk_out, + idx_sort_out, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + if (sorted) { + if (select_min) { + cub::DeviceSegmentedSort::StableSortPairs(sort_temp_storage, + temp_storage_bytes, + sort_in, + out, + sort_in_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } else { + cub::DeviceSegmentedSort::StableSortPairsDescending(sort_temp_storage, + temp_storage_bytes, + sort_in, + out, + sort_in_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } + } +} + +template +void standalone_stable_radix_topk_one_block_(void* buf, + size_t& buf_size, + T const* in, + IdxT const* in_idx, + int batch_size, + IdxT len, + IdxT k, + T* out, + IdxT* out_idx, + bool select_min, + cudaStream_t stream, + bool sorted = false) { + static_assert(air_topk_stable::calc_num_passes() > 1); + + char* bufs = nullptr; + void* sort_temp_storage = nullptr; + T* topk_out = nullptr; + IdxT* topk_out_idx = nullptr; + T* sort_in = nullptr; + IdxT* sort_in_idx = nullptr; + + size_t temp_storage_bytes = 0; + size_t temp_storage_bytes_sort = 0; + const IdxT buf_len = air_topk_stable::calc_buf_len(len); + + air_topk_stable::ComputeOffset computeoffset(k); + thrust::counting_iterator counting_iter(0); + thrust::transform_iterator, + thrust::counting_iterator> + transform_iter(counting_iter, computeoffset); + + cub::DeviceSegmentedSort::SortPairs(NULL, + temp_storage_bytes, + out_idx, + out_idx, + out, + out, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + if (sorted) { + if (select_min) { + cub::DeviceSegmentedSort::StableSortPairs(NULL, + temp_storage_bytes_sort, + out, + out, + out_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } else { + cub::DeviceSegmentedSort::StableSortPairsDescending( + NULL, + temp_storage_bytes_sort, + out, + out, + out_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } + } + + temp_storage_bytes = max(temp_storage_bytes, temp_storage_bytes_sort); + { + size_t total_size = 0; + size_t sort_buffer_size = 0; + if (sorted) { + sort_buffer_size = k * batch_size; + } + std::vector sizes = { + buf_len * 2 * (sizeof(T) + sizeof(IdxT)) * batch_size, + temp_storage_bytes, + sizeof(*topk_out) * k * batch_size, + sizeof(*topk_out_idx) * k * batch_size, + sizeof(*sort_in) * sort_buffer_size, + sizeof(*sort_in_idx) * sort_buffer_size}; + total_size = calc_aligned_size(sizes); + + if (!buf) { + buf_size = total_size; + return; + } + + std::vector aligned_pointers = calc_aligned_pointers(buf, sizes); + bufs = static_cast(aligned_pointers[0]); + sort_temp_storage = aligned_pointers[1]; + topk_out = static_cast(aligned_pointers[2]); + topk_out_idx = static_cast(aligned_pointers[3]); + if (sorted) { + sort_in = static_cast(aligned_pointers[4]); + sort_in_idx = static_cast(aligned_pointers[5]); + } + } + + air_topk_stable:: + radix_topk_one_block_kernel + <<>>( + in, in_idx, len, k, topk_out, topk_out_idx, select_min, bufs); + + T* idx_sort_out = sorted ? sort_in : out; + IdxT* idx_sort_out_idx = sorted ? sort_in_idx : out_idx; + cub::DeviceSegmentedSort::SortPairs(sort_temp_storage, + temp_storage_bytes, + topk_out_idx, + idx_sort_out_idx, + topk_out, + idx_sort_out, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + if (sorted) { + if (select_min) { + cub::DeviceSegmentedSort::StableSortPairs(sort_temp_storage, + temp_storage_bytes, + sort_in, + out, + sort_in_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } else { + cub::DeviceSegmentedSort::StableSortPairsDescending(sort_temp_storage, + temp_storage_bytes, + sort_in, + out, + sort_in_idx, + out_idx, + k * batch_size, + batch_size, + transform_iter, + transform_iter + 1, + stream); + } + } +} + +template +void standalone_stable_radix_11bits(void* buf, + size_t& buf_size, + T const* in, + int batch_size, + IdxT len, + IdxT k, + T* out, + IdxT* out_idx, + bool greater, + cudaStream_t stream = 0) { + constexpr int items_per_thread = 32; + constexpr int block_dim = 512; + constexpr bool fused_last_filter = false; + if (len <= block_dim * items_per_thread) { + standalone_stable_radix_topk_one_block_( + buf, + buf_size, + in, + static_cast(nullptr), + batch_size, + len, + k, + out, + out_idx, + !greater, + stream, + sorted); + } else { + int32_t sm_cnt = xllm::Device::sm_count(); + unsigned grid_dim = air_topk_stable::calc_grid_dim( + batch_size, len, sm_cnt); + + if (grid_dim == 1) { + standalone_stable_radix_topk_one_block_( + buf, + buf_size, + in, + static_cast(nullptr), + batch_size, + len, + k, + out, + out_idx, + !greater, + stream, + sorted); + } else { + standalone_stable_radix_topk_( + buf, + buf_size, + in, + static_cast(nullptr), + batch_size, + len, + k, + out, + out_idx, + !greater, + fused_last_filter, + grid_dim, + stream, + sorted); + } + } +} + +inline int nextPowerOfTwo(int num) { + if (num <= 0) { + return 1; // Handle invalid input + } + int power = 1; + while (power < num) { + // Check for overflow before shifting + if (power > INT_MAX / 2) { + return power; + } + power <<= 1; + } + return power; +} + +template +void moe_reduce_topk(T const* in, + int batch_size, + IdxT len, + IdxT k, + T* out, + IdxT* out_idx, + bool greater, + cudaStream_t stream = 0) { + using InputT = T; + using OutputT = T; + const uint32_t max_num_blocks = 1024; + const uint32_t num_blocks = std::min( + static_cast((batch_size - 1) / moe_topk::kWARPS_PER_BLOCK + 1), + max_num_blocks); + + uint32_t max_len = nextPowerOfTwo(len) < 32 ? 32 : nextPowerOfTwo(len); + uint32_t moe_topk = nextPowerOfTwo(k); + + auto* kernel_instance = + &moe_topk::moe_topk_kernel; + + switch (max_len) { + case 32: + switch (moe_topk) { + case 1: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 2: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 4: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 8: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + default: + kernel_instance = nullptr; + break; + } + break; + case 64: + switch (moe_topk) { + case 1: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 2: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 4: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 8: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + default: + kernel_instance = nullptr; + break; + } + break; + case 96: + switch (moe_topk) { + case 1: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 2: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 4: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 8: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + default: + kernel_instance = nullptr; + break; + } + break; + case 128: + switch (moe_topk) { + case 1: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 2: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 4: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + case 8: + kernel_instance = + &moe_topk::moe_topk_kernel; + break; + default: + kernel_instance = nullptr; + break; + } + break; + default: + kernel_instance = nullptr; + break; + } + + dim3 moe_topk_grid_dim(num_blocks); + dim3 moe_topk_block_dim(moe_topk::kBLOCK_SIZE); + + kernel_instance<<>>( + in, out, out_idx, batch_size, len, k); +} +#endif + +/////////////// + +template +size_t invokeComputeTopkLastDimWorkspaceSize(SizeType32 batchSize, + SizeType32 inputLength, + SizeType32 k, + bool is_largest, + bool sorted) { + using IdxT = SizeType32; + + size_t buf_size = 0; + void* workspace = nullptr; + T const* in = nullptr; + T* out_val = nullptr; + IdxT* out_idx = nullptr; + + constexpr int block_dim = 512; + constexpr bool fused_last_filter = false; + int32_t sm_cnt = xllm::Device::sm_count(); + unsigned grid_dim = air_topk_stable::calc_grid_dim( + batchSize, inputLength, sm_cnt); + + if (sorted) { + standalone_stable_radix_topk_( + workspace, + buf_size, + in, + static_cast(nullptr), + batchSize, + inputLength, + k, + out_val, + out_idx, + !is_largest, + fused_last_filter, + grid_dim, + 0, + true); + } else { + standalone_stable_radix_topk_( + workspace, + buf_size, + in, + static_cast(nullptr), + batchSize, + inputLength, + k, + out_val, + out_idx, + !is_largest, + fused_last_filter, + grid_dim, + 0, + false); + } + return buf_size; +} + +template +size_t invokeComputeTopkLastDimWorkspaceSize(SizeType32 batchSize, + SizeType32 inputLength, + SizeType32 k, + bool is_largest) { + return invokeComputeTopkLastDimWorkspaceSize( + batchSize, inputLength, k, is_largest, true); +} + +#define INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE(T) \ + template size_t invokeComputeTopkLastDimWorkspaceSize( \ + SizeType32 batchSize, \ + SizeType32 inputLength, \ + SizeType32 k, \ + bool is_largest) + +INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE(int); +INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE(float); +INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE(half); +#ifdef ENABLE_BF16 +INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE(__nv_bfloat16); +#endif +#undef INSTANTIATE_COMPUTE_TOPK_LastDim_WORKSPACE_SIZE_DATA_TYPE + +// Might need FP8 in the future. + +/////////////// + +template +void invokeTopkLastDim(SizeType32 batchSize, + SizeType32 inputLength, + SizeType32 k, + bool is_largest, + void const* __restrict__ input, + void* __restrict__ out_val, + void* __restrict__ out_idx, + void* workspace, + cudaStream_t stream, + bool sorted) { + size_t buf_size = 0; // will be overwritten by the kernel + T const* in = reinterpret_cast(input); + T* out_val_ = reinterpret_cast(out_val); + SizeType32* out_idx_ = reinterpret_cast(out_idx); + if (inputLength <= 128 && k <= 8 && is_largest == true) { + // This method does not require a buffer, but since the implementation may + // vary in different cases, we still allocate the buffer in case AIR TopK is + // used instead. + moe_reduce_topk( + in, batchSize, inputLength, k, out_val_, out_idx_, !is_largest, stream); + } else { + if (sorted) { + standalone_stable_radix_11bits(workspace, + buf_size, + in, + batchSize, + inputLength, + k, + out_val_, + out_idx_, + is_largest, + stream); + } else { + standalone_stable_radix_11bits(workspace, + buf_size, + in, + batchSize, + inputLength, + k, + out_val_, + out_idx_, + is_largest, + stream); + } + } +} + +template +void invokeTopkLastDim(SizeType32 batchSize, + SizeType32 inputLength, + SizeType32 k, + bool is_largest, + void const* __restrict__ input, + void* __restrict__ out_val, + void* __restrict__ out_idx, + void* workspace, + cudaStream_t stream) { + invokeTopkLastDim(batchSize, + inputLength, + k, + is_largest, + input, + out_val, + out_idx, + workspace, + stream, + true); +} + +#define INSTANTIATE_TOPK_LastDim_DATA_TYPE(T) \ + template void invokeTopkLastDim(SizeType32 batchSize, \ + SizeType32 inputLength, \ + SizeType32 k, \ + bool is_largest, \ + void const* __restrict__ input, \ + void* __restrict__ out_val, \ + void* __restrict__ out_idx, \ + void* workspace, \ + cudaStream_t stream) + +INSTANTIATE_TOPK_LastDim_DATA_TYPE(int); +INSTANTIATE_TOPK_LastDim_DATA_TYPE(float); +INSTANTIATE_TOPK_LastDim_DATA_TYPE(half); +#ifdef ENABLE_BF16 +INSTANTIATE_TOPK_LastDim_DATA_TYPE(__nv_bfloat16); +#endif +#undef INSTANTIATE_TOPK_LastDim_DATA_TYPE + +} // namespace reduce_topk + +} // namespace xllm::kernel::cuda \ No newline at end of file diff --git a/ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp b/ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp new file mode 100644 index 00000000..f2ac239c --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/beam_search.cpp @@ -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 +#include +#include +#include + +#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(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 \ No newline at end of file diff --git a/ex_engine/xllm_kernels/cuda/xattention/cache_select.cu b/ex_engine/xllm_kernels/cuda/xattention/cache_select.cu new file mode 100644 index 00000000..db273ef8 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/cache_select.cu @@ -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 +#include +#include +#include +#include +#include + +#include +#include + +#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 +__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(blockIdx.x); + const int32_t kv = static_cast(blockIdx.y); + const int32_t layer = static_cast(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(static_cast(k_ptrs_i64[layer])); + scalar_t* __restrict__ v_cache = + reinterpret_cast(static_cast(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(req) * Beam; + const int64_t step_kv_stride = static_cast(Kv) * D; + const int64_t kv_d_base = static_cast(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(threadIdx.x); d < D; + d += static_cast(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(threadIdx.x); d < D; + d += static_cast(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(B64); + const int32_t Beam = static_cast(Beam64); + const int32_t Kv = static_cast(Kv64); + const int32_t MaxStep = static_cast(MaxStep64); + const int32_t D = static_cast(D64); + const int32_t MaxReq = static_cast(MaxReq64); + const int32_t Layer = static_cast(Layer64); + const int32_t decode_step_i32 = static_cast(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(threads_per_block), 1, 1); + + CHECK_LE(Kv64, static_cast(UINT32_MAX)) << "Kv too large for grid.y"; + CHECK_LE(Layer64, 65535) << "layer_num too large for grid.z"; + dim3 grid_dim(static_cast(B), + static_cast(Kv), + static_cast(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 + <<>>( + k_ptrs_i64.data_ptr(), + v_ptrs_i64.data_ptr(), + beam_index_i32.data_ptr(), + block_table_i32.data_ptr(), + 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& unshared_k_cache, + std::vector& 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(unshared_k_cache.size()), layer_num) + << "unshared_k_cache length mismatch"; + CHECK_EQ(static_cast(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 k_ptrs_host(static_cast(layer_num)); + std::vector v_ptrs_host(static_cast(layer_num)); + + for (int64_t layer = 0; layer < layer_num; ++layer) { + auto k = unshared_k_cache[static_cast(layer)]; + auto v = unshared_v_cache[static_cast(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(layer)] = + static_cast(reinterpret_cast(k.data_ptr())); + v_ptrs_host[static_cast(layer)] = + static_cast(reinterpret_cast(v.data_ptr())); + } + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + C10_CUDA_CHECK( + cudaMemcpyAsync(k_ptrs_i64.data_ptr(), + k_ptrs_host.data(), + static_cast(layer_num) * sizeof(int64_t), + cudaMemcpyHostToDevice, + stream)); + C10_CUDA_CHECK( + cudaMemcpyAsync(v_ptrs_i64.data_ptr(), + v_ptrs_host.data(), + static_cast(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 diff --git a/ex_engine/xllm_kernels/cuda/xattention/decoder_reshape_and_cache.cu b/ex_engine/xllm_kernels/cuda/xattention/decoder_reshape_and_cache.cu new file mode 100644 index 00000000..2d2d0c09 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/decoder_reshape_and_cache.cu @@ -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 +#include +#include + +#include +#include + +#include "kernels/cuda/utils.h" +#include "xattention_ops_api.h" + +namespace { + +template +struct VecType; + +template <> +struct VecType { + using type = uint4; // 8 elements * 2 bytes = 16 bytes + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + using type = uint4; // 8 elements * 2 bytes = 16 bytes + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + 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 +__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::type; + constexpr int32_t VEC_WIDTH = VecType::vec_width; + + const int64_t token_idx = static_cast(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(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(threadIdx.x); + linear_idx < total_vecs; + linear_idx += static_cast(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( + proj_k + k_token_base + head_idx * head_dim + vec_offset); + const auto* v_src_vec = reinterpret_cast( + proj_v + v_token_base + head_idx * head_dim + vec_offset); + auto* k_dst_vec = + reinterpret_cast(unshared_k_cache + dst_token_base + + head_idx * cache_stride3 + vec_offset); + auto* v_dst_vec = + reinterpret_cast(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(total_tokens), 1); + + DISPATCH_FLOATING_TYPES( + proj_k.scalar_type(), "decoder_reshape_and_cache_kernel", [&] { + constexpr int32_t VEC_WIDTH = (std::is_same_v || + std::is_same_v) + ? 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( + 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(proj_k.data_ptr()); + const auto proj_v_ptr = + reinterpret_cast(proj_v.data_ptr()); + const auto k_cache_ptr = reinterpret_cast( + unshared_k_cache.data_ptr()); + const auto v_cache_ptr = reinterpret_cast( + unshared_v_cache.data_ptr()); + 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(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 + <<>>( + proj_k.data_ptr(), + proj_v.data_ptr(), + unshared_k_cache.data_ptr(), + unshared_v_cache.data_ptr(), + step.data_ptr(), + 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 diff --git a/ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu b/ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu new file mode 100644 index 00000000..572601f2 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu @@ -0,0 +1,168 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include +#include +#include + +#include + +#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 +__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(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(shared_o[base_idx + d]); + const float unshared_val = static_cast(unshared_o[base_idx + d]); + const float combined = w_shared * shared_val + w_unshared * unshared_val; + output[base_idx + d] = static_cast(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(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 + <<>>( + output.data_ptr(), + shared_o.data_ptr(), + shared_lse.data_ptr(), + unshared_o.data_ptr(), + unshared_lse.data_ptr(), + B, + H, + D); + }); + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace xllm::kernel::cuda \ No newline at end of file diff --git a/ex_engine/xllm_kernels/cuda/xattention/prefill_reshape_and_cache.cu b/ex_engine/xllm_kernels/cuda/xattention/prefill_reshape_and_cache.cu new file mode 100644 index 00000000..07fb0609 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/prefill_reshape_and_cache.cu @@ -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 +#include +#include +#include + +#include +#include + +#include "kernels/cuda/cuda_ops_api.h" +#include "kernels/cuda/utils.h" +using at::device_of; + +namespace { + +template +struct VecType; + +template <> +struct VecType { + using type = uint4; // 8 elements * 2 bytes = 16 bytes + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + using type = uint4; // 8 elements * 2 bytes = 16 bytes + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + using type = float4; // 4 elements * 4 bytes = 16 bytes + static constexpr int32_t vec_width = 4; +}; + +template +__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::type; + constexpr int32_t VEC_WIDTH = VecType::vec_width; + const int64_t token_idx = static_cast(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( + proj_k + k_token_base + head_offset + vec_offset); + const auto* v_src_vec = reinterpret_cast( + proj_v + v_token_base + head_idx * v_stride1 + vec_offset); + auto* k_dst_vec = reinterpret_cast( + shared_k_cache + dst_token_base + head_offset + vec_offset); + auto* v_dst_vec = reinterpret_cast( + 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(shared_len), 1); + + DISPATCH_FLOATING_TYPES( + proj_k.scalar_type(), "prefill_reshape_and_cache_kernel", [&] { + constexpr int32_t VEC_WIDTH = (std::is_same_v || + std::is_same_v) + ? 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( + 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(proj_k.data_ptr()); + const auto proj_v_ptr = + reinterpret_cast(proj_v.data_ptr()); + const auto k_cache_ptr = reinterpret_cast( + shared_k_cache.data_ptr()); + const auto v_cache_ptr = reinterpret_cast( + shared_v_cache.data_ptr()); + + 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(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 + <<>>( + proj_k.data_ptr(), + proj_v.data_ptr(), + shared_k_cache.data_ptr(), + shared_v_cache.data_ptr(), + shared_len, + kv_heads, + head_dim, + k_stride0, + v_stride0, + v_stride1); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h b/ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h new file mode 100644 index 00000000..5b84f495 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/xattention/xattention_ops_api.h @@ -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 + +#include + +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& unshared_k_cache, + std::vector& 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 diff --git a/ex_engine/xllm_kernels/ilu/CMakeLists.txt b/ex_engine/xllm_kernels/ilu/CMakeLists.txt new file mode 100644 index 00000000..fa26c886 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/CMakeLists.txt @@ -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 +) diff --git a/upstream_ref/sgemm_cuda/01_naive.cuh b/upstream_ref/sgemm_cuda/01_naive.cuh new file mode 100644 index 00000000..04ed8e8c --- /dev/null +++ b/upstream_ref/sgemm_cuda/01_naive.cuh @@ -0,0 +1,37 @@ +# pragma once + +#include +#include +#include +#include + +/* + +Matrix sizes: +MxK * KxN = MxN + +*/ + +__global__ void sgemm_naive( + int M, int N, int K, // sizes + float alpha, const float *A, const float *B, float beta, float *C // pointers used to point to matrices +) { + // compute position in C that this thread is responsible for + // "which block" * "width of block" to get to start of block + "which thread" + const uint x = blockIdx.x * blockDim.x + threadIdx.x; // "which row?" (inverted from graphical intuition, confusingly) + const uint y = blockIdx.y * blockDim.y + threadIdx.y; // "which column?" + + // if M or N are not multiples of 32, there will be "extra"/"remainder" threads on the last block in x/y. + // we don't want those leftover threads to do anything (tile quantisation) + if (x < M && y < N) { + float tmp = 0.0; + for (int i = 0; i < K; ++i) { // K is the size of the row in A, col in B i.e. the dot product + // A: x * K gives the start of relevant row, i enumerates across the row (col by col) + // B: y gives the relevant column, i * N enumerates down the column, (row by row) + tmp += A[x * K + i] * B[i * N + y]; + } + // C = alpha*(A@B) + beta*C + // x * N takes to start of relevant row, y moves across to the relevant column + C[x * N + y] = alpha * tmp + beta * C[x * N + y]; + } +} diff --git a/upstream_ref/sgemm_cuda/02_kernel_global_mem_coalesce.cuh b/upstream_ref/sgemm_cuda/02_kernel_global_mem_coalesce.cuh new file mode 100644 index 00000000..2f6f4504 --- /dev/null +++ b/upstream_ref/sgemm_cuda/02_kernel_global_mem_coalesce.cuh @@ -0,0 +1,25 @@ +#pragma once + +#include +#include +#include +#include +#include + +template +// __global__ is used to specify that the function is run on GPU, called by host (CPU) +__global__ void sgemm_global_mem_coalesce(int M, int N, int K, float alpha, + const float *A, const float *B, + float beta, float *C) { + const int cRow = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE); // note that blockDim is now 1-dimensional + const int cCol = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE); + + // if statement is necessary to make things work under tile quantization + if (cRow < M && cCol < N) { + float tmp = 0.0; + for (int i = 0; i < K; ++i) { + tmp += A[cRow * K + i] * B[i * N + cCol]; + } + C[cRow * N + cCol] = alpha * tmp + beta * C[cRow * N + cCol]; + } +} diff --git a/upstream_ref/sgemm_cuda/03_kernel_shared_mem_blocking.cuh b/upstream_ref/sgemm_cuda/03_kernel_shared_mem_blocking.cuh new file mode 100644 index 00000000..975cc47f --- /dev/null +++ b/upstream_ref/sgemm_cuda/03_kernel_shared_mem_blocking.cuh @@ -0,0 +1,54 @@ +#pragma once + +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void sgemm_shared_mem_block(int M, int N, int K, float alpha, + const float *A, const float *B, + float beta, float *C) { + // output C block we want to compute with this threadBlock + const uint cRow = blockIdx.x; + const uint cCol = blockIdx.y; + + // allocate buffer for current block in fast SMEM (shared between all threads in block) + __shared__ float As[BLOCKSIZE * BLOCKSIZE]; + __shared__ float Bs[BLOCKSIZE * BLOCKSIZE]; + + // the inner row and col that we are accessing in this specific thread + const uint threadRow = threadIdx.x / BLOCKSIZE; // note similarity to previous kernel + const uint threadCol = threadIdx.x % BLOCKSIZE; + + // advance pointers to the starting positions (they are input as pointers to first elements in the matrices) + A += cRow * BLOCKSIZE * K; // row=cRow, col=0 (the start of the relevant row) + B += cCol * BLOCKSIZE; // row=0, col=cCol (top of relevant col) + C += cRow * BLOCKSIZE * N + cCol * BLOCKSIZE; // row=cRow, col=cCol + + float tmp = 0.0; + for (int bkIdx=0; bkIdx < K; bkIdx+=BLOCKSIZE) { // shifting the whole block along the row of A and col of B + // have each thread load one of the elements in A and B + // make the threadCol (=threadIdx.x) the consecutive index + // to allow GMEM access coalescing + As[threadRow * BLOCKSIZE + threadCol] = A[threadRow * K + threadCol]; + Bs[threadRow * BLOCKSIZE + threadCol] = B[threadRow * N + threadCol]; + + // ensure cache is fully populated + __syncthreads(); + A += BLOCKSIZE; // for next iteration + B += BLOCKSIZE * N; + + // execute the dotproduct on the currently cached block + for (int dotIdx = 0; dotIdx < BLOCKSIZE; ++dotIdx) { + tmp += As[threadRow * BLOCKSIZE + dotIdx] * + Bs[dotIdx * BLOCKSIZE + threadCol]; + } + // sync so faster threads don't fetch the next block into cache + _syncthreads(); + } + C[threadRow * N + threadCol] = alpha * tmp + beta * C[threadRow * N + threadCol]; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/04_kernel_1D_blocktiling.cuh b/upstream_ref/sgemm_cuda/04_kernel_1D_blocktiling.cuh new file mode 100644 index 00000000..12c9c5d7 --- /dev/null +++ b/upstream_ref/sgemm_cuda/04_kernel_1D_blocktiling.cuh @@ -0,0 +1,80 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void sgemm1DBlocktiling(int M, int N, int K, float alpha, + const float *A, const float *B, float beta, + float *C) { + // If we flip x and y here we get ~30% less performance for large matrices. + // The current, 30% faster configuration ensures that blocks with sequential + // blockIDs access columns of B sequentially, while sharing the same row of A. + // The slower configuration would share columns of A, but access into B would + // be non-sequential. So the faster configuration has better spatial locality + // and hence a greater L2 hit rate. + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // each warp will calculate 32*TM elements, with 32 being the columnar dim. + const int threadCol = threadIdx.x % BN; + const int threadRow = threadIdx.x / BN; + + // allocate space for the current blocktile in SMEM + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // todo: adjust this to each thread to load multiple entries and + // better exploit the cache sizes + assert(BM * BK == blockDim.x); + assert(BN * BK == blockDim.x); + const uint innerColA = threadIdx.x % BK; // warp-level GMEM coalescing + const uint innerRowA = threadIdx.x / BK; + const uint innerColB = threadIdx.x % BN; // warp-level GMEM coalescing + const uint innerRowB = threadIdx.x / BN; + + // allocate thread-local cache for results in registerfile + float threadResults[TM] = {0.0}; + + // outer loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + As[innerRowA * BK + innerColA] = A[innerRowA * K + innerColA]; + Bs[innerRowB * BN + innerColB] = B[innerRowB * N + innerColB]; + __syncthreads(); + + // advance blocktile + A += BK; + B += BK * N; + + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // we make the dotproduct loop the outside loop, which facilitates + // reuse of the Bs entry, which we can cache in a tmp var. + float tmpB = Bs[dotIdx * BN + threadCol]; + for (uint resIdx = 0; resIdx < TM; ++resIdx) { + threadResults[resIdx] += + As[(threadRow * TM + resIdx) * BK + dotIdx] * tmpB; + } + } + __syncthreads(); + } + + // write out the results + for (uint resIdx = 0; resIdx < TM; ++resIdx) { + C[(threadRow * TM + resIdx) * N + threadCol] = + alpha * threadResults[resIdx] + + beta * C[(threadRow * TM + resIdx) * N + threadCol]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/05_kernel_2D_blocktiling.cuh b/upstream_ref/sgemm_cuda/05_kernel_2D_blocktiling.cuh new file mode 100644 index 00000000..0b361f7c --- /dev/null +++ b/upstream_ref/sgemm_cuda/05_kernel_2D_blocktiling.cuh @@ -0,0 +1,102 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void __launch_bounds__((BM * BN) / (TM * TN), 1) + sgemm2DBlocktiling(int M, int N, int K, float alpha, const float *A, + const float *B, float beta, float *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + const uint totalResultsBlocktile = BM * BN; + // A thread is responsible for calculating TM*TN elements in the blocktile + const uint numThreadsBlocktile = totalResultsBlocktile / (TM * TN); + + // ResultsPerBlock / ResultsPerThread == ThreadsPerBlock + assert(numThreadsBlocktile == blockDim.x); + + // BN/TN are the number of threads to span a column + const int threadCol = threadIdx.x % (BN / TN); + const int threadRow = threadIdx.x / (BN / TN); + + // allocate space for the current blocktile in smem + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // calculating the indices that this thread will load into SMEM + const uint innerRowA = threadIdx.x / BK; + const uint innerColA = threadIdx.x % BK; + // calculates the number of rows of As that are being loaded in a single step + // by a single block + const uint strideA = numThreadsBlocktile / BK; + const uint innerRowB = threadIdx.x / BN; + const uint innerColB = threadIdx.x % BN; + // for both As and Bs we want each load to span the full column-width, for + // better GMEM coalescing (as opposed to spanning full row-width and iterating + // across columns) + const uint strideB = numThreadsBlocktile / BN; + + // allocate thread-local cache for results in registerfile + float threadResults[TM * TN] = {0.0}; + // register caches for As and Bs + float regM[TM] = {0.0}; + float regN[TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) { + As[(innerRowA + loadOffset) * BK + innerColA] = + A[(innerRowA + loadOffset) * K + innerColA]; + } + for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) { + Bs[(innerRowB + loadOffset) * BN + innerColB] = + B[(innerRowB + loadOffset) * N + innerColB]; + } + __syncthreads(); + + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // block into registers + for (uint i = 0; i < TM; ++i) { + regM[i] = As[(threadRow * TM + i) * BK + dotIdx]; + } + for (uint i = 0; i < TN; ++i) { + regN[i] = Bs[dotIdx * BN + threadCol * TN + i]; + } + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[resIdxM * TN + resIdxN] += + regM[resIdxM] * regN[resIdxN]; + } + } + } + __syncthreads(); + } + + // write out the results + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] = + alpha * threadResults[resIdxM * TN + resIdxN] + + beta * C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN]; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/06_kernel_vectorize.cuh b/upstream_ref/sgemm_cuda/06_kernel_vectorize.cuh new file mode 100644 index 00000000..665d4e50 --- /dev/null +++ b/upstream_ref/sgemm_cuda/06_kernel_vectorize.cuh @@ -0,0 +1,98 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void sgemmVectorize(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // BN/TN are the number of threads to span a column + const int threadCol = threadIdx.x % (BN / TN); + const int threadRow = threadIdx.x / (BN / TN); + + // allocate space for the current blocktile in smem + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // calculating the indices that this thread will load into SMEM + // we'll load 128bit / 32bit = 4 elements per thread at each step + const uint innerRowA = threadIdx.x / (BK / 4); + const uint innerColA = threadIdx.x % (BK / 4); + const uint innerRowB = threadIdx.x / (BN / 4); + const uint innerColB = threadIdx.x % (BN / 4); + + // allocate thread-local cache for results in registerfile + float threadResults[TM * TN] = {0.0}; + float regM[TM] = {0.0}; + float regN[TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + // transpose A while loading it + float4 tmp = + reinterpret_cast(&A[innerRowA * K + innerColA * 4])[0]; + As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x; + As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y; + As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z; + As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w; + + reinterpret_cast(&Bs[innerRowB * BN + innerColB * 4])[0] = + reinterpret_cast(&B[innerRowB * N + innerColB * 4])[0]; + __syncthreads(); + + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // block into registers + for (uint i = 0; i < TM; ++i) { + regM[i] = As[dotIdx * BM + threadRow * TM + i]; + } + for (uint i = 0; i < TN; ++i) { + regN[i] = Bs[dotIdx * BN + threadCol * TN + i]; + } + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[resIdxM * TN + resIdxN] += + regM[resIdxM] * regN[resIdxN]; + } + } + } + __syncthreads(); + } + + // write out the results + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) { + // load C vector into registers + float4 tmp = reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0]; + // perform GEMM update in reg + tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x; + tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y; + tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z; + tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w; + // write back + reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/07_kernel_resolve_bank_conflicts.cuh b/upstream_ref/sgemm_cuda/07_kernel_resolve_bank_conflicts.cuh new file mode 100644 index 00000000..d571008e --- /dev/null +++ b/upstream_ref/sgemm_cuda/07_kernel_resolve_bank_conflicts.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void sgemmResolveBankConflicts(int M, int N, int K, float alpha, + float *A, float *B, float beta, + float *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // BN/TN are the number of threads to span a column + const int threadCol = threadIdx.x % (BN / TN); + const int threadRow = threadIdx.x / (BN / TN); + + // allocate space for the current blocktile in smem + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // calculating the indices that this thread will load into SMEM + // we'll load 128bit / 32bit = 4 elements per thread at each step + const uint innerRowA = threadIdx.x / (BK / 4); + const uint innerColA = threadIdx.x % (BK / 4); + const uint innerRowB = threadIdx.x / (BN / 4); + const uint innerColB = threadIdx.x % (BN / 4); + + // allocate thread-local cache for results in registerfile + float threadResults[TM * TN] = {0.0}; + float regM[TM] = {0.0}; + float regN[TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + // transpose A while loading it + float4 tmp = + reinterpret_cast(&A[innerRowA * K + innerColA * 4])[0]; + As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x; + As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y; + As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z; + As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w; + + // "linearize" Bs while storing it + tmp = reinterpret_cast(&B[innerRowB * N + innerColB * 4])[0]; + Bs[((innerColB % 2) * 4 + innerRowB * 8 + 0) * 16 + innerColB / 2] = tmp.x; + Bs[((innerColB % 2) * 4 + innerRowB * 8 + 1) * 16 + innerColB / 2] = tmp.y; + Bs[((innerColB % 2) * 4 + innerRowB * 8 + 2) * 16 + innerColB / 2] = tmp.z; + Bs[((innerColB % 2) * 4 + innerRowB * 8 + 3) * 16 + innerColB / 2] = tmp.w; + __syncthreads(); + + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // block into registers + for (uint i = 0; i < TM; ++i) { + regM[i] = As[dotIdx * BM + threadRow * TM + i]; + } + for (uint i = 0; i < TN; ++i) { + regN[i] = Bs[(dotIdx * 8 + i) * 16 + threadCol]; + } + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[resIdxM * TN + resIdxN] += + regM[resIdxM] * regN[resIdxN]; + } + } + } + __syncthreads(); + } + + // write out the results + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) { + // load C vector into registers + float4 tmp = reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0]; + // perform GEMM update in reg + tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x; + tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y; + tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z; + tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w; + // write back + reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/08_kernel_bank_extra_col.cuh b/upstream_ref/sgemm_cuda/08_kernel_bank_extra_col.cuh new file mode 100644 index 00000000..c362cba8 --- /dev/null +++ b/upstream_ref/sgemm_cuda/08_kernel_bank_extra_col.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__global__ void sgemmResolveBankExtraCol(int M, int N, int K, float alpha, + float *A, float *B, float beta, + float *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // BN/TN are the number of threads to span a column + const int threadCol = threadIdx.x % (BN / TN); + const int threadRow = threadIdx.x / (BN / TN); + + // allocate space for the current blocktile in smem + __shared__ float As[BM * BK]; + const int extraCols = 5; + __shared__ float Bs[BK * (BN + extraCols)]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // calculating the indices that this thread will load into SMEM + // we'll load 128bit / 32bit = 4 elements per thread at each step + const uint innerRowA = threadIdx.x / (BK / 4); + const uint innerColA = threadIdx.x % (BK / 4); + const uint innerRowB = threadIdx.x / (BN / 4); + const uint innerColB = threadIdx.x % (BN / 4); + + // allocate thread-local cache for results in registerfile + float threadResults[TM * TN] = {0.0}; + float regM[TM] = {0.0}; + float regN[TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + // transpose A while loading it + float4 tmp = + reinterpret_cast(&A[innerRowA * K + innerColA * 4])[0]; + As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x; + As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y; + As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z; + As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w; + + tmp = reinterpret_cast(&B[innerRowB * N + innerColB * 4])[0]; + Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 0] = tmp.x; + Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 1] = tmp.y; + Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 2] = tmp.z; + Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 3] = tmp.w; + __syncthreads(); + + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // block into registers + for (uint i = 0; i < TM; ++i) { + regM[i] = As[dotIdx * BM + threadRow * TM + i]; + } + for (uint i = 0; i < TN; ++i) { + regN[i] = Bs[dotIdx * (BN + extraCols) + threadCol * TN + i]; + } + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[resIdxM * TN + resIdxN] += + regM[resIdxM] * regN[resIdxN]; + } + } + } + __syncthreads(); + } + + // write out the results + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) { + // load C vector into registers + float4 tmp = reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0]; + // perform GEMM update in reg + tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x; + tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y; + tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z; + tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w; + // write back + reinterpret_cast( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/09_kernel_autotuned.cuh b/upstream_ref/sgemm_cuda/09_kernel_autotuned.cuh new file mode 100644 index 00000000..6841952b --- /dev/null +++ b/upstream_ref/sgemm_cuda/09_kernel_autotuned.cuh @@ -0,0 +1,127 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +const int K9_NUM_THREADS = 256; + +template +__global__ void __launch_bounds__(K9_NUM_THREADS) + sgemmAutotuned(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // size of warptile + constexpr int WM = TM * 16; + constexpr int WN = TN * 16; + // iterations of warptile + constexpr int WMITER = CEIL_DIV(BM, WM); + constexpr int WNITER = CEIL_DIV(BN, WN); + + // Placement of the thread in the warptile + const int threadCol = threadIdx.x % (WN / TN); + const int threadRow = threadIdx.x / (WN / TN); + + // allocate space for the current blocktile in smem + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + C += cRow * BM * N + cCol * BN; + + // calculating the indices that this thread will load into SMEM + // we'll load 128bit / 32bit = 4 elements per thread at each step + const uint innerRowA = threadIdx.x / (BK / 4); + const uint innerColA = threadIdx.x % (BK / 4); + constexpr uint rowStrideA = (K9_NUM_THREADS * 4) / BK; + const uint innerRowB = threadIdx.x / (BN / 4); + const uint innerColB = threadIdx.x % (BN / 4); + constexpr uint rowStrideB = K9_NUM_THREADS / (BN / 4); + + // allocate thread-local cache for results in registerfile + float threadResults[WMITER * WNITER * TM * TN] = {0.0}; + float regM[TM] = {0.0}; + float regN[TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) { + float4 tmp = reinterpret_cast( + &A[(innerRowA + offset) * K + innerColA * 4])[0]; + // transpose A while storing it + As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x; + As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y; + As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z; + As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w; + } + + for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) { + reinterpret_cast( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &B[(innerRowB + offset) * N + innerColB * 4])[0]; + } + __syncthreads(); + + for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) { + for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) { + // calculate per-thread results + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // block into registers + for (uint i = 0; i < TM; ++i) { + regM[i] = As[dotIdx * BM + (wmIdx * WM) + threadRow * TM + i]; + } + for (uint i = 0; i < TN; ++i) { + regN[i] = Bs[dotIdx * BN + (wnIdx * WN) + threadCol * TN + i]; + } + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[(wmIdx * TM + resIdxM) * (WNITER * TN) + + wnIdx * TN + resIdxN] += + regM[resIdxM] * regN[resIdxN]; + } + } + } + } + } + __syncthreads(); + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + } + + // write out the results + for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) { + for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) { + float *C_interim = C + (wmIdx * WM * N) + (wnIdx * WN); + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) { + // load C vector into registers + float4 tmp = reinterpret_cast( + &C_interim[(threadRow * TM + resIdxM) * N + threadCol * TN + + resIdxN])[0]; + // perform GEMM update in reg + const int i = + (wmIdx * TM + resIdxM) * (WNITER * TN) + wnIdx * TN + resIdxN; + tmp.x = alpha * threadResults[i + 0] + beta * tmp.x; + tmp.y = alpha * threadResults[i + 1] + beta * tmp.y; + tmp.z = alpha * threadResults[i + 2] + beta * tmp.z; + tmp.w = alpha * threadResults[i + 3] + beta * tmp.w; + // write back + reinterpret_cast(&C_interim[(threadRow * TM + resIdxM) * N + + threadCol * TN + resIdxN])[0] = + tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/CMakeLists_edtallison.txt b/upstream_ref/sgemm_cuda/CMakeLists_edtallison.txt new file mode 100644 index 00000000..fde93413 --- /dev/null +++ b/upstream_ref/sgemm_cuda/CMakeLists_edtallison.txt @@ -0,0 +1,36 @@ +cmake_minimum_required(VERSION 3.19) +project(NVIDIA_SGEMM_PRACTICE LANGUAGES CXX CUDA) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +find_package(CUDA REQUIRED) + +# ensure cuda is available +include(CheckLanguage) +check_language(CUDA) + +set(CMAKE_CXX_STANDARD 20) +set(CUDA_COMPUTE_CAPABILITY 75) + +# in debug mode, add debug symbols to device code +# this disables most optimizations and kills performance +add_compile_options("$<$,$>:-G;-src-in-ptx>") +# add_compile_options("--ptxas-options=-v") + +# Configure header file search paths +include_directories(${CUDA_INCLUDE_DIRS}) +include_directories(${PROJECT_SOURCE_DIR}/src) +# Configure the source file path to be compiled +aux_source_directory(${PROJECT_SOURCE_DIR}/src SRC) + +# generate executable +add_executable(sgemm sgemm.cu ${SRC}) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(sgemm ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES}) + +add_executable(cuBLAS_sgemm cuBLAS_sgemm.cu ) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(cuBLAS_sgemm ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES}) + +add_executable(simplest_kernel simplest_kernel.cu) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(simplest_kernel ${CUDA_LIBRARIES}) \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/CMakeLists_siboehm.txt b/upstream_ref/sgemm_cuda/CMakeLists_siboehm.txt new file mode 100644 index 00000000..7e344af8 --- /dev/null +++ b/upstream_ref/sgemm_cuda/CMakeLists_siboehm.txt @@ -0,0 +1,36 @@ +cmake_minimum_required(VERSION 3.19) +project(NVIDIA_SGEMM_PRACTICE LANGUAGES CXX CUDA) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +find_package(CUDA REQUIRED) + +# ensure cuda is available +include(CheckLanguage) +check_language(CUDA) + +set(CMAKE_CXX_STANDARD 20) +set(CUDA_COMPUTE_CAPABILITY 86) + +# in debug mode, add debug symbols to device code +# this disables most optimizations and kills performance +add_compile_options("$<$,$>:-G;-src-in-ptx>") +# add_compile_options("--ptxas-options=-v") + +# Configure header file search paths +include_directories(${CUDA_INCLUDE_DIRS}) +include_directories(${PROJECT_SOURCE_DIR}/src) +# Configure the source file path to be compiled +aux_source_directory(${PROJECT_SOURCE_DIR}/src SRC) + +# generate executable +add_executable(sgemm sgemm.cu ${SRC}) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(sgemm ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES}) + +add_executable(cuBLAS_sgemm cuBLAS_sgemm.cu ) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(cuBLAS_sgemm ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES}) + +add_executable(simplest_kernel simplest_kernel.cu) +set_target_properties(sgemm PROPERTIES CUDA_ARCHITECTURES ${CUDA_COMPUTE_CAPABILITY}) +target_link_libraries(simplest_kernel ${CUDA_LIBRARIES}) \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/Makefile_edtallison b/upstream_ref/sgemm_cuda/Makefile_edtallison new file mode 100644 index 00000000..715f1967 --- /dev/null +++ b/upstream_ref/sgemm_cuda/Makefile_edtallison @@ -0,0 +1,34 @@ +.PHONY: all build debug clean profile bench cuobjdump + +CMAKE := cmake + +BUILD_DIR := build +BENCHMARK_DIR := benchmark_results + +all: build + +build: + @mkdir -p $(BUILD_DIR) + @cd $(BUILD_DIR) && $(CMAKE) -DCMAKE_BUILD_TYPE=Release .. + @$(MAKE) -C $(BUILD_DIR) + +debug: + @mkdir -p $(BUILD_DIR) + @cd $(BUILD_DIR) && $(CMAKE) -DCMAKE_BUILD_TYPE=Debug .. + @$(MAKE) -C $(BUILD_DIR) + +clean: + @rm -rf $(BUILD_DIR) + +FUNCTION := $$(cuobjdump -symbols build/sgemm | grep -i Warptiling | awk '{print $$NF}') + +cuobjdump: build + @cuobjdump -arch sm_86 -sass -fun $(FUNCTION) build/sgemm | c++filt > build/cuobjdump.sass + @cuobjdump -arch sm_86 -ptx -fun $(FUNCTION) build/sgemm | c++filt > build/cuobjdump.ptx + +# Usage: make profile KERNEL= PREFIX= +profile: build + @ncu --set full --export $(BENCHMARK_DIR)/$(PREFIX)kernel_$(KERNEL) --force-overwrite $(BUILD_DIR)/sgemm $(KERNEL) + +bench: build + @bash gen_benchmark_results.sh diff --git a/upstream_ref/sgemm_cuda/cuBLAS_sgemm.cu b/upstream_ref/sgemm_cuda/cuBLAS_sgemm.cu new file mode 100644 index 00000000..c6062184 --- /dev/null +++ b/upstream_ref/sgemm_cuda/cuBLAS_sgemm.cu @@ -0,0 +1,108 @@ +#include +#include +#include + +/* + * A stand-alone script to invoke & benchmark standard cuBLAS SGEMM performance + */ + +int main(int argc, char *argv[]) { + int m = 2; + int k = 3; + int n = 4; + int print = 1; + cudaError_t cudaStat; // cudaMalloc status + cublasStatus_t stat; // cuBLAS functions status + cublasHandle_t handle; // cuBLAS context + + int i, j; + + float *a, *b, *c; + + // malloc for a,b,c... + a = (float *)malloc(m * k * sizeof(float)); + b = (float *)malloc(k * n * sizeof(float)); + c = (float *)malloc(m * n * sizeof(float)); + + int ind = 11; + for (j = 0; j < m * k; j++) { + a[j] = (float)ind++; + } + + ind = 11; + for (j = 0; j < k * n; j++) { + b[j] = (float)ind++; + } + + ind = 11; + for (j = 0; j < m * n; j++) { + c[j] = (float)ind++; + } + + // DEVICE + float *d_a, *d_b, *d_c; + + // cudaMalloc for d_a, d_b, d_c... + cudaMalloc((void **)&d_a, m * k * sizeof(float)); + cudaMalloc((void **)&d_b, k * n * sizeof(float)); + cudaMalloc((void **)&d_c, m * n * sizeof(float)); + + stat = cublasCreate(&handle); // initialize CUBLAS context + + cudaMemcpy(d_a, a, m * k * sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(d_b, b, k * n * sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(d_c, c, m * n * sizeof(float), cudaMemcpyHostToDevice); + + float alpha = 1.0f; + float beta = 0.5f; + + if (print == 1) { + printf("alpha = %4.0f, beta = %4.0f\n", alpha, beta); + printf("A = (mxk: %d x %d)\n", m, k); + for (i = 0; i < m; i++) { + for (j = 0; j < k; j++) { + printf("%4.1f ", a[i * m + j]); + } + printf("\n"); + } + printf("B = (kxn: %d x %d)\n", k, n); + for (i = 0; i < k; i++) { + for (j = 0; j < n; j++) { + printf("%4.1f ", b[i * n + j]); + } + printf("\n"); + } + printf("C = (mxn: %d x %d)\n", m, n); + for (i = 0; i < m; i++) { + for (j = 0; j < n; j++) { + printf("%4.1f ", c[i * n + j]); + } + printf("\n"); + } + } + + stat = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, d_b, n, + d_a, k, &beta, d_c, n); + + cudaMemcpy(c, d_c, m * n * sizeof(float), cudaMemcpyDeviceToHost); + + if (print == 1) { + printf("\nC after SGEMM = \n"); + for (i = 0; i < m; i++) { + for (j = 0; j < n; j++) { + printf("%4.1f ", c[i * n + j]); + } + printf("\n"); + } + } + + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_c); + cublasDestroy(handle); // destroy CUBLAS context + free(a); + free(b); + free(c); + + return EXIT_SUCCESS; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_1.cuh b/upstream_ref/sgemm_cuda/kernel_1.cuh new file mode 100644 index 00000000..94ce9a77 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_1.cuh @@ -0,0 +1,19 @@ +#pragma once + +#include +#include +#include +#include + +__global__ __launch_bounds__(1024) void +mysgemm_v1(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + + int gx = blockIdx.x * blockDim.x + threadIdx.x; // 全局x + int gy = blockIdx.y * blockDim.y + threadIdx.y; // 全局y + + float tmp = 0.; + for (int i = 0; i < K; i++) { + tmp += A[gy * K + i] * B[i * N + gx]; // 两次全局内存访问和一次FMA（累加乘） + } + C[gy * N + gx] = alpha * tmp + beta * C[gy * N + gx]; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_2.cuh b/upstream_ref/sgemm_cuda/kernel_2.cuh new file mode 100644 index 00000000..73bde1a1 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_2.cuh @@ -0,0 +1,45 @@ +#pragma once + +#include +#include +#include +#include + +template +__global__ void mysgemm_v2(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + + const int BM = BLOCK_SIZE; + const int BN = BLOCK_SIZE; + const int BK = BLOCK_SIZE; + + int tx = threadIdx.x % BN; + int ty = threadIdx.x / BN; + + // 申请共享内存空间 + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + + float tmp = 0.; + for (int k = 0; k < K; k += BK) { + // 缓存A_tile和B_tile + As[ty * BK + tx] = A[ty * K + tx]; + Bs[ty * BN + tx] = B[ty * N + tx]; + // 同步所有线程缓存完成 + __syncthreads(); + A += BK; + B += BK * N; + for (int i = 0; i < BK; i++) { + tmp += As[ty * BK + i] * Bs[i * BN + tx]; + } + // FMA计算需要读取缓存数据，在新一轮写入缓存前进行同步，确保所有线程计算完成 + __syncthreads(); + } + C[ty * N + tx] = alpha * tmp + beta * C[ty * N + tx]; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_3.cuh b/upstream_ref/sgemm_cuda/kernel_3.cuh new file mode 100644 index 00000000..f55a32ce --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_3.cuh @@ -0,0 +1,71 @@ +#pragma once + +#include +#include +#include +#include + +template +__global__ void mysgemm_v3(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + int thread_num = BM * BN / TM; // 一个线程负责block中计算TM个元素 + + int tx = threadIdx.x % BN; + int ty = threadIdx.x / BN * TM; + + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + + /* + 当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列 + a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存； + + 若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运; + 若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运; + */ + int a_tile_row = threadIdx.x / BK; + int a_tile_col = threadIdx.x % BK; + int a_tile_stride = thread_num / BK; + + int b_tile_row = threadIdx.x / BN; + int b_tile_col = threadIdx.x % BN; + int b_tile_stride = thread_num / BN; + + float tmp[TM + 1] = {0.}; // 每个线程负责TM个元素，则需要申请TM个寄存器保存累加值，额外的一个寄存器用于缓存； + #pragma unroll + for (int k = 0; k < K; k += BK) { + #pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col]; + } + #pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col]; + } + __syncthreads(); + A += BK; + B += BK * N; + #pragma unroll + for (int i = 0; i < BK; i++) { + tmp[TM] = Bs[tx + i * BN]; // 额外的一个寄存器，避免反复从共享内存中读取Bs[tx + i * BN] + #pragma unroll // 循环展开，增加指令并行度 + for (int j = 0; j < TM; j++) { + tmp[j] += As[(ty + j) * BK + i] * tmp[TM]; + } + } + __syncthreads(); + } + #pragma unroll + for (int j = 0; j < TM; j++) { + C[(ty + j) * N + tx] = alpha * tmp[j] + beta * C[(ty + j) * N + tx]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_4.cuh b/upstream_ref/sgemm_cuda/kernel_4.cuh new file mode 100644 index 00000000..bb2214fa --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_4.cuh @@ -0,0 +1,76 @@ +#pragma once + +#include +#include +#include +#include + +template +__global__ void mysgemm_v4(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + + int block_row_thread = BN / TN; + int block_col_thread = BM / TM; + int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素 + + int tx = (threadIdx.x % block_row_thread) * TN; + int ty = (threadIdx.x / block_row_thread) * TM; + + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + + /* + 当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列 + a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存； + + 若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运; + 若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运; + */ + int a_tile_row = threadIdx.x / BK; + int a_tile_col = threadIdx.x % BK; + int a_tile_stride = thread_num / BK; + + int b_tile_row = threadIdx.x / BN; + int b_tile_col = threadIdx.x % BN; + int b_tile_stride = thread_num / BN; + + float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存； + #pragma unroll + for (int k = 0; k < K; k += BK) { + #pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col]; + } + #pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col]; + } + __syncthreads(); + A += BK; + B += BK * N; + #pragma unroll + for (int i = 0; i < BK; i++) { + #pragma unroll // 循环展开，增加指令并行度 + for (int j = 0; j < TM; j++) { + for (int l = 0; l < TN; l++) + tmp[j][l] += As[(ty + j) * BK + i] * Bs[tx + l + i * BN]; + } + } + __syncthreads(); + } + #pragma unroll + for (int j = 0; j < TM; j++) { + for (int l = 0; l < TN; l++) + C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_5.cuh b/upstream_ref/sgemm_cuda/kernel_5.cuh new file mode 100644 index 00000000..2fcd0587 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_5.cuh @@ -0,0 +1,88 @@ +#pragma once + +#include +#include +#include +#include + +template +__global__ void mysgemm_v5(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + + int block_row_thread = BN / TN; + int block_col_thread = BM / TM; + int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素 + + int tx = (threadIdx.x % block_row_thread) * TN; + int ty = (threadIdx.x / block_row_thread) * TM; + + __shared__ float As[BM * BK]; + __shared__ float Bs[BK * BN]; + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + + /* + 当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列 + a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存； + + 若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运; + 若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运; + */ + int a_tile_row = threadIdx.x / BK; + int a_tile_col = threadIdx.x % BK; + int a_tile_stride = thread_num / BK; + + int b_tile_row = threadIdx.x / BN; + int b_tile_col = threadIdx.x % BN; + int b_tile_stride = thread_num / BN; + + float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存； + float a_frag[TM] = {0.}; + float b_frag[TN] = {0.}; + + #pragma unroll + for (int k = 0; k < K; k += BK) { + #pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col]; + } + #pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col]; + } + __syncthreads(); + A += BK; + B += BK * N; + #pragma unroll + for (int i = 0; i < BK; i++) { + #pragma unroll + for (int j = 0; j < TM; j++) { + a_frag[j] = As[(ty + j) * BK + i]; + } + #pragma unroll + for (int l = 0; l < TN; l++) { + b_frag[l] = Bs[tx + l + i * BN]; + } + #pragma unroll + for (int j = 0; j < TM; j++) { + #pragma unroll + for (int l = 0; l < TN; l++) + tmp[j][l] += a_frag[j] * b_frag[l]; + } + } + __syncthreads(); + } + #pragma unroll + for (int j = 0; j < TM; j++) { + for (int l = 0; l < TN; l++) + C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_6.cuh b/upstream_ref/sgemm_cuda/kernel_6.cuh new file mode 100644 index 00000000..f40c2501 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_6.cuh @@ -0,0 +1,110 @@ +#pragma once + +#include +#include +#include +#include + +#define OFFSET(row, col, ld) ((row)*(ld)+(col)) +#define FETCH_FLOAT4(pointer) (reinterpret_cast(&(pointer))[0]) + +template +__global__ void mysgemm_v6(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + + const int block_row_thread = BN / TN; + const int block_col_thread = BM / TM; + const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素 + + // 当前线程对应thread tile的左上角元素在block中的位置 + int tx = (threadIdx.x % block_row_thread) * TN; + int ty = (threadIdx.x / block_row_thread) * TM; + + __shared__ float As[BK * BM]; + __shared__ float Bs[BK * BN]; + + + const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮 + const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮 + + int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运 + int a_tile_col = threadIdx.x % (BK / 4) * 4; + int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行 + + int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运 + int b_tile_col = threadIdx.x % (BN / 4) * 4; + int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行 + + float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存； + + // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小 + float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵 + + float a_frag[TM]; // 缓存As共享内存 + float b_frag[TN]; // 缓存Bs共享内存 + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + +#pragma unroll + for (int k = 0; k < K; k += BK) { +#pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + int ldg_index = i / a_tile_stride * 4; // 第ldg_index轮 + FETCH_FLOAT4(ldg_a_reg[ldg_index]) = + FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]); + // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取 + As[OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index]; + As[OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1]; + As[OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2]; + As[OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3]; + } +#pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + FETCH_FLOAT4(Bs[OFFSET(b_tile_row + i, b_tile_col, BN)]) = + FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置 + } + __syncthreads(); + A += BK; + B += BK * N; +#pragma unroll + for (int i = 0; i < BK; i++) { +#pragma unroll + for (int m = 0; m < TM; m += 4) { + FETCH_FLOAT4(a_frag[m]) = FETCH_FLOAT4(As[OFFSET(i, ty + m, BM)]); // 偏移到当前thread tile + } +#pragma unroll + for (int n = 0; n < TN; n += 4) { + FETCH_FLOAT4(b_frag[n]) = FETCH_FLOAT4(Bs[OFFSET(i, tx + n, BN)]); // 偏移到当前thread tile + } +#pragma unroll + for (int m = 0; m < TM; m++) { +#pragma unroll + for (int n = 0; n < TN; n++) { + accum[m][n] += a_frag[m] * b_frag[n]; + } + } + } + __syncthreads(); + } +#pragma unroll + for (int m = 0; m < TM; m++) { +#pragma unroll + for (int n = 0; n < TN; n += 4) { + float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]); + //float4 atmp = FETCH_FLOAT4(accum[m][n]); + ctmp.x = alpha * accum[m][n] + beta * ctmp.x; + ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y; + ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z; + ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w; + FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_7.cuh b/upstream_ref/sgemm_cuda/kernel_7.cuh new file mode 100644 index 00000000..53639644 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_7.cuh @@ -0,0 +1,180 @@ +#pragma once + +#include +#include +#include +#include + +#define OFFSET(row, col, ld) ((row)*(ld)+(col)) +#define FETCH_FLOAT4(pointer) (reinterpret_cast(&(pointer))[0]) + +template +__global__ void mysgemm_v7(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + int bx = blockIdx.x; + int by = blockIdx.y; + + const int block_row_thread = BN / TN; + const int block_col_thread = BM / TM; + const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素 + + // 当前线程对应thread tile的左上角元素在block中的位置 + int tx = (threadIdx.x % block_row_thread) * TN; + int ty = (threadIdx.x / block_row_thread) * TM; + + __shared__ float As[2][BK * BM]; // 增加一倍共享内存大小用于缓存 + __shared__ float Bs[2][BK * BN]; + + + const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮 + const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮 + + int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运 + int a_tile_col = threadIdx.x % (BK / 4) * 4; + int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行 + + int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运 + int b_tile_col = threadIdx.x % (BN / 4) * 4; + int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行 + + float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存； + + // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小 + float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵 + float ldg_b_reg[4 * ldg_b_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵 + + float a_frag[2][TM]; // 缓存As共享内存,增加一倍寄存器大小用于缓存 + float b_frag[2][TN]; // 缓存Bs共享内存,增加一倍寄存器大小用于缓存 + + // 移动到当前block + A = &A[by * BM * K]; + B = &B[bx * BN]; + C = &C[by * BM * N + bx * BN]; + + // first global to shared +#pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + int ldg_index = i / a_tile_stride * 4; // 第ldg_index轮 + FETCH_FLOAT4(ldg_a_reg[ldg_index]) = + FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]); + // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取 + As[0][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index]; + As[0][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1]; + As[0][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2]; + As[0][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3]; + } +#pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + FETCH_FLOAT4(Bs[0][OFFSET(b_tile_row + i, b_tile_col, BN)]) = + FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置 + } + __syncthreads(); + + // first shared to frag +#pragma unroll + for (int m = 0; m < TM; m += 4) { + FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4(As[0][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile + } +#pragma unroll + for (int n = 0; n < TN; n += 4) { + FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4(Bs[0][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile + } + + + int write_index = 1; + int load_index; + int k = 0; + do { + k += BK; + // load global to reg + if (k < K) { +#pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + int ldg_index = i / a_tile_stride * 4; // 第ldg_index轮 + FETCH_FLOAT4(ldg_a_reg[ldg_index]) = + FETCH_FLOAT4(A[OFFSET(a_tile_row + i, k + a_tile_col, K)]); + } +#pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + int ldg_index = i / b_tile_stride * 4; // 第ldg_index轮 + FETCH_FLOAT4(ldg_b_reg[ldg_index]) = + FETCH_FLOAT4(B[OFFSET(k + b_tile_row + i, b_tile_col, N)]); + } + } + + load_index = write_index ^ 1; +#pragma unroll + for (int bk = 0; bk < BK - 1; bk++) { + for (int m = 0; m < TM; m += 4) { + FETCH_FLOAT4(a_frag[(bk + 1) % 2][m]) = FETCH_FLOAT4( + As[load_index][OFFSET(bk + 1, ty + m, BM)]); // 偏移到当前thread tile + } +#pragma unroll + for (int n = 0; n < TN; n += 4) { + FETCH_FLOAT4(b_frag[(bk + 1) % 2][n]) = FETCH_FLOAT4( + Bs[load_index][OFFSET(bk + 1, tx + n, BN)]); // 偏移到当前thread tile + } +#pragma unroll + for (int m = 0; m < TM; m++) { + for (int n = 0; n < TN; n++) { + accum[m][n] += a_frag[bk % 2][m] * b_frag[bk % 2][n]; + } + } + } + if (k < K) { +#pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + int ldg_index = i / a_tile_stride * 4; + As[write_index][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index]; + As[write_index][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1]; + As[write_index][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2]; + As[write_index][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3]; + } +#pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + int ldg_index = i / b_tile_stride * 4; + FETCH_FLOAT4(Bs[write_index][OFFSET(b_tile_row + i, b_tile_col, BN)]) = + FETCH_FLOAT4(ldg_b_reg[ldg_index]); + } + __syncthreads(); +#pragma unroll + for (int m = 0; m < TM; m += 4) { + FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4( + As[write_index][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile + } +#pragma unroll + for (int n = 0; n < TN; n += 4) { + FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4( + Bs[write_index][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile + } + + write_index ^= 1; + } +#pragma unroll + for (int m = 0; m < TM; m++) { +#pragma unroll + for (int n = 0; n < TN; n++) { + accum[m][n] += a_frag[(BK - 1) % 2][m] * b_frag[(BK - 1) % 2][n]; + } + } + + + } while (k < K); + + // C = alpha*AB+C +#pragma unroll + for (int m = 0; m < TM; m++) { +#pragma unroll + for (int n = 0; n < TN; n += 4) { + float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]); + ctmp.x = alpha * accum[m][n] + beta * ctmp.x; + ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y; + ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z; + ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w; + FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernel_wangzyon.cuh b/upstream_ref/sgemm_cuda/kernel_wangzyon.cuh new file mode 100644 index 00000000..cde4381c --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernel_wangzyon.cuh @@ -0,0 +1,9 @@ +#pragma once + +#include "kernel/kernel_1.cuh" +#include "kernel/kernel_2.cuh" +#include "kernel/kernel_3.cuh" +#include "kernel/kernel_4.cuh" +#include "kernel/kernel_5.cuh" +#include "kernel/kernel_6.cuh" +#include "kernel/kernel_7.cuh" \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/kernels.cuh b/upstream_ref/sgemm_cuda/kernels.cuh new file mode 100644 index 00000000..7691b297 --- /dev/null +++ b/upstream_ref/sgemm_cuda/kernels.cuh @@ -0,0 +1,14 @@ +#pragma once + +#include "kernels/10_kernel_warptiling.cuh" +#include "kernels/11_kernel_double_buffering.cuh" +#include "kernels/12_kernel_double_buffering.cuh" +#include "kernels/1_naive.cuh" +#include "kernels/2_kernel_global_mem_coalesce.cuh" +#include "kernels/3_kernel_shared_mem_blocking.cuh" +#include "kernels/4_kernel_1D_blocktiling.cuh" +#include "kernels/5_kernel_2D_blocktiling.cuh" +#include "kernels/6_kernel_vectorize.cuh" +#include "kernels/7_kernel_resolve_bank_conflicts.cuh" +#include "kernels/8_kernel_bank_extra_col.cuh" +#include "kernels/9_kernel_autotuned.cuh" \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/runner.cu b/upstream_ref/sgemm_cuda/runner.cu new file mode 100644 index 00000000..b3b731a2 --- /dev/null +++ b/upstream_ref/sgemm_cuda/runner.cu @@ -0,0 +1,549 @@ +#include "kernels.cuh" +#include "runner.cuh" +#include +#include +#include +#include + +float get_sec() { + struct timeval time; + gettimeofday(&time, NULL); + return (1e6 * time.tv_sec + time.tv_usec); +} + +float cpu_elapsed_time(float &beg, float &end) { return 1.0e-6 * (end - beg); } + +void cudaCheck(cudaError_t error, const char *file, int line) { + if (error != cudaSuccess) { + printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, + cudaGetErrorString(error)); + exit(EXIT_FAILURE); + } +}; + +void CudaDeviceInfo() { + int deviceId; + + cudaGetDevice(&deviceId); + + cudaDeviceProp props{}; + cudaGetDeviceProperties(&props, deviceId); + + printf("Device ID: %d\n\ + Name: %s\n\ + Compute Capability: %d.%d\n\ + memoryBusWidth: %d\n\ + maxThreadsPerBlock: %d\n\ + maxThreadsPerMultiProcessor: %d\n\ + maxRegsPerBlock: %d\n\ + maxRegsPerMultiProcessor: %d\n\ + totalGlobalMem: %zuMB\n\ + sharedMemPerBlock: %zuKB\n\ + sharedMemPerMultiprocessor: %zuKB\n\ + totalConstMem: %zuKB\n\ + multiProcessorCount: %d\n\ + Warp Size: %d\n", + deviceId, props.name, props.major, props.minor, props.memoryBusWidth, + props.maxThreadsPerBlock, props.maxThreadsPerMultiProcessor, + props.regsPerBlock, props.regsPerMultiprocessor, + props.totalGlobalMem / 1024 / 1024, props.sharedMemPerBlock / 1024, + props.sharedMemPerMultiprocessor / 1024, props.totalConstMem / 1024, + props.multiProcessorCount, props.warpSize); +}; + +void randomize_matrix(float *mat, int N) { + // NOTICE: Use gettimeofday instead of srand((unsigned)time(NULL)); the time + // precision is too low and the same random number is generated. + struct timeval time {}; + gettimeofday(&time, nullptr); + srand(time.tv_usec); + for (int i = 0; i < N; i++) { + float tmp = (float)(rand() % 5) + 0.01 * (rand() % 5); + tmp = (rand() % 2 == 0) ? tmp : tmp * (-1.); + mat[i] = tmp; + } +} + +void range_init_matrix(float *mat, int N) { + for (int i = 0; i < N; i++) { + mat[i] = i; + } +} + +void zero_init_matrix(float *mat, int N) { + for (int i = 0; i < N; i++) { + mat[i] = 0.0; + } +} + +void copy_matrix(const float *src, float *dest, int N) { + int i; + for (i = 0; src + i && dest + i && i < N; i++) + *(dest + i) = *(src + i); + if (i != N) + printf("copy failed at %d while there are %d elements in total.\n", i, N); +} + +void print_matrix(const float *A, int M, int N, std::ofstream &fs) { + int i; + fs << std::setprecision(2) + << std::fixed; // Set floating-point precision and fixed notation + fs << "["; + for (i = 0; i < M * N; i++) { + if ((i + 1) % N == 0) + fs << std::setw(5) << A[i]; // Set field width and write the value + else + fs << std::setw(5) << A[i] << ", "; + if ((i + 1) % N == 0) { + if (i + 1 < M * N) + fs << ";\n"; + } + } + fs << "]\n"; +} + +bool verify_matrix(float *matRef, float *matOut, int N) { + double diff = 0.0; + int i; + for (i = 0; i < N; i++) { + diff = std::fabs(matRef[i] - matOut[i]); + if (isnan(diff) || diff > 0.01) { + printf("Divergence! Should %5.2f, Is %5.2f (Diff %5.2f) at %d\n", + matRef[i], matOut[i], diff, i); + return false; + } + } + return true; +} + +int div_ceil(int numerator, int denominator) { + std::div_t res = std::div(numerator, denominator); + return res.rem ? (res.quot + 1) : res.quot; +} + +void runCublasFP32(cublasHandle_t handle, int M, int N, int K, float alpha, + float *A, float *B, float beta, float *C) { + // cuBLAS uses column-major order. So we change the order of our row-major A & + // B, since (B^T*A^T)^T = (A*B) + // This runs cuBLAS in full fp32 mode + cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F, + N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N, CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +} + +void runCublasBF16(cublasHandle_t handle, int M, int N, int K, float alpha, + float *A, float *B, float beta, float *C) { + // This runs cuBLAS with mixed precision (performing the mul with operands + // downcast to bf16), which is ~4x faster + cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F, + N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N, + CUBLAS_COMPUTE_32F_FAST_16BF, CUBLAS_GEMM_DEFAULT_TENSOR_OP); +} + +void runCublasTF32(cublasHandle_t handle, int M, int N, int K, float alpha, + float *A, float *B, float beta, float *C) { + // This runs cuBLAS with mixed precision (performing the mul with operands + // downcast to bf16), which is ~4x faster + cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F, + N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N, + CUBLAS_COMPUTE_32F_FAST_TF32, CUBLAS_GEMM_DEFAULT_TENSOR_OP); +} + +void run_sgemm_naive(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32)); + dim3 blockDim(32, 32); + sgemm_naive<<>>(M, N, K, alpha, A, B, beta, C); +} + +void run_sgemm_coalesce(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32)); + dim3 blockDim(32 * 32); + sgemm_global_mem_coalesce<32> + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void run_sgemm_shared_mem_block(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32)); + dim3 blockDim(32 * 32); + // L1 cache becomes useless, since we access GMEM only via SMEM, so we carve + // out all of L1 to SMEM. This doesn't currently make a difference, since + // occupancy is limited by reg and thread count, but it's good to do anyway. + cudaFuncSetAttribute(sgemm_shared_mem_block<32>, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared); + sgemm_shared_mem_block<32> + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void runSgemm1DBlocktiling(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + const uint BM = 64; + const uint BN = 64; + const uint BK = 8; + const uint TM = 8; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / TM); + sgemm1DBlocktiling + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void runSgemm2DBlocktiling(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + const uint BK = 8; + const uint TM = 8; + const uint TN = 8; + if (M >= 128 and N >= 128) { + const uint BM = 128; + const uint BN = 128; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemm2DBlocktiling + <<>>(M, N, K, alpha, A, B, beta, C); + } else { + // this is a hacky solution to the underlying problem + // of not having proper bounds checking in the kernel + const uint BM = 64; + const uint BN = 64; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemm2DBlocktiling + <<>>(M, N, K, alpha, A, B, beta, C); + } +} + +void runSgemmVectorize(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + const uint BK = 8; + const uint TM = 8; + const uint TN = 8; + if (M >= 128 and N >= 128) { + const uint BM = 128; + const uint BN = 128; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmVectorize + <<>>(M, N, K, alpha, A, B, beta, C); + } else { + // this is a hacky solution to the underlying problem + // of not having proper bounds checking in the kernel + const uint BM = 64; + const uint BN = 64; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmVectorize + <<>>(M, N, K, alpha, A, B, beta, C); + } +} + +void runSgemmResolveBankConflicts(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + const uint BK = 8; + const uint TM = 8; + const uint TN = 8; + if (M >= 128 and N >= 128) { + const uint BM = 128; + const uint BN = 128; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmResolveBankConflicts + <<>>(M, N, K, alpha, A, B, beta, C); + } else { + // this is a hacky solution to the underlying problem + // of not having proper bounds checking in the kernel + const uint BM = 64; + const uint BN = 64; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmResolveBankConflicts + <<>>(M, N, K, alpha, A, B, beta, C); + } +} + +void runSgemmResolveBankExtraCol(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + const uint BK = 8; + const uint TM = 8; + const uint TN = 8; + if (M >= 128 and N >= 128) { + const uint BM = 128; + const uint BN = 128; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmResolveBankExtraCol + <<>>(M, N, K, alpha, A, B, beta, C); + } else { + // this is a hacky solution to the underlying problem + // of not having proper bounds checking in the kernel + const uint BM = 64; + const uint BN = 64; + dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 blockDim((BM * BN) / (TM * TN)); + sgemmResolveBankExtraCol + <<>>(M, N, K, alpha, A, B, beta, C); + } +} + +void runSgemmAutotuned(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + // A100 + // const uint K9_BK = 16; + // const uint K9_TM = 4; + // const uint K9_TN = 4; + // const uint K9_BM = 64; + // const uint K9_BN = 64; + // A6000 + const uint K9_BK = 16; + const uint K9_TM = 8; + const uint K9_TN = 8; + const uint K9_BM = 128; + const uint K9_BN = 128; + dim3 blockDim(K9_NUM_THREADS); + + static_assert( + (K9_NUM_THREADS * 4) % K9_BK == 0, + "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization issues " + "during GMEM->SMEM tiling (loading only parts of the final row of Bs " + "during each iteraion)"); + static_assert( + (K9_NUM_THREADS * 4) % K9_BN == 0, + "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization issues " + "during GMEM->SMEM tiling (loading only parts of the final row of As " + "during each iteration)"); + static_assert( + K9_BN % (16 * K9_TN) == 0, + "K9_BN must be a multiple of 16*K9_TN to avoid quantization effects"); + static_assert( + K9_BM % (16 * K9_TM) == 0, + "K9_BM must be a multiple of 16*K9_TM to avoid quantization effects"); + static_assert((K9_BM * K9_BK) % (4 * K9_NUM_THREADS) == 0, + "K9_BM*K9_BK must be a multiple of 4*256 to vectorize loads"); + static_assert((K9_BN * K9_BK) % (4 * K9_NUM_THREADS) == 0, + "K9_BN*K9_BK must be a multiple of 4*256 to vectorize loads"); + + dim3 gridDim(CEIL_DIV(N, K9_BN), CEIL_DIV(M, K9_BM)); + sgemmAutotuned + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void runSgemmWarptiling(int M, int N, int K, float alpha, float *A, float *B, + float beta, float *C) { + // Settings for A100 + // const uint K10_NUM_THREADS = 128; + // const uint K10_BN = 128; + // const uint K10_BM = 64; + // const uint K10_BK = 16; + // const uint K10_WN = 64; + // const uint K10_WM = 32; + // const uint K10_WNITER = 1; + // const uint K10_TN = 4; + // const uint K10_TM = 4; + // Settings for A6000 + const uint K10_NUM_THREADS = 128; + const uint K10_BN = 128; + const uint K10_BM = 128; + const uint K10_BK = 16; + const uint K10_WN = 64; + const uint K10_WM = 64; + const uint K10_WNITER = 4; + const uint K10_TN = 4; + const uint K10_TM = 8; + dim3 blockDim(K10_NUM_THREADS); + + constexpr uint NUM_WARPS = K10_NUM_THREADS / 32; + + // warptile in threadblocktile + static_assert((K10_BN % K10_WN == 0) and (K10_BM % K10_WM == 0)); + static_assert((K10_BN / K10_WN) * (K10_BM / K10_WM) == NUM_WARPS); + + // threads in warpsubtile + static_assert((K10_WM * K10_WN) % (WARPSIZE * K10_TM * K10_TN * K10_WNITER) == + 0); + constexpr uint K10_WMITER = + (K10_WM * K10_WN) / (32 * K10_TM * K10_TN * K10_WNITER); + // warpsubtile in warptile + static_assert((K10_WM % K10_WMITER == 0) and (K10_WN % K10_WNITER == 0)); + + static_assert((K10_NUM_THREADS * 4) % K10_BK == 0, + "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of Bs during each iteraion)"); + static_assert((K10_NUM_THREADS * 4) % K10_BN == 0, + "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of As during each iteration)"); + static_assert(K10_BN % (16 * K10_TN) == 0, + "BN must be a multiple of 16*TN to avoid quantization effects"); + static_assert(K10_BM % (16 * K10_TM) == 0, + "BM must be a multiple of 16*TM to avoid quantization effects"); + static_assert((K10_BM * K10_BK) % (4 * K10_NUM_THREADS) == 0, + "BM*BK must be a multiple of 4*256 to vectorize loads"); + static_assert((K10_BN * K10_BK) % (4 * K10_NUM_THREADS) == 0, + "BN*BK must be a multiple of 4*256 to vectorize loads"); + + dim3 gridDim(CEIL_DIV(N, K10_BN), CEIL_DIV(M, K10_BM)); + sgemmWarptiling + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void runSgemmDoubleBuffering(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + // Settings for A100 + // const uint K11_NUM_THREADS = 256; + // const uint K11_BN = 128; + // const uint K11_BM = 64; + // const uint K11_BK = 16; + // const uint K11_WN = 32; + // const uint K11_WM = 32; + // const uint K11_WNITER = 2; + // const uint K11_TN = 4; + // const uint K11_TM = 4; + // Settings for A6000 + const uint K11_NUM_THREADS = 256; + const uint K11_BN = 256; + const uint K11_BM = 128; + const uint K11_BK = 16; + const uint K11_WN = 32; + const uint K11_WM = 128; + const uint K11_WNITER = 1; + const uint K11_TN = 8; + const uint K11_TM = 8; + dim3 blockDim(K11_NUM_THREADS); + + constexpr uint NUM_WARPS = K11_NUM_THREADS / 32; + + // warptile in threadblocktile + static_assert((K11_BN % K11_WN == 0) and (K11_BM % K11_WM == 0)); + static_assert((K11_BN / K11_WN) * (K11_BM / K11_WM) == NUM_WARPS); + + // threads in warpsubtile + static_assert((K11_WM * K11_WN) % (WARPSIZE * K11_TM * K11_TN * K11_WNITER) == + 0); + constexpr uint K11_WMITER = + (K11_WM * K11_WN) / (32 * K11_TM * K11_TN * K11_WNITER); + // warpsubtile in warptile + static_assert((K11_WM % K11_WMITER == 0) and (K11_WN % K11_WNITER == 0)); + + static_assert((K11_NUM_THREADS / 2 * 4) % K11_BK == 0, + "NUM_THREADS*4 must be multiple of BK to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of Bs during each iteraion)"); + static_assert((K11_NUM_THREADS / 2 * 4) % K11_BN == 0, + "NUM_THREADS*4 must be multiple of BN to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of As during each iteration)"); + static_assert(K11_BN % (16 * K11_TN) == 0, + "BN must be a multiple of 16*TN to avoid quantization effects"); + static_assert(K11_BM % (16 * K11_TM) == 0, + "BM must be a multiple of 16*TM to avoid quantization effects"); + static_assert((K11_BM * K11_BK) % (4 * K11_NUM_THREADS / 2) == 0, + "BM*BK must be a multiple of 4*256 to vectorize loads"); + static_assert((K11_BN * K11_BK) % (4 * K11_NUM_THREADS / 2) == 0, + "BN*BK must be a multiple of 4*256 to vectorize loads"); + + dim3 gridDim(CEIL_DIV(N, K11_BN), CEIL_DIV(M, K11_BM)); + sgemmDoubleBuffering + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void runSgemmDoubleBuffering2(int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C) { + // Settings for A6000 + const uint K12_NUM_THREADS = 128; + const uint K12_BN = 128; + const uint K12_BM = 128; + const uint K12_BK = 16; + const uint K12_WN = 64; + const uint K12_WM = 64; + const uint K12_WNITER = 4; + const uint K12_TN = 4; + const uint K12_TM = 8; + dim3 blockDim(K12_NUM_THREADS); + + constexpr uint NUM_WARPS = K12_NUM_THREADS / 32; + + // warptile in threadblocktile + static_assert((K12_BN % K12_WN == 0) and (K12_BM % K12_WM == 0)); + static_assert((K12_BN / K12_WN) * (K12_BM / K12_WM) == NUM_WARPS); + + // threads in warpsubtile + static_assert((K12_WM * K12_WN) % (WARPSIZE * K12_TM * K12_TN * K12_WNITER) == + 0); + constexpr uint K12_WMITER = + (K12_WM * K12_WN) / (32 * K12_TM * K12_TN * K12_WNITER); + // warpsubtile in warptile + static_assert((K12_WM % K12_WMITER == 0) and (K12_WN % K12_WNITER == 0)); + + static_assert((K12_NUM_THREADS * 4) % K12_BK == 0, + "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of Bs during each iteraion)"); + static_assert((K12_NUM_THREADS * 4) % K12_BN == 0, + "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of As during each iteration)"); + static_assert(K12_BN % (16 * K12_TN) == 0, + "BN must be a multiple of 16*TN to avoid quantization effects"); + static_assert(K12_BM % (16 * K12_TM) == 0, + "BM must be a multiple of 16*TM to avoid quantization effects"); + static_assert((K12_BM * K12_BK) % (4 * K12_NUM_THREADS) == 0, + "BM*BK must be a multiple of 4*256 to vectorize loads"); + static_assert((K12_BN * K12_BK) % (4 * K12_NUM_THREADS) == 0, + "BN*BK must be a multiple of 4*256 to vectorize loads"); + + dim3 gridDim(CEIL_DIV(N, K12_BN), CEIL_DIV(M, K12_BM)); + runSgemmDoubleBuffering2 + <<>>(M, N, K, alpha, A, B, beta, C); +} + +void run_kernel(int kernel_num, int M, int N, int K, float alpha, float *A, + float *B, float beta, float *C, cublasHandle_t handle) { + switch (kernel_num) { + case 0: + runCublasFP32(handle, M, N, K, alpha, A, B, beta, C); + break; + case 1: + run_sgemm_naive(M, N, K, alpha, A, B, beta, C); + break; + case 2: + run_sgemm_coalesce(M, N, K, alpha, A, B, beta, C); + break; + case 3: + run_sgemm_shared_mem_block(M, N, K, alpha, A, B, beta, C); + break; + case 4: + runSgemm1DBlocktiling(M, N, K, alpha, A, B, beta, C); + break; + case 5: + runSgemm2DBlocktiling(M, N, K, alpha, A, B, beta, C); + break; + case 6: + runSgemmVectorize(M, N, K, alpha, A, B, beta, C); + break; + case 7: + runSgemmResolveBankConflicts(M, N, K, alpha, A, B, beta, C); + break; + case 8: + runSgemmResolveBankExtraCol(M, N, K, alpha, A, B, beta, C); + break; + case 9: + runSgemmAutotuned(M, N, K, alpha, A, B, beta, C); + break; + case 10: + runSgemmWarptiling(M, N, K, alpha, A, B, beta, C); + break; + case 11: + runSgemmDoubleBuffering(M, N, K, alpha, A, B, beta, C); + break; + case 12: + runSgemmDoubleBuffering2(M, N, K, alpha, A, B, beta, C); + break; + default: + throw std::invalid_argument("Unknown kernel number"); + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/runner.cuh b/upstream_ref/sgemm_cuda/runner.cuh new file mode 100644 index 00000000..3ac967e6 --- /dev/null +++ b/upstream_ref/sgemm_cuda/runner.cuh @@ -0,0 +1,26 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +void cudaCheck(cudaError_t error, const char *file, + int line); // CUDA error check +void CudaDeviceInfo(); // print CUDA information + +void range_init_matrix(float *mat, int N); +void randomize_matrix(float *mat, int N); +void zero_init_matrix(float *mat, int N); +void copy_matrix(const float *src, float *dest, int N); +void print_matrix(const float *A, int M, int N, std::ofstream &fs); +bool verify_matrix(float *mat1, float *mat2, int N); + +float get_current_sec(); // Get the current moment +float cpu_elapsed_time(float &beg, float &end); // Calculate time difference + +void run_kernel(int kernel_num, int m, int n, int k, float alpha, float *A, + float *B, float beta, float *C, cublasHandle_t handle); \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/sgemm_siboehm.cu b/upstream_ref/sgemm_cuda/sgemm_siboehm.cu new file mode 100644 index 00000000..bb6dd38c --- /dev/null +++ b/upstream_ref/sgemm_cuda/sgemm_siboehm.cu @@ -0,0 +1,168 @@ +#include +#include +#include +#include +#include +#include +#include + +#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) + +const std::string errLogFile = "matrixValidationFailure.txt"; + +int main(int argc, char **argv) { + if (argc != 2) { + std::cerr << "Please select a kernel (range 0 - 12, 0 for NVIDIA cuBLAS)" + << std::endl; + exit(EXIT_FAILURE); + } + + // get kernel number + int kernel_num = std::stoi(argv[1]); + if (kernel_num < 0 || kernel_num > 12) { + std::cerr << "Please enter a valid kernel number (0-12)" << std::endl; + exit(EXIT_FAILURE); + } + + // get environment variable for device + int deviceIdx = 0; + if (getenv("DEVICE") != NULL) { + deviceIdx = atoi(getenv("DEVICE")); + } + cudaCheck(cudaSetDevice(deviceIdx)); + + printf("Running kernel %d on device %d.\n", kernel_num, deviceIdx); + + // print some device info + // CudaDeviceInfo(); + + // Declare the handle, create the handle, cublasCreate will return a value of + // type cublasStatus_t to determine whether the handle was created + // successfully (the value is 0) + cublasHandle_t handle; + if (cublasCreate(&handle)) { + std::cerr << "Create cublas handle error." << std::endl; + exit(EXIT_FAILURE); + }; + + // Using cudaEvent for gpu stream timing, cudaEvent is equivalent to + // publishing event tasks in the target stream + float elapsed_time; + cudaEvent_t beg, end; + cudaEventCreate(&beg); + cudaEventCreate(&end); + + // cuBLAS FLOPs ceiling is reached at 8192 + std::vector SIZE = {128, 256, 512, 1024, 2048, 4096}; + + long m, n, k, max_size; + max_size = SIZE[SIZE.size() - 1]; + std::cout << "Max size: " << max_size << std::endl; + + float alpha = 0.5, beta = 3.0; // GEMM input parameters, C=α*AB+β*C + + float *A = nullptr, *B = nullptr, *C = nullptr, + *C_ref = nullptr; // host matrices + float *dA = nullptr, *dB = nullptr, *dC = nullptr, + *dC_ref = nullptr; // device matrices + + A = (float *)malloc(sizeof(float) * max_size * max_size); + B = (float *)malloc(sizeof(float) * max_size * max_size); + C = (float *)malloc(sizeof(float) * max_size * max_size); + C_ref = (float *)malloc(sizeof(float) * max_size * max_size); + + randomize_matrix(A, max_size * max_size); + randomize_matrix(B, max_size * max_size); + randomize_matrix(C, max_size * max_size); + + cudaCheck(cudaMalloc((void **)&dA, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dB, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dC, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dC_ref, sizeof(float) * max_size * max_size)); + + cudaCheck(cudaMemcpy(dA, A, sizeof(float) * max_size * max_size, + cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dB, B, sizeof(float) * max_size * max_size, + cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dC, C, sizeof(float) * max_size * max_size, + cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dC_ref, C, sizeof(float) * max_size * max_size, + cudaMemcpyHostToDevice)); + + int repeat_times = 50; + for (int size : SIZE) { + m = n = k = size; + + std::cout << "dimensions(m=n=k) " << m << ", alpha: " << alpha + << ", beta: " << beta << std::endl; + // Verify the correctness of the calculation, and execute it once before the + // kernel function timing to avoid cold start errors + if (kernel_num != 0) { + run_kernel(0, m, n, k, alpha, dA, dB, beta, dC_ref, + handle); // cuBLAS + run_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, + handle); // Executes the kernel, modifies the result matrix + cudaCheck(cudaDeviceSynchronize()); + cudaCheck(cudaGetLastError()); // Check for async errors during kernel run + cudaMemcpy(C, dC, sizeof(float) * m * n, cudaMemcpyDeviceToHost); + cudaMemcpy(C_ref, dC_ref, sizeof(float) * m * n, cudaMemcpyDeviceToHost); + + if (!verify_matrix(C_ref, C, m * n)) { + std::cout + << "Failed to pass the correctness verification against NVIDIA " + "cuBLAS." + << std::endl; + if (m <= 128) { + std::cout << " Logging faulty output into " << errLogFile << "\n"; + std::ofstream fs; + fs.open(errLogFile); + fs << "A:\n"; + print_matrix(A, m, n, fs); + fs << "B:\n"; + print_matrix(B, m, n, fs); + fs << "C:\n"; + print_matrix(C, m, n, fs); + fs << "Should:\n"; + print_matrix(C_ref, m, n, fs); + } + exit(EXIT_FAILURE); + } + } + + cudaEventRecord(beg); + for (int j = 0; j < repeat_times; j++) { + // We don't reset dC between runs to save time + run_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, handle); + } + cudaEventRecord(end); + cudaEventSynchronize(beg); + cudaEventSynchronize(end); + cudaEventElapsedTime(&elapsed_time, beg, end); + elapsed_time /= 1000.; // Convert to seconds + + long flops = 2 * m * n * k; + printf( + "Average elapsed time: (%7.6f) s, performance: (%7.1f) GFLOPS. size: " + "(%ld).\n", + elapsed_time / repeat_times, + (repeat_times * flops * 1e-9) / elapsed_time, m); + fflush(stdout); + // make dC and dC_ref equal again (we modified dC while calling our kernel + // for benchmarking) + cudaCheck(cudaMemcpy(dC, dC_ref, sizeof(float) * m * n, + cudaMemcpyDeviceToDevice)); + } + + // Free up CPU and GPU space + free(A); + free(B); + free(C); + free(C_ref); + cudaFree(dA); + cudaFree(dB); + cudaFree(dC); + cudaFree(dC_ref); + cublasDestroy(handle); + + return 0; +}; \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/sgemm_wangzyon.cu b/upstream_ref/sgemm_cuda/sgemm_wangzyon.cu new file mode 100644 index 00000000..b93c59d3 --- /dev/null +++ b/upstream_ref/sgemm_cuda/sgemm_wangzyon.cu @@ -0,0 +1,119 @@ +#include +#include +#include +#include + +#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) + +int main(int argc, char **argv) { + if (argc != 2) { + printf("Please select a kernel (range 0 - 11, here 0 is for NVIDIA cuBLAS).\n"); + exit(EXIT_FAILURE); + } + + // cuda kernel num + int kernel_num = atoi(argv[1]); + if (kernel_num < 0 || kernel_num > 11) { + printf("Please enter a valid kernel number (0-11).\n"); + exit(EXIT_FAILURE); + } else { + printf("Select kernel %d.\n", kernel_num); + }; + + // 申明句柄，创建句柄, cublasCreate会返回一个cublasStatus_t类型的值，用来判断句柄是否创建成功(值为0) + cublasHandle_t handle; + if (cublasCreate(&handle)) { + printf("Create cublas handle error.\n"); + exit(EXIT_FAILURE); + }; + + // 采用cudaEvent进行gpu流计时，cudaEvent相当于在目标流中发布事件任务 + float elapsed_time; + cudaEvent_t beg, end; + cudaEventCreate(&beg); + cudaEventCreate(&end); + + // matrix size + int size_len = 24; + int SIZE[size_len]; + for (int i = 0; i < size_len; i++) + SIZE[i] = 256 * (i + 1); + + int m, n, k, max_size; + max_size = SIZE[size_len - 1]; + printf("max_size=%d\n", max_size); + + float alpha = 1.0, beta = 0.; //two arbitary input parameters，C=α*AB+β*C + + float *A = NULL, *B = NULL, *C = NULL, *C_ref = NULL; //host matrices + float *dA = NULL, *dB = NULL, *dC = NULL, *dC_ref = NULL; //device matrices + + A = (float *) malloc(sizeof(float) * max_size * max_size); + B = (float *) malloc(sizeof(float) * max_size * max_size); + C = (float *) malloc(sizeof(float) * max_size * max_size); + C_ref = (float *) malloc(sizeof(float) * max_size * max_size); + + randomize_matrix(A, max_size * max_size); + randomize_matrix(B, max_size * max_size); + randomize_matrix(C, max_size * max_size); + copy_matrix(C, C_ref, max_size * max_size); + + cudaCheck(cudaMalloc((void **) &dA, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **) &dB, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **) &dC, sizeof(float) * max_size * max_size)); + cudaCheck(cudaMalloc((void **) &dC_ref, sizeof(float) * max_size * max_size)); + + cudaCheck(cudaMemcpy(dA, A, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dB, B, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dC, C, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dC_ref, C_ref, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice)); + + int repeat_times = 10; + for (int i = 0; i < size_len; i++) { + m = n = k = SIZE[i]; + + printf("m=n=k=%d\n", m); + // 验证计算正确性，同时在核函数计时前预先执行一次，避免冷启动误差 + if (kernel_num != 0) { + test_kernel(0, m, n, k, alpha, dA, dB, beta, dC_ref, handle); // cuBLAS + test_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, handle); // user define + cudaDeviceSynchronize(); + cudaMemcpy(C, dC, sizeof(float) * m * n, cudaMemcpyDeviceToHost); + cudaMemcpy(C_ref, dC_ref, sizeof(float) * m * n, cudaMemcpyDeviceToHost); + cudaDeviceSynchronize(); + + if (!verify_matrix(C_ref, C, m * n)) { + printf("Failed to pass the correctness verification against NVIDIA cuBLAS. Exited.\n"); + exit(EXIT_FAILURE); + } + } + cudaDeviceSynchronize(); + + cudaEventRecord(beg); + for (int j = 0; j < repeat_times; j++) { + test_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, handle); + } + cudaEventRecord(end); + cudaEventSynchronize(beg); + cudaEventSynchronize(end); + cudaEventElapsedTime(&elapsed_time, beg, end); + elapsed_time /= 1000.; //换算成秒 + + printf("Average elasped time: (%f) second, performance: (%f) GFLOPS. size: (%d).\n", + elapsed_time / repeat_times, 2. * 1e-9 * repeat_times * m * n * k / elapsed_time, m); + fflush(stdout); + copy_matrix(C_ref, C, m * n); //sync C with cuBLAS to prepare for the next run + } + + // 释放CPU和GPU空间 + free(A); + free(B); + free(C); + free(C_ref); + cudaFree(dA); + cudaFree(dB); + cudaFree(dC); + cudaFree(dC_ref); + + return 0; +}; diff --git a/upstream_ref/sgemm_cuda/simplest_kernel.cu b/upstream_ref/sgemm_cuda/simplest_kernel.cu new file mode 100644 index 00000000..2fc15ccc --- /dev/null +++ b/upstream_ref/sgemm_cuda/simplest_kernel.cu @@ -0,0 +1,46 @@ +#include +#include +#include + +__global__ void kernel(uint *A, uint *B, int row) { + auto x = threadIdx.x / 4; + auto y = threadIdx.x % 4; + A[x * row + y] = x; + B[x * row + y] = y; +} + +int main(int argc, char **argv) { + uint *Xs, *Ys; + uint *Xs_d, *Ys_d; + + uint SIZE = 4; + + Xs = (uint *)malloc(SIZE * SIZE * sizeof(uint)); + Ys = (uint *)malloc(SIZE * SIZE * sizeof(uint)); + + cudaMalloc((void **)&Xs_d, SIZE * SIZE * sizeof(uint)); + cudaMalloc((void **)&Ys_d, SIZE * SIZE * sizeof(uint)); + + dim3 grid_size(1, 1, 1); + dim3 block_size(4 * 4); + + kernel<<>>(Xs_d, Ys_d, 4); + + cudaMemcpy(Xs, Xs_d, SIZE * SIZE * sizeof(uint), cudaMemcpyDeviceToHost); + cudaMemcpy(Ys, Ys_d, SIZE * SIZE * sizeof(uint), cudaMemcpyDeviceToHost); + + cudaDeviceSynchronize(); + + for (int row = 0; row < SIZE; ++row) { + for (int col = 0; col < SIZE; ++col) { + std::cout << "[" << Xs[row * SIZE + col] << "|" << Ys[row * SIZE + col] + << "] "; + } + std::cout << "\n"; + } + + cudaFree(Xs_d); + cudaFree(Ys_d); + free(Xs); + free(Ys); +} diff --git a/upstream_ref/sgemm_cuda/utils.cu b/upstream_ref/sgemm_cuda/utils.cu new file mode 100644 index 00000000..bd94ef8b --- /dev/null +++ b/upstream_ref/sgemm_cuda/utils.cu @@ -0,0 +1,199 @@ +#include +#include "utils.cuh" +#include "kernel.cuh" + +float get_sec() { + struct timeval time; + gettimeofday(&time, NULL); + return (1e6 * time.tv_sec + time.tv_usec); +} + +float cpu_elapsed_time(float &beg, float &end) { + return 1.0e-6 * (end - beg); +} + +void cudaCheck(cudaError_t error, const char *file, int line) { + if (error != cudaSuccess) { + printf("[CUDA ERROR] at file %s(line %d):\n%s\n", file, line, cudaGetErrorString(error)); + exit(EXIT_FAILURE); + } + return; +}; + +void CudaDeviceInfo() { + int deviceId; + + cudaGetDevice(&deviceId); + + cudaDeviceProp props; + cudaGetDeviceProperties(&props, deviceId); + + /* + * There should be no need to modify the output string below. + */ + + printf("Device ID: %d\n\ + *Number of SMs: %d\n\ + Compute Capability Major: %d\n\ + Compute Capability Minor: %d\n\ + memoryBusWidth: %d\n\ + *maxThreadsPerBlock: %d\n\ + maxThreadsPerMultiProcessor: %d\n\ + *totalGlobalMem: %zuM\n\ + sharedMemPerBlock: %zuKB\n\ + *sharedMemPerMultiprocessor: %zuKB\n\ + totalConstMem: %zuKB\n\ + *multiProcessorCount: %d\n\ + *Warp Size: %d\n", + deviceId, + props.multiProcessorCount, + props.major, + props.minor, + props.memoryBusWidth, + props.maxThreadsPerBlock, + props.maxThreadsPerMultiProcessor, + props.totalGlobalMem / 1024 / 1024, + props.sharedMemPerBlock / 1024, + props.sharedMemPerMultiprocessor / 1024, + props.totalConstMem / 1024, + props.multiProcessorCount, + props.warpSize); +}; + +void randomize_matrix(float *mat, int N) { + // NOTICE: 使用gettimeofdays替代srand((unsigned)time(NULL));time精度过低，产生相同随机数 + struct timeval time; + gettimeofday(&time, NULL); + srand(time.tv_usec); + for (int i = 0; i < N; i++) { + float tmp = (float) (rand() % 5) + 0.01 * (rand() % 5); + tmp = (rand() % 2 == 0) ? tmp : tmp * (-1.); + mat[i] = tmp; + } +} + +void copy_matrix(float *src, float *dest, int N) { + int i; + for (i = 0; src + i && dest + i && i < N; i++) + *(dest + i) = *(src + i); + if (i != N) + printf("copy failed at %d while there are %d elements in total.\n", i, N); +} + +void print_matrix(const float *A, int M, int N) { + int i; + printf("["); + for (i = 0; i < M * N; i++) { + if ((i + 1) % N == 0) + printf("%5.2f ", A[i]); + else + printf("%5.2f, ", A[i]); + if ((i + 1) % N == 0) { + if (i + 1 < M * N) + printf(";\n"); + } + } + printf("]\n"); +} + +bool verify_matrix(float *mat1, float *mat2, int N) { + double diff = 0.0; + int i; + for (i = 0; mat1 + i && mat2 + i && i < N; i++) { + diff = fabs((double) mat1[i] - (double) mat2[i]); + if (diff > 1e-2) { + printf("error. %5.2f,%5.2f,%d\n", mat1[i], mat2[i], i); + return false; + } + } + return true; +} + +#define CEIL_DIV(M, N) ((M) + (N)-1) / (N) + +void test_cublas(cublasHandle_t handle, int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + //cublas列主序计算：https://www.cnblogs.com/cuancuancuanhao/p/7763256.html + cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, N, A, K, &beta, C, N); +} + +void test_mysgemm_v1(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(32, 32); + dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32)); + mysgemm_v1<<>>(M, N, K, alpha, A, B, beta, C); +} + +void test_mysgemm_v2(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(1024); + dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32)); + mysgemm_v2<32><<>>(M, N, K, alpha, A, B, beta, C); +} + +void test_mysgemm_v3(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(512); + dim3 gridDim(CEIL_DIV(M, 64), CEIL_DIV(N, 64)); + mysgemm_v3<64, 64, 8, 8><<>>(M, N, K, alpha, A, B, beta, C); +} + +void test_mysgemm_v4(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(256); + dim3 gridDim(CEIL_DIV(M, 128), CEIL_DIV(N, 128)); + mysgemm_v4<128, 128, 8, 8, 8><<>>(M, N, K, alpha, A, B, beta, C); +} + +void test_mysgemm_v5(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(256); + dim3 gridDim(CEIL_DIV(M, 128), CEIL_DIV(N, 128)); + mysgemm_v5<128, 128, 8, 8, 8><<>>(M, N, K, alpha, A, B, beta, C); +} + +//void test_mysgemm_v6(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { +// dim3 blockDim(4); +// dim3 gridDim(CEIL_DIV(M, 8), CEIL_DIV(N, 8)); +// mysgemm_v6<8, 8, 4, 4, 4><<>>(M, N, K, alpha, A, B, beta, C); +//} + +void test_mysgemm_v6(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(256); + dim3 gridDim(CEIL_DIV(M, 128), CEIL_DIV(N, 128)); + mysgemm_v6<128, 128, 8, 8, 8><<>>(M, N, K, alpha, A, B, beta, C); +} + +void test_mysgemm_v7(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) { + dim3 blockDim(256); + dim3 gridDim(CEIL_DIV(M, 128), CEIL_DIV(N, 128)); + mysgemm_v7<128, 128, 8, 8, 8><<>>(M, N, K, alpha, A, B, beta, C); +} + + + +void test_kernel(int kernel_num, int M, int N, int K, float alpha, float *A, float *B, float beta, float *C, + cublasHandle_t handle) { + switch (kernel_num) { + case 0: + test_cublas(handle, M, N, K, alpha, A, B, beta, C); + break; + case 1: + test_mysgemm_v1(M, N, K, alpha, A, B, beta, C); + break; + case 2: + test_mysgemm_v2(M, N, K, alpha, A, B, beta, C); + break; + case 3: + test_mysgemm_v3(M, N, K, alpha, A, B, beta, C); + break; + case 4: + test_mysgemm_v4(M, N, K, alpha, A, B, beta, C); + break; + case 5: + test_mysgemm_v5(M, N, K, alpha, A, B, beta, C); + break; + case 6: + test_mysgemm_v6(M, N, K, alpha, A, B, beta, C); + break; + case 7: + test_mysgemm_v7(M, N, K, alpha, A, B, beta, C); + break; + default: + break; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/utils.cuh b/upstream_ref/sgemm_cuda/utils.cuh new file mode 100644 index 00000000..fa423199 --- /dev/null +++ b/upstream_ref/sgemm_cuda/utils.cuh @@ -0,0 +1,42 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +/* +===================================== +CUDA操作 +===================================== +*/ +void cudaCheck(cudaError_t error, const char *file, int line); //CUDA错误检查 +void CudaDeviceInfo(); // 打印CUDA信息 + +/* +===================================== +矩阵操作 +===================================== +*/ +void randomize_matrix(float *mat, int N); // 随机初始化矩阵 +void copy_matrix(float *src, float *dest, int N); // 复制矩阵 +void print_matrix(const float *A, int M, int N); // 打印矩阵 +bool verify_matrix(float *mat1, float *mat2, int N); // 验证矩阵 + +/* +===================================== +计时操作 +===================================== +*/ +float get_current_sec(); // 获取当前时刻 +float cpu_elapsed_time(float &beg, float &end); // 计算时间差 + +/* +===================================== +kernel操作 +===================================== +*/ +//调用指定核函数计算矩阵乘法 +void test_kernel(int kernel_num, int m, int n, int k, float alpha, float *A, float *B, float beta, float *C, cublasHandle_t handle); \ No newline at end of file