diff --git a/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu b/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu index 537d9124..8477a923 100644 --- a/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu +++ b/ex_engine/xllm_kernels/cuda/hgemm_blocktiling.cu @@ -1,262 +1,167 @@ -// hgemm_blocktiling.cu — FP16 GEMM kernel for BI-V100 (ivcore10) +// hgemm_blocktiling.cu — FP16 GEMM for BI-V100 // -// 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. +// 1:1 from siboehm/SGEMM_CUDA kernel 6 (sgemmVectorize). +// Changes: float→__half, float4→load 4 halfs, FP32 accumulator. +// No WARPSIZE usage. No cooperative_groups. CUDA 10.2 safe. #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; +template +__global__ void hgemmVectorize(int M, int N, int K, float alpha, + const __half *A, const __half *B, + float beta, __half *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; - const int block_row_thread = BN / TN; - const int block_col_thread = BM / TM; - const int thread_num = block_row_thread * block_col_thread; + // 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); - int tx = (threadIdx.x % block_row_thread) * TN; - int ty = (threadIdx.x / block_row_thread) * TM; + // allocate space for the current blocktile in smem + // A stored transposed: As[BK][BM], B normal: Bs[BK][BN] + __shared__ __half As[BM * BK]; + __shared__ __half Bs[BK * BN]; - // 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] + // 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; - // 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; + // calculating the indices that this thread will load into SMEM + // FP16: load 4 halfs (8 bytes) per step. 4 halfs per thread. + // siboehm: float4 = 4 floats = 128bit. We do 4 halfs = 64bit. + 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); - 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; + // allocate thread-local cache for results in registerfile + // FP32 accumulation to avoid FP16 precision loss + float threadResults[TM * TN] = {0.0f}; + __half regM[TM]; + __half regN[TN]; - 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; + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // populate the SMEM caches + // transpose A while loading it (same as siboehm) + // Load 4 halfs from A + __half a0 = A[innerRowA * K + innerColA * 4 + 0]; + __half a1 = A[innerRowA * K + innerColA * 4 + 1]; + __half a2 = A[innerRowA * K + innerColA * 4 + 2]; + __half a3 = A[innerRowA * K + innerColA * 4 + 3]; + As[(innerColA * 4 + 0) * BM + innerRowA] = a0; + As[(innerColA * 4 + 1) * BM + innerRowA] = a1; + As[(innerColA * 4 + 2) * BM + innerRowA] = a2; + As[(innerColA * 4 + 3) * BM + innerRowA] = a3; - // FP32 accumulators to avoid precision loss - float accum[TM][TN] = {0.0f}; + // Load 4 halfs from B (no transpose) + Bs[innerRowB * BN + innerColB * 4 + 0] = B[innerRowB * N + innerColB * 4 + 0]; + Bs[innerRowB * BN + innerColB * 4 + 1] = B[innerRowB * N + innerColB * 4 + 1]; + Bs[innerRowB * BN + innerColB * 4 + 2] = B[innerRowB * N + innerColB * 4 + 2]; + Bs[innerRowB * BN + innerColB * 4 + 3] = B[innerRowB * N + innerColB * 4 + 3]; + __syncthreads(); - // Register cache for A transpose - __half ldg_a_reg[4 * ldg_a_num]; + // advance blocktile + A += BK; // move BK columns to right + B += BK * N; // move BK rows down - // 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; - } - } + // 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]; + } + // FP32 accumulation + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + float aVal = __half2float(regM[resIdxM]); + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[resIdxM * TN + resIdxN] += + aVal * __half2float(regN[resIdxN]); } - - // 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(); + } } + __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); - } + // write out the results + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 1) { + uint row = cRow * BM + threadRow * TM + resIdxM; + uint col = cCol * BN + threadCol * TN + resIdxN; + if (row < M && col < N) { + float c_old = __half2float(C[(threadRow * TM + resIdxM) * N + + threadCol * TN + resIdxN]); + C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] = + __float2half(alpha * threadResults[resIdxM * TN + resIdxN] + + beta * c_old); + } } + } } // ============================================================================ -// Launch wrapper +// Launch wrapper — matches siboehm runSgemmVectorize // ============================================================================ void launch_hgemm_blocktiling( int M, int N, int K, - const __half* alpha, + const __half* alpha_ptr, const __half* A, int lda, const __half* B, int ldb, - const __half* beta, + const __half* beta_ptr, __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 + cudaStream_t stream) +{ 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 + // 256 threads — same as siboehm + constexpr int NUM_THREADS = (BM * BN) / (TM * TN); dim3 grid(CEIL_DIV(N, BN), CEIL_DIV(M, BM)); - dim3 block(thread_num); + dim3 block(NUM_THREADS); - hgemm_blocktiling_v6 - <<>>(M, N, K, *alpha, A, B, *beta, C); + float alpha = 1.0f, beta = 0.0f; + if (alpha_ptr) alpha = __half2float(*alpha_ptr); + if (beta_ptr) beta = __half2float(*beta_ptr); + + hgemmVectorize + <<>>(M, N, K, alpha, A, B, beta, C); } // ============================================================================ -// MoE expert GEMM: loop over experts, each with different token count +// MoE expert GEMM — C++ loop over experts (replaces Python for-loop) // ============================================================================ -// 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 int* expert_counts, // host, [num_experts] + const int* expert_offsets, // host, [num_experts] + int N, int K, const __half* input, // (total_tokens, K) - const __half* weights, // (num_experts, N, K) — each expert weight + const __half* weights, // (num_experts, N, K) __half* output, // (total_tokens, N) - cudaStream_t stream -) { - __half alpha = __float2half(1.0f); - __half beta = __float2half(0.0f); - + cudaStream_t stream) +{ for (int e = 0; e < num_experts; e++) { - int M = expert_counts[e]; - if (M == 0) continue; + int M_e = expert_counts[e]; + if (M_e == 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) + int off = expert_offsets[e]; + const __half* A = input + off * K; + const __half* B = weights + (long long)e * N * K; + __half* C_e = output + off * 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); + launch_hgemm_blocktiling(M_e, N, K, + nullptr, A, K, B, N, nullptr, C_e, N, stream); } }