diff --git a/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh b/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh new file mode 100644 index 000000000..804e44cdf --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh @@ -0,0 +1,631 @@ +/* + * Adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu + * + * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include + +#include + +#include +#include +#include + +namespace { + +using bf16_t = __nv_bfloat16; + +__device__ void hmma_16_8_16_f32acc_bf16ab( + float (&d_reg)[4], const bf16_t (&a_reg)[8], const bf16_t (&b_reg)[4], float const (&c_reg)[4]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t a0 = *reinterpret_cast(a_reg + 0); + uint32_t a1 = *reinterpret_cast(a_reg + 2); + uint32_t a2 = *reinterpret_cast(a_reg + 4); + uint32_t a3 = *reinterpret_cast(a_reg + 6); + uint32_t b0 = *reinterpret_cast(b_reg + 0); + uint32_t b1 = *reinterpret_cast(b_reg + 2); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3]) + : "r"(a0), + "r"(a1), + "r"(a2), + "r"(a3), + "r"(b0), + "r"(b1), + "f"(d_reg[0]), + "f"(d_reg[1]), + "f"(d_reg[2]), + "f"(d_reg[3])); +#endif +} + +extern "C" { +__device__ uint32_t __nvvm_get_smem_pointer(void*); +} + +__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if (pred) { + uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr); + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(smemPtrAsUint32), "l"(gPtr), "n"(16)); + } +#endif +} + +__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3]) + : "r"(__nvvm_get_smem_pointer(smem_ptr))); +#endif +} + +template +__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) { + uint32_t row_idx = *reinterpret_cast(&row_idx_); + uint32_t col_idx = *reinterpret_cast(&col_idx_); + row_idx = row_idx % 8; + row_idx = row_idx * (16 / sizeof(Type)); + col_idx = col_idx ^ row_idx; + return *reinterpret_cast(&col_idx); +} + +__device__ void initialize_barrier(uint64_t* smem_barrier, int thread_count = 1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr), "r"(thread_count)); +#endif +} + +__device__ void wait_barrier(uint64_t* smem_barrier, int phase_bit) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra DONE;\n" + "bra LAB_WAIT;\n" + "DONE:\n" + "}\n" ::"r"(smem_int_ptr), + "r"(phase_bit)); +#endif +} + +__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t wait_complete; + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr); + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t" + "selp.b32 %0, 1, 0, P1; \n\t" + "}" + : "=r"(wait_complete) + : "r"(smem_int_ptr), "r"(phase_bit)); + return static_cast(wait_complete); +#endif + return false; +} + +__device__ void arrive_barrier(uint64_t* smem_barrier) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile( + "{\n" + ".reg .b64 state; \n" + "mbarrier.arrive.shared::cta.b64 state, [%0];\n" + "}\n" ::"r"(smem_int_ptr)); +#endif +} + +__device__ void ldgsts_arrive(uint64_t* smem_barrier) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); + asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];" : : "r"(smem_int_ptr)); +#endif +} + +template +struct GmemLoaderA { + static constexpr int elem_bytes = 2; + static constexpr int vec_bytes = 16; + static constexpr int vec_elems = vec_bytes / elem_bytes; + static constexpr int thread_cnt = 64; + static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0); + static constexpr int a_inst_cnt_per_iter = (tile_m * tile_k) / (vec_elems * thread_cnt); + static_assert(gemm_k % tile_k == 0); + static constexpr int k_iter_cnt = gemm_k / tile_k; + + static constexpr int mma_warp_cnt = 4; + static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt; + static constexpr int k_each_chunk = gemm_k / mma_warp_cnt; + + private: + __device__ int k_project(int tile_k_idx) { + return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k); + } + + public: + __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_, uint64_t* smem_barrier_) + : gmem_a(gmem_a_local_), smem_a(smem_a_), smem_barrier(smem_barrier_), local_tid(threadIdx.x % thread_cnt) {} + + __device__ void prepare() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#pragma unroll + for (int i = 0; i < a_inst_cnt_per_iter; i++) { + int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; + int m_idx = linear_idx / tile_k; + int k_idx = linear_idx % tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); + a_smem_offsets[i] = m_idx * tile_k + k_idx; + } +#endif + } + + __device__ void issue_mainloop() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#pragma unroll 1 + for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { + if (need_wait) { + wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit); + } + int next_stage_idx = stage_idx + 1; + int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; + next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx; + if (loop_idx != k_iter_cnt - 1) { + need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit); + } + +#pragma unroll + for (int i = 0; i < a_inst_cnt_per_iter; i++) { + int smem_offset = a_smem_offsets[i]; + bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; + int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; + int m_idx = linear_idx / tile_k; + int k_idx = linear_idx % tile_k; + int gmem_offset = m_idx * gemm_k + k_project(k_idx); + bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset; + ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true); + } + ldgsts_arrive(smem_barrier + stage_idx * 2); + + stage_idx = next_stage_idx; + phase_bit = next_phase_bit; + gmem_a += per_mma_warp_k; + } +#endif + } + + bf16_t const* gmem_a; + bf16_t* smem_a; + uint64_t* smem_barrier; + int local_tid; + int stage_idx = 0; + int phase_bit = 1; + bool need_wait = true; + + int a_smem_offsets[a_inst_cnt_per_iter]; +}; + +template +struct GmemLoaderB { + static constexpr int elem_bytes = 2; + static constexpr int vec_bytes = 16; + static constexpr int vec_elems = vec_bytes / elem_bytes; + static constexpr int thread_cnt = 64; + static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0); + static constexpr int b_inst_cnt_per_iter = (tile_n * tile_k) / (vec_elems * thread_cnt); + static_assert(gemm_k % tile_k == 0); + static constexpr int k_iter_cnt = gemm_k / tile_k; + + static constexpr int mma_warp_cnt = 4; + static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt; + static constexpr int k_each_chunk = gemm_k / mma_warp_cnt; + + private: + __device__ int k_project(int tile_k_idx) { + return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k); + } + + public: + __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_, uint64_t* smem_barrier_, int gemm_n_) + : gmem_b(gmem_b_local_), + smem_b(smem_b_), + smem_barrier(smem_barrier_), + gemm_n(gemm_n_), + local_tid(threadIdx.x % thread_cnt) {} + + __device__ void prepare() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#pragma unroll + for (int i = 0; i < b_inst_cnt_per_iter; i++) { + int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; + int n_idx = linear_idx / tile_k; + int k_idx = linear_idx % tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(n_idx, k_idx); + b_smem_offsets[i] = n_idx * tile_k + k_idx; + preds[i] = n_idx < gemm_n; + } +#endif + } + + __device__ void issue_mainloop() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#pragma unroll 1 + for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { + if (need_wait) { + wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit); + } + int next_stage_idx = stage_idx + 1; + int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; + next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx; + if (loop_idx != k_iter_cnt - 1) { + need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit); + } +#pragma unroll + for (int i = 0; i < b_inst_cnt_per_iter; i++) { + int smem_offset = b_smem_offsets[i]; + bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset; + int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; + int n_idx = linear_idx / tile_k; + int k_idx = linear_idx % tile_k; + int gmem_offset = n_idx * gemm_k + k_project(k_idx); + bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset; + ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]); + } + ldgsts_arrive(smem_barrier + stage_idx * 2); + + stage_idx = next_stage_idx; + phase_bit = next_phase_bit; + gmem_b += per_mma_warp_k; + } +#endif + } + + bf16_t const* gmem_b; + bf16_t* smem_b; + uint64_t* smem_barrier; + int gemm_n; + int local_tid; + int stage_idx = 0; + int phase_bit = 1; + bool need_wait = true; + + int b_smem_offsets[b_inst_cnt_per_iter]; + uint32_t preds[b_inst_cnt_per_iter]; +}; + +template +struct MmaComputer { + static constexpr int elem_bytes = 2; + static constexpr int thread_cnt = 128; + static_assert(gemm_k % tile_k == 0); + static_assert(tile_k % (thread_cnt / 32) == 0); + static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32); + static constexpr int k_iter_cnt = gemm_k / tile_k; + static constexpr int k_phase_cnt = per_warp_tile_k / 16; + static constexpr int m_iter_cnt = (tile_m + 15) / 16; + static constexpr int n_iter_cnt = (tile_n + 7) / 8; + static_assert(m_iter_cnt == 1); + static_assert(n_iter_cnt == 1 || n_iter_cnt == 2); + + __device__ MmaComputer( + bf16_t* gmem_c_local_, bf16_t* smem_a_, bf16_t* smem_b_, uint64_t* smem_barrier_, int warp_idx_, int gemm_n_) + : gmem_c(gmem_c_local_), + smem_a(smem_a_), + smem_b(smem_b_), + smem_barrier(smem_barrier_), + warp_idx(warp_idx_ - (thread_cnt / 32)), + gemm_n(gemm_n_) {} + + private: + __device__ constexpr int internal_b_atom_func(int tid) { + if constexpr (tile_n < 8) { + return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n; + } else { + return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8)); + } + } + + public: + __device__ void prepare() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#pragma unroll + for (int i = 0; i < k_phase_cnt; i++) { + int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; + int m_idx = linear_idx % tile_m; + int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); + a_smem_offsets[0][i] = m_idx * tile_k + k_idx; + } +#pragma unroll + for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { +#pragma unroll + for (int i = 0; i < k_phase_cnt; i += 2) { + int linear_idx = internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8; + int n_idx = linear_idx % tile_n; + int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(n_idx, k_idx); + b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx; + } + } +#endif + } + + __device__ void issue_mainloop() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#pragma unroll 1 + for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { + wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit); + +#pragma unroll + for (int i = 0; i < k_phase_cnt; i++) { + int smem_offset = a_smem_offsets[0][i]; + bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; + ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[0][i])); + } + +#pragma unroll + for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { +#pragma unroll + for (int i = 0; i < k_phase_cnt; i += 2) { + int smem_offset = b_smem_offsets[n_iter_idx][i]; + bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset; + ldsm_x4(smem_ptr_this_iter, reinterpret_cast(b_reg[n_iter_idx][i])); + } + } + +#pragma unroll + for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) { +#pragma unroll + for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { + hmma_16_8_16_f32acc_bf16ab( + acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx], b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]); + } + } + ::arrive_barrier(smem_barrier + 1 + stage_idx * 2); + stage_idx += 1; + phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; + stage_idx = stage_idx == stage_cnt ? 0 : stage_idx; + } +#endif + } + + __device__ void epi() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); + constexpr int thread_m = 2; + constexpr int thread_n = 2 * n_iter_cnt; + constexpr int cta_mma_n = n_iter_cnt * 8; + float acc_reg_reorg[thread_m][thread_n]; + + for (int i = 0; i < thread_m; i++) { + for (int j = 0; j < thread_n; j++) { + acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)]; + } + } + + float* smem_c = reinterpret_cast(smem_a); + auto smem_c_index_func = [&](int m_idx, int n_idx) { + int group_rows = 32 / cta_mma_n; + int group_cnt = 2; + return (m_idx % group_rows * cta_mma_n) + (m_idx / group_rows * (32 + group_cnt)) + n_idx; + }; + constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2); + +#pragma unroll + for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) { +#pragma unroll + for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) { + int m_idx = (lane_idx / 4) + m_idx_thread * 8; + int n_idx = ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8; + smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] = acc_reg_reorg[m_idx_thread][n_idx_thread]; + } + } + asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); + + if (warp_idx == 0) { + constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32; + float acc_final[final_acc_reg_cnt]{}; + +#pragma unroll + for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) { + int linear_idx = reg_idx * 32 + lane_idx; + int m_idx = linear_idx % tile_m; + int n_idx = linear_idx / tile_m; + acc_final[reg_idx] += smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] + + smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] + + smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] + + smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c]; + } + +#pragma unroll + for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) { + int linear_idx = reg_idx * 32 + lane_idx; + int m_idx = linear_idx % tile_m; + int n_idx = linear_idx / tile_m; + if (m_idx < tile_m && n_idx < gemm_n) { + gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx]; + } + } + } +#endif + } + + bf16_t* gmem_c; + bf16_t* smem_a; + bf16_t* smem_b; + uint64_t* smem_barrier; + int warp_idx; + int gemm_n; + int stage_idx = 0; + int phase_bit = 0; + int lane_idx = threadIdx.x % 32; + int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k; + + int a_smem_offsets[m_iter_cnt][k_phase_cnt]; + int b_smem_offsets[n_iter_cnt][k_phase_cnt]; + + bf16_t a_reg[m_iter_cnt][k_phase_cnt][8]; + bf16_t b_reg[n_iter_cnt][k_phase_cnt][4]; + float acc_reg[m_iter_cnt][n_iter_cnt][4]{}; +}; + +template +__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( + bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + constexpr int load_thread_cnt = 128; + constexpr int compute_thread_cnt = 128; + constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt; + (void)thread_cnt; + static_assert(gemm_m % 16 == 0); + static_assert(gemm_k % tile_k == 0); + static_assert(gemm_m % tile_m == 0); + static_assert(tile_k == 128 || tile_k == 256 || tile_k == 512 || tile_k == 1024); + static_assert(tile_m == 16); + constexpr int g2s_vec_bytes = 16; + constexpr int a_elem_bytes = 2; + constexpr int b_elem_bytes = 2; + static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k * stage_cnt <= 225 * 1024); + static_assert((tile_m * tile_k * a_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0); + static_assert((tile_n * tile_k * b_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0); + + extern __shared__ char smem[]; + uint64_t* smem_barrier = reinterpret_cast(smem); + bf16_t* smem_a = reinterpret_cast(smem + (stage_cnt * 8 * 2 + 1024) / 1024 * 1024); + bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt; + + int cta_m_idx = tile_m * blockIdx.x; + int cta_n_idx = tile_n * blockIdx.y; + bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k; + bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k; + bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx; + + int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + + if (warp_idx == 4) { + for (int i = 0; i < stage_cnt; i++) { + initialize_barrier(smem_barrier + i * 2 + 0, load_thread_cnt); + initialize_barrier(smem_barrier + i * 2 + 1, compute_thread_cnt); + } + } + __syncthreads(); + + if (warp_idx < 2) { + GmemLoaderA a_loader(gmem_a_local, smem_a, smem_barrier); + a_loader.prepare(); + a_loader.issue_mainloop(); + } else if (warp_idx < 4) { + GmemLoaderB b_loader(gmem_b_local, smem_b, smem_barrier, gemm_n); + b_loader.prepare(); + b_loader.issue_mainloop(); + } else { + MmaComputer mma_computer( + gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n); + mma_computer.prepare(); + mma_computer.issue_mainloop(); + mma_computer.epi(); + } + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, DLDevice device) { + constexpr int gemm_m = kHdOut; // 2112 + int const gemm_n = num_tokens; // 16 + constexpr int gemm_k = kHdIn; // 7168 + constexpr int batch_size = 1; + std::swap(mat_a, mat_b); + constexpr int tile_m = 16; + constexpr int tile_n = kTileN; // 8 or 16 + constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 +#if defined(SGL_CUDA_ARCH) && SGL_CUDA_ARCH >= 1200 + constexpr int smem_stage_budget = 96 * 1024; +#else + constexpr int smem_stage_budget = 192 * 1024; +#endif + constexpr int max_stage_cnt = smem_stage_budget / ((tile_m + tile_n) * tile_k * sizeof(bf16_t)); + constexpr int k_iter_cnt = gemm_k / tile_k; + constexpr int stage_cnt = k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt; + int cta_m_cnt = gemm_m / tile_m; + int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n; + constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024; // 4096 + constexpr int smem_bytes = ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) / 1024 * 1024; + + dim3 grid(cta_m_cnt, cta_n_cnt, 1); + dim3 block_size(256); + + auto kernel = fused_a_gemm_kernel; + if (smem_bytes >= (48 * 1024)) { + host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)); + } + host::LaunchKernel(grid, block_size, device, smem_bytes).enable_pdl(kUsePDL)(kernel, output, mat_a, mat_b, gemm_n); +} + +template +struct DSV3FusedAGemmKernel { + static void + run(const tvm::ffi::TensorView mat_a, const tvm::ffi::TensorView mat_b, const tvm::ffi::TensorView output) { + using namespace host; + + auto M = SymbolicSize{"num_tokens"}; + auto K = SymbolicSize{"hd_in"}; + auto N = SymbolicSize{"hd_out"}; + auto device = SymbolicDevice{}; + K.set_value(kHdIn); + N.set_value(kHdOut); + device.set_options(); + + // mat_a: [num_tokens, hd_in] row-major; output: [num_tokens, hd_out] row-major; + // mat_b: [hd_in, hd_out] column-major (weight.T), i.e. strides {1, hd_in}. + TensorMatcher({M, K}).with_dtype().with_device(device).verify(mat_a); + TensorMatcher({K, N}).with_dtype().with_device(device).with_strides({1, kHdIn}).verify(mat_b); + TensorMatcher({M, N}).with_dtype().with_device(device).verify(output); + + const int num_tokens = static_cast(M.unwrap()); + RuntimeCheck( + num_tokens >= 1 && num_tokens <= 16, "dsv3_fused_a_gemm: num_tokens must be in [1, 16], got ", num_tokens); + + const DLDevice dev = device.unwrap(); + auto* out_ptr = static_cast(output.data_ptr()); + auto* a_ptr = static_cast(mat_a.data_ptr()); + auto* b_ptr = static_cast(mat_b.data_ptr()); + + if (num_tokens <= 8) { + invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); + } else { + invokeFusedAGemm(out_ptr, a_ptr, b_ptr, num_tokens, dev); + } + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/cutedsl_dsv3_fused_a_gemm.py b/python/sglang/jit_kernel/cutedsl_dsv3_fused_a_gemm.py new file mode 100644 index 000000000..7e89986d2 --- /dev/null +++ b/python/sglang/jit_kernel/cutedsl_dsv3_fused_a_gemm.py @@ -0,0 +1,384 @@ +# Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. +# Copyright 2023-2024 SGLang Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CuTe DSL DeepSeek-V3 fused-A GEMM (sm90+): out[M, N] = mat_a[M, K] @ weight, +N in {2112, 6144}, M = num_tokens in [1, 16], K any multiple of 1024, bf16. + +Adapted from NVIDIA TensorRT-LLM dsv3FusedAGemm.cu +(cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu), reimplemented +in the CuTe DSL: AB-swap, warp-specialized 4-way split-K, cp.async + mbarrier +pipeline, ldmatrix + mma.sync.m16n8k16, 3-4-3 swizzle. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.utils import get_smem_capacity_in_bytes + +from sglang.kernel_api_logging import debug_kernel_api +from sglang.srt.utils import get_device_sm +from sglang.srt.utils.common import direct_register_custom_op + +TILE_M = 16 +TILE_K = 256 +SPLITK = 4 +LOAD_WARPS = 4 +MAX_NSTAGE = 16 +PWK = TILE_K // SPLITK +KSTEPS = PWK // 16 +COMPUTE_THREADS = SPLITK * 32 +LOADER_THREADS = LOAD_WARPS * 32 +NTHREADS = COMPUTE_THREADS + LOADER_THREADS +KI = TILE_K // 2 + +_BAR_I32 = 2 * MAX_NSTAGE * 2 + + +def _stage_i32(tile_n: int) -> int: + return (TILE_M + tile_n) * KI + + +def _cp_async_16b(smem_ptr, gmem_ptr): + llvm.inline_asm( + None, + [smem_ptr.toint().ir_value(), gmem_ptr.toint().ir_value()], + "{ .reg .u32 sa; cvt.u32.u64 sa, $0; cp.async.cg.shared.global.L2::128B [sa], [$1], 16; }", + "l,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=0, + ) + + +def _cp_async_16b_pred(smem_ptr, gmem_ptr, pred_i32): + llvm.inline_asm( + None, + [smem_ptr.toint().ir_value(), gmem_ptr.toint().ir_value(), pred_i32.ir_value()], + "{ .reg .pred p; .reg .u32 sa; setp.ne.s32 p, $2, 0; cvt.u32.u64 sa, $0; " + "@p cp.async.cg.shared.global.L2::128B [sa], [$1], 16; }", + "l,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=0, + ) + + +def _ldmatrix_x4(smem_ptr): + i32 = ir.IntegerType.get_signless(32) + res = llvm.inline_asm( + llvm.StructType.get_literal([i32, i32, i32, i32]), + [smem_ptr.toint().ir_value()], + "{ .reg .u32 sa; cvt.u32.u64 sa, $4; " + "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {$0,$1,$2,$3}, [sa]; }", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=0, + ) + return [llvm.extractvalue(i32, res, [i]) for i in range(4)] + + +def _ldmatrix_x2(smem_ptr): + i32 = ir.IntegerType.get_signless(32) + res = llvm.inline_asm( + llvm.StructType.get_literal([i32, i32]), + [smem_ptr.toint().ir_value()], + "{ .reg .u32 sa; cvt.u32.u64 sa, $2; " + "ldmatrix.sync.aligned.x2.m8n8.shared.b16 {$0,$1}, [sa]; }", + "=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=0, + ) + return [llvm.extractvalue(i32, res, [i]) for i in range(2)] + + +def _mma_m16n8k16(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3): + f32 = ir.F32Type.get() + res = llvm.inline_asm( + llvm.StructType.get_literal([f32, f32, f32, f32]), + [a0, a1, a2, a3, b0, b1, c0, c1, c2, c3], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0,$1,$2,$3}, {$4,$5,$6,$7}, {$8,$9}, {$10,$11,$12,$13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=0, + ) + return [llvm.extractvalue(f32, res, [i]) for i in range(4)] + + +def _swizzle_343(row, col): + return col ^ ((row % 8) * 4) + + +def _global_k_col(tile, col, kgi): + kp_warp = KI // SPLITK + kp_chunk = kgi // SPLITK + return (col // kp_warp) * kp_chunk + tile * kp_warp + (col % kp_warp) + + +def _load_stage(ltid, feat0, sa, sb, mW, mA, tile, buf, M, kgi, tile_n): + for it in range(TILE_M * KI // (LOADER_THREADS * 4)): + idx = (it * LOADER_THREADS + ltid) * 4 + row, col = idx // KI, idx % KI + _cp_async_16b( + sa.iterator + (buf * TILE_M * KI + row * KI + _swizzle_343(row, col)), + mW.iterator + ((feat0 + row) * kgi + _global_k_col(tile, col, kgi)), + ) + for it in range(tile_n * KI // (LOADER_THREADS * 4)): + idx = (it * LOADER_THREADS + ltid) * 4 + row, col = idx // KI, idx % KI + pred = (row < M).to(cutlass.Int32) + _cp_async_16b_pred( + sb.iterator + (buf * tile_n * KI + row * KI + _swizzle_343(row, col)), + mA.iterator + (row * pred * kgi + _global_k_col(tile, col, kgi)), + pred, + ) + + +@cute.kernel +def _dsv3_fused_a_gemm_kernel( + mW: cute.Tensor, + mA: cute.Tensor, + mOut: cute.Tensor, + M: cutlass.Int32, + num_kt: cutlass.Constexpr, + nstage: cutlass.Constexpr, + tile_n: cutlass.Constexpr, +): + NB = tile_n // 8 + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + warp, lane = tid // 32, tid % 32 + r0, cc = lane // 4, lane % 4 + feat0 = bid * TILE_M + kgi = num_kt * KI + + base = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) + bar = cute.recast_ptr(base, dtype=cutlass.Int64) + full, empty = bar, bar + MAX_NSTAGE + sa_off = _BAR_I32 + sc_warp_stride = TILE_M * tile_n + 2 + sC = cute.make_tensor( + cute.recast_ptr(base + sa_off, dtype=cutlass.Float32), + cute.make_layout((SPLITK, TILE_M, tile_n), stride=(sc_warp_stride, tile_n, 1)), + ) + sA = cute.make_tensor( + base + sa_off, + cute.make_layout((nstage, TILE_M, KI), stride=(TILE_M * KI, KI, 1)), + ) + sB = cute.make_tensor( + base + sa_off + nstage * TILE_M * KI, + cute.make_layout((nstage, tile_n, KI), stride=(tile_n * KI, KI, 1)), + ) + + if tid == 0: + for s in range(nstage): + cute.arch.mbarrier_init(full + s, LOADER_THREADS) + cute.arch.mbarrier_init(empty + s, COMPUTE_THREADS) + cute.arch.barrier() + + if warp >= SPLITK: + cute.arch.griddepcontrol_wait() + ltid = tid - COMPUTE_THREADS + for kt in cutlass.range_constexpr(num_kt): + st = kt % nstage + if kt >= nstage: + cute.arch.mbarrier_wait(empty + st, ((kt // nstage) & 1) ^ 1) + _load_stage(ltid, feat0, sA, sB, mW, mA, kt, st, M, kgi, tile_n) + cute.arch.cp_async_mbarrier_arrive_noinc(full + st) + else: + acc = [[cutlass.Float32(0.0) for _ in range(4)] for _ in range(NB)] + ph = 0 + for kt in cutlass.range_constexpr(num_kt): + buf = kt % nstage + cute.arch.mbarrier_wait(full + buf, ph) + brow_lo = lane % 8 + boff = ((lane // 8) & 1) * 4 + for step in cutlass.range_constexpr(KSTEPS): + kbh = warp * (PWK // 2) + step * 8 + arow, aoff = lane % 16, (lane // 16) * 4 + a = _ldmatrix_x4( + sA.iterator + + (buf * TILE_M * KI + arow * KI + _swizzle_343(arow, kbh + aoff)) + ) + for nb in cutlass.range_constexpr(NB): + brow = nb * 8 + brow_lo + bb = _ldmatrix_x2( + sB.iterator + + ( + buf * tile_n * KI + + brow * KI + + _swizzle_343(brow, kbh + boff) + ) + ) + d = _mma_m16n8k16( + a[0], + a[1], + a[2], + a[3], + bb[0], + bb[1], + acc[nb][0].ir_value(), + acc[nb][1].ir_value(), + acc[nb][2].ir_value(), + acc[nb][3].ir_value(), + ) + for i in cutlass.range_constexpr(4): + acc[nb][i] = cutlass.Float32(d[i]) + cute.arch.mbarrier_arrive(empty + buf) + ph = (ph ^ 1) if buf == nstage - 1 else ph + for nb in cutlass.range_constexpr(NB): + for i in cutlass.range_constexpr(4): + m = r0 + (8 if i >= 2 else 0) + n = nb * 8 + cc * 2 + (i % 2) + sC[warp, m, n] = acc[nb][i] + + cute.arch.barrier(barrier_id=1, number_of_threads=COMPUTE_THREADS) + nred = TILE_M * tile_n + for it in cutlass.range_constexpr( + (nred + COMPUTE_THREADS - 1) // COMPUTE_THREADS + ): + e = it * COMPUTE_THREADS + tid + if e < nred: + m, n = e // tile_n, e % tile_n + s = sC[0, m, n] + for w in cutlass.range_constexpr(1, SPLITK): + s = s + sC[w, m, n] + if n < M: + mOut[n, feat0 + m] = s.to(cutlass.BFloat16) + + cute.arch.griddepcontrol_launch_dependents() + + +@cute.jit +def _dsv3_fused_a_gemm_host( + mW: cute.Tensor, + mA: cute.Tensor, + mOut: cute.Tensor, + M: cutlass.Int32, + stream: cuda.CUstream, + num_kt: cutlass.Constexpr, + gemm_m: cutlass.Constexpr, + smem_bytes: cutlass.Constexpr, + nstage: cutlass.Constexpr, + tile_n: cutlass.Constexpr, +): + _dsv3_fused_a_gemm_kernel(mW, mA, mOut, M, num_kt, nstage, tile_n).launch( + grid=[gemm_m // TILE_M, 1, 1], + block=[NTHREADS, 1, 1], + max_number_threads=[NTHREADS, 1, 1], + min_blocks_per_mp=1, + smem=smem_bytes, + use_pdl=True, + stream=stream, + ) + + +_compiled: dict[tuple[int, int, int], object] = {} + + +def _pick_nstage(num_kt: int, tile_n: int) -> int: + nstage = (get_smem_capacity_in_bytes() // 4 - _BAR_I32) // _stage_i32(tile_n) + return min(nstage, MAX_NSTAGE, num_kt) + + +def _pick_tile_n(num_tokens: int) -> int: + return 8 if num_tokens <= 8 else 16 + + +def _compiled_kernel(num_kt: int, gemm_m: int, tile_n: int): + if get_device_sm() < 90: + raise RuntimeError("dsv3_fused_a_gemm requires SM90 (Hopper) or later") + if (num_kt, gemm_m, tile_n) not in _compiled: + nstage = _pick_nstage(num_kt, tile_n) + smem_bytes = (_BAR_I32 + nstage * _stage_i32(tile_n)) * 4 + k = num_kt * TILE_K + w = torch.empty(gemm_m, k, dtype=torch.bfloat16, device="cuda") + a = torch.empty(16, k, dtype=torch.bfloat16, device="cuda") + o = torch.empty(16, gemm_m, dtype=torch.bfloat16, device="cuda") + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + _compiled[(num_kt, gemm_m, tile_n)] = cute.compile( + _dsv3_fused_a_gemm_host, + from_dlpack(w.view(torch.int32)), + from_dlpack(a.view(torch.int32)), + from_dlpack(o), + cutlass.Int32(16), + stream, + num_kt, + gemm_m, + smem_bytes, + nstage, + tile_n, + ) + return _compiled[(num_kt, gemm_m, tile_n)] + + +def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + M, K = mat_a.shape + N = mat_b.shape[1] + assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16 + assert K % 1024 == 0, f"K must be a multiple of 1024, got {K}" + assert N % TILE_M == 0, f"N must be a multiple of {TILE_M}, got {N}" + assert ( + tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1 + ), "mat_b must be [K, N] column-major" + assert 1 <= M <= 16, "num_tokens must be in [1, 16]" + assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]" + + weight = mat_b.t() + out = torch.empty(M, N, dtype=torch.bfloat16, device=mat_a.device) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + _compiled_kernel(K // TILE_K, N, _pick_tile_n(M))( + from_dlpack(weight.view(torch.int32)), + from_dlpack(mat_a.view(torch.int32)), + from_dlpack(out), + M, + stream, + ) + return out + + +def _dsv3_fused_a_gemm_fake(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + return mat_a.new_empty((mat_a.shape[0], mat_b.shape[1]), dtype=torch.bfloat16) + + +direct_register_custom_op( + op_name="cutedsl_dsv3_fused_a_gemm", + op_func=_dsv3_fused_a_gemm_run, + mutates_args=[], + fake_impl=_dsv3_fused_a_gemm_fake, +) + + +@debug_kernel_api +def dsv3_fused_a_gemm( + mat_a: torch.Tensor, mat_b: torch.Tensor, output: torch.Tensor | None = None +) -> torch.Tensor: + """out[M, N] = mat_a[M, K] @ mat_b, with mat_a row-major [M, K] (M in [1, 16]), + mat_b column-major [K, N] (the weight, stride(0) == 1), N a multiple of 16 + (e.g. 2112, 6144), K a multiple of 1024.""" + result = torch.ops.sglang.cutedsl_dsv3_fused_a_gemm(mat_a, mat_b) + if output is not None: + output.copy_(result) + return output + return result diff --git a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py new file mode 100644 index 000000000..c718a5131 --- /dev/null +++ b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py @@ -0,0 +1,90 @@ +""" +JIT kernel for DeepSeek V3 fused QKV-A GEMM (min-latency). + +Replaces the AOT sgl_kernel.dsv3_fused_a_gemm for SM90+ (Hopper) GPUs. +Shapes: hd_in a multiple of 256, hd_out a multiple of 16, num_tokens 1-16, bfloat16. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.kernel_api_logging import debug_kernel_api +from sglang.srt.utils.common import direct_register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_dsv3_fused_a_gemm_module(hd_in: int, hd_out: int, use_pdl: bool) -> Module: + args = make_cpp_args(hd_in, hd_out, use_pdl) + return load_jit( + "dsv3_fused_a_gemm", + *args, + cuda_files=["gemm/dsv3_fused_a_gemm.cuh"], + cuda_wrappers=[ + ("dsv3_fused_a_gemm", f"DSV3FusedAGemmKernel<{args}>::run"), + ], + ) + + +def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]" + output = torch.empty( + (mat_a.shape[0], mat_b.shape[1]), + device=mat_a.device, + dtype=mat_a.dtype, + ) + module = _jit_dsv3_fused_a_gemm_module( + mat_a.shape[1], mat_b.shape[1], is_arch_support_pdl() + ) + module.dsv3_fused_a_gemm(mat_a, mat_b, output) + return output + + +def _dsv3_fused_a_gemm_fake(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + return mat_a.new_empty((mat_a.shape[0], mat_b.shape[1]), dtype=torch.bfloat16) + + +direct_register_custom_op( + op_name="jit_dsv3_fused_a_gemm", + op_func=_dsv3_fused_a_gemm_run, + mutates_args=[], + fake_impl=_dsv3_fused_a_gemm_fake, +) + + +@debug_kernel_api +def dsv3_fused_a_gemm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + DeepSeek V3 fused QKV-A GEMM kernel (JIT variant). + + Args: + mat_a: Input tensor of shape [num_tokens, hd_in], bfloat16, row-major. + hd_in must be a multiple of 256 and num_tokens in [1, 16]. + mat_b: Weight tensor of shape [hd_in, hd_out], bfloat16, column-major + (i.e. ``weight.T`` of a row-major [hd_out, hd_in] weight). + hd_out must be a multiple of 16. + output: Optional pre-allocated output tensor of shape [num_tokens, hd_out]. + + Returns: + Output tensor of shape [num_tokens, hd_out]. + """ + result = torch.ops.sglang.jit_dsv3_fused_a_gemm(mat_a, mat_b) + if output is not None: + output.copy_(result) + return output + return result diff --git a/python/sglang/jit_kernel/fused_a_gemm.py b/python/sglang/jit_kernel/fused_a_gemm.py new file mode 100644 index 000000000..7ab37f31c --- /dev/null +++ b/python/sglang/jit_kernel/fused_a_gemm.py @@ -0,0 +1,52 @@ +"""Unified entry point for the DeepSeek-V3 fused QKV-A GEMM. + +Dispatches to one of three interchangeable implementations via ``backend``: + +- ``"aot"``: prebuilt ``sgl_kernel.dsv3_fused_a_gemm`` (CUDA C++). +- ``"jit"``: runtime-compiled CUDA C++ (``sglang.jit_kernel.dsv3_fused_a_gemm``). +- ``"cutedsl"``: CuTe DSL (``sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm``). +- ``"auto"``: CuTe DSL on SM120+, otherwise the JIT kernel. + +All backends share the signature ``(mat_a, mat_b, output=None) -> Tensor`` with +``mat_a`` row-major ``[M, K]`` (M in [1, 16], bf16) and ``mat_b`` the column-major +weight ``[K, N]`` (``weight.T``). +""" + +from enum import Enum + +import torch + +from sglang.srt.utils.common import is_sm120_supported + + +class FusedAGemmBackend(str, Enum): + AUTO = "auto" + AOT = "aot" + JIT = "jit" + CUTEDSL = "cutedsl" + + +_AUTO_BACKEND = ( + FusedAGemmBackend.CUTEDSL if is_sm120_supported() else FusedAGemmBackend.JIT +) + + +def dsv3_fused_a_gemm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + output: torch.Tensor | None = None, + backend: FusedAGemmBackend | str = FusedAGemmBackend.AUTO, +) -> torch.Tensor: + backend = FusedAGemmBackend(backend) + if backend == FusedAGemmBackend.AUTO: + backend = _AUTO_BACKEND + + if backend == FusedAGemmBackend.AOT: + from sgl_kernel import dsv3_fused_a_gemm as impl + elif backend == FusedAGemmBackend.JIT: + from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm as impl + else: + from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import ( + dsv3_fused_a_gemm as impl, + ) + return impl(mat_a, mat_b, output) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index eda36be57..7131bcbbd 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -201,11 +201,10 @@ if _use_aiter: pass if _is_cuda: - from sgl_kernel import dsv3_fused_a_gemm - from sglang.jit_kernel.dsv3_router_gemm import ( dsv3_router_gemm as _jit_dsv3_router_gemm, ) + from sglang.jit_kernel.fused_a_gemm import dsv3_fused_a_gemm elif _is_npu: from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import ( forward_dsa_core_npu, @@ -1769,11 +1768,12 @@ class DeepseekV2AttentionMLA( self.has_fused_proj and not self.is_packed_weight and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.bfloat16 - and self.fused_qkv_a_proj_with_mqa.weight.shape[0] == 2112 - and self.fused_qkv_a_proj_with_mqa.weight.shape[1] == 7168 + and self.fused_qkv_a_proj_with_mqa.weight.shape[0] % 16 == 0 + and self.fused_qkv_a_proj_with_mqa.weight.shape[1] % 256 == 0 and _is_cuda - and 90 <= _device_sm < 120 + and _device_sm >= 90 ) + self.fused_a_gemm_backend = "auto" self.init_mha_forward() self.init_mla_forward() @@ -1987,7 +1987,9 @@ class DeepseekV2AttentionMLA( and not lora_active ): qkv_latent = dsv3_fused_a_gemm( - hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T + hidden_states, + self.fused_qkv_a_proj_with_mqa.weight.T, + backend=self.fused_a_gemm_backend, ) else: qkv_latent = self.fused_qkv_a_proj_with_mqa(hidden_states)[0] diff --git a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py new file mode 100644 index 000000000..b1d34cc2f --- /dev/null +++ b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py @@ -0,0 +1,94 @@ +"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs +sgl_kernel AOT vs torch. + +Run on SM90+ (Hopper or later): + python test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py +""" + +import torch +import torch.nn.functional as F +import triton.testing +from sgl_kernel import dsv3_fused_a_gemm as sgl_kernel_dsv3_fused_a_gemm + +from sglang.jit_kernel.benchmark import marker +from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import ( + dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm, +) +from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm +from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.srt.utils.common import is_sm120_supported +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.utils import is_in_ci + +register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large") + +IS_CI = is_in_ci() + +DTYPE = torch.bfloat16 +DEVICE = "cuda" +HD_OUT = 2112 +HD_IN_LIST = [6144, 7168] + +AOT_HD_IN = 7168 +HAS_AOT = not is_sm120_supported() + +NUM_TOKENS_LIST = [1, 8, 16] if IS_CI else list(range(1, 17)) + +LINE_VALS = ["cutedsl", "jit", "sgl_kernel", "torch"] +LINE_NAMES = ["CuTe DSL", "CUDA JIT", "sgl_kernel AOT", "torch F.linear"] +STYLES = [("blue", "-"), ("orange", "--"), ("red", ":"), ("green", "-.")] + + +def _median_us(fn, *args) -> float: + result = marker.do_bench( + fn, + input_args=args, + use_cuda_graph=True, + metrics=(0.5,), + disable_log_bandwidth=True, + ) + return result.times[0] * 1e6 + + +def _bench(num_tokens, provider, hd_in): + if provider == "sgl_kernel" and not (HAS_AOT and hd_in == AOT_HD_IN): + return float("nan") + + mat_a = torch.randn((num_tokens, hd_in), dtype=DTYPE, device=DEVICE) + mat_b = torch.randn((HD_OUT, hd_in), dtype=DTYPE, device=DEVICE).transpose(0, 1) + fn_map = { + "cutedsl": cutedsl_dsv3_fused_a_gemm, + "jit": dsv3_fused_a_gemm, + "sgl_kernel": sgl_kernel_dsv3_fused_a_gemm, + "torch": lambda a, b: F.linear(a, b.T), + } + return _median_us(fn_map[provider], mat_a, mat_b) + + +@triton.testing.perf_report( + [ + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=NUM_TOKENS_LIST, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name=f"dsv3-fused-a-gemm-bf16-K{hd_in}-N{HD_OUT}", + args={"hd_in": hd_in}, + ) + for hd_in in HD_IN_LIST + ] +) +def benchmark(num_tokens, provider, hd_in): + return _bench(num_tokens, provider, hd_in) + + +if __name__ == "__main__": + if is_hip_runtime() or get_jit_cuda_arch().major < 9: + print( + "dsv3_fused_a_gemm JIT kernel requires SM90+ (Hopper). Skipping benchmark." + ) + else: + benchmark.run(print_data=True) diff --git a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py b/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py new file mode 100644 index 000000000..83162b52d --- /dev/null +++ b/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py @@ -0,0 +1,42 @@ +"""Tests for the CuTe DSL DeepSeek-V3 fused-A GEMM kernel.""" + +import sys + +import pytest +import torch + +from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import dsv3_fused_a_gemm +from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + +# hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes. +HD_INS = [6144, 7168] +# hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants. +HD_OUTS = [2112, 2624] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("hd_out", HD_OUTS) +@pytest.mark.parametrize("hd_in", HD_INS) +@pytest.mark.parametrize("num_tokens", list(range(1, 17))) +def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out): + if is_hip_runtime() or get_jit_cuda_arch().major < 9: + pytest.skip("SM90+ required") + + torch.manual_seed(num_tokens) + weight = torch.randn(hd_out, hd_in, dtype=torch.bfloat16, device="cuda") + mat_a = torch.randn(num_tokens, hd_in, dtype=torch.bfloat16, device="cuda") + mat_b = weight.t() + + out = dsv3_fused_a_gemm(mat_a, mat_b) + assert out.shape == (num_tokens, hd_out) + assert out.dtype == torch.bfloat16 + + ref = (mat_a.float() @ weight.float().T).bfloat16() + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2.5) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/test_dsv3_fused_a_gemm.py b/test/registered/jit/test_dsv3_fused_a_gemm.py new file mode 100644 index 000000000..5c34eef58 --- /dev/null +++ b/test/registered/jit/test_dsv3_fused_a_gemm.py @@ -0,0 +1,43 @@ +"""Tests for JIT dsv3_fused_a_gemm kernel.""" + +import sys + +import pytest +import torch +import torch.nn.functional as F + +from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm +from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + +# hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes. +HD_INS = [6144, 7168] +# hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants. +HD_OUTS = [2112, 2624] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("hd_out", HD_OUTS) +@pytest.mark.parametrize("hd_in", HD_INS) +@pytest.mark.parametrize("num_tokens", list(range(1, 17))) +def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out): + if is_hip_runtime() or get_jit_cuda_arch().major < 9: + pytest.skip("SM90+ required") + + mat_a = torch.randn(num_tokens, hd_in, dtype=torch.bfloat16, device="cuda") + mat_b = torch.randn(hd_out, hd_in, dtype=torch.bfloat16, device="cuda").transpose( + 0, 1 + ) + + ref = F.linear(mat_a, mat_b.T) + out = dsv3_fused_a_gemm(mat_a, mat_b) + + assert out.shape == (num_tokens, hd_out) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))