diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 17ac44803aa..b83eaf669aa 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -13,8 +13,9 @@ steps: - tests/test_config - tests/test_logger - tests/test_vllm_port + - tests/test_jit_monitor.py commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py - label: Engine (1 GPU) key: engine-1-gpu diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index 3e3c7921798..2f56099c66f 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -1473,6 +1473,12 @@ async def main() -> None: "(for example: --warmup-percentages=0%%,50%%)", ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code when loading the tokenizer.", + ) + args = parser.parse_args() logger.info(args) @@ -1515,7 +1521,9 @@ async def main() -> None: np.random.seed(args.seed) logger.info("Loading tokenizer") - tokenizer = AutoTokenizer.from_pretrained(args.model) + tokenizer = AutoTokenizer.from_pretrained( + args.model, trust_remote_code=args.trust_remote_code + ) await get_server_info(args.url) diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 18afe4b7925..4750dd78838 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata( isa = cpu_attention::ISA::NEON; } else if (isa_hint == "vxe") { isa = cpu_attention::ISA::VXE; + } else if (isa_hint == "vsx") { + isa = cpu_attention::ISA::VSX; } else { TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint); } @@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache( return cpu_attention::ISA::NEON; } else if (isa == "vxe") { return cpu_attention::ISA::VXE; + } else if (isa == "vsx") { + return cpu_attention::ISA::VSX; } else { TORCH_CHECK(false, "Invalid ISA type: " + isa); } diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index f5b473bd262..b9987fb26c1 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -12,7 +12,7 @@ #include "cpu/utils.hpp" namespace cpu_attention { -enum class ISA { AMX, VEC, VEC16, NEON, VXE }; +enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX }; // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { @@ -164,6 +164,9 @@ struct AttentionMetadata { case ISA::VXE: ss << "VXE, "; break; + case ISA::VSX: + ss << "VSX, "; + break; } ss << "workitem_group_num: " << workitem_group_num << ", reduction_item_num: " << reduction_item_num diff --git a/csrc/cpu/cpu_attn_vec.hpp b/csrc/cpu/cpu_attn_vec.hpp index 61cae12d67d..c3983e0578a 100644 --- a/csrc/cpu/cpu_attn_vec.hpp +++ b/csrc/cpu/cpu_attn_vec.hpp @@ -27,8 +27,8 @@ FORCE_INLINE std::pair load_b_pair_vec( return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)}; } else { using load_vec_t = typename VecTypeTrait::vec_t; - return {vec_op::FP32Vec16(load_vec_t(ptr)), - vec_op::FP32Vec16(load_vec_t(ptr + 16))}; + return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)), + vec_op::FP32Vec16(load_vec_t(ptr + 16))); } } diff --git a/csrc/cpu/cpu_attn_vsx.hpp b/csrc/cpu/cpu_attn_vsx.hpp new file mode 100644 index 00000000000..c7e1502bcb0 --- /dev/null +++ b/csrc/cpu/cpu_attn_vsx.hpp @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#ifndef CPU_ATTN_VSX_HPP +#define CPU_ATTN_VSX_HPP + +#include "cpu_attn_impl.hpp" +#include +#include + +namespace cpu_attention { + +namespace { + +// ppc64le Vector = 16 bytes (128 bits) +#define BLOCK_SIZE_ALIGNMENT 32 +#define HEAD_SIZE_ALIGNMENT 32 +#define MAX_Q_HEAD_NUM_PER_ITER 16 + +template +FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0, + __vector float& b1); + +// [1] Float Specialization +template <> +FORCE_INLINE void load_row8_B_as_f32(const float* p, __vector float& b0, + __vector float& b1) { + b0 = vec_xl(0, const_cast(p)); + b1 = vec_xl(0, const_cast(p + 4)); +} + +// [2] BFloat16 Specialization (Little Endian ppc64le) +// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32. +// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)] +// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed. +// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]} +// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair: +// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs +// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits +// Cast to float32 gives exactly (bf16_bits << 16) per element. +template <> +FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p, + __vector float& b0, + __vector float& b1) { + __vector unsigned short raw = vec_xl( + 0, reinterpret_cast(const_cast(p))); + __vector unsigned short zeros = vec_splat_u16(0); + + // LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32 + b0 = (__vector float)vec_mergeh(zeros, raw); + b1 = (__vector float)vec_mergel(zeros, raw); +} + +// Note: c10::Half (FP16) is not supported on PowerPC architecture + +template +FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4( + const float* __restrict A, // [M x K] + const kv_cache_t* __restrict B, // [K x 8] + float* __restrict C, // [M x 8] + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) { + static_assert(1 <= M && M <= 8, "M must be in [1,8]"); + +#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7) +#define IF_M(i) if constexpr (M > (i)) + + // 1. Define A pointers +#define DECL_A(i) const float* a##i = A + (i) * lda; + ROWS_APPLY(DECL_A) +#undef DECL_A + + // 2. Define Accumulators (2 vectors covers 8 columns) +#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1; + ROWS_APPLY(DECL_ACC) +#undef DECL_ACC + + // 3. Initialize Accumulators (Load C or Zero) +#define INIT_ACC(i) \ + IF_M(i) { \ + if (accumulate) { \ + acc##i##_0 = vec_xl(0, const_cast(C + (i) * ldc + 0)); \ + acc##i##_1 = vec_xl(0, const_cast(C + (i) * ldc + 4)); \ + } else { \ + acc##i##_0 = vec_splats(0.0f); \ + acc##i##_1 = vec_splats(0.0f); \ + } \ + } + ROWS_APPLY(INIT_ACC) +#undef INIT_ACC + + int32_t k = 0; + + for (; k + 3 < K; k += 4) { + // Load 4 values of A for each Row M: A[k...k+3] +#define LOAD_A4(i) \ + __vector float a##i##v; \ + IF_M(i) a##i##v = vec_xl(0, const_cast(a##i + k)); + ROWS_APPLY(LOAD_A4) +#undef LOAD_A4 + + // FMA for specific lane L of A + // ppc64le: vec_madd(b, vec_splat(a, lane), acc) +#define FMAS_LANE(i, aiv, L) \ + IF_M(i) { \ + __vector float a_broad = vec_splat(aiv, L); \ + acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \ + acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \ + } + + // Unroll K=0..3 + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 0) * ldb, b0, b1); +#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0) + ROWS_APPLY(STEP_K0) +#undef STEP_K0 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 1) * ldb, b0, b1); +#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1) + ROWS_APPLY(STEP_K1) +#undef STEP_K1 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 2) * ldb, b0, b1); +#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2) + ROWS_APPLY(STEP_K2) +#undef STEP_K2 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 3) * ldb, b0, b1); +#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3) + ROWS_APPLY(STEP_K3) +#undef STEP_K3 + } +#undef FMAS_LANE + } + + for (; k < K; ++k) { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)k * ldb, b0, b1); +#define TAIL_ROW(i) \ + IF_M(i) { \ + __vector float ai = vec_splats(*(a##i + k)); \ + acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \ + acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \ + } + ROWS_APPLY(TAIL_ROW) +#undef TAIL_ROW + } + +#define STORE_ROW(i) \ + IF_M(i) { \ + vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \ + vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \ + } + ROWS_APPLY(STORE_ROW) +#undef STORE_ROW + +#undef ROWS_APPLY +#undef IF_M +} + +template +FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A, + const kv_cache_t* __restrict B, + float* __restrict C, int32_t M, + int32_t K, int64_t lda, + int64_t ldb, int64_t ldc, + bool accumulate) { + static_assert(N % 8 == 0, "N must be a multiple of 8"); + for (int32_t m = 0; m < M;) { + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1; + const float* Ab = A + m * lda; + float* Cb = C + m * ldc; + + for (int32_t n = 0; n < N; n += 8) { + const kv_cache_t* Bn = B + n; + float* Cn = Cb + n; + switch (mb) { + case 8: + gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 4: + gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 2: + gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + default: + gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + } + } + m += mb; + } +} + +template +class TileGemmPPC64 { + public: + template + FORCE_INLINE static void gemm(const int32_t m_size, + float* __restrict__ a_tile, + kv_cache_t* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + const int64_t ldb, const int64_t ldc, + const int32_t block_size, + const int32_t dynamic_k_size, + const bool accum_c) { + if constexpr (phase == AttentionGemmPhase::QK) { + gemm_macro_ppc64le_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c); + } else { + gemm_macro_ppc64le_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc, + accum_c); + } + } +}; + +} // namespace + +template +class AttentionImpl { + public: + using query_t = scalar_t; + using q_buffer_t = float; + using kv_cache_t = scalar_t; + using logits_buffer_t = float; + using partial_output_buffer_t = float; + using prob_buffer_t = float; + + constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT; + constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT; + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER; + constexpr static int64_t HeadDim = head_dim; + constexpr static ISA ISAType = ISA::VSX; + constexpr static bool scale_on_logits = + false; // Scale is applied to Q during copy + + public: + AttentionImpl() {} + + template