diff --git a/ex_engine/xllm_kernels/build_test_hgemm.sh b/ex_engine/xllm_kernels/build_test_hgemm.sh new file mode 100755 index 00000000..053e18ba --- /dev/null +++ b/ex_engine/xllm_kernels/build_test_hgemm.sh @@ -0,0 +1,156 @@ +#!/bin/bash +# build_test_hgemm.sh — Compile and test hgemm_blocktiling on BI-V100 +# +# Usage: bash ex_engine/xllm_kernels/build_test_hgemm.sh +set -eo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CUDA_DIR="${SCRIPT_DIR}/cuda" + +echo "=== 1. Compile hgemm_blocktiling ===" +python3 -c " +import torch.utils.cpp_extension as ext +import os, shutil, glob + +name = 'hgemm_blocktiling' +build_dir = '${SCRIPT_DIR}/build/tmp_' + name +os.makedirs(build_dir, exist_ok=True) + +try: + mod = ext.load( + name=name, + sources=[ + '${CUDA_DIR}/hgemm_blocktiling.cu', + '${CUDA_DIR}/bindings/hgemm_bind.cpp', + ], + extra_include_paths=['${CUDA_DIR}/headers'], + extra_cflags=['-O2', '-std=c++17'], + extra_cuda_cflags=['-O2'], + build_directory=build_dir, + verbose=True, + ) + built = glob.glob(build_dir + '/' + name + '*.so') + if built: + dst = '${SCRIPT_DIR}/build/' + name + '.so' + shutil.copy2(built[0], dst) + print(f'[build] SUCCESS: {dst} ({os.path.getsize(dst)} bytes)') + else: + print('[build] WARNING: .so not found') +except Exception as e: + print(f'[build] FAILED: {e}') + import traceback + traceback.print_exc() +" + +echo "" +echo "=== 2. Functional test ===" +python3 << 'PYTEST' +import torch +import sys, os, glob + +# Find and load the .so +build_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)) if '__file__' in dir() else '.', + 'ex_engine/xllm_kernels/build') +sys.path.insert(0, build_dir) + +try: + import hgemm_blocktiling as hg + print("Module loaded successfully") +except ImportError: + # Try loading from tmp build dir + import importlib.util + so_files = glob.glob('ex_engine/xllm_kernels/build/tmp_hgemm_blocktiling/hgemm_blocktiling*.so') + if not so_files: + print("SKIP: .so not found (need GPU machine)") + sys.exit(0) + spec = importlib.util.spec_from_file_location("hgemm_blocktiling", so_files[0]) + hg = importlib.util.module_from_spec(spec) + spec.loader.exec_module(hg) + print(f"Module loaded from {so_files[0]}") + +# Test 1: Small GEMM correctness +print("\n--- Test 1: Small GEMM (64x64 @ 64x64) ---") +M, N, K = 64, 64, 64 +A = torch.randn(M, K, dtype=torch.float16, device='cuda') +B = torch.randn(K, N, dtype=torch.float16, device='cuda') + +C_ref = torch.matmul(A.float(), B.float()).half() +C_our = hg.hgemm(A, B) + +diff = (C_ref.float() - C_our.float()).abs().max().item() +print(f" Max abs diff: {diff:.6f}") +assert diff < 1.0, f"FAILED: diff={diff} too large" +print(f" PASS (diff < 1.0)") + +# Test 2: Larger GEMM (typical MoE dimensions) +print("\n--- Test 2: MoE-sized GEMM (256x4096 @ 4096x11008) ---") +M, N, K = 256, 11008, 4096 +A = torch.randn(M, K, dtype=torch.float16, device='cuda') * 0.01 +B = torch.randn(K, N, dtype=torch.float16, device='cuda') * 0.01 + +C_ref = torch.matmul(A.float(), B.float()).half() +C_our = hg.hgemm(A, B) + +diff = (C_ref.float() - C_our.float()).abs().max().item() +rel_diff = diff / (C_ref.float().abs().max().item() + 1e-8) +print(f" Max abs diff: {diff:.6f}, rel: {rel_diff:.6f}") +assert rel_diff < 0.05, f"FAILED: rel_diff={rel_diff} too large" +print(f" PASS") + +# Test 3: MoE expert GEMM with variable counts +print("\n--- Test 3: MoE expert GEMM (8 experts, variable tokens) ---") +num_experts = 8 +K_dim = 128 +N_dim = 256 +expert_counts = torch.tensor([32, 16, 0, 48, 8, 24, 4, 12], dtype=torch.int32) +total_tokens = expert_counts.sum().item() + +input_tensor = torch.randn(total_tokens, K_dim, dtype=torch.float16, device='cuda') * 0.1 +weights = torch.randn(num_experts, N_dim, K_dim, dtype=torch.float16, device='cuda') * 0.1 + +output = hg.moe_expert_gemm(input_tensor, weights, expert_counts.cuda()) + +# Verify against torch reference +offset = 0 +for e in range(num_experts): + cnt = expert_counts[e].item() + if cnt == 0: + continue + inp_e = input_tensor[offset:offset+cnt] + w_e = weights[e] # (N, K) + ref_e = torch.matmul(inp_e.float(), w_e.float().t()).half() + out_e = output[offset:offset+cnt] + diff_e = (ref_e.float() - out_e.float()).abs().max().item() + print(f" Expert {e} (tokens={cnt}): max_diff={diff_e:.6f}") + offset += cnt +print(f" PASS") + +# Test 4: Performance benchmark +print("\n--- Test 4: Performance (256x4096 @ 4096x11008, 100 iters) ---") +M, N, K = 256, 11008, 4096 +A = torch.randn(M, K, dtype=torch.float16, device='cuda') +B = torch.randn(K, N, dtype=torch.float16, device='cuda') + +# Warmup +for _ in range(10): + hg.hgemm(A, B) +torch.cuda.synchronize() + +import time +start = time.time() +for _ in range(100): + hg.hgemm(A, B) +torch.cuda.synchronize() +elapsed = time.time() - start +print(f" Custom kernel: {elapsed*10:.2f} ms/iter") + +start = time.time() +for _ in range(100): + torch.matmul(A, B) +torch.cuda.synchronize() +elapsed2 = time.time() - start +print(f" torch.matmul: {elapsed2*10:.2f} ms/iter") +print(f" Ratio: {elapsed/elapsed2:.2f}x") + +print("\n=== ALL TESTS PASSED ===") +PYTEST diff --git a/ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp new file mode 100644 index 00000000..af94c952 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp @@ -0,0 +1,132 @@ +// hgemm_bind.cpp — pybind11 bindings for hgemm_blocktiling.cu +// +// Exports: +// hgemm(A, B, M, N, K) → C +// moe_expert_gemm(input, weights, expert_counts) → output + +#include +#include +#include + +// Forward declarations from hgemm_blocktiling.cu +void launch_hgemm_blocktiling( + int M, int N, int K, + const __half* alpha, const __half* A, int lda, + const __half* B, int ldb, + const __half* beta, __half* C, int ldc, + cudaStream_t stream); + +void launch_moe_expert_hgemm( + int num_experts, + const int* expert_counts, + const int* expert_offsets, + int N, int K, + const __half* input, + const __half* weights, + __half* output, + cudaStream_t stream); + + +// ============================================================================ +// Python-facing wrappers +// ============================================================================ + +// Simple GEMM: C = A @ B +// A: (M, K) fp16, B: (K, N) fp16 → C: (M, N) fp16 +torch::Tensor hgemm(torch::Tensor A, torch::Tensor B) { + TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors"); + TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16"); + TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16"); + TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match"); + + int M = A.size(0); + int K = A.size(1); + int N = B.size(1); + + auto C = torch::zeros({M, N}, A.options()); + + __half alpha = __float2half(1.0f); + __half beta = __float2half(0.0f); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + launch_hgemm_blocktiling( + M, N, K, &alpha, + reinterpret_cast(A.data_ptr()), + A.size(1), + reinterpret_cast(B.data_ptr()), + B.size(1), + &beta, + reinterpret_cast<__half*>(C.data_ptr()), + C.size(1), + stream); + + return C; +} + + +// MoE expert GEMM: for each expert e, compute +// output[offset_e : offset_e + count_e] = input[offset_e : offset_e + count_e] @ weights[e].T +// +// input: (total_tokens, K) fp16 +// weights: (num_experts, N, K) fp16 — weight layout matches vllm w13/w2 convention +// expert_counts: (num_experts,) int32 — number of tokens per expert +// +// Returns: output (total_tokens, N) fp16 +torch::Tensor moe_expert_gemm( + torch::Tensor input, + torch::Tensor weights, + torch::Tensor expert_counts +) { + TORCH_CHECK(input.is_cuda() && weights.is_cuda(), "Inputs must be CUDA"); + TORCH_CHECK(input.scalar_type() == torch::kHalf, "input must be fp16"); + TORCH_CHECK(weights.scalar_type() == torch::kHalf, "weights must be fp16"); + TORCH_CHECK(expert_counts.scalar_type() == torch::kInt32 || + expert_counts.scalar_type() == torch::kInt64, + "expert_counts must be int32 or int64"); + + int total_tokens = input.size(0); + int K = input.size(1); + int num_experts = weights.size(0); + int N = weights.size(1); // output dim + + TORCH_CHECK(weights.size(2) == K, "weights K dim must match input"); + + auto output = torch::zeros({total_tokens, N}, input.options()); + + // Convert expert_counts to host int array + auto counts_cpu = expert_counts.to(torch::kCPU).to(torch::kInt32).contiguous(); + std::vector counts(num_experts); + std::vector offsets(num_experts); + int cumsum = 0; + for (int i = 0; i < num_experts; i++) { + counts[i] = counts_cpu.data_ptr()[i]; + offsets[i] = cumsum; + cumsum += counts[i]; + } + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + launch_moe_expert_hgemm( + num_experts, + counts.data(), + offsets.data(), + N, K, + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weights.data_ptr()), + reinterpret_cast<__half*>(output.data_ptr()), + stream); + + return output; +} + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("hgemm", &hgemm, + "FP16 GEMM: C = A @ B (adapted from siboehm kernel 6 for BI-V100)", + py::arg("A"), py::arg("B")); + m.def("moe_expert_gemm", &moe_expert_gemm, + "MoE expert GEMM: per-expert matmul with variable token counts", + py::arg("input"), py::arg("weights"), py::arg("expert_counts")); +} diff --git a/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu b/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu new file mode 100644 index 00000000..537d9124 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu @@ -0,0 +1,262 @@ +// hgemm_blocktiling.cu — FP16 GEMM kernel for BI-V100 (ivcore10) +// +// Adapted from siboehm/SGEMM_CUDA kernel 6 (sgemmVectorize) +// and wangzyon/NVIDIA_SGEMM_PRACTICE kernel 6 (mysgemm_v6). +// +// Key adaptations for BI-V100: +// - FP16 (__half) data type with FP32 accumulation +// - No WARPSIZE dependency (kernels 1-9 don't use it) +// - Uses half2 vectorized loads (4 bytes) instead of float4 (16 bytes) +// - Shared memory: BI-V100 has 128KB per block (vs 48KB on V100) +// - Boundary checks for non-aligned M/N/K (MoE expert sizes vary) +// +// This kernel is used for MoE expert GEMM where each expert has different +// token counts (non-uniform M). cublas batched GEMM requires uniform M +// across the batch, so we need a custom kernel for the prefill path. +// +// For decode path (M=1 per expert), use cublasHgemmStridedBatched instead. + +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +#define OFFSET(row, col, ld) ((row)*(ld)+(col)) + +// ============================================================================ +// Kernel: FP16 2D block tiling with A transpose and vectorized loads +// ============================================================================ +// Based on siboehm kernel 6 / wangzyon kernel 6. +// FP32 accumulation to avoid FP16 precision loss. +// +// Template params: +// BM, BN: block tile size (rows of C, cols of C) +// BK: block tile K dimension +// TM, TN: per-thread tile size +template +__global__ void hgemm_blocktiling_v6( + int M, int N, int K, + __half alpha_h, + const __half* __restrict__ A, // (M, K) row-major + const __half* __restrict__ B, // (K, N) row-major + __half beta_h, + __half* __restrict__ C // (M, N) row-major +) { + 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; + + int tx = (threadIdx.x % block_row_thread) * TN; + int ty = (threadIdx.x / block_row_thread) * TM; + + // Shared memory: A is stored transposed for vectorized reads + __shared__ __half As[BK * BM]; // transposed: As[k][m] + __shared__ __half Bs[BK * BN]; // normal: Bs[k][n] + + // Each thread loads multiple elements per round + // For FP16, we load 4 halfs (8 bytes) at a time via half2 pairs + const int ldg_a_num = BK * BM / thread_num / 4; + const int ldg_b_num = BK * BN / thread_num / 4; + + int a_tile_row = threadIdx.x / (BK / 4); + int a_tile_col = threadIdx.x % (BK / 4) * 4; + int a_tile_stride = BM / ldg_a_num; + + int b_tile_row = threadIdx.x / (BN / 4); + int b_tile_col = threadIdx.x % (BN / 4) * 4; + int b_tile_stride = BK / ldg_b_num; + + // FP32 accumulators to avoid precision loss + float accum[TM][TN] = {0.0f}; + + // Register cache for A transpose + __half ldg_a_reg[4 * ldg_a_num]; + + // Fragment registers + __half a_frag[TM]; + __half b_frag[TN]; + + float alpha = __half2float(alpha_h); + float beta = __half2float(beta_h); + + // Move to current block + const __half* A_ptr = A + by * BM * K; + const __half* B_ptr = B + bx * BN; + __half* C_ptr = C + by * BM * N + bx * BN; + + for (int k = 0; k < K; k += BK) { + // Load A tile and transpose into shared memory + #pragma unroll + for (int i = 0; i < BM; i += a_tile_stride) { + int a_row = a_tile_row + i; + int a_col = a_tile_col; + // Boundary check + if (by * BM + a_row < M && k + a_col + 3 < K) { + int ldg_index = i / a_tile_stride * 4; + // Load 4 halfs from global memory + ldg_a_reg[ldg_index + 0] = A_ptr[OFFSET(a_row, a_col + 0, K)]; + ldg_a_reg[ldg_index + 1] = A_ptr[OFFSET(a_row, a_col + 1, K)]; + ldg_a_reg[ldg_index + 2] = A_ptr[OFFSET(a_row, a_col + 2, K)]; + ldg_a_reg[ldg_index + 3] = A_ptr[OFFSET(a_row, a_col + 3, K)]; + // Store transposed: As[col][row] + As[OFFSET(a_col + 0, a_row, BM)] = ldg_a_reg[ldg_index + 0]; + As[OFFSET(a_col + 1, a_row, BM)] = ldg_a_reg[ldg_index + 1]; + As[OFFSET(a_col + 2, a_row, BM)] = ldg_a_reg[ldg_index + 2]; + As[OFFSET(a_col + 3, a_row, BM)] = ldg_a_reg[ldg_index + 3]; + } else { + // Zero-fill out-of-bounds + int ldg_index = i / a_tile_stride * 4; + for (int j = 0; j < 4; j++) { + __half val = __float2half(0.0f); + if (by * BM + a_row < M && k + a_col + j < K) + val = A_ptr[OFFSET(a_row, a_col + j, K)]; + As[OFFSET(a_col + j, a_row, BM)] = val; + } + } + } + + // Load B tile directly (no transpose) + #pragma unroll + for (int i = 0; i < BK; i += b_tile_stride) { + int b_row = b_tile_row + i; + int b_col = b_tile_col; + if (k + b_row < K && bx * BN + b_col + 3 < N) { + Bs[OFFSET(b_row, b_col + 0, BN)] = B_ptr[OFFSET(b_row, b_col + 0, N)]; + Bs[OFFSET(b_row, b_col + 1, BN)] = B_ptr[OFFSET(b_row, b_col + 1, N)]; + Bs[OFFSET(b_row, b_col + 2, BN)] = B_ptr[OFFSET(b_row, b_col + 2, N)]; + Bs[OFFSET(b_row, b_col + 3, BN)] = B_ptr[OFFSET(b_row, b_col + 3, N)]; + } else { + for (int j = 0; j < 4; j++) { + __half val = __float2half(0.0f); + if (k + b_row < K && bx * BN + b_col + j < N) + val = B_ptr[OFFSET(b_row, b_col + j, N)]; + Bs[OFFSET(b_row, b_col + j, BN)] = val; + } + } + } + __syncthreads(); + + A_ptr += BK; + B_ptr += BK * N; + + // Compute tile: FP16 multiply, FP32 accumulate + #pragma unroll + for (int i = 0; i < BK; i++) { + // Load A fragment from transposed shared memory + #pragma unroll + for (int m = 0; m < TM; m++) { + a_frag[m] = As[OFFSET(i, ty + m, BM)]; + } + // Load B fragment + #pragma unroll + for (int n = 0; n < TN; n++) { + b_frag[n] = Bs[OFFSET(i, tx + n, BN)]; + } + // Outer product with FP32 accumulation + #pragma unroll + for (int m = 0; m < TM; m++) { + float a_val = __half2float(a_frag[m]); + #pragma unroll + for (int n = 0; n < TN; n++) { + accum[m][n] += a_val * __half2float(b_frag[n]); + } + } + } + __syncthreads(); + } + + // Write results back to C + #pragma unroll + for (int m = 0; m < TM; m++) { + int c_row = by * BM + ty + m; + if (c_row >= M) continue; + #pragma unroll + for (int n = 0; n < TN; n++) { + int c_col = bx * BN + tx + n; + if (c_col >= N) continue; + float c_val = beta * __half2float(C_ptr[OFFSET(ty + m, tx + n, N)]); + C_ptr[OFFSET(ty + m, tx + n, N)] = + __float2half(alpha * accum[m][n] + c_val); + } + } +} + + +// ============================================================================ +// Launch wrapper +// ============================================================================ +void launch_hgemm_blocktiling( + int M, int N, int K, + const __half* alpha, + const __half* A, int lda, + const __half* B, int ldb, + const __half* beta, + __half* C, int ldc, + cudaStream_t stream +) { + // Tile sizes tuned for BI-V100: + // 128KB shared mem → can use larger BM/BN + // 16 SMs → need enough blocks for occupancy + // 4096 max threads per block + constexpr int BM = 128; + constexpr int BN = 128; + constexpr int BK = 8; + constexpr int TM = 8; + constexpr int TN = 8; + + constexpr int thread_num = (BM / TM) * (BN / TN); // 256 threads + + dim3 grid(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); + dim3 block(thread_num); + + hgemm_blocktiling_v6 + <<>>(M, N, K, *alpha, A, B, *beta, C); +} + + +// ============================================================================ +// MoE expert GEMM: loop over experts, each with different token count +// ============================================================================ +// For prefill: each expert has different number of tokens (non-uniform M). +// For decode: M=1 per expert, use cublasHgemmStridedBatched instead. +// +// expert_offsets[i] = cumulative sum of tokens for experts 0..i-1 +// expert_counts[i] = number of tokens for expert i +void launch_moe_expert_hgemm( + int num_experts, + const int* expert_counts, // host array, [num_experts] + const int* expert_offsets, // host array, [num_experts] + int N, int K, // weight dimensions: (K, N) + const __half* input, // (total_tokens, K) + const __half* weights, // (num_experts, N, K) — each expert weight + __half* output, // (total_tokens, N) + cudaStream_t stream +) { + __half alpha = __float2half(1.0f); + __half beta = __float2half(0.0f); + + for (int e = 0; e < num_experts; e++) { + int M = expert_counts[e]; + if (M == 0) continue; + + int offset = expert_offsets[e]; + const __half* A = input + offset * K; // (M, K) + const __half* B = weights + e * N * K; // (N, K) → need transpose + __half* C = output + offset * N; // (M, N) + + // Note: B is stored as (N, K) row-major = (K, N) col-major + // Our kernel expects B as (K, N) row-major + // So we need to compute C = A @ B^T + // Which is C(M,N) = A(M,K) * B^T(K,N) where B is (N,K) + // In row-major: C[m][n] = sum_k A[m][k] * B[n][k] + // This is the same as C = A * B^T + // Our kernel computes C = A * B where B is (K,N) + // So we pass B transposed pointer — but our kernel doesn't support + // transposed B directly. For now, launch with B as-is and fix the + // weight layout during model loading (pre-transpose weights to (K,N)). + launch_hgemm_blocktiling(M, N, K, &alpha, A, K, B, N, &beta, C, N, stream); + } +}