data: complete SGEMM upstream from 3 repos (siboehm+wangzyon+edtallison) + xllm fused_qknorm_rope + xattention kernels

SGEMM repos (upstream_ref/sgemm_cuda/, 41 files):
  siboehm/SGEMM_CUDA: kernel 1-12, runner, CMake, cuBLAS benchmark
  wangzyon/NVIDIA_SGEMM_PRACTICE: kernel 1-7 (Chinese comments), utils
  edtallison/sgemm-cuda: kernel 01-09 (learning notes), Makefile

xllm kernels (ex_engine/xllm_kernels/cuda/):
  fused_qknorm_rope.cu + bind — saves 128 kernel launches/fwd
  xattention/ — 6 files from upstream xllm
  headers: corex_compat_utils.h, topk_last_dim.cuh
  ilu/CMakeLists.txt

SO_BUILD_MANIFEST.md — complete .so inventory and call chain analysis

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'
+```

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)"

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 <torch/extension.h>
+
+namespace xllm::kernel::cuda {
+void fused_qk_norm_rope(
+    torch::Tensor& qkv,
+    int64_t num_heads_q,
+    int64_t num_heads_k,
+    int64_t num_heads_v,
+    int64_t head_dim,
+    double eps,
+    const torch::Tensor& q_weight,
+    const torch::Tensor& k_weight,
+    const torch::Tensor& cos_sin_cache,
+    bool interleaved,
+    const torch::Tensor& position_ids);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fused_qk_norm_rope",
+          &xllm::kernel::cuda::fused_qk_norm_rope,
+          "Fused QK-Norm + RoPE (xllm CUDA kernel)",
+          py::arg("qkv"),
+          py::arg("num_heads_q"),
+          py::arg("num_heads_k"),
+          py::arg("num_heads_v"),
+          py::arg("head_dim"),
+          py::arg("eps") = 1e-6,
+          py::arg("q_weight"),
+          py::arg("k_weight"),
+          py::arg("cos_sin_cache"),
+          py::arg("interleaved") = false,
+          py::arg("position_ids"));
+}

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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <torch/cuda.h>
+
+#include <cmath>
+#include <type_traits>
+
+#include "cuda_ops_api.h"
+#include "type_convert.cuh"
+#include "utils.h"
+
+using at::device_of;
+
+// Borrowed from:
+// https://github.com/vllm-project/vllm/blob/022f3cea5327cc720a325c50931e1edcfdf2d32b/csrc/fused_qknorm_rope_kernel.cu
+
+constexpr uint32_t kFinalMask = 0xffffffffu;
+
+namespace {
+
+using namespace xllm::kernel::cuda;
+
+template <typename T, int num>
+struct packed_as;
+// Specialization for packed_as used in this kernel.
+template <>
+struct packed_as<uint, 1> {
+  using type = uint;
+};
+
+template <>
+struct packed_as<uint, 2> {
+  using type = uint2;
+};
+
+template <>
+struct packed_as<uint, 4> {
+  using type = uint4;
+};
+
+template <typename T>
+__inline__ __device__ T warp_reduce_sum(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(kFinalMask, val, mask, 32);
+  return val;
+}
+
+template <typename T>
+inline __device__ __host__ T div_up(T m, T n) {
+  return (m + n - 1) / n;
+}
+
+// Perform per-head QK Norm and RoPE in a single kernel.
+// scalar_t_in: data type of QKV and RMSNorm weights
+// scalar_t_cache: data type of cos/sin cache
+// head_dim: the dimension of each head
+// interleave: interleave=!is_neox.
+template <typename scalar_t_in,
+          typename scalar_t_cache,
+          int head_dim,
+          bool interleave>
+__global__ void fused_qknorm_rope_kernel(
+    void* qkv_void,                  // Combined QKV tensor
+    int const num_heads_q,           // Number of query heads
+    int const num_heads_k,           // Number of key heads
+    int const num_heads_v,           // Number of value heads
+    float const eps,                 // Epsilon for RMS normalization
+    void const* q_weight_void,       // RMSNorm weights for query
+    void const* k_weight_void,       // RMSNorm weights for key
+    void const* cos_sin_cache_void,  // Pre-computed cos/sin cache
+    int64_t const* position_ids,     // Position IDs for RoPE
+    int const num_tokens,            // Number of tokens
+    int const rotary_dim             // Dimension for RoPE
+) {
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
+  if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
+                std::is_same_v<scalar_t_cache, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+
+    using Converter = _typeConvert<scalar_t_in>;
+    static_assert(Converter::exists,
+                  "Input QKV data type is not supported for this CUDA "
+                  "architecture or toolkit version.");
+    using T_in = typename Converter::hip_type;
+    using T2_in = typename Converter::packed_hip_type;
+
+    using CacheConverter = _typeConvert<scalar_t_cache>;
+    static_assert(CacheConverter::exists,
+                  "Cache data type is not supported for this CUDA architecture "
+                  "or toolkit version.");
+    using T_cache = typename CacheConverter::hip_type;
+
+    T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
+    T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
+    T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
+    T_cache const* cos_sin_cache =
+        reinterpret_cast<T_cache const*>(cos_sin_cache_void);
+
+    int const warpsPerBlock = blockDim.x / 32;
+    int const warpId = threadIdx.x / 32;
+    int const laneId = threadIdx.x % 32;
+
+    // Calculate global warp index to determine which head/token this warp
+    // processes
+    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
+
+    // Total number of attention heads (Q and K)
+    int const total_qk_heads = num_heads_q + num_heads_k;
+
+    // Determine which token and head type (Q or K) this warp processes
+    int const tokenIdx = globalWarpIdx / total_qk_heads;
+    int const localHeadIdx = globalWarpIdx % total_qk_heads;
+
+    // Skip if this warp is assigned beyond the number of tokens
+    if (tokenIdx >= num_tokens) return;
+
+    bool const isQ = localHeadIdx < num_heads_q;
+    int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
+
+    int const num_heads = num_heads_q + num_heads_k + num_heads_v;
+
+    static_assert(head_dim % (32 * 2) == 0,
+                  "head_dim must be divisible by 64 (each warp processes one "
+                  "head, and each thread gets even number of "
+                  "elements)");
+    constexpr int numElemsPerThread = head_dim / 32;
+    float elements[numElemsPerThread];
+    constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
+    static_assert(elemSizeBytes % 4 == 0,
+                  "numSizeBytes must be a multiple of 4");
+    constexpr int vecSize =
+        elemSizeBytes /
+        4;  // Use packed_as<uint, vecSize> to perform loading/saving.
+    using vec_T = typename packed_as<uint, vecSize>::type;
+
+    int offsetWarp;  // Offset for the warp
+    if (isQ) {
+      // Q segment: token offset + head offset within Q segment
+      offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
+    } else {
+      // K segment: token offset + entire Q segment + head offset within K
+      // segment
+      offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
+                   headIdx * head_dim;
+    }
+    int offsetThread = offsetWarp + laneId * numElemsPerThread;
+
+    // Sum of squares for RMSNorm
+    float sumOfSquares = 0.0f;
+
+    // Load.
+    {
+      vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
+      constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
+#pragma unroll
+      for (int i = 0; i < num_packed_elems; i++) {
+        // Interpret the generic vector chunk as the specific packed type
+        T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
+        // Convert to float2 for computation
+        float2 vals = Converter::convert(packed_val);
+        sumOfSquares += vals.x * vals.x;
+        sumOfSquares += vals.y * vals.y;
+
+        elements[2 * i] = vals.x;
+        elements[2 * i + 1] = vals.y;
+      }
+    }
+
+    // Reduce sum across warp using the utility function
+    sumOfSquares = warp_reduce_sum(sumOfSquares);
+
+    // Compute RMS normalization factor
+    float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
+
+    // Normalize elements
+#pragma unroll
+    for (int i = 0; i < numElemsPerThread; i++) {
+      int dim = laneId * numElemsPerThread + i;
+      float weight = isQ ? Converter::convert(q_weight[dim])
+                         : Converter::convert(k_weight[dim]);
+      elements[i] *= rms_rcp * weight;
+    }
+
+    // Apply RoPE to normalized elements
+    float elements2[numElemsPerThread];  // Additional buffer required for RoPE.
+
+    int64_t pos_id = position_ids[tokenIdx];
+
+    // Calculate cache pointer for this position - similar to
+    // pos_encoding_kernels.cu
+    T_cache const* cache_ptr = cos_sin_cache + pos_id * rotary_dim;
+    int const embed_dim = rotary_dim / 2;
+    T_cache const* cos_ptr = cache_ptr;
+    T_cache const* sin_ptr = cache_ptr + embed_dim;
+    int const rotary_lanes = rotary_dim / numElemsPerThread;  // rotary range
+    if (laneId < rotary_lanes) {
+      if constexpr (interleave) {
+        // Perform interleaving. Use pre-computed cos/sin values.
+#pragma unroll
+        for (int i = 0; i < numElemsPerThread / 2; ++i) {
+          int const idx0 = 2 * i;
+          int const idx1 = 2 * i + 1;
+          // Global dimension index in the head
+          int const dim_idx = laneId * numElemsPerThread + idx0;
+
+          float const val0 = elements[idx0];
+          float const val1 = elements[idx1];
+
+          int const half_dim = dim_idx / 2;
+          float const cos_val =
+              CacheConverter::convert(__ldg(cos_ptr + half_dim));
+          float const sin_val =
+              CacheConverter::convert(__ldg(sin_ptr + half_dim));
+
+          elements[idx0] = val0 * cos_val - val1 * sin_val;
+          elements[idx1] = val0 * sin_val + val1 * cos_val;
+        }
+      } else {
+        // Before data exchange with in warp, we need to sync.
+        __syncwarp();
+        int pairOffset = (rotary_dim / 2) / numElemsPerThread;
+        // Get the data from the other half of the warp. Use pre-computed
+        // cos/sin values.
+#pragma unroll
+        for (int i = 0; i < numElemsPerThread; i++) {
+          elements2[i] = __shfl_xor_sync(kFinalMask, elements[i], pairOffset);
+
+          if (laneId < pairOffset) {
+            elements2[i] = -elements2[i];
+          }
+          int dim_idx = laneId * numElemsPerThread + i;
+
+          dim_idx = (dim_idx * 2) % rotary_dim;
+          int half_dim = dim_idx / 2;
+          float cos_val = CacheConverter::convert(__ldg(cos_ptr + half_dim));
+          float sin_val = CacheConverter::convert(__ldg(sin_ptr + half_dim));
+
+          elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
+        }
+        // __shfl_xor_sync does not provide memfence. Need to sync again.
+        __syncwarp();
+      }
+    }
+    // Store.
+    {
+      vec_T vec;
+      constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
+#pragma unroll
+      for (int i = 0; i < num_packed_elems; i++) {
+        // Convert from float2 back to the specific packed type
+        float2 vals = {elements[2 * i], elements[2 * i + 1]};
+        T2_in packed_val = Converter::convert(vals);
+        // Place it into the generic vector
+        *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
+      }
+      *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
+    }
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
+  }
+#endif
+}
+
+// Borrowed from
+// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
+#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
+  if (interleave) {                                      \
+    const bool INTERLEAVE = true;                        \
+    __VA_ARGS__                                          \
+  } else {                                               \
+    const bool INTERLEAVE = false;                       \
+    __VA_ARGS__                                          \
+  }
+
+template <typename scalar_t_in, typename scalar_t_cache>
+void launch_fused_qknorm_rope(void* qkv,
+                              int const num_tokens,
+                              int const num_heads_q,
+                              int const num_heads_k,
+                              int const num_heads_v,
+                              int const head_dim,
+                              int const rotary_dim,
+                              float const eps,
+                              void const* q_weight,
+                              void const* k_weight,
+                              void const* cos_sin_cache,
+                              bool const interleave,
+                              int64_t const* position_ids,
+                              cudaStream_t stream) {
+  constexpr int blockSize = 256;
+
+  int const warpsPerBlock = blockSize / 32;
+  int const totalQKHeads = num_heads_q + num_heads_k;
+  int const totalWarps = num_tokens * totalQKHeads;
+
+  int const gridSize = div_up(totalWarps, warpsPerBlock);
+  dim3 gridDim(gridSize);
+  dim3 blockDim(blockSize);
+
+  switch (head_dim) {
+    case 64:
+      DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
+        fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
+            <<<gridDim, blockDim, 0, stream>>>(qkv,
+                                               num_heads_q,
+                                               num_heads_k,
+                                               num_heads_v,
+                                               eps,
+                                               q_weight,
+                                               k_weight,
+                                               cos_sin_cache,
+                                               position_ids,
+                                               num_tokens,
+                                               rotary_dim);
+      });
+      break;
+    case 128:
+      DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
+        fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
+            <<<gridDim, blockDim, 0, stream>>>(qkv,
+                                               num_heads_q,
+                                               num_heads_k,
+                                               num_heads_v,
+                                               eps,
+                                               q_weight,
+                                               k_weight,
+                                               cos_sin_cache,
+                                               position_ids,
+                                               num_tokens,
+                                               rotary_dim);
+      });
+      break;
+    case 256:
+      DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
+        fused_qknorm_rope_kernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
+            <<<gridDim, blockDim, 0, stream>>>(qkv,
+                                               num_heads_q,
+                                               num_heads_k,
+                                               num_heads_v,
+                                               eps,
+                                               q_weight,
+                                               k_weight,
+                                               cos_sin_cache,
+                                               position_ids,
+                                               num_tokens,
+                                               rotary_dim);
+      });
+      break;
+    default:
+      CHECK(false) << "Unsupported head dimension for fusedQKNormRope: "
+                   << head_dim;
+  }
+}
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+void fused_qk_norm_rope(
+    torch::Tensor& qkv,   // Combined QKV tensor [num_tokens,
+                          // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
+    int64_t num_heads_q,  // Number of query heads
+    int64_t num_heads_k,  // Number of key heads
+    int64_t num_heads_v,  // Number of value heads
+    int64_t head_dim,     // Dimension per head
+    double eps,           // Epsilon for RMS normalization
+    const torch::Tensor& q_weight,  // RMSNorm weights for query [head_dim]
+    const torch::Tensor& k_weight,  // RMSNorm weights for key [head_dim]
+    const torch::Tensor&
+        cos_sin_cache,  // Cos/sin cache [max_position, rotary_dim]
+    bool interleaved,   // Whether RoPE is applied in interleaved style
+    const torch::Tensor& position_ids  // Position IDs for RoPE [num_tokens]
+) {
+  // Input validation
+  CHECK(qkv.is_cuda()) << "qkv must be a CUDA tensor";
+  CHECK(qkv.is_contiguous()) << "qkv must be contiguous";
+  CHECK(position_ids.is_cuda()) << "position_ids must be a CUDA tensor";
+  CHECK(position_ids.is_contiguous()) << "position_ids must be contiguous";
+  CHECK(q_weight.is_cuda()) << "q_weight must be a CUDA tensor";
+  CHECK(q_weight.is_contiguous()) << "q_weight must be contiguous";
+  CHECK(k_weight.is_cuda()) << "k_weight must be a CUDA tensor";
+  CHECK(k_weight.is_contiguous()) << "k_weight must be contiguous";
+  CHECK(cos_sin_cache.is_cuda()) << "cos_sin_cache must be a CUDA tensor";
+  CHECK(cos_sin_cache.is_contiguous()) << "cos_sin_cache must be contiguous";
+  CHECK(position_ids.scalar_type() == torch::kInt64)
+      << "position_ids dtype is " << position_ids.scalar_type()
+      << ", while Int64 is expected";
+
+  CHECK(qkv.dim() == 2) << "QKV tensor must be 2D: [num_tokens, "
+                        << "(num_heads_q+num_heads_k+num_heads_v)*head_dim]";
+  CHECK(position_ids.dim() == 1) << "Position IDs must be 1D: [num_tokens]";
+  CHECK(q_weight.dim() == 1) << "Query weights must be 1D: [head_dim]";
+  CHECK(k_weight.dim() == 1) << "Key weights must be 1D: [head_dim]";
+  CHECK(cos_sin_cache.dim() == 2)
+      << "Cos/sin cache must be 2D: [max_position, rotary_dim]";
+  CHECK(q_weight.size(0) == head_dim)
+      << "Query weights size must match head dimension";
+  CHECK(k_weight.size(0) == head_dim)
+      << "Key weights size must match head dimension";
+
+  CHECK(cos_sin_cache.size(1) % 2 == 0) << "rotary_dim must be even";
+  CHECK(cos_sin_cache.size(1) <= head_dim)
+      << "rotary_dim must be less than or equal to head_dim";
+
+  CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
+        qkv.scalar_type() == k_weight.scalar_type())
+      << "qkv, q_weight and k_weight must have the same dtype";
+
+  int64_t num_tokens = qkv.size(0);
+  CHECK(position_ids.size(0) == num_tokens)
+      << "Number of tokens in position_ids must match QKV";
+
+  int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
+  CHECK(qkv.size(1) == total_heads * head_dim)
+      << "QKV tensor size must match total number of heads and head dimension";
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(qkv));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
+    using qkv_scalar_t = scalar_t;
+    DISPATCH_FLOATING_TYPES(
+        cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
+          using cache_scalar_t = scalar_t;
+          launch_fused_qknorm_rope<qkv_scalar_t, cache_scalar_t>(
+              qkv.data_ptr(),
+              static_cast<int>(num_tokens),
+              static_cast<int>(num_heads_q),
+              static_cast<int>(num_heads_k),
+              static_cast<int>(num_heads_v),
+              static_cast<int>(head_dim),
+              static_cast<int>(cos_sin_cache.size(1)),
+              static_cast<float>(eps),
+              q_weight.data_ptr(),
+              k_weight.data_ptr(),
+              cos_sin_cache.data_ptr(),
+              interleaved,
+              reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
+              stream);
+        });
+  });
+}
+
+}  // namespace xllm::kernel::cuda

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 <torch/torch.h>
+#include <c10/cuda/CUDAGuard.h>
+
+// Replace glog CHECK with TORCH_CHECK
+#ifndef CHECK
+#define CHECK(cond) TORCH_CHECK(cond)
+#endif
+
+#ifndef CHECK_EQ
+#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b))
+#endif
+
+#ifndef CHECK_GE
+#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b))
+#endif
+
+// Include device_utils for DISPATCH_HALF_TYPES etc
+#include "device_utils.cuh"
+
+// ffi namespace stub (some headers reference it)
+namespace ffi {
+template <typename T>
+using Array = std::vector<T>;
+}
+
+// HOST_DEVICE_INLINE
+#if defined(__CUDACC__) || defined(_NVHPC_CUDA)
+#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
+#else
+#define HOST_DEVICE_INLINE inline
+#endif

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 <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+#include <torch/script.h>
+#include <torch/torch.h>
+
+#include "cuda.h"
+
+namespace xllm::kernel::cuda {
+
+void beam_search(torch::Tensor acc_logprob,
+                 torch::Tensor in_sequence_group,
+                 torch::Tensor top_tokens,
+                 torch::Tensor top_logprobs,
+                 torch::Tensor out_acc_logprob,
+                 torch::Tensor out_token_ids,
+                 torch::Tensor out_token_index,
+                 torch::Tensor out_beam_count_prefix_sums,
+                 torch::Tensor out_sequence_group,
+                 uint32_t batch_size,
+                 uint32_t current_step) {
+  torch::Device device = acc_logprob.device();
+
+  uint32_t beam_size = in_sequence_group.size(1);
+
+  uint32_t top_k = top_tokens.size(1);
+  uint32_t total_rounds = in_sequence_group.size(2);
+
+  CHECK_EQ(beam_size, top_k) << "beam_size must be equal with top_k.";
+
+  if (current_step == 0) {
+    auto tokens_view =
+        top_tokens.view({batch_size, top_k}).slice(1, 0, beam_size);
+    auto init_probs_view =
+        top_logprobs.view({batch_size, top_k}).slice(1, 0, beam_size);
+
+    out_token_ids.view({batch_size, beam_size}).copy_(tokens_view);
+    out_acc_logprob.view({batch_size, beam_size}).copy_(init_probs_view);
+
+    auto indices =
+        torch::arange(
+            beam_size,
+            torch::TensorOptions().dtype(torch::kInt32).device(device))
+            .unsqueeze(0)
+            .expand({batch_size, -1})
+            .reshape({-1, 1});
+    out_token_index.copy_(indices);
+
+    auto sequence_view =
+        out_sequence_group.view({batch_size, beam_size, total_rounds});
+    sequence_view.slice(2, 0, 1).squeeze(2).copy_(tokens_view);
+
+  } else {
+    auto combined_probs =
+        (acc_logprob + top_logprobs).view({batch_size, beam_size * top_k});
+
+    auto topk_result = torch::topk(combined_probs, beam_size, -1);
+    auto new_probs = std::get<0>(topk_result);    // [batch_size, beam_size]
+    auto new_indices = std::get<1>(topk_result);  // [batch_size, beam_size]
+
+    auto ordered_indices = new_indices.argsort(static_cast<int64_t>(1), false);
+    // Reorder new_probs (and corresponding new_indices) by ordered_indices to
+    // keep alignment.
+    if (current_step < total_rounds - 1) {
+      new_probs = new_probs.gather(1, ordered_indices);
+      new_indices = new_indices.gather(1, ordered_indices);
+    }
+
+    auto parent_beam = (new_indices / top_k).to(torch::kLong);
+    auto token_in_beam = (new_indices % top_k).to(torch::kLong);
+
+    auto top_tokens_reshaped = top_tokens.view({batch_size, beam_size, top_k});
+
+    auto batch_idx =
+        torch::arange(batch_size,
+                      torch::TensorOptions().dtype(torch::kLong).device(device))
+            .unsqueeze(1)
+            .expand_as(parent_beam);
+
+    using torch::indexing::TensorIndex;
+    auto new_tokens = top_tokens_reshaped.index({TensorIndex(batch_idx),
+                                                 TensorIndex(parent_beam),
+                                                 TensorIndex(token_in_beam)});
+
+    out_acc_logprob.view({batch_size, beam_size}).copy_(new_probs);
+    out_token_index.view({batch_size, beam_size})
+        .copy_(new_indices.to(torch::kInt32));
+    out_token_ids.view({batch_size, beam_size}).copy_(new_tokens);
+
+    auto batch_range =
+        torch::arange(
+            batch_size,
+            torch::TensorOptions().dtype(torch::kInt32).device(device))
+            .unsqueeze(1)
+            .expand({-1, beam_size});
+    auto beam_range =
+        torch::arange(
+            beam_size,
+            torch::TensorOptions().dtype(torch::kInt32).device(device))
+            .unsqueeze(0)
+            .expand({batch_size, -1});
+
+    using torch::indexing::Slice;
+    using torch::indexing::TensorIndex;
+    out_sequence_group.slice(2, 0, current_step) =
+        in_sequence_group.index({TensorIndex(batch_range),
+                                 TensorIndex(parent_beam.to(torch::kInt32)),
+                                 Slice(0, current_step)});
+
+    out_sequence_group.slice(2, current_step, current_step + 1) =
+        new_tokens.unsqueeze(2);
+  }
+}
+
+}  // namespace xllm::kernel::cuda

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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <vector>
+
+#include "xattention_ops_api.h"
+
+namespace {
+
+// In-place cache selection kernel for Xattention.
+// Reorders KV cache entries based on beam search results. After beam search,
+// the beam indices may have changed, and this kernel copies KV cache data from
+// old beam positions to new beam positions to maintain consistency.
+// Inputs:
+//   k_ptrs_i64      : [Layer] - pointers to K cache tensors for each layer
+//   v_ptrs_i64      : [Layer] - pointers to V cache tensors for each layer
+//   beam_index      : [B*Beam] - mapping from new beam index to old beam index
+//   block_table     : [B] - request ID per batch item (extracted from [B*Beam,
+//   1]) B               : batch size (actual batch size, not batch_size *
+//   beam_size) Beam            : beam width Kv              : number of KV
+//   heads MaxStep         : maximum decode steps D               : head
+//   dimension MaxReq          : maximum number of requests Layer           :
+//   number of transformer layers decode_step     : current decode step
+//   (0-indexed)
+// Cache layout: [MaxReq, Beam, MaxStep, Kv, D]
+// The kernel performs two passes to avoid overwriting data:
+//   pass-1: copy from old_beam > new_beam (increasing new_beam)
+//   pass-2: copy from old_beam < new_beam (decreasing new_beam)
+template <typename scalar_t>
+__global__ void cache_select_inplace_ptrs_kernel(
+    const int64_t* __restrict__ k_ptrs_i64,   // [Layer]
+    const int64_t* __restrict__ v_ptrs_i64,   // [Layer]
+    const int32_t* __restrict__ beam_index,   // [B*Beam]
+    const int32_t* __restrict__ block_table,  // [B]
+    int32_t B,
+    int32_t Beam,
+    int32_t Kv,
+    int32_t MaxStep,
+    int32_t D,
+    int32_t MaxReq,
+    int32_t Layer,
+    int32_t decode_step) {
+  const int32_t b = static_cast<int32_t>(blockIdx.x);
+  const int32_t kv = static_cast<int32_t>(blockIdx.y);
+  const int32_t layer = static_cast<int32_t>(blockIdx.z);
+
+  if (b >= B || kv >= Kv || layer >= Layer) {
+    return;
+  }
+
+  const int32_t step_end =
+      decode_step < (MaxStep - 1) ? decode_step : (MaxStep - 1);
+
+  const int32_t req = block_table[b];
+  if (req < 0 || req >= MaxReq) {
+    return;
+  }
+
+  scalar_t* __restrict__ k_cache =
+      reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(k_ptrs_i64[layer]));
+  scalar_t* __restrict__ v_cache =
+      reinterpret_cast<scalar_t*>(static_cast<uintptr_t>(v_ptrs_i64[layer]));
+
+  // base(req, beam, s, kv, d) = ((((req*Beam + beam)*MaxStep + s)*Kv + kv) * D
+  // + d)
+  const int64_t req_base = static_cast<int64_t>(req) * Beam;
+  const int64_t step_kv_stride = static_cast<int64_t>(Kv) * D;
+  const int64_t kv_d_base = static_cast<int64_t>(kv) * D;
+
+  // grid_step is typically small; loop over s in-kernel to reduce launch
+  // blocks.
+  for (int32_t s = 0; s <= step_end; ++s) {
+    // pass-1: new_beam increasing, copy if old_beam > new_beam
+    for (int32_t new_beam = 0; new_beam < Beam; ++new_beam) {
+      const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
+      if (old_beam >= 0 && old_beam < Beam && old_beam > new_beam) {
+        const int64_t dst_base =
+            ((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
+        const int64_t src_base =
+            ((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
+        for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
+             d += static_cast<int32_t>(blockDim.x)) {
+          k_cache[dst_base + d] = k_cache[src_base + d];
+          v_cache[dst_base + d] = v_cache[src_base + d];
+        }
+      }
+    }
+
+    // pass-2: new_beam decreasing, copy if old_beam < new_beam
+    for (int32_t new_beam = Beam - 1; new_beam >= 0; --new_beam) {
+      const int32_t old_beam = beam_index[b * Beam + new_beam] / Beam;
+      if (old_beam >= 0 && old_beam < Beam && old_beam < new_beam) {
+        const int64_t dst_base =
+            ((req_base + new_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
+        const int64_t src_base =
+            ((req_base + old_beam) * MaxStep + s) * step_kv_stride + kv_d_base;
+        for (int32_t d = static_cast<int32_t>(threadIdx.x); d < D;
+             d += static_cast<int32_t>(blockDim.x)) {
+          k_cache[dst_base + d] = k_cache[src_base + d];
+          v_cache[dst_base + d] = v_cache[src_base + d];
+        }
+      }
+    }
+  }
+}
+
+void cache_select_cuda_launch_ptrs(
+    torch::Tensor k0,
+    torch::Tensor v0,
+    torch::Tensor k_ptrs_i64,       // [Layer] int64 (CUDA)
+    torch::Tensor v_ptrs_i64,       // [Layer] int64 (CUDA)
+    torch::Tensor beam_index_i32,   // [B*Beam, 1] int32
+    torch::Tensor block_table_i32,  // [B] int32
+    int64_t decode_step,
+    int64_t layer_num) {
+  CHECK(k_ptrs_i64.is_cuda() && v_ptrs_i64.is_cuda())
+      << "k_ptrs_i64/v_ptrs_i64 must be CUDA";
+  CHECK_EQ(k_ptrs_i64.scalar_type(), torch::kInt64)
+      << "k_ptrs_i64/v_ptrs_i64 must be int64";
+  CHECK_EQ(v_ptrs_i64.scalar_type(), torch::kInt64)
+      << "k_ptrs_i64/v_ptrs_i64 must be int64";
+  CHECK(k_ptrs_i64.is_contiguous() && v_ptrs_i64.is_contiguous())
+      << "k_ptrs_i64/v_ptrs_i64 must be contiguous";
+
+  const int64_t B64 = block_table_i32.size(0);
+  const int64_t Beam64 = k0.size(1);
+  const int64_t MaxStep64 = k0.size(2);
+  const int64_t Kv64 = k0.size(3);
+  const int64_t D64 = k0.size(4);
+  const int64_t MaxReq64 = k0.size(0);
+  const int64_t Layer64 = layer_num;
+
+  const int32_t B = static_cast<int32_t>(B64);
+  const int32_t Beam = static_cast<int32_t>(Beam64);
+  const int32_t Kv = static_cast<int32_t>(Kv64);
+  const int32_t MaxStep = static_cast<int32_t>(MaxStep64);
+  const int32_t D = static_cast<int32_t>(D64);
+  const int32_t MaxReq = static_cast<int32_t>(MaxReq64);
+  const int32_t Layer = static_cast<int32_t>(Layer64);
+  const int32_t decode_step_i32 = static_cast<int32_t>(decode_step);
+
+  // Warp-aligned threads, capped to keep occupancy reasonable.
+  int threads_per_block = ((D + 31) / 32) * 32;
+  if (threads_per_block < 32) {
+    threads_per_block = 32;
+  }
+  if (threads_per_block > 256) {
+    threads_per_block = 256;
+  }
+  dim3 block_dim(static_cast<unsigned int>(threads_per_block), 1, 1);
+
+  CHECK_LE(Kv64, static_cast<int64_t>(UINT32_MAX)) << "Kv too large for grid.y";
+  CHECK_LE(Layer64, 65535) << "layer_num too large for grid.z";
+  dim3 grid_dim(static_cast<unsigned int>(B),
+                static_cast<unsigned int>(Kv),
+                static_cast<unsigned int>(Layer));
+
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(torch::ScalarType::Half,
+                                  torch::ScalarType::BFloat16,
+                                  k0.scalar_type(),
+                                  "cache_select_inplace_ptrs_kernel",
+                                  [&] {
+                                    cache_select_inplace_ptrs_kernel<scalar_t>
+                                        <<<grid_dim, block_dim, 0, stream>>>(
+                                            k_ptrs_i64.data_ptr<int64_t>(),
+                                            v_ptrs_i64.data_ptr<int64_t>(),
+                                            beam_index_i32.data_ptr<int32_t>(),
+                                            block_table_i32.data_ptr<int32_t>(),
+                                            B,
+                                            Beam,
+                                            Kv,
+                                            MaxStep,
+                                            D,
+                                            MaxReq,
+                                            Layer,
+                                            decode_step_i32);
+                                  });
+
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+void cache_select(const torch::Tensor& beam_index,  // [B*Beam, 1]
+                  std::vector<torch::Tensor>& unshared_k_cache,
+                  std::vector<torch::Tensor>& unshared_v_cache,
+                  const torch::Tensor& block_table,  // [B*Beam, 1]
+                  int64_t decode_step,
+                  int64_t beam_size,
+                  int64_t layer_num) {
+  CHECK_GE(layer_num, 0) << "layer_num must be >= 0";
+  if (layer_num == 0) {
+    return;
+  }
+  CHECK_EQ(static_cast<int64_t>(unshared_k_cache.size()), layer_num)
+      << "unshared_k_cache length mismatch";
+  CHECK_EQ(static_cast<int64_t>(unshared_v_cache.size()), layer_num)
+      << "unshared_v_cache length mismatch";
+
+  CHECK(beam_index.is_cuda()) << "beam_index must be CUDA";
+  CHECK(block_table.is_cuda()) << "block_table must be CUDA";
+  CHECK_EQ(block_table.dim(), 2) << "block_table must be [B*Beam, 1]";
+  CHECK_EQ(block_table.size(1), 1) << "block_table must be [B*Beam, 1]";
+  CHECK_EQ(beam_index.dim(), 2) << "beam_index must be [B*Beam, 1]";
+  CHECK_EQ(beam_index.size(1), 1) << "beam_index must be [B*Beam, 1]";
+  CHECK_GE(decode_step, 0) << "decode_step must be >= 0";
+  CHECK_GT(beam_size, 0) << "beam_size must be > 0";
+
+  // block_table is [B*Beam, 1] with sequential values [0,1,2,3,...]
+  // Infer actual batch_size
+  CHECK_EQ(block_table.size(0) % beam_size, 0)
+      << "block_table.size(0) must be divisible by beam_size";
+  const int64_t B = block_table.size(0) / beam_size;
+  CHECK_EQ(beam_index.size(0), B * beam_size)
+      << "beam_index size mismatch with B*beam_size";
+
+  // Prepare indices (int32, contiguous).
+  auto beam_index_i32 = beam_index.to(torch::kInt32).contiguous();
+  auto block_table_i32 = torch::arange(
+      0,
+      B,
+      torch::TensorOptions().dtype(torch::kInt32).device(block_table.device()));
+  // Validate shapes/dtypes against layer 0.
+  const auto& k0 = unshared_k_cache[0];
+  const auto& v0 = unshared_v_cache[0];
+  CHECK(k0.is_cuda() && v0.is_cuda()) << "cache must be CUDA";
+  CHECK(k0.is_contiguous() && v0.is_contiguous()) << "cache must be contiguous";
+  CHECK_EQ(k0.dim(), 5) << "cache must be 5D [MaxReq, Beam, MaxStep, Kv, D]";
+  CHECK_EQ(v0.sizes(), k0.sizes()) << "k/v cache shapes must match";
+  CHECK_EQ(k0.size(1), beam_size) << "beam_size mismatch with cache";
+  CHECK_LT(decode_step, k0.size(2)) << "decode_step must be < max_decode_step";
+
+  // Pack layer pointers into CUDA int64 tensors so we can launch once.
+  // Note: pointer values are produced on host (data_ptr()), then copied to GPU.
+  c10::cuda::CUDAGuard device_guard(k0.device());
+  auto ptr_cuda_opts =
+      torch::TensorOptions().dtype(torch::kInt64).device(k0.device());
+  auto k_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
+  auto v_ptrs_i64 = torch::empty({layer_num}, ptr_cuda_opts);
+  std::vector<int64_t> k_ptrs_host(static_cast<size_t>(layer_num));
+  std::vector<int64_t> v_ptrs_host(static_cast<size_t>(layer_num));
+
+  for (int64_t layer = 0; layer < layer_num; ++layer) {
+    auto k = unshared_k_cache[static_cast<size_t>(layer)];
+    auto v = unshared_v_cache[static_cast<size_t>(layer)];
+    CHECK(k.is_cuda() && v.is_cuda()) << "cache must be CUDA";
+    CHECK(k.is_contiguous() && v.is_contiguous()) << "cache must be contiguous";
+    CHECK_EQ(k.sizes(), k0.sizes()) << "all layers must have same cache shape";
+    CHECK_EQ(v.sizes(), k0.sizes()) << "all layers must have same cache shape";
+    CHECK_EQ(k.scalar_type(), k0.scalar_type())
+        << "all layers must have same dtype";
+    CHECK_EQ(v.scalar_type(), k0.scalar_type())
+        << "all layers must have same dtype";
+    CHECK_EQ(k.get_device(), k0.get_device())
+        << "all layers must be on the same CUDA device";
+    CHECK_EQ(v.get_device(), k0.get_device())
+        << "all layers must be on the same CUDA device";
+
+    k_ptrs_host[static_cast<size_t>(layer)] =
+        static_cast<int64_t>(reinterpret_cast<uintptr_t>(k.data_ptr()));
+    v_ptrs_host[static_cast<size_t>(layer)] =
+        static_cast<int64_t>(reinterpret_cast<uintptr_t>(v.data_ptr()));
+  }
+
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  C10_CUDA_CHECK(
+      cudaMemcpyAsync(k_ptrs_i64.data_ptr<int64_t>(),
+                      k_ptrs_host.data(),
+                      static_cast<size_t>(layer_num) * sizeof(int64_t),
+                      cudaMemcpyHostToDevice,
+                      stream));
+  C10_CUDA_CHECK(
+      cudaMemcpyAsync(v_ptrs_i64.data_ptr<int64_t>(),
+                      v_ptrs_host.data(),
+                      static_cast<size_t>(layer_num) * sizeof(int64_t),
+                      cudaMemcpyHostToDevice,
+                      stream));
+
+  cache_select_cuda_launch_ptrs(k0,
+                                v0,
+                                k_ptrs_i64,
+                                v_ptrs_i64,
+                                beam_index_i32,
+                                block_table_i32,
+                                decode_step,
+                                layer_num);
+}
+
+}  // namespace xllm::kernel::cuda

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 <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <torch/cuda.h>
+
+#include <cstdint>
+#include <type_traits>
+
+#include "kernels/cuda/utils.h"
+#include "xattention_ops_api.h"
+
+namespace {
+
+template <typename scalar_t>
+struct VecType;
+
+template <>
+struct VecType<c10::Half> {
+  using type = uint4;  // 8 elements * 2 bytes = 16 bytes
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<c10::BFloat16> {
+  using type = uint4;  // 8 elements * 2 bytes = 16 bytes
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<float> {
+  using type = float4;  // 4 elements * 4 bytes = 16 bytes
+  static constexpr int32_t vec_width = 4;
+};
+
+// decoder reshape and cache kernel.
+// Copies proj_k and proj_v into unshared_k_cache / unshared_v_cache.
+// Inputs:
+//   proj_k           : [batch_size, beam_size, kv_heads, head_dim]
+//   proj_v           : [batch_size, beam_size, kv_heads, head_dim]
+//   step             : [1] - current decode step
+//   batch_size       : batch size
+//   beam_size        : beam size
+//   kv_heads         : number of kv heads
+//   head_dim         : head dimension
+//   k_stride0        : proj_k.stride(0)
+//   k_stride1        : proj_k.stride(1)
+//   v_stride0        : proj_v.stride(0)
+//   v_stride1        : proj_v.stride(1)
+//   cache_stride0    : unshared_k_cache.stride(0)
+//   cache_stride1    : unshared_k_cache.stride(1)
+//   cache_stride2    : unshared_k_cache.stride(2)
+//   cache_stride3    : unshared_k_cache.stride(3)
+// Outputs:
+//   unshared_k_cache : [max_batch_size, beam_size, max_step, kv_heads,
+//   head_dim]
+//   unshared_v_cache : [max_batch_size, beam_size, max_step, kv_heads,
+//   head_dim]
+
+template <typename scalar_t>
+__global__ void decoder_reshape_and_cache_kernel(
+    const scalar_t* __restrict__ proj_k,
+    const scalar_t* __restrict__ proj_v,
+    scalar_t* __restrict__ unshared_k_cache,
+    scalar_t* __restrict__ unshared_v_cache,
+    const int32_t* __restrict__ step,
+    const int64_t batch_size,
+    const int64_t beam_size,
+    const int64_t kv_heads,
+    const int64_t head_dim,
+    const int64_t k_stride0,
+    const int64_t k_stride1,
+    const int64_t v_stride0,
+    const int64_t v_stride1,
+    const int64_t cache_stride0,
+    const int64_t cache_stride1,
+    const int64_t cache_stride2,
+    const int64_t cache_stride3) {
+  using VecTypeT = typename VecType<scalar_t>::type;
+  constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
+
+  const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
+  const int64_t total_tokens = batch_size * beam_size;
+  if (token_idx >= total_tokens) {
+    return;
+  }
+
+  const int64_t batch_idx = token_idx / beam_size;
+  const int64_t beam_idx = token_idx - batch_idx * beam_size;
+
+  __shared__ int32_t current_step_s;
+  if (threadIdx.x == 0) {
+    current_step_s = __ldg(step);
+  }
+  __syncthreads();
+  const int64_t current_step = static_cast<int64_t>(current_step_s);
+
+  const int64_t vecs_per_head = head_dim / VEC_WIDTH;
+  const int64_t total_vecs = kv_heads * vecs_per_head;
+
+  const int64_t k_token_base = batch_idx * k_stride0 + beam_idx * k_stride1;
+  const int64_t v_token_base = batch_idx * v_stride0 + beam_idx * v_stride1;
+  const int64_t dst_token_base = batch_idx * cache_stride0 +
+                                 beam_idx * cache_stride1 +
+                                 current_step * cache_stride2;
+
+  for (int64_t linear_idx = static_cast<int64_t>(threadIdx.x);
+       linear_idx < total_vecs;
+       linear_idx += static_cast<int64_t>(blockDim.x)) {
+    const int64_t head_idx = linear_idx / vecs_per_head;
+    const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
+    const int64_t vec_offset = vec_idx * VEC_WIDTH;
+
+    const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
+        proj_k + k_token_base + head_idx * head_dim + vec_offset);
+    const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
+        proj_v + v_token_base + head_idx * head_dim + vec_offset);
+    auto* k_dst_vec =
+        reinterpret_cast<VecTypeT*>(unshared_k_cache + dst_token_base +
+                                    head_idx * cache_stride3 + vec_offset);
+    auto* v_dst_vec =
+        reinterpret_cast<VecTypeT*>(unshared_v_cache + dst_token_base +
+                                    head_idx * cache_stride3 + vec_offset);
+
+    *k_dst_vec = *k_src_vec;
+    *v_dst_vec = *v_src_vec;
+  }
+}
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+
+void decoder_reshape_and_cache(torch::Tensor proj_k,
+                               torch::Tensor proj_v,
+                               torch::Tensor unshared_k_cache,
+                               torch::Tensor unshared_v_cache,
+                               torch::Tensor step) {
+  CHECK_EQ(proj_k.dim(), 4) << "proj_k must be 4-dimensional";
+  CHECK_EQ(proj_v.dim(), 4) << "proj_v must be 4-dimensional";
+  CHECK_EQ(unshared_k_cache.dim(), 5)
+      << "unshared_k_cache must be 5-dimensional";
+  CHECK_EQ(unshared_v_cache.dim(), 5)
+      << "unshared_v_cache must be 5-dimensional";
+  CHECK(proj_k.is_cuda() && proj_v.is_cuda() && unshared_k_cache.is_cuda() &&
+        unshared_v_cache.is_cuda() && step.is_cuda())
+      << "all tensors must be CUDA tensors";
+  CHECK_EQ(step.dim(), 1) << "step must be 1-dimensional";
+  CHECK_EQ(step.size(0), 1) << "step must have shape [1]";
+  CHECK_EQ(step.scalar_type(), at::ScalarType::Int)
+      << "step must be int32 (torch::kInt32)";
+
+  const int64_t batch_size = proj_k.size(0);
+  const int64_t beam_size = proj_k.size(1);
+  const int64_t kv_heads = proj_k.size(2);
+  const int64_t head_dim = proj_k.size(3);
+
+  CHECK_EQ(proj_v.sizes(), proj_k.sizes())
+      << "proj_v and proj_k must have same shape";
+  CHECK_EQ(unshared_k_cache.size(3), kv_heads)
+      << "unshared_k_cache kv_heads mismatch";
+  CHECK_EQ(unshared_k_cache.size(4), head_dim)
+      << "unshared_k_cache head_dim mismatch";
+  CHECK(unshared_v_cache.sizes() == unshared_k_cache.sizes())
+      << "unshared_v_cache and unshared_k_cache must have same shape";
+
+  // This kernel is specialized for qkv-slice layouts:
+  // last dim contiguous and kv head stride tightly packed by head_dim.
+  CHECK_EQ(proj_k.stride(3), 1) << "proj_k must satisfy stride(3)=1";
+  CHECK_EQ(proj_v.stride(3), 1) << "proj_v must satisfy stride(3)=1";
+  CHECK_EQ(proj_k.stride(2), head_dim)
+      << "proj_k must satisfy stride(2)=head_dim";
+  CHECK_EQ(proj_v.stride(2), head_dim)
+      << "proj_v must satisfy stride(2)=head_dim";
+  CHECK_EQ(unshared_k_cache.stride(4), 1)
+      << "unshared_k_cache must satisfy stride(4)=1";
+  CHECK_EQ(unshared_v_cache.stride(4), 1)
+      << "unshared_v_cache must satisfy stride(4)=1";
+  CHECK_EQ(unshared_k_cache.stride(3), head_dim)
+      << "unshared_k_cache must satisfy stride(3)=head_dim";
+  CHECK_EQ(unshared_v_cache.stride(3), head_dim)
+      << "unshared_v_cache must satisfy stride(3)=head_dim";
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  const int64_t k_stride0 = proj_k.stride(0);
+  const int64_t k_stride1 = proj_k.stride(1);
+  const int64_t v_stride0 = proj_v.stride(0);
+  const int64_t v_stride1 = proj_v.stride(1);
+  const int64_t cache_stride0 = unshared_k_cache.stride(0);
+  const int64_t cache_stride1 = unshared_k_cache.stride(1);
+  const int64_t cache_stride2 = unshared_k_cache.stride(2);
+  const int64_t cache_stride3 = unshared_k_cache.stride(3);
+
+  // Launch kernel: one block per (batch, beam), threads cover
+  // kv_heads*head_dim.
+  const int64_t total_tokens = batch_size * beam_size;
+  dim3 grid_dim(1, static_cast<unsigned int>(total_tokens), 1);
+
+  DISPATCH_FLOATING_TYPES(
+      proj_k.scalar_type(), "decoder_reshape_and_cache_kernel", [&] {
+        constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
+                                       std::is_same_v<scalar_t, c10::BFloat16>)
+                                          ? 8
+                                          : 4;  // FP16/BF16: 8, Float: 4
+        constexpr int32_t kWarpSize = 32;
+        constexpr int32_t kMaxThreadsPerBlock = 256;
+        constexpr int32_t kAlignmentBytes = 16;  // 128-bit alignment
+
+        CHECK(head_dim % VEC_WIDTH == 0)
+            << "head_dim must be divisible by vector width: " << VEC_WIDTH;
+        const int64_t vecs_per_head = head_dim / VEC_WIDTH;
+        const int64_t total_vecs = kv_heads * vecs_per_head;
+        CHECK(total_vecs > 0) << "total_vecs must be > 0";
+
+        int32_t threads_per_block = static_cast<int32_t>(
+            total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
+                                             : total_vecs);
+        threads_per_block =
+            ((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
+        if (threads_per_block < kWarpSize) {
+          threads_per_block = kWarpSize;
+        }
+        dim3 block_dim(threads_per_block, 1, 1);
+
+        const auto proj_k_ptr =
+            reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
+        const auto proj_v_ptr =
+            reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
+        const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
+            unshared_k_cache.data_ptr<scalar_t>());
+        const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
+            unshared_v_cache.data_ptr<scalar_t>());
+        CHECK(proj_k_ptr % kAlignmentBytes == 0)
+            << "proj_k data_ptr must be 16-byte aligned";
+        CHECK(proj_v_ptr % kAlignmentBytes == 0)
+            << "proj_v data_ptr must be 16-byte aligned";
+        CHECK(k_cache_ptr % kAlignmentBytes == 0)
+            << "unshared_k_cache data_ptr must be 16-byte aligned";
+        CHECK(v_cache_ptr % kAlignmentBytes == 0)
+            << "unshared_v_cache data_ptr must be 16-byte aligned";
+
+        const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
+        CHECK((k_stride0 * scalar_bytes) % kAlignmentBytes == 0)
+            << "proj_k stride(0) bytes must be 16-byte aligned";
+        CHECK((k_stride1 * scalar_bytes) % kAlignmentBytes == 0)
+            << "proj_k stride(1) bytes must be 16-byte aligned";
+        CHECK((v_stride0 * scalar_bytes) % kAlignmentBytes == 0)
+            << "proj_v stride(0) bytes must be 16-byte aligned";
+        CHECK((v_stride1 * scalar_bytes) % kAlignmentBytes == 0)
+            << "proj_v stride(1) bytes must be 16-byte aligned";
+        CHECK((cache_stride0 * scalar_bytes) % kAlignmentBytes == 0)
+            << "cache stride(0) bytes must be 16-byte aligned";
+        CHECK((cache_stride1 * scalar_bytes) % kAlignmentBytes == 0)
+            << "cache stride(1) bytes must be 16-byte aligned";
+        CHECK((cache_stride2 * scalar_bytes) % kAlignmentBytes == 0)
+            << "cache stride(2) bytes must be 16-byte aligned";
+        CHECK((cache_stride3 * scalar_bytes) % kAlignmentBytes == 0)
+            << "cache stride(3) bytes must be 16-byte aligned";
+
+        decoder_reshape_and_cache_kernel<scalar_t>
+            <<<grid_dim, block_dim, 0, stream>>>(
+                proj_k.data_ptr<scalar_t>(),
+                proj_v.data_ptr<scalar_t>(),
+                unshared_k_cache.data_ptr<scalar_t>(),
+                unshared_v_cache.data_ptr<scalar_t>(),
+                step.data_ptr<int32_t>(),
+                batch_size,
+                beam_size,
+                kv_heads,
+                head_dim,
+                k_stride0,
+                k_stride1,
+                v_stride0,
+                v_stride1,
+                cache_stride0,
+                cache_stride1,
+                cache_stride2,
+                cache_stride3);
+      });
+
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+}  // namespace xllm::kernel::cuda

ex_engine/xllm_kernels/cuda/xattention/lse_combine.cu

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

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 <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/cuda.h>
+
+#include <cstdint>
+#include <type_traits>
+
+#include "kernels/cuda/cuda_ops_api.h"
+#include "kernels/cuda/utils.h"
+using at::device_of;
+
+namespace {
+
+template <typename scalar_t>
+struct VecType;
+
+template <>
+struct VecType<c10::Half> {
+  using type = uint4;  // 8 elements * 2 bytes = 16 bytes
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<c10::BFloat16> {
+  using type = uint4;  // 8 elements * 2 bytes = 16 bytes
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<float> {
+  using type = float4;  // 4 elements * 4 bytes = 16 bytes
+  static constexpr int32_t vec_width = 4;
+};
+
+template <typename scalar_t>
+__global__ void prefill_reshape_and_cache_kernel(
+    const scalar_t* __restrict__ proj_k,    // [shared_len, kv_heads, head_dim]
+    const scalar_t* __restrict__ proj_v,    // [shared_len, kv_heads, head_dim]
+    scalar_t* __restrict__ shared_k_cache,  // [shared_len, kv_heads, head_dim]
+    scalar_t* __restrict__ shared_v_cache,  // [shared_len, kv_heads, head_dim]
+    const int64_t shared_len,
+    const int64_t kv_heads,
+    const int64_t head_dim,
+    const int64_t k_stride0,    // proj_k.stride(0)
+    const int64_t v_stride0,    // proj_v.stride(0)
+    const int64_t v_stride1) {  // proj_v.stride(1), same as head_dim
+  using VecTypeT = typename VecType<scalar_t>::type;
+  constexpr int32_t VEC_WIDTH = VecType<scalar_t>::vec_width;
+  const int64_t token_idx = static_cast<int64_t>(blockIdx.y);
+  if (token_idx >= shared_len) {
+    return;
+  }
+
+  const int64_t vecs_per_head = head_dim / VEC_WIDTH;
+  const int64_t total_vecs = kv_heads * vecs_per_head;
+  const int64_t k_token_base = token_idx * k_stride0;
+  const int64_t v_token_base = token_idx * v_stride0;
+  const int64_t dst_token_base = token_idx * kv_heads * head_dim;
+
+  for (int64_t linear_idx = threadIdx.x; linear_idx < total_vecs;
+       linear_idx += blockDim.x) {
+    const int64_t head_idx = linear_idx / vecs_per_head;
+    const int64_t vec_idx = linear_idx - head_idx * vecs_per_head;
+    const int64_t head_offset = head_idx * head_dim;
+    const int64_t vec_offset = vec_idx * VEC_WIDTH;
+
+    const auto* k_src_vec = reinterpret_cast<const VecTypeT*>(
+        proj_k + k_token_base + head_offset + vec_offset);
+    const auto* v_src_vec = reinterpret_cast<const VecTypeT*>(
+        proj_v + v_token_base + head_idx * v_stride1 + vec_offset);
+    auto* k_dst_vec = reinterpret_cast<VecTypeT*>(
+        shared_k_cache + dst_token_base + head_offset + vec_offset);
+    auto* v_dst_vec = reinterpret_cast<VecTypeT*>(
+        shared_v_cache + dst_token_base + head_offset + vec_offset);
+
+    *k_dst_vec = *k_src_vec;
+    *v_dst_vec = *v_src_vec;
+  }
+}
+
+}  // namespace
+
+namespace xllm::kernel::cuda {
+void prefill_reshape_and_cache(
+    torch::Tensor proj_k,  // [shared_len, kv_heads, head_dim]
+    torch::Tensor proj_v,  // [shared_len, kv_heads, head_dim]
+    torch::Tensor
+        shared_k_cache,  // [num_shared_kv_seq_len, kv_heads, head_dim]
+    torch::Tensor shared_v_cache) {
+  CHECK(proj_k.dim() == 3) << "proj_k must be 3-dimensional";
+  CHECK(proj_v.dim() == 3) << "proj_v must be 3-dimensional";
+  CHECK(shared_k_cache.dim() == 3) << "shared_k_cache must be 3-dimensional";
+  CHECK(shared_v_cache.dim() == 3) << "shared_v_cache must be 3-dimensional";
+  CHECK(proj_k.is_cuda() && proj_v.is_cuda() && shared_k_cache.is_cuda() &&
+        shared_v_cache.is_cuda())
+      << "all tensors must be CUDA tensors";
+
+  const int64_t shared_len = proj_k.size(0);
+  const int64_t kv_heads = proj_k.size(1);
+  const int64_t head_dim = proj_k.size(2);
+  CHECK(proj_v.sizes() == proj_k.sizes())
+      << "proj_v and proj_k must have same shape";
+  CHECK(shared_k_cache.size(0) >= shared_len &&
+        shared_k_cache.size(1) == kv_heads &&
+        shared_k_cache.size(2) == head_dim)
+      << "shared_k_cache shape mismatch";
+  CHECK(shared_v_cache.size(0) >= shared_len &&
+        shared_v_cache.size(1) == kv_heads &&
+        shared_v_cache.size(2) == head_dim)
+      << "shared_v_cache shape mismatch";
+
+  shared_k_cache = shared_k_cache.slice(0, 0, shared_len);
+  shared_v_cache = shared_v_cache.slice(0, 0, shared_len);
+
+  // This kernel is specialized for qkv-slice layouts:
+  // last dim contiguous and head stride tightly packed by head_dim.
+  CHECK(proj_k.stride(2) == 1 && proj_v.stride(2) == 1)
+      << "proj_k/proj_v must be contiguous on head_dim (stride(2)=1)";
+  CHECK(proj_k.stride(1) == head_dim && proj_v.stride(1) == head_dim)
+      << "proj_k/proj_v must satisfy stride(1)=head_dim for qkv-slice layout";
+  CHECK(shared_k_cache.stride(2) == 1 && shared_v_cache.stride(2) == 1)
+      << "shared caches must be contiguous on head_dim (stride(2)=1)";
+  CHECK(shared_k_cache.stride(1) == head_dim &&
+        shared_v_cache.stride(1) == head_dim)
+      << "shared caches must satisfy stride(1)=head_dim";
+  CHECK(shared_k_cache.stride(0) == kv_heads * head_dim &&
+        shared_v_cache.stride(0) == kv_heads * head_dim)
+      << "shared caches must be contiguous on token stride";
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(proj_k));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  const int64_t k_stride0 = proj_k.stride(0);
+  const int64_t v_stride0 = proj_v.stride(0);
+  const int64_t v_stride1 = proj_v.stride(1);
+  dim3 grid_dim(1, static_cast<unsigned int>(shared_len), 1);
+
+  DISPATCH_FLOATING_TYPES(
+      proj_k.scalar_type(), "prefill_reshape_and_cache_kernel", [&] {
+        constexpr int32_t VEC_WIDTH = (std::is_same_v<scalar_t, c10::Half> ||
+                                       std::is_same_v<scalar_t, c10::BFloat16>)
+                                          ? 8
+                                          : 4;  // FP16/BF16: 8, Float: 4
+        constexpr int32_t kWarpSize = 32;
+        constexpr int32_t kMaxThreadsPerBlock = 256;
+
+        CHECK(head_dim % VEC_WIDTH == 0)
+            << "head_dim must be divisible by vector width: " << VEC_WIDTH;
+        const int64_t vecs_per_head = head_dim / VEC_WIDTH;
+        const int64_t total_vecs = kv_heads * vecs_per_head;
+        CHECK(total_vecs > 0) << "total_vecs must be > 0";
+
+        int32_t threads_per_block = static_cast<int32_t>(
+            total_vecs > kMaxThreadsPerBlock ? kMaxThreadsPerBlock
+                                             : total_vecs);
+        threads_per_block =
+            ((threads_per_block + kWarpSize - 1) / kWarpSize) * kWarpSize;
+        if (threads_per_block < kWarpSize) {
+          threads_per_block = kWarpSize;
+        }
+        dim3 block_dim(threads_per_block, 1, 1);
+
+        const auto proj_k_ptr =
+            reinterpret_cast<std::uintptr_t>(proj_k.data_ptr<scalar_t>());
+        const auto proj_v_ptr =
+            reinterpret_cast<std::uintptr_t>(proj_v.data_ptr<scalar_t>());
+        const auto k_cache_ptr = reinterpret_cast<std::uintptr_t>(
+            shared_k_cache.data_ptr<scalar_t>());
+        const auto v_cache_ptr = reinterpret_cast<std::uintptr_t>(
+            shared_v_cache.data_ptr<scalar_t>());
+
+        constexpr int32_t alignment_bytes = 16;  // 128-bit alignment
+        CHECK(proj_k_ptr % alignment_bytes == 0)
+            << "proj_k data_ptr must be 16-byte aligned";
+        CHECK(proj_v_ptr % alignment_bytes == 0)
+            << "proj_v data_ptr must be 16-byte aligned";
+        CHECK(k_cache_ptr % alignment_bytes == 0)
+            << "shared_k_cache data_ptr must be 16-byte aligned";
+        CHECK(v_cache_ptr % alignment_bytes == 0)
+            << "shared_v_cache data_ptr must be 16-byte aligned";
+
+        const int64_t scalar_bytes = static_cast<int64_t>(sizeof(scalar_t));
+        CHECK((k_stride0 * scalar_bytes) % alignment_bytes == 0)
+            << "proj_k stride(0) bytes must be 16-byte aligned";
+        CHECK((v_stride0 * scalar_bytes) % alignment_bytes == 0)
+            << "proj_v stride(0) bytes must be 16-byte aligned";
+        CHECK((v_stride1 * scalar_bytes) % alignment_bytes == 0)
+            << "proj_v stride(1) bytes must be 16-byte aligned";
+
+        prefill_reshape_and_cache_kernel<scalar_t>
+            <<<grid_dim, block_dim, 0, stream>>>(
+                proj_k.data_ptr<scalar_t>(),
+                proj_v.data_ptr<scalar_t>(),
+                shared_k_cache.data_ptr<scalar_t>(),
+                shared_v_cache.data_ptr<scalar_t>(),
+                shared_len,
+                kv_heads,
+                head_dim,
+                k_stride0,
+                v_stride0,
+                v_stride1);
+      });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+}  // namespace xllm::kernel::cuda

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 <torch/torch.h>
+
+#include <vector>
+
+namespace xllm::kernel::cuda {
+
+void decoder_reshape_and_cache(torch::Tensor proj_k,
+                               torch::Tensor proj_v,
+                               torch::Tensor unshared_k_cache,
+                               torch::Tensor unshared_v_cache,
+                               torch::Tensor step);
+
+void cache_select(const torch::Tensor& beam_index,
+                  std::vector<torch::Tensor>& unshared_k_cache,
+                  std::vector<torch::Tensor>& unshared_v_cache,
+                  const torch::Tensor& block_table,
+                  int64_t decode_step,
+                  int64_t beam_size,
+                  int64_t layer_num);
+
+void lse_combine(torch::Tensor output,
+                 torch::Tensor shared_o,
+                 torch::Tensor shared_lse,
+                 torch::Tensor unshared_o,
+                 torch::Tensor unshared_lse);
+
+void prefill_reshape_and_cache(
+    torch::Tensor proj_k,  // [shared_len, kv_heads, head_dim]
+    torch::Tensor proj_v,  // [shared_len, kv_heads, head_dim]
+    torch::Tensor
+        shared_k_cache,  // [num_shared_kv_seq_len, kv_heads, head_dim]
+    torch::Tensor shared_v_cache);
+
+void beam_search(torch::Tensor acc_logprob,
+                 torch::Tensor in_sequence_group,
+                 torch::Tensor top_tokens,
+                 torch::Tensor top_logprobs,
+                 torch::Tensor out_acc_logprob,
+                 torch::Tensor out_token_ids,
+                 torch::Tensor out_token_index,
+                 torch::Tensor out_beam_count_prefix_sums,
+                 torch::Tensor out_sequence_group,
+                 uint32_t batch_size,
+                 uint32_t current_step);
+
+}  // namespace xllm::kernel::cuda

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
+)

upstream_ref/sgemm_cuda/01_naive.cuh

@@ -0,0 +1,37 @@
+# pragma once
+
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+/*
+
+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];
+    }
+}

upstream_ref/sgemm_cuda/02_kernel_global_mem_coalesce.cuh

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+template <const uint BLOCKSIZE>
+// __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];
+    }
+}

upstream_ref/sgemm_cuda/03_kernel_shared_mem_blocking.cuh

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BLOCKSIZE>
+__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];
+}

upstream_ref/sgemm_cuda/04_kernel_1D_blocktiling.cuh

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM>
+__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];
+  }
+}

upstream_ref/sgemm_cuda/05_kernel_2D_blocktiling.cuh

@@ -0,0 +1,102 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__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];
+    }
+  }
+}

upstream_ref/sgemm_cuda/06_kernel_vectorize.cuh

@@ -0,0 +1,98 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__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<float4 *>(&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<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
+        reinterpret_cast<float4 *>(&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<float4 *>(
+          &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<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/07_kernel_resolve_bank_conflicts.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__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<float4 *>(&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<float4 *>(&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<float4 *>(
+          &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<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/08_kernel_bank_extra_col.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__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<float4 *>(&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<float4 *>(&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<float4 *>(
+          &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<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/09_kernel_autotuned.cuh

@@ -0,0 +1,127 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+const int K9_NUM_THREADS = 256;
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__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<float4 *>(
+          &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<float4 *>(
+          &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+          reinterpret_cast<float4 *>(
+              &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<float4 *>(
+              &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<float4 *>(&C_interim[(threadRow * TM + resIdxM) * N +
+                                                threadCol * TN + resIdxN])[0] =
+              tmp;
+        }
+      }
+    }
+  }
+}

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("$<$<AND:$<CONFIG:Debug>,$<COMPILE_LANGUAGE:CUDA>>:-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})

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("$<$<AND:$<CONFIG:Debug>,$<COMPILE_LANGUAGE:CUDA>>:-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})

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=<integer> PREFIX=<optional string>
+profile: build
+	@ncu --set full --export $(BENCHMARK_DIR)/$(PREFIX)kernel_$(KERNEL) --force-overwrite $(BUILD_DIR)/sgemm $(KERNEL)
+
+bench: build
+	@bash gen_benchmark_results.sh

upstream_ref/sgemm_cuda/cuBLAS_sgemm.cu

@@ -0,0 +1,108 @@
+#include <cstdio>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+/*
+ * 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;
+}

upstream_ref/sgemm_cuda/kernel_1.cuh

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+__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];
+}

upstream_ref/sgemm_cuda/kernel_2.cuh

@@ -0,0 +1,45 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BLOCK_SIZE>
+__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];
+}

upstream_ref/sgemm_cuda/kernel_3.cuh

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM>
+__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];
+    }
+}

upstream_ref/sgemm_cuda/kernel_4.cuh

@@ -0,0 +1,76 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__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];
+    }
+}

upstream_ref/sgemm_cuda/kernel_5.cuh

@@ -0,0 +1,88 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__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];
+    }
+}

upstream_ref/sgemm_cuda/kernel_6.cuh

@@ -0,0 +1,110 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__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;
+        }
+    }
+}

upstream_ref/sgemm_cuda/kernel_7.cuh

@@ -0,0 +1,180 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__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;
+        }
+    }
+}

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"

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"

upstream_ref/sgemm_cuda/runner.cu

@@ -0,0 +1,549 @@
+#include "kernels.cuh"
+#include "runner.cuh"
+#include <cmath>
+#include <cstdio>
+#include <fstream>
+#include <iomanip>
+
+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<<<gridDim, blockDim>>>(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>
+      <<<gridDim, blockDim>>>(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>
+      <<<gridDim, blockDim>>>(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<BM, BN, BK, TM>
+      <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(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<K9_BM, K9_BN, K9_BK, K9_TM, K9_TN>
+      <<<gridDim, blockDim>>>(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<K10_BM, K10_BN, K10_BK, K10_WM, K10_WN, K10_WNITER, K10_TM,
+                  K10_TN, K10_NUM_THREADS>
+      <<<gridDim, blockDim>>>(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<K11_BM, K11_BN, K11_BK, K11_WM, K11_WN, K11_WNITER,
+                       K11_TM, K11_TN, K11_NUM_THREADS>
+      <<<gridDim, blockDim>>>(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<K12_BM, K12_BN, K12_BK, K12_WM, K12_WN, K12_WNITER,
+                           K12_TM, K12_TN, K12_NUM_THREADS>
+      <<<gridDim, blockDim>>>(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");
+  }
+}

upstream_ref/sgemm_cuda/runner.cuh

@@ -0,0 +1,26 @@
+#pragma once
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+#include <fstream>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/time.h>
+#include <time.h>
+#include <unistd.h>
+
+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);

upstream_ref/sgemm_cuda/sgemm_siboehm.cu

@@ -0,0 +1,168 @@
+#include <cstdio>
+#include <cstdlib>
+#include <ctime>
+#include <fstream>
+#include <iostream>
+#include <runner.cuh>
+#include <vector>
+
+#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<int> 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;
+};

upstream_ref/sgemm_cuda/sgemm_wangzyon.cu

@@ -0,0 +1,119 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/time.h>
+#include <utils.cuh>
+
+#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;
+};

upstream_ref/sgemm_cuda/simplest_kernel.cu

@@ -0,0 +1,46 @@
+#include <cuda_runtime.h>
+#include <iostream>
+#include <vector>
+
+__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<<<grid_size, block_size>>>(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);
+}

upstream_ref/sgemm_cuda/utils.cu

@@ -0,0 +1,199 @@
+#include <stdio.h>
+#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<<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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><<<gridDim, blockDim>>>(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;
+    }
+}

upstream_ref/sgemm_cuda/utils.cuh

@@ -0,0 +1,42 @@
+#pragma once
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <unistd.h>
+#include <sys/time.h>
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+
+/*
+=====================================
+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);