diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 71c98e881..0833a68c2 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -1208,14 +1208,6 @@ class BailingMoELinearForCausalLM(nn.Module): ) if _is_hip: self_attn.w_scale *= 2.0 - # TODO: remove this after adding FP8 support in bmm cpu kernel - if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: - self_attn.w_kc = ( - self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale - ) - self_attn.w_vc = ( - self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale - ) else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 9eca43ce7..fe1b3c966 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.fp8_kernel import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.deepseek_common.utils import ( FORWARD_ABSORB_CORE_ATTENTION_BACKENDS, + _is_cpu, _is_cublas_ge_129, _is_cuda, _is_gfx95_supported, @@ -268,18 +269,24 @@ class DeepseekMLAForwardMixin: ) elif self.w_kc.dtype == torch.float8_e4m3fn: - # fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612 - q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( - q_nope.transpose(0, 1), - ( - torch.zeros((1,), dtype=torch.float32, device=q_nope.device) - if _is_cublas_ge_129 - else zero_allocator.allocate(1) - ), - ) - q_nope_out = bmm_fp8( - q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 - ) + if _is_cpu: + q_nope_out = torch.bmm( + q_nope.to(torch.bfloat16).transpose(0, 1), + self.w_kc.to(torch.bfloat16) * self.w_scale, + ) + else: + # fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612 + q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( + q_nope.transpose(0, 1), + ( + torch.zeros((1,), dtype=torch.float32, device=q_nope.device) + if _is_cublas_ge_129 + else zero_allocator.allocate(1) + ), + ) + q_nope_out = bmm_fp8( + q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 + ) else: q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) @@ -455,22 +462,31 @@ class DeepseekMLAForwardMixin: attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) elif self.w_vc.dtype == torch.float8_e4m3fn: - attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8( - attn_output.transpose(0, 1), - ( - torch.zeros((1,), dtype=torch.float32, device=attn_output.device) - if _is_cublas_ge_129 - else zero_allocator.allocate(1) - ), - ) - attn_bmm_output = bmm_fp8( - attn_output_val, - self.w_vc, - attn_output_scale, - self.w_scale, - torch.bfloat16, - ) - attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + if _is_cpu: + attn_bmm_output = torch.bmm( + attn_output.to(torch.bfloat16).transpose(0, 1), + self.w_vc.to(torch.bfloat16) * self.w_scale, + ) + attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + else: + attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8( + attn_output.transpose(0, 1), + ( + torch.zeros( + (1,), dtype=torch.float32, device=attn_output.device + ) + if _is_cublas_ge_129 + else zero_allocator.allocate(1) + ), + ) + attn_bmm_output = bmm_fp8( + attn_output_val, + self.w_vc, + attn_output_scale, + self.w_scale, + torch.bfloat16, + ) + attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) else: if is_in_piecewise_cuda_graph(): # torch dynamo requires out= op was called where output tensor was non-contiguous diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py index 43a5eed37..b6c76df09 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py @@ -100,6 +100,7 @@ class DeepseekMLACpuForwardMixin: else None ) ), + self.w_scale, True, # is_vnni self.weight_block_size, self.q_lora_rank, @@ -144,7 +145,7 @@ class DeepseekMLACpuForwardMixin: attn_output.transpose(0, 1), self.w_vc, True, # is_vnni - None, # scale + self.w_scale, # scale ) attn_output = output output, _ = self.o_proj(attn_output) diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index c54ab358d..12ce382ed 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -46,8 +46,6 @@ from sglang.srt.model_loader.utils import ( ) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_common.utils import ( - _is_cpu, - _is_cpu_amx_available, _is_cuda, _is_fp8_fnuz, _is_hip, @@ -583,14 +581,6 @@ class DeepseekV2WeightLoaderMixin: ) if _is_hip: self_attn.w_scale *= 2.0 - # TODO: remove this after adding FP8 support in bmm cpu kernel - if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: - self_attn.w_kc = ( - self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale - ) - self_attn.w_vc = ( - self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale - ) else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index c55c741ad..028148a00 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -776,18 +776,6 @@ class LongcatFlashForCausalLM(nn.Module): ) if _is_hip: self_attn.w_scale *= 2.0 - # TODO: remove this after adding FP8 support in bmm cpu kernel - if ( - _is_cpu - and _is_cpu_amx_available - and w.dtype == torch.float8_e4m3fn - ): - self_attn.w_kc = ( - self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale - ) - self_attn.w_vc = ( - self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale - ) else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index abaf27855..d3dc20f36 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -426,10 +426,6 @@ class LongcatFlashForCausalLMNextN(LongcatFlashForCausalLM): ) if _is_hip: self_attn.w_scale *= 2.0 - # TODO: remove this after adding FP8 support in bmm cpu kernel - if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: - self_attn.w_kc = self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale - self_attn.w_vc = self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] diff --git a/sgl-kernel/csrc/cpu/bmm.cpp b/sgl-kernel/csrc/cpu/bmm.cpp index 9e809a464..48d515ef4 100644 --- a/sgl-kernel/csrc/cpu/bmm.cpp +++ b/sgl-kernel/csrc/cpu/bmm.cpp @@ -4,11 +4,11 @@ namespace { -template +template void bmm_kernel_impl( scalar_t* __restrict__ out, const scalar_t* __restrict__ mat1, - const scalar_t* __restrict__ mat2, + const packed_t* __restrict__ mat2, int64_t B, int64_t M, int64_t N, @@ -67,6 +67,67 @@ void bmm_kernel_impl( }); } +template <> +void bmm_kernel_impl( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ mat1, + const at::Float8_e4m3fn* __restrict__ mat2, + int64_t B, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideB, + int64_t mat1_strideM, + int64_t out_strideB, + int64_t out_strideM, + float scale) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + // mat2 contiguous in [B, N, K] + int64_t mat2_strideB = N * K; + int64_t mat2_strideN = K; + + const bool use_brgemm = can_use_brgemm(M); + + // parallel on [B, MB, NB] + parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + // for brgemm, use float32 for accumulate + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + // for brgemm when mat2 is float8_e4m3 + alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t bs = mb / MB; + int64_t mb_start = (mb % MB) * BLOCK_M; + int64_t mb_size = std::min(M - mb_start, BLOCK_M); + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(N - nb_start, BLOCK_N); + + tinygemm_kernel( + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM, + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */, + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, + /* Btmp*/ Btmp, + /* Ctmp*/ Ctmp, + /*scale*/ scale, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + }); + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + } // anonymous namespace // mat1 : [B, M, K] @@ -94,7 +155,7 @@ void bmm_cpu( int64_t N = mat2.size(1); int64_t K = mat1.size(2); - TORCH_CHECK(!scale.has_value(), "bmm: do not support fp8 weight for now.") + const bool use_fp8_w8a16 = scale.has_value(); TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x."); int64_t mat1_strideB = mat1.stride(0); @@ -105,12 +166,32 @@ void bmm_cpu( // check shapes TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!"); TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!"); + if (!use_fp8_w8a16) { + AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] { + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + B, + M, + N, + K, + mat1_strideB, + mat1_strideM, + out_strideB, + out_strideM); + }); + } else { // fp8 bmm + float scale_val = 0.f; - AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] { - bmm_kernel_impl( - out.data_ptr(), - mat1.data_ptr(), - packed_w.data_ptr(), + auto scale_tensor = scale.value(); + TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor."); + scale_val = scale_tensor.item(); + + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), B, M, N, @@ -118,6 +199,7 @@ void bmm_cpu( mat1_strideB, mat1_strideM, out_strideB, - out_strideM); - }); + out_strideM, + scale_val); + } } diff --git a/sgl-kernel/csrc/cpu/gemm.h b/sgl-kernel/csrc/cpu/gemm.h index cbed7edea..e11b224fe 100644 --- a/sgl-kernel/csrc/cpu/gemm.h +++ b/sgl-kernel/csrc/cpu/gemm.h @@ -232,6 +232,7 @@ void tinygemm_kernel( int64_t ldc, bool brg); +// block quantization template void tinygemm_kernel( const scalar_t* __restrict__ A, @@ -250,6 +251,23 @@ void tinygemm_kernel( int64_t block_size_K, bool do_unpack = true); +// per tensor quantization +template +void tinygemm_kernel( + const scalar_t* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + scalar_t* __restrict__ C, + scalar_t* __restrict__ Btmp, + float* __restrict__ Ctmp, + float scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + bool brg); + template void tinygemm_kernel( scalar_t* C, diff --git a/sgl-kernel/csrc/cpu/gemm_fp8.cpp b/sgl-kernel/csrc/cpu/gemm_fp8.cpp index b2821982a..15bd44434 100644 --- a/sgl-kernel/csrc/cpu/gemm_fp8.cpp +++ b/sgl-kernel/csrc/cpu/gemm_fp8.cpp @@ -42,6 +42,25 @@ inline void copy_add_stub( out[d] = static_cast(input[d] + bias[d]); } } +template +inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int size, float scale) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec vscale = fVec(scale); + + int d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + fVec data0 = fVec::loadu(input + d) * vscale; + fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale; + bVec out_vec = convert_from_float_ext(data0, data1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] * scale); + } +} inline void unpack_B( at::BFloat16* __restrict__ Btmp, @@ -100,6 +119,41 @@ inline void unpack_B( #endif } +inline void unpack_B( + at::BFloat16* __restrict__ Btmp, + const at::Float8_e4m3fn* __restrict__ packed_B, + int N, + int K, + int ldb, + int ldb_tmp) { +#if defined(CPU_CAPABILITY_AVX512) + // [K/2, N, 2] + const int K2 = K >> 1; + const int ldb2 = ldb; // ldb * 2 >> 1; + const uint16_t* b_ptr = reinterpret_cast(packed_B); + + // prefetch distance + constexpr int PREFETCH_SIZE_K = 64; +#pragma GCC unroll 4 + for (int k = 0; k < K2; ++k) { + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2); + if constexpr (PREFETCH_SIZE_K > 0) { + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0); + } + + __m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0); + __m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1); + + __m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0); + __m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1); + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)bf16_0); + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)bf16_1); + } +#else + TORCH_CHECK(false, "unpack_B: scalar path not implemented!"); +#endif +} + template struct tinygemm_kernel_nn { static inline void apply( @@ -117,6 +171,20 @@ struct tinygemm_kernel_nn { } }; +template +struct tinygemm_kernel_nn2 { + static inline void apply( + const scalar_t* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + scalar_t* __restrict__ C, + float scale, + int K, + int lda, + int ldb, + int ldc) { + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); + } +}; #if defined(CPU_CAPABILITY_AVX512) template struct tinygemm_kernel_nn { @@ -221,6 +289,76 @@ struct tinygemm_kernel_nn{}(storec); } }; + +template +struct tinygemm_kernel_nn2 { + static inline void apply( + const at::BFloat16* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + at::BFloat16* __restrict__ C, + float scale, + int K, + int lda, + int ldb, + int ldc) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N / 16; + + // prefetch distance + constexpr int PREFETCH_SIZE_K = 64; + + __m512bh va; + __m512bh vb[COLS]; + __m512 vc[ROWS * COLS]; + + const __m512 vscale = _mm512_set1_ps(scale); + + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + Unroll{}(loadc); + + const int K2 = K >> 1; + const int lda2 = lda >> 1; + const int ldb2 = ldb; // ldb * 2 >> 1; + const float* a_ptr = reinterpret_cast(A); + const uint16_t* b_ptr = reinterpret_cast(B); + + auto compute = [&](auto i, int k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k])); + } + if constexpr (row == 0) { + if constexpr (col % 2 == 0) { + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2 + col * 16); + if constexpr (PREFETCH_SIZE_K > 0) { + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0); + } + vb[col + 0] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 0)); + vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1)); + } + } + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + }; + for (int k = 0; k < K2; ++k) { + Unroll{}(compute, k); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + // for COLS = 2, 4 use 512bit store + if constexpr (col % 2 == 0) { + __m512 vc0 = _mm512_mul_ps(vc[row * COLS + col + 0], vscale); + __m512 vc1 = _mm512_mul_ps(vc[row * COLS + col + 1], vscale); + _mm512_storeu_si512( + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0))); + } + }; + Unroll{}(storec); + } +}; #endif #define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \ @@ -236,6 +374,10 @@ struct tinygemm_kernel_nn::apply( \ + A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, scale, K, lda, ldb, ldc); + template struct brgemm { static inline void apply( @@ -256,6 +398,8 @@ struct brgemm { TORCH_CHECK(false, "struct brgemm: primary template not implemented!"); } }; +template +struct brgemm2 {}; template struct brgemm { @@ -301,6 +445,42 @@ struct brgemm { } }; +template <> +struct brgemm2 { + static inline void apply( + const at::BFloat16* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + at::BFloat16* __restrict__ C, + at::BFloat16* __restrict__ Btmp, + float* __restrict__ Ctmp, + float scale, + int M, + int N, + int K, + int lda, + int ldb, + int ldc) { + constexpr int BLOCK_N = block_size_n(); + + // [BLOCK_K, BLOCK_N] -> [BLOCK_K / 2, BLOCK_N * 2] + const int ldb_tmp = block_size_n(); + + // accumulate across K per BLOCK_K + for (int k = 0; k < K; k += BLOCK_K) { + int kb_size = std::min(BLOCK_K, K - k); + unpack_B(Btmp, B + k * ldb, N, kb_size, ldb, ldb_tmp); + + const bool add_C = (k != 0); + at::native::cpublas::brgemm(M, N, kb_size, lda, ldb_tmp, BLOCK_N, add_C, A + k, Btmp, Ctmp); + } + + // copy from Ctmp to C and mul scale + for (int m = 0; m < M; ++m) { + copy_mul_stub(C + m * ldc, Ctmp + m * BLOCK_N, N, scale); + } + } +}; + template void tinygemm_kernel( const scalar_t* __restrict__ A, @@ -356,7 +536,103 @@ void tinygemm_kernel( } } } +template +void tinygemm_kernel2( + const scalar_t* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + scalar_t* __restrict__ C, + scalar_t* __restrict__ Btmp, + float* __restrict__ Ctmp, + float scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + bool brg) { + if (brg) { + brgemm2::apply(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc); + return; + } + // pattern: 1-8-8 + if (M == 1) { + constexpr int64_t BLOCK_N = 128; + const int64_t NB = div_up(N, BLOCK_N); + int64_t mb_start = 0; + + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (nb_size >> 4) { + case 2: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 32); + break; + case 4: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 64); + break; + case 6: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 96); + break; + case 8: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 128); + break; + default: + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + } + } + return; + } + + // pattern: 1-4-16 + constexpr int64_t BLOCK_M = 4; + constexpr int64_t BLOCK_N = 64; + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + for (int64_t mb = 0; mb < MB; ++mb) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(BLOCK_M, M - mb_start); + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size >> 4) { + // mb_size = 1 + case 0x12: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 32); + break; + case 0x14: + LAUNCH_TINYGEMM_KERNEL_NN2(1, 64); + break; + // mb_size = 2 + case 0x22: + LAUNCH_TINYGEMM_KERNEL_NN2(2, 32); + break; + case 0x24: + LAUNCH_TINYGEMM_KERNEL_NN2(2, 64); + break; + // mb_size = 3 + case 0x32: + LAUNCH_TINYGEMM_KERNEL_NN2(3, 32); + break; + case 0x34: + LAUNCH_TINYGEMM_KERNEL_NN2(3, 64); + break; + // mb_size = 4 + case 0x42: + LAUNCH_TINYGEMM_KERNEL_NN2(4, 32); + break; + case 0x44: + LAUNCH_TINYGEMM_KERNEL_NN2(4, 64); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + } + } + } +} template void fp8_scaled_mm_kernel_impl( scalar_t* __restrict__ out, @@ -450,7 +726,23 @@ void tinygemm_kernel( tinygemm_kernel( A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack); } - +template +void tinygemm_kernel( + const scalar_t* __restrict__ A, + const at::Float8_e4m3fn* __restrict__ B, + scalar_t* __restrict__ C, + scalar_t* __restrict__ Btmp, + float* __restrict__ Ctmp, + float scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + bool brg) { + tinygemm_kernel2(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg); +} #define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \ template void tinygemm_kernel( \ const TYPE* __restrict__ A, \ @@ -469,8 +761,25 @@ void tinygemm_kernel( int64_t block_size_K, \ bool do_unpack) +#define INSTANTIATE_TINYGEMM_TEMPLATE2(TYPE) \ + template void tinygemm_kernel( \ + const TYPE* __restrict__ A, \ + const at::Float8_e4m3fn* __restrict__ B, \ + TYPE* __restrict__ C, \ + TYPE* __restrict__ Btmp, \ + float* __restrict__ Ctmp, \ + float scale, \ + int64_t M, \ + int64_t N, \ + int64_t K, \ + int64_t lda, \ + int64_t ldb, \ + int64_t ldc, \ + bool brg) + INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16); INSTANTIATE_TINYGEMM_TEMPLATE(at::Half); +INSTANTIATE_TINYGEMM_TEMPLATE2(at::BFloat16); at::Tensor fp8_scaled_mm_cpu( at::Tensor& mat1, diff --git a/sgl-kernel/csrc/cpu/qkv_proj.cpp b/sgl-kernel/csrc/cpu/qkv_proj.cpp index b3e2072e8..ba0e80d7e 100644 --- a/sgl-kernel/csrc/cpu/qkv_proj.cpp +++ b/sgl-kernel/csrc/cpu/qkv_proj.cpp @@ -434,6 +434,7 @@ std::tuple qkv_proj_with_rope( std::optional q_a_proj_scale, std::optional q_b_proj_scale, std::optional kv_a_proj_scale, + std::optional w_scale, bool is_vnni, std::optional> block_size) { RECORD_FUNCTION( @@ -601,10 +602,9 @@ std::tuple qkv_proj_with_rope( qb.as_strided_({num_seqs, num_heads, qk_head_dim}, {num_heads * qk_head_dim, qk_head_dim, 1}); // stage 4: bmm - std::optional scale; auto q_nope = qb.narrow(2, 0, qk_nope_head_dim).transpose_(0, 1); auto q_nope_out = q_input.narrow(2, 0, kv_lora_rank).transpose_(0, 1); - bmm_cpu(q_nope_out, q_nope, w_kc, is_vnni, scale); + bmm_cpu(q_nope_out, q_nope, w_kc, is_vnni, w_scale); // stage 5: rope AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rotary_emb_kernel_impl", [&] { @@ -643,6 +643,7 @@ std::tuple qkv_proj_with_rope_fused_weight( bool use_fp8_w8a16, std::optional qkv_a_proj_scale, std::optional q_b_proj_scale, + std::optional w_scale, bool is_vnni, std::optional> block_size, int64_t q_lora_rank, @@ -696,6 +697,7 @@ std::tuple qkv_proj_with_rope_fused_weight( q_a_proj_s, q_b_proj_scale, kv_a_proj_s, + w_scale, is_vnni, block_size); } diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 4f42e23bc..67f68a569 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -245,6 +245,7 @@ std::tuple qkv_proj_with_rope( std::optional q_a_proj_scale, std::optional q_b_proj_scale, std::optional kv_a_proj_scale, + std::optional w_scale, bool is_vnni, std::optional> block_size); @@ -262,6 +263,7 @@ std::tuple qkv_proj_with_rope_fused_weight( bool use_fp8_w8a16, std::optional qkv_a_proj_scale, std::optional q_b_proj_scale, + std::optional w_scale, bool is_vnni, std::optional> block_size, int64_t q_lora_rank, @@ -515,14 +517,14 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor " "kv_a_proj_weight, Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, " "Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? q_a_proj_scale, Tensor? " - "q_b_proj_scale, Tensor? " - "kv_a_proj_scale, bool is_vnni, int[]? block_size) -> (Tensor, Tensor, Tensor)"); + "q_b_proj_scale, Tensor? kv_a_proj_scale, Tensor? w_scale, " + "bool is_vnni, int[]? block_size) -> (Tensor, Tensor, Tensor)"); m.impl("qkv_proj_with_rope", torch::kCPU, &qkv_proj_with_rope); m.def( "qkv_proj_with_rope_fused_weight(Tensor hidden_states, Tensor qkv_a_proj_weight, Tensor q_b_proj_weight, " "Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, " "Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? qkv_a_proj_scale, Tensor? " - "q_b_proj_scale," + "q_b_proj_scale, Tensor? w_scale," "bool is_vnni, int[]? block_size, int q_lora_rank, int kv_lora_rank," "int qk_rope_head_dim) -> (Tensor, Tensor, Tensor)"); m.impl("qkv_proj_with_rope_fused_weight", torch::kCPU, &qkv_proj_with_rope_fused_weight); diff --git a/test/srt/cpu/test_bmm.py b/test/srt/cpu/test_bmm.py new file mode 100644 index 000000000..013e69f0b --- /dev/null +++ b/test/srt/cpu/test_bmm.py @@ -0,0 +1,95 @@ +import itertools +import unittest + +# TODO: use interface in cpu.py +import torch +import torch.nn as nn +from utils import precision + +from sglang.srt.layers.quantization.fp8_utils import input_to_float8 +from sglang.test.test_utils import CustomTestCase + +torch.manual_seed(1234) + + +class Mod(nn.Module): + def __init__(self, input_channel, output_channel, has_bias): + super(Mod, self).__init__() + self.linear = torch.nn.Linear(input_channel, output_channel, has_bias) + + def forward(self, x): + return self.linear(x) + + +class TestBmm(CustomTestCase): + M = [1, 2, 11, 111] + N = [128 + 32, 512] + K = [512 + 32, 128 + 32] + B = [1, 16, 17] + chunk = [True, False] + + def _get_bmm_inputs(self, B, M, N, K, chunk, dtype): + if chunk: + mat1 = ( + torch.randn(M, B, K + 64, dtype=dtype).narrow(2, 0, K).transpose_(0, 1) + ) + mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2) + mat3 = ( + torch.randn(M, B, N + 64, dtype=dtype).narrow(2, 0, N).transpose_(0, 1) + ) + else: + mat1 = torch.randn(M, B, K, dtype=dtype).transpose_(0, 1) + mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2) + mat3 = torch.randn(M, B, N, dtype=dtype).transpose_(0, 1) + return mat1, mat2, mat3 + + def _bf16_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16): + mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype) + ref = torch.bmm(mat1, mat2) + mat2_t = mat2.transpose_(1, 2) + mat3.zero_() + torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2, False, None) + atol = rtol = precision[ref.dtype] + torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol) + + packed_B = torch.ops.sgl_kernel.convert_weight_packed(mat2_t) + mat3.zero_() + torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B, True, None) + torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol) + + def _fp8_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16): + mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype) + mat2_q, mat2_s = input_to_float8(mat2) + ref = torch.bmm(mat1, mat2_q.to(torch.bfloat16)) * mat2_s + mat2_q_t = mat2_q.transpose_(1, 2).contiguous() + mat3.zero_() + atol = rtol = precision[ref.dtype] + torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2_q_t, False, mat2_s) + torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol) + + packed_B_q = torch.ops.sgl_kernel.convert_weight_packed(mat2_q_t) + mat3.zero_() + torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B_q, True, mat2_s) + torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol) + + def test_bmm(self): + for params in itertools.product( + self.B, + self.M, + self.N, + self.K, + self.chunk, + ): + with self.subTest( + B=params[0], + M=params[1], + N=params[2], + K=params[3], + chunk=params[4], + ): + self._bf16_bmm(*params) + self._fp8_bmm(*params) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/cpu/test_qkv_proj_with_rope.py b/test/srt/cpu/test_qkv_proj_with_rope.py index 366043dc9..7015d4504 100644 --- a/test/srt/cpu/test_qkv_proj_with_rope.py +++ b/test/srt/cpu/test_qkv_proj_with_rope.py @@ -8,6 +8,7 @@ from utils import ( precision, ) +from sglang.srt.layers.quantization.fp8_utils import input_to_float8 from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.test.test_utils import CustomTestCase @@ -186,6 +187,7 @@ class TestQKVProjWithROPE(CustomTestCase): None, None, None, + None, True, None, ) @@ -203,6 +205,7 @@ class TestQKVProjWithROPE(CustomTestCase): False, None, None, + None, True, None, q_lora_rank, @@ -274,6 +277,7 @@ class TestQKVProjWithROPE(CustomTestCase): w1_s, w2_s, w3_s, + None, True, None, ) @@ -294,6 +298,7 @@ class TestQKVProjWithROPE(CustomTestCase): False, fused_weight_s, w2_s, + None, True, None, q_lora_rank, @@ -320,6 +325,7 @@ class TestQKVProjWithROPE(CustomTestCase): torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1 ) w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1 + w_kc_q, w_kc_s = input_to_float8(w_kc) kv_a_proj_weight = ( torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1 ) @@ -350,13 +356,14 @@ class TestQKVProjWithROPE(CustomTestCase): ) = convert_weight( kv_a_proj_weight, [scale_block_size_N, scale_block_size_K], torch.bfloat16 ) + w_kc_dq = w_kc_q.to(torch.bfloat16) * w_kc_s q_ref, k_ref, v_ref = native_torch( q_input, hidden_states, q_a_proj_weight_dq, norm_weight1, q_b_proj_weight_dq, - w_kc.transpose(1, 2), + w_kc_dq.transpose(1, 2), kv_a_proj_with_mqa_weight_dq, norm_weight2, pos, @@ -367,13 +374,13 @@ class TestQKVProjWithROPE(CustomTestCase): fp8_kv_a_proj_with_mqa_weight_packed = convert_weight_packed( fp8_kv_a_proj_with_mqa_weight ) - w_kc = convert_weight_packed(w_kc) + w_kc_q = convert_weight_packed(w_kc_q) q_out, k_out, v_out = qkv_proj_with_rope( hidden_states, fp8_q_a_proj_weight_packed, fp8_q_b_proj_weight_packed, fp8_kv_a_proj_with_mqa_weight_packed, - w_kc, + w_kc_q, norm_weight1, norm_weight2, pos, @@ -384,6 +391,7 @@ class TestQKVProjWithROPE(CustomTestCase): q_a_proj_weight_scale_inv.float(), q_b_proj_weight_scale_inv.float(), kv_a_proj_with_mqa_weight_scale_inv.float(), + w_kc_s, True, [scale_block_size_N, scale_block_size_K], ) @@ -399,7 +407,7 @@ class TestQKVProjWithROPE(CustomTestCase): hidden_states, fused_weight_packed, fp8_q_b_proj_weight_packed, - w_kc, + w_kc_q, norm_weight1, norm_weight2, pos, @@ -409,6 +417,7 @@ class TestQKVProjWithROPE(CustomTestCase): True, fused_weight_s.float(), q_b_proj_weight_scale_inv.float(), + w_kc_s, True, [scale_block_size_N, scale_block_size_K], q_lora_rank, diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index be6877844..5ed743aa0 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -49,6 +49,7 @@ suite_xeon = { "per-commit-cpu": [ TestFile("cpu/test_activation.py"), TestFile("cpu/test_binding.py"), + TestFile("cpu/test_bmm.py"), TestFile("cpu/test_causal_conv1d.py"), TestFile("cpu/test_cpu_graph.py"), TestFile("cpu/test_decode.py"),