diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 458bef9518..ef2b4ba073 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -98,7 +98,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pip install dpctl && - pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py --ignore=v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py && + pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index e9ddaacea5..4ade47eff3 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -956,6 +956,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col, std::printf("%s", ss.str().c_str()); } +// This allows us to store probabilities in a packed format +// for GEMM implementations that need packed A. +template +struct ProbabilityTokenStore { + using prob_buffer_t = typename attention_impl_t::prob_buffer_t; + static constexpr int32_t TokenStride = 1; + + FORCE_INLINE static void store_probabilities( + prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values, + const int32_t, const int64_t) { + using prob_buffer_vec_t = typename VecTypeTrait::vec_t; + prob_buffer_vec_t output_vec(values); + output_vec.save(probability); + } +}; + +template +struct ProbabilityTokenStore< + attention_impl_t, + std::void_t> + : attention_impl_t::ProbabilityTokenStore {}; + template class AttentionMainLoop { public: @@ -966,6 +988,7 @@ class AttentionMainLoop { using partial_output_buffer_t = typename attention_impl_t::partial_output_buffer_t; using prob_buffer_t = typename attention_impl_t::prob_buffer_t; + using probability_token_store_t = ProbabilityTokenStore; static constexpr int64_t max_q_head_num_per_iter = attention_impl_t::MaxQHeadNumPerIteration; @@ -974,6 +997,8 @@ class AttentionMainLoop { static constexpr int64_t headdim_alignment = attention_impl_t::HeadDimAlignment; static constexpr int64_t head_dim = attention_impl_t::HeadDim; + static constexpr int32_t probability_token_stride = + probability_token_store_t::TokenStride; static constexpr ISA ISAType = attention_impl_t::ISAType; static constexpr bool scale_on_logits = attention_impl_t::scale_on_logits; // apply scale on logits, otherwise @@ -1152,8 +1177,6 @@ class AttentionMainLoop { ? kv_tile_token_num : kv_tile_token_num * (sizeof(logits_buffer_t) / sizeof(prob_buffer_t)); - constexpr int64_t prob_buffer_elem_size = - prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t); partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer; bool accum_c = !is_first_iter; for (int32_t block_idx = start_block_idx; block_idx < end_block_idx; @@ -1187,9 +1210,7 @@ class AttentionMainLoop { remaining_group_num -= curr_group_num_in_block; curr_group_offset = 0; curr_group_num_in_block = token_group_num_per_block; - curr_prob_buffer = reinterpret_cast( - reinterpret_cast(curr_prob_buffer) + - curr_token_num * prob_buffer_elem_size); + curr_prob_buffer += curr_token_num * probability_token_stride; curr_partial_q_buffer = partial_q_buffer; accum_c = true; } @@ -1262,7 +1283,6 @@ class AttentionMainLoop { std::is_same_v; #endif - using prob_buffer_vec_t = typename VecTypeTrait::vec_t; constexpr bool prequantize_probabilities = []() { if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) { return tile_gemm_t::prequantize_probabilities; @@ -1331,12 +1351,6 @@ class AttentionMainLoop { vec = fast_exp(vec); } - if constexpr (prequantize_probabilities) { - vec.save(curr_logits_buffer_iter); - } else { - prob_buffer_vec_t output_vec(vec); - output_vec.save(curr_prob_buffer_iter); - } #else vec.save(curr_logits_buffer_iter); for (int32_t k = 0; k < 16; ++k) { @@ -1345,10 +1359,19 @@ class AttentionMainLoop { vec = vec_op::FP32Vec16(curr_logits_buffer_iter); #endif + if constexpr (prequantize_probabilities) { + vec.save(curr_logits_buffer_iter); + } else { + probability_token_store_t::store_probabilities( + curr_prob_buffer_iter, vec, i, + logits_buffer_stride * sizeof(logits_buffer_t) / + sizeof(prob_buffer_t)); + } + sum_vec = sum_vec + vec; curr_logits_buffer_iter += 16; - curr_prob_buffer_iter += 16; + curr_prob_buffer_iter += 16 * probability_token_stride; } if constexpr (prequantize_probabilities) { tile_gemm_t::quantize_probability_row( diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp index db4c5df2e8..eed7499444 100644 --- a/csrc/cpu/cpu_attn_neon.hpp +++ b/csrc/cpu/cpu_attn_neon.hpp @@ -11,7 +11,7 @@ namespace cpu_attention { namespace { -#define BLOCK_SIZE_ALIGNMENT 32 +#define BLOCK_SIZE_ALIGNMENT 16 #define HEAD_SIZE_ALIGNMENT 32 #define MAX_Q_HEAD_NUM_PER_ITER 16 @@ -387,7 +387,7 @@ class AttentionImpl { }; #ifdef ARM_BF16_SUPPORT -// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment. +// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment. template class AttentionImpl : public AttentionImplNEONBFMMLA #include +#include +#include #include -#include namespace cpu_attention { +class BfmmlaGemm { + public: + static constexpr int32_t KTile = 4; + static constexpr int32_t NTile = 8; + static constexpr int32_t MaxRows = 8; + + FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a, + const c10::BFloat16* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t n, const int32_t k, + const int64_t a_pair_stride, + const int64_t b_n_group_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const auto* a_ptr = reinterpret_cast(a); + const auto* b_ptr = reinterpret_cast(b); + + for (int32_t n_idx = 0; n_idx < n; n_idx += 16) { + const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride; + float* c_panel = c + n_idx; + + // Preserve this range so the inactive row pair is optimized away. + if (m <= 2) { + gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride, + b_n_group_stride, b_k_group_stride, ldc, accumulate); + } else if (m <= 4) { + gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride, + b_n_group_stride, b_k_group_stride, ldc, accumulate); + } else { + gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride, + ldc, accumulate); + gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k, + a_pair_stride, b_k_group_stride, ldc, accumulate); + } + } + } + + private: + FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a, + const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); + } + + FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a, + const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); + } + + FORCE_INLINE static void init_accumulators( + float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45, + float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc, + const int32_t rows, const bool accumulate) { + if (!accumulate || rows == 0) { + acc01 = vdupq_n_f32(0.0f); + acc23 = vdupq_n_f32(0.0f); + acc45 = vdupq_n_f32(0.0f); + acc67 = vdupq_n_f32(0.0f); + return; + } + + const float32x4_t row0_0123 = vld1q_f32(c); + const float32x4_t row0_4567 = vld1q_f32(c + 4); + const float32x4_t row1_0123 = + (rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f); + const float32x4_t row1_4567 = + (rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f); + + acc01 = zip_low_pairs(row0_0123, row1_0123); + acc23 = zip_high_pairs(row0_0123, row1_0123); + acc45 = zip_low_pairs(row0_4567, row1_4567); + acc67 = zip_high_pairs(row0_4567, row1_4567); + } + + FORCE_INLINE static void store_accumulators( + const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45, + const float32x4_t acc67, float* __restrict__ c, const int64_t ldc, + const int32_t rows) { + if (rows == 0) { + return; + } + + vst1q_f32(c, zip_low_pairs(acc01, acc23)); + vst1q_f32(c + 4, zip_low_pairs(acc45, acc67)); + if (rows == 2) { + vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23)); + vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67)); + } + } + + FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a, + const int32_t rows) { + if (rows == 0) { + return vdupq_n_bf16(bfloat16_t{}); + } + // Packed A reserves both rows for an M tail. + return vld1q_bf16(a); + } + + FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a, + const bfloat16_t* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t k, + const int64_t a_pair_stride, + const int64_t b_n_group_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const int32_t rows01 = std::min(2, std::max(0, m)); + const int32_t rows23 = std::min(2, std::max(0, m - 2)); + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc0189, acc011011, acc011213, acc011415; + float32x4_t acc2389, acc231011, acc231213, acc231415; + init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01, + accumulate); + init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23, accumulate); + init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc, + rows01, accumulate); + init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8, + ldc, rows23, accumulate); + + const bfloat16_t* a01 = a; + const bfloat16_t* a23 = a + a_pair_stride; + const bfloat16_t* b0 = b; + const bfloat16_t* b1 = b + b_n_group_stride; + +#pragma GCC unroll 4 + for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) { + const bfloat16x8_t av01 = load_a_pair(a01, rows01); + const bfloat16x8_t av23 = load_a_pair(a23, rows23); + const bfloat16x8_t b01 = vld1q_bf16(b0); + const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile); + const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile); + const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile); + const bfloat16x8_t b89 = vld1q_bf16(b1); + const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile); + const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile); + const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile); + + acc0101 = vbfmmlaq_f32(acc0101, av01, b01); + acc2301 = vbfmmlaq_f32(acc2301, av23, b01); + acc0123 = vbfmmlaq_f32(acc0123, av01, b23); + acc2323 = vbfmmlaq_f32(acc2323, av23, b23); + acc0145 = vbfmmlaq_f32(acc0145, av01, b45); + acc2345 = vbfmmlaq_f32(acc2345, av23, b45); + acc0167 = vbfmmlaq_f32(acc0167, av01, b67); + acc2367 = vbfmmlaq_f32(acc2367, av23, b67); + acc0189 = vbfmmlaq_f32(acc0189, av01, b89); + acc2389 = vbfmmlaq_f32(acc2389, av23, b89); + acc011011 = vbfmmlaq_f32(acc011011, av01, b1011); + acc231011 = vbfmmlaq_f32(acc231011, av23, b1011); + acc011213 = vbfmmlaq_f32(acc011213, av01, b1213); + acc231213 = vbfmmlaq_f32(acc231213, av23, b1213); + acc011415 = vbfmmlaq_f32(acc011415, av01, b1415); + acc231415 = vbfmmlaq_f32(acc231415, av23, b1415); + + a01 += 2 * KTile; + a23 += 2 * KTile; + b0 += b_k_group_stride; + b1 += b_k_group_stride; + } + + store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01); + store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23); + store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc, + rows01); + store_accumulators(acc2389, acc231011, acc231213, acc231415, + c + 2 * ldc + 8, ldc, rows23); + } + + FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a, + const bfloat16_t* __restrict__ b, + float* __restrict__ c, const int32_t m, + const int32_t k, + const int64_t a_pair_stride, + const int64_t b_k_group_stride, + const int64_t ldc, const bool accumulate) { + const int32_t rows01 = std::min(2, std::max(0, m)); + const int32_t rows23 = std::min(2, std::max(0, m - 2)); + const int32_t rows45 = std::min(2, std::max(0, m - 4)); + const int32_t rows67 = std::min(2, std::max(0, m - 6)); + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc4501, acc4523, acc4545, acc4567; + float32x4_t acc6701, acc6723, acc6745, acc6767; + init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01, + accumulate); + init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23, accumulate); + init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc, + rows45, accumulate); + init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc, + rows67, accumulate); + + const bfloat16_t* a01 = a; + const bfloat16_t* a23 = a + a_pair_stride; + const bfloat16_t* a45 = a + 2 * a_pair_stride; + const bfloat16_t* a67 = a + 3 * a_pair_stride; + const bfloat16_t* b_ptr = b; + +#pragma GCC unroll 4 + for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) { + const bfloat16x8_t av01 = load_a_pair(a01, rows01); + const bfloat16x8_t av23 = load_a_pair(a23, rows23); + const bfloat16x8_t av45 = load_a_pair(a45, rows45); + const bfloat16x8_t av67 = load_a_pair(a67, rows67); + const bfloat16x8_t b01 = vld1q_bf16(b_ptr); + const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile); + const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile); + const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile); + + acc0101 = vbfmmlaq_f32(acc0101, av01, b01); + acc2301 = vbfmmlaq_f32(acc2301, av23, b01); + acc4501 = vbfmmlaq_f32(acc4501, av45, b01); + acc6701 = vbfmmlaq_f32(acc6701, av67, b01); + acc0123 = vbfmmlaq_f32(acc0123, av01, b23); + acc2323 = vbfmmlaq_f32(acc2323, av23, b23); + acc4523 = vbfmmlaq_f32(acc4523, av45, b23); + acc6723 = vbfmmlaq_f32(acc6723, av67, b23); + acc0145 = vbfmmlaq_f32(acc0145, av01, b45); + acc2345 = vbfmmlaq_f32(acc2345, av23, b45); + acc4545 = vbfmmlaq_f32(acc4545, av45, b45); + acc6745 = vbfmmlaq_f32(acc6745, av67, b45); + acc0167 = vbfmmlaq_f32(acc0167, av01, b67); + acc2367 = vbfmmlaq_f32(acc2367, av23, b67); + acc4567 = vbfmmlaq_f32(acc4567, av45, b67); + acc6767 = vbfmmlaq_f32(acc6767, av67, b67); + + a01 += 2 * KTile; + a23 += 2 * KTile; + a45 += 2 * KTile; + a67 += 2 * KTile; + b_ptr += b_k_group_stride; + } + + store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01); + store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc, + rows23); + store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc, + rows45); + store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc, + rows67); + } +}; + namespace { -// BFMMLA tile dimensions -constexpr int32_t TILE_ROWS = 2; // M dimension -constexpr int32_t TILE_K = 4; // K reduction -constexpr int32_t TILE_COLS = 2; // N dimension (column-pair) - -// Derived constants -constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs -constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache -constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block -constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K; -constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS; -constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing - -// Matrix Packing and Accumulator -// Reshape two rows of Q into BFMMLA-friendly interleaved -// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3] -// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7] -// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk -FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0, - const c10::BFloat16* __restrict r1, - c10::BFloat16* __restrict dst, - int32_t K) { - const uint16_t* s0 = reinterpret_cast(r0); - const uint16_t* s1 = reinterpret_cast(r1); - uint16_t* d = reinterpret_cast(dst); - - // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output) - int32_t k = 0; - for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) { - vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k))); - } - - // Handle K tail: pack remaining elements with zero-padding - const int32_t tail = K - k; - if (tail > 0) { - // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1], - // pad] - for (int32_t t = 0; t < tail; ++t) { - d[t] = s0[k + t]; - d[t + TILE_K] = s1[k + t]; - } - // Zero-pad the rest - for (int32_t t = tail; t < TILE_K; ++t) { - d[t] = 0; - d[t + TILE_K] = 0; - } - } -} - -// 2x2 accumulator load/store with compile-time row count -template -FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) { - static_assert(m_rows == 1 || m_rows == 2); - float32x2_t row0 = vld1_f32(base + col_off); - float32x2_t row1 = - (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f); - return vcombine_f32(row0, row1); -} - -template -FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc, - int col_off) { - static_assert(m_rows == 1 || m_rows == 2); - vst1_f32(base + col_off, vget_low_f32(acc)); - if constexpr (m_rows == 2) { - vst1_f32(base + ldc + col_off, vget_high_f32(acc)); - } -} - -// Initialize 4 column-pair accumulators for 2 rows (8 columns total) -#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \ - do { \ - if (accum) { \ - if (m_rows == 2) { \ - a0 = load_acc_2x2<2>(Crow, ldc, 0); \ - a1 = load_acc_2x2<2>(Crow, ldc, 2); \ - a2 = load_acc_2x2<2>(Crow, ldc, 4); \ - a3 = load_acc_2x2<2>(Crow, ldc, 6); \ - } else { \ - a0 = load_acc_2x2<1>(Crow, ldc, 0); \ - a1 = load_acc_2x2<1>(Crow, ldc, 2); \ - a2 = load_acc_2x2<1>(Crow, ldc, 4); \ - a3 = load_acc_2x2<1>(Crow, ldc, 6); \ - } \ - } else { \ - a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \ - } \ - } while (0) - -// Store 4 column-pair accumulators back to C matrix -#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \ - do { \ - if (m_rows == 2) { \ - store_acc_2x2<2>(a0, Crow, ldc, 0); \ - store_acc_2x2<2>(a1, Crow, ldc, 2); \ - store_acc_2x2<2>(a2, Crow, ldc, 4); \ - store_acc_2x2<2>(a3, Crow, ldc, 6); \ - } else { \ - store_acc_2x2<1>(a0, Crow, ldc, 0); \ - store_acc_2x2<1>(a1, Crow, ldc, 2); \ - store_acc_2x2<1>(a2, Crow, ldc, 4); \ - store_acc_2x2<1>(a3, Crow, ldc, 6); \ - } \ - } while (0) - -// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs -#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \ - do { \ - r0 = vbfmmlaq_f32(r0, a, b0); \ - r1 = vbfmmlaq_f32(r1, a, b1); \ - r2 = vbfmmlaq_f32(r2, a, b2); \ - r3 = vbfmmlaq_f32(r3, a, b3); \ - } while (0) - -// Micro-kernel: updates a small fixed tile using BFMMLA. -// RP = number of row-pairs (1,2,4) -// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B. -// A_packed interleaves RP row-pairs; B layout is driven by the attention phase: -// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T) -// - AttentionGemmPhase::PV -> token-row layout (P @ V) -// K_static < 0 enables runtime K (PV only) -template -FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon( - const bfloat16_t* const* __restrict A_packed_rp, - const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk, - float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride, - int32_t K_runtime = 0) { - static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4"); - static_assert(K_static < 0 || K_static % TILE_K == 0, - "K must be divisible by TILE_K"); - static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, - "Runtime K only supported for PV"); - - constexpr bool runtime_k = (K_static < 0); - const int32_t K_iters = - runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K); - const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0; - - if (!runtime_k) { - // Help the compiler fold away unused K_runtime when K is compile-time - (void)K_runtime; - } - - auto* C_al = C; - const auto* B_al = B_blk; - - // Setup A pointers - const bfloat16_t* a_ptr[4] = { - A_packed_rp[0], - (RP >= 2) ? A_packed_rp[1] : nullptr, - (RP >= 4) ? A_packed_rp[2] : nullptr, - (RP >= 4) ? A_packed_rp[3] : nullptr, - }; - - // Setup B pointers based on layout - const bfloat16_t* b_ptr[4]; - if constexpr (phase == AttentionGemmPhase::PV) { - b_ptr[0] = B_blk + 0 * b_stride; - b_ptr[1] = B_blk + 1 * b_stride; - b_ptr[2] = B_blk + 2 * b_stride; - b_ptr[3] = B_blk + 3 * b_stride; - } - - float32x4_t acc[4][4]; - -// Initialize accumulators -#define INIT_RP(rp) \ - if constexpr (RP > rp) { \ - INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ - C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \ - } - INIT_RP(0); - INIT_RP(1); - INIT_RP(2); - INIT_RP(3); -#undef INIT_RP - - // Main compute loop - for (int32_t ki = 0; ki < K_iters; ++ki) { - bfloat16x8_t b0, b1, b2, b3; - if constexpr (phase == AttentionGemmPhase::PV) { - b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE); - b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE); - b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE); - b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE); - } else { - const bfloat16_t* b_base = B_al + ki * b_stride; - b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE); - b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE); - b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE); - b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE); - } - -#define COMPUTE_RP(rp) \ - if constexpr (RP > rp) { \ - bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \ - BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \ - b1, b2, b3); \ - } - COMPUTE_RP(0); - COMPUTE_RP(1); - COMPUTE_RP(2); - COMPUTE_RP(3); -#undef COMPUTE_RP - } - - // K tail for runtime PV: fallback path - if constexpr (runtime_k) { - if (K_tail > 0) { - const int32_t tail_offset = K_iters * V_INNER_STRIDE; - const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK; - for (int32_t kt = 0; kt < K_tail; ++kt) { - float32x4_t b_vecs[4]; - for (int32_t p = 0; p < 4; ++p) { - const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS; - const float b0 = vcvtah_f32_bf16(bp[0]); - const float b1 = vcvtah_f32_bf16(bp[1]); - const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1); - b_vecs[p] = vcombine_f32(b_pair, b_pair); - } - -#define TAIL_RP(rp) \ - if constexpr (RP > rp) { \ - const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \ - float a_row0 = vcvtah_f32_bf16(ap[kt]); \ - float a_row1 = \ - (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \ - const float32x4_t a_vec = \ - vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \ - for (int32_t p = 0; p < 4; ++p) { \ - acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \ - } \ - } - TAIL_RP(0); - TAIL_RP(1); - TAIL_RP(2); - TAIL_RP(3); -#undef TAIL_RP - } - } - } - - // Store results -#define STORE_RP(rp) \ - if constexpr (RP > rp) { \ - STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ - C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \ - } - STORE_RP(0); - STORE_RP(1); - STORE_RP(2); - STORE_RP(3); -#undef STORE_RP -} - -// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the -// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables -// runtime K (PV only). -template -FORCE_INLINE void gemm_packA_compute_MB_xN( - const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, - float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc, - int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { - static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]"); - static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, - "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); - static_assert(K_static < 0 || K_static % TILE_K == 0, - "K must be divisible by TILE_K"); - static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, - "Runtime K only supported for PV"); - - constexpr bool runtime_k = (K_static < 0); - const int32_t K_val = runtime_k ? K_runtime : K_static; - - // Keep small packs on-stack to avoid heap churn - constexpr int32_t STACK_PACK_STRIDE = - (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; - - constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS; - const int32_t pack_stride = - runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK - : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; - - alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE]; - std::vector A_packed_heap; - c10::BFloat16* A_packed = - (pack_stride <= STACK_PACK_STRIDE) - ? A_packed_stack - : (A_packed_heap.resize(ROW_PAIRS * pack_stride), - A_packed_heap.data()); - - for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) { - const int32_t m = rp * TILE_ROWS; - const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1; - const c10::BFloat16* A0 = A + m * lda; - const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0; - reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val); - } - - for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) { - const c10::BFloat16* B_blk_c10 = - (phase == AttentionGemmPhase::PV) - ? (B + (n / TILE_COLS) * b_layout_stride) - : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride); - const bfloat16_t* B_blk = reinterpret_cast(B_blk_c10); - - // Process row-pairs in groups of 4, 2, then 1 - int32_t row_pair_idx = 0; - -#define PROCESS_RP_GROUP(group_size) \ - for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \ - row_pair_idx += group_size) { \ - const bfloat16_t* Ap[group_size]; \ - int32_t mr[group_size]; \ - for (int32_t i = 0; i < group_size; ++i) { \ - Ap[i] = reinterpret_cast( \ - A_packed + (row_pair_idx + i) * pack_stride); \ - mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \ - } \ - float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \ - if constexpr (runtime_k) { \ - gemm_rowpairs_x8_bfmmla_neon( \ - Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \ - } else { \ - gemm_rowpairs_x8_bfmmla_neon( \ - Ap, mr, B_blk, C_blk, ldc, accumulate, \ - (phase == AttentionGemmPhase::PV) ? b_layout_stride \ - : b_reduction_stride); \ - } \ - } - - PROCESS_RP_GROUP(4); - PROCESS_RP_GROUP(2); - PROCESS_RP_GROUP(1); -#undef PROCESS_RP_GROUP - } -} - -// Macro-kernel: iterates over M in MB={8,4,2,1} chunks. -// Supports compile-time K specialization when K >= 0; otherwise uses runtime K -// (runtime K path is only supported for PV). -template -FORCE_INLINE void gemm_macro_neon_bfmmla( - const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, - float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc, - int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { - static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, - "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); - - if constexpr (K >= 0) { - static_assert(K % TILE_K == 0, "K must be divisible by TILE_K"); - for (int32_t m = 0; m < M;) { - const int32_t rem = M - m; - const c10::BFloat16* A_blk = A + m * lda; - float* C_blk = C + m * ldc; - -#define DISPATCH_MB(mb) \ - gemm_packA_compute_MB_xN(A_blk, B, C_blk, 0, lda, ldc, \ - b_layout_stride, \ - b_reduction_stride, accumulate) - - if (rem >= 8) { - DISPATCH_MB(8); - m += 8; - } else if (rem >= 4) { - DISPATCH_MB(4); - m += 4; - } else if (rem >= 2) { - DISPATCH_MB(2); - m += 2; - } else { - DISPATCH_MB(1); - m += 1; - } -#undef DISPATCH_MB - } - } else { - static_assert(phase == AttentionGemmPhase::PV, - "Runtime K specialization only supported for PV."); - const int32_t K_val = K_runtime; - - for (int32_t m = 0; m < M;) { - const int32_t rem = M - m; - const c10::BFloat16* A_blk = A + m * lda; - float* C_blk = C + m * ldc; - -#define DISPATCH_MB_RUNTIME(mb) \ - gemm_packA_compute_MB_xN(A_blk, B, C_blk, K_val, lda, ldc, \ - b_layout_stride, \ - b_reduction_stride, accumulate) - - if (rem >= 8) { - DISPATCH_MB_RUNTIME(8); - m += 8; - } else if (rem >= 4) { - DISPATCH_MB_RUNTIME(4); - m += 4; - } else if (rem >= 2) { - DISPATCH_MB_RUNTIME(2); - m += 2; - } else { - DISPATCH_MB_RUNTIME(1); - m += 1; - } -#undef DISPATCH_MB_RUNTIME - } - } -} - -#undef INIT_ACC_ROWPAIR_4 -#undef STORE_ACC_ROWPAIR_4 -#undef BFMMLA_COMPUTE_4 +constexpr int32_t TILE_K = BfmmlaGemm::KTile; +constexpr int32_t TILE_COLS = 2; +constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile; +constexpr int32_t K_TOKENS_PER_GROUP = 8; +constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; +constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K; +constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS; } // namespace -// TileGemm Adapter for Attention - template class TileGemmNEONBFMMLA { public: @@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA { [[maybe_unused]] const int32_t block_size, [[maybe_unused]] const int32_t dynamic_k_size, const bool accum_c) { - static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0); - // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for - // API parity with other tile_gemm implementations. + static_assert(BlockTokens % 16 == 0); if constexpr (head_dim_ct >= 0) { - static_assert(head_dim_ct == HeadDim, - "BFMMLA expects head_dim_ct to match HeadDim; PV passes " - "-1 for API parity."); + static_assert(head_dim_ct == HeadDim); } + const auto* a = reinterpret_cast(a_tile); if constexpr (phase == AttentionGemmPhase::QK) { - const int64_t b_reduction_stride = K_INNER_STRIDE; - const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE; + constexpr int64_t b_n_group_stride = + (HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE; - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride, - accum_c); + for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) { + const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row); + BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m, + BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride, + K_CACHE_K_GROUP_STRIDE, ldc, accum_c); + } } else { - const int64_t b_pair_stride = - (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + const int64_t b_n_group_stride = + (block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE; - // PV gemm with runtime K specialization - switch (dynamic_k_size) { - case 32: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 32, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 128: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 128, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 256: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 256, lda, ldc, b_pair_stride, 0, accum_c); - break; - default: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c); - break; + for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) { + const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row); + BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m, + HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride, + K_CACHE_K_GROUP_STRIDE, ldc, accum_c); } } } @@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA { static constexpr int64_t HeadDim = head_dim; static constexpr ISA ISAType = isa_type; static constexpr bool scale_on_logits = false; + static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK; + static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K; - static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0); - static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0); + static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0); + static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0); static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K"); public: @@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA { attention_iteration(CPU_ATTENTION_PARAMS); } + struct ProbabilityTokenStore { + static constexpr int32_t TokenStride = 2; + + FORCE_INLINE static void store_probabilities( + c10::BFloat16* __restrict__ probability, + const vec_op::FP32Vec16& values, const int32_t row, + const int64_t row_stride) { + const int32_t row_in_pair = row & 1; + auto* dst = reinterpret_cast( + probability + row_in_pair * (BfmmlaGemm::KTile - row_stride)); + vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0])); + vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1])); + vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2])); + vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3])); + } + }; + // Key cache stride per token group (TokenColumn layout; QK) static constexpr int64_t k_cache_token_group_stride( [[maybe_unused]] const int32_t block_size) { static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0); return (BlockSizeAlignment / K_TOKENS_PER_GROUP) * - ((head_dim / TILE_K) * K_INNER_STRIDE); + ((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE); } // Value cache stride per token group (TokenRow layout; PV) static constexpr int64_t v_cache_token_group_stride( [[maybe_unused]] const int32_t block_size) { static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0); - return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride; } // The stride to move to the "next" head_dim group @@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA { return head_dim * block_size; } - // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics + // Scale Q and write row pairs in BFMMLA reduction order. static void copy_q_heads_tile(c10::BFloat16* __restrict__ src, c10::BFloat16* __restrict__ q_buffer, const int32_t q_num, @@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA { const int64_t q_head_stride, float scale) { constexpr int32_t dim = static_cast(head_dim); const float32x4_t scale_vec = vdupq_n_f32(scale); + const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{}); + const int32_t row_num = q_num * q_heads_per_kv; - for (int32_t qi = 0; qi < q_num; ++qi) { - for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) { - c10::BFloat16* __restrict__ curr_q = - src + qi * q_num_stride + hi * q_head_stride; - c10::BFloat16* __restrict__ dst = - q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim; + for (int32_t row = 0; row < row_num; row += 2) { + const int32_t q0 = row / q_heads_per_kv; + const int32_t h0 = row % q_heads_per_kv; + const auto* row0 = reinterpret_cast( + src + q0 * q_num_stride + h0 * q_head_stride); + const bool has_row1 = row + 1 < row_num; + const int32_t q1 = (row + 1) / q_heads_per_kv; + const int32_t h1 = (row + 1) % q_heads_per_kv; + const auto* row1 = has_row1 + ? reinterpret_cast( + src + q1 * q_num_stride + h1 * q_head_stride) + : nullptr; + auto* dst = reinterpret_cast(q_buffer + row * head_dim); - for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) { - bfloat16x8_t in8 = - vld1q_bf16(reinterpret_cast(curr_q + i)); - float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec); - float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec); + for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) { + const bfloat16x8_t in0 = vld1q_bf16(row0 + k); + const bfloat16x4_t out0_lo = + vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec)); + const bfloat16x4_t out0_hi = + vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec)); - bfloat16x4_t lo_b = vcvt_bf16_f32(lo); - bfloat16x4_t hi_b = vcvt_bf16_f32(hi); - bfloat16x8_t out = vcombine_bf16(lo_b, hi_b); - vst1q_bf16(reinterpret_cast(dst + i), out); + bfloat16x4_t out1_lo = zero; + bfloat16x4_t out1_hi = zero; + if (has_row1) { + const bfloat16x8_t in1 = vld1q_bf16(row1 + k); + out1_lo = + vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec)); + out1_hi = + vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec)); } + + vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo)); + vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi)); } } } @@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA { public: // Reshape and cache K/V into BFMMLA-optimized layouts // K cache: - // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE] + // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K] + // [K_CACHE_K_GROUP_STRIDE] // - TokenColumn // V cache: - // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE] - // - TokenRows + // [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK] + // [VCacheKGroupStride] static void reshape_and_cache( const c10::BFloat16* __restrict__ key, const c10::BFloat16* __restrict__ value, @@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA { const int64_t block_size, [[maybe_unused]] const int64_t block_size_stride, const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) { - const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE; - const int64_t v_pair_stride = - (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE; + const int64_t v_n_group_stride = + (block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride; -#pragma omp parallel for - for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { - for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { +#pragma omp parallel for collapse(2) + for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { + for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { const int64_t pos = slot_mapping[token_idx]; if (pos < 0) continue; @@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA { key_base + block_in_block * k_block_stride; for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) { - uint16_t* dst_u16 = reinterpret_cast( - block_base + hd4 * K_INNER_STRIDE + - pair_in_block * V_INNER_STRIDE + lane_base); - const uint16_t* src_u16 = - reinterpret_cast(key_src + hd4 * TILE_K); - vst1_u16(dst_u16, vld1_u16(src_u16)); + bfloat16_t* dst = reinterpret_cast( + block_base + hd4 * K_CACHE_K_GROUP_STRIDE + + pair_in_block * B_COL_PAIR_STRIDE + lane_base); + const bfloat16_t* src = + reinterpret_cast(key_src + hd4 * TILE_K); + vst1_bf16(dst, vld1_bf16(src)); } } @@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA { value_cache + block_idx * num_blocks_stride + head_idx * cache_head_num_stride; - const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK; + const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK; const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1); - c10::BFloat16* __restrict row_block_base = - value_base + row_block * V_INNER_STRIDE; + const auto* src = reinterpret_cast(value_src); + auto* dst = reinterpret_cast(value_base); + for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) { + const bfloat16x8_t values = + vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK); + const bfloat16x4_t low = vget_low_bf16(values); + const bfloat16x4_t high = vget_high_bf16(values); + bfloat16_t* group = + dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride; - for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) { - c10::BFloat16* __restrict dst_val = - row_block_base + hd2 * v_pair_stride; - - const uint16_t* src_u16 = - reinterpret_cast(value_src); - uint16_t* dst_u16 = reinterpret_cast(dst_val); - dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0]; - dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] = - src_u16[hd2 * TILE_COLS + 1]; + vst1_lane_bf16(group + lane, low, 0); + vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1); + vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2); + vst1_lane_bf16( + group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low, + 3); + vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0); + vst1_lane_bf16( + group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, + high, 1); + vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2); + vst1_lane_bf16( + group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, + high, 3); } } } @@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA { } // namespace cpu_attention -#endif // CPU_ATTN_ASIMD_BFMMLA_HPP +#endif // CPU_ATTN_NEON_BFMMLA_HPP diff --git a/examples/disaggregated/disaggregated_encoder/README.md b/examples/disaggregated/disaggregated_encoder/README.md index 71a0d0e72e..ec10f23fc0 100644 --- a/examples/disaggregated/disaggregated_encoder/README.md +++ b/examples/disaggregated/disaggregated_encoder/README.md @@ -106,6 +106,82 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. | | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). | +### Dynamic registration + +Alternatively, let the external launcher register ready instances over HTTP. +No vLLM configuration changes or worker registration threads are required. +Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration +and removal. These are trusted control-plane APIs; do not expose them publicly. + +```bash +export ADMIN_API_KEY="your-admin-key" +python disagg_epd_proxy.py --port 8000 --dynamic-registration +``` + +Omit the static server URL flags and register every stage. The roles determine +the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` + +`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`, +even when D registers first or all P instances go offline; otherwise requests +return `503`. The proxy rejects mixing combined PD with standalone P/D, +including unhealthy instances still in the registry. Explicitly remove the +old topology's P/D or PD registrations before switching topologies. + +After each instance is ready, the launcher registers its reachable HTTP URL: + +```bash +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"encode","url":"http://e-host:8001"}' +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"prefill_decode","url":"http://pd-host:8002"}' +``` + +For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`. +Example and NIXL EC connectors need no additional registration fields. +For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode` +or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port +layout: supply the TP-rank-0 address for each DP replica, in DP-rank order; +replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers +the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`: + +```json +{ + "role": "prefill_decode", + "url": "http://pd-host:8002", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"] +} +``` + +The proxy selects one consumer replica and uses it for both the encoder push +and the consumer HTTP request. Standalone D does not need EC control addresses. +Port allocation and avoiding collisions remain the launcher's responsibility. + +Inspect or remove instances without restarting the proxy: + +```bash +curl http://proxy-host:8000/instances +curl --fail-with-body -X DELETE \ + 'http://proxy-host:8000/instances?url=http://e-host:8001' \ + -H "X-API-Key: $ADMIN_API_KEY" +``` + +Registration is idempotent. The proxy probes registered instances every +`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe +(default 2). After `--fail-threshold` consecutive failures (default 3), it stops +sending new requests to that instance. Healthy instances rejoin automatically; +unreachable ones are forgotten after `--evicted-ttl` seconds (default 900; +`0` retains them indefinitely). Removal stops new routing; already routed +requests retain their selected endpoints, so drain requests before stopping +the instance. The launcher must re-register instances after a proxy restart. + +The example launch scripts also support this flow: export `ADMIN_API_KEY` and +set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or +`disagg_1e1p1d_example.sh`. Their default remains static routing. + +### Static configuration + The proxy batches images from the same user request assigned to the same encoder. Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each @@ -137,6 +213,6 @@ For E + P + D setup: ```bash $ python disagg_encoder_proxy.py \ --encode-servers-urls "http://e1:8001,http://e2:8001" \ - --prefill-servers-urls "http://p1:8003,http://p2:8004" \ + --prefill-servers-urls "http://p1:8003,http://p2:8004" \ --decode-servers-urls "http://d1:8005,http://d2:8006" ``` diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh index cb51333e90..57392a0408 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh @@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --prefill-servers-urls "http://localhost:$PREFILL_PORT" + --decode-servers-urls "http://localhost:$DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "http://localhost:$PREFILL_PORT" \ - --decode-servers-urls "http://localhost:$DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh index 4bf60ea66d..44e43b8d7f 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh @@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "disable" \ - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 3317b091fb..abae7eb65d 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -1,12 +1,13 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""disagg_encoder_proxy.py +"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation. -Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two -clusters: - • encode (multimodal feature extraction) - • decode (language-model inference) +Use static server URLs, or --dynamic-registration for launcher-managed HTTP +registration of E, P/PD and D instances. Both modes route OpenAI-compatible +``/v1/chat/completions`` requests through encoder and inference clusters and +share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming +forwarding. For MM input we: 1. Extract *every* image/audio/video item. @@ -20,22 +21,27 @@ from __future__ import annotations import argparse import asyncio +import enum import hashlib import itertools import json import logging import os import random +import secrets import time import uuid from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field from typing import Any import aiohttp import msgspec import uvicorn -from fastapi import FastAPI, HTTPException, Request +from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import AnyHttpUrl, BaseModel, Field, model_validator ############################################################################### # FastAPI app & global state @@ -44,6 +50,340 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") logger = logging.getLogger("proxy") +DEFAULT_PROBE_INTERVAL = 5.0 +DEFAULT_PROBE_TIMEOUT = 2.0 +DEFAULT_FAIL_THRESHOLD = 3 +# Stop probing an instance that has been down this long. 0 probes forever, +# which is what a cluster that restarts instances in place wants. +DEFAULT_EVICTED_TTL = 900.0 + + +class InstanceRole(str, enum.Enum): + ENCODE = "encode" + PREFILL = "prefill" + DECODE = "decode" + PREFILL_DECODE = "prefill_decode" + + +@dataclass +class InstanceRecord: + """One registered instance. + + Attributes: + role: Which stage this instance serves. + url: Base OpenAI-compatible URL, e.g. ``http://host:8000``. + ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica, + supplied by the launcher from the consumer's fixed port config. + dp_size: Data-parallel replicas behind `url`, so the proxy can pick + a replica and name the same one to both halves of a request. + + """ + + role: InstanceRole + url: str + ec_zmq_addrs: list[str] = field(default_factory=list) + dp_size: int = 1 + + +class InstanceRegistry: + def __init__( + self, + probe_interval: float = DEFAULT_PROBE_INTERVAL, + probe_timeout: float = DEFAULT_PROBE_TIMEOUT, + fail_threshold: int = DEFAULT_FAIL_THRESHOLD, + evicted_ttl: float = DEFAULT_EVICTED_TTL, + ): + self._probe_interval = probe_interval + self._probe_timeout = probe_timeout + self._fail_threshold = fail_threshold + self._evicted_ttl = evicted_ttl + + self._live: dict[str, InstanceRecord] = {} + self._evicted: dict[str, InstanceRecord] = {} + self._evicted_since: dict[str, float] = {} + self._fail_counts: dict[str, int] = {} + # One cursor per role, only ever incremented. Rebuilding it whenever + # the roster changes -- what an `itertools.cycle` over a mutable list + # forces -- restarts every fan-out at the first instance and hot-spots + # it after each registration. + self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole} + self._replica_cursors: dict[str, int] = {} + self._probe_task: asyncio.Task | None = None + + def register(self, record: InstanceRecord) -> bool: + """Add or refresh an instance without overriding its health status.""" + roles = { + other.role + for other in itertools.chain(self._live.values(), self._evicted.values()) + } + split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE} + if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or ( + record.role in split_roles and InstanceRole.PREFILL_DECODE in roles + ): + raise ValueError("Cannot mix prefill_decode with standalone prefill/decode") + key = record.url + previous = self._live.get(key) or self._evicted.get(key) + if previous is not None: + if previous == record: + return False + if previous.role is not record.role: + raise ValueError("Unregister the instance before changing its role") + target = self._live if key in self._live else self._evicted + target[key] = record + return False + self._live[key] = record + logger.info("Registered instance %s: %s", record.role.value, record.url) + return True + + def unregister(self, url: str) -> bool: + """Drop an instance for good, so a probe cannot bring it back.""" + key = url + if key not in self._live and key not in self._evicted: + return False + self._live.pop(key, None) + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + logger.info("Unregistered instance: %s", url) + return True + + def instances(self, role: InstanceRole) -> list[InstanceRecord]: + return [record for record in self._live.values() if record.role is role] + + def urls(self, role: InstanceRole) -> list[str]: + return [record.url for record in self.instances(role)] + + def pick(self, role: InstanceRole) -> InstanceRecord | None: + """Take the next instance of `role` in round-robin order.""" + picked = self.pick_many(role, 1) + return picked[0] if picked else None + + def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]: + """Take `count` instances, continuing the rotation across calls. + + A multimodal request fans out one encoder request per item, so the + assignment has to be contiguous with the previous request's rather + than restart at the first instance every time. + """ + alive = self.instances(role) + if not alive or count <= 0: + return [] + start = self._cursors[role] + self._cursors[role] = start + count + return [alive[(start + offset) % len(alive)] for offset in range(count)] + + def next_replica(self, record: InstanceRecord) -> int: + """Take the next data-parallel replica of `record`, round-robin. + + The encoder pushes to one replica's receive channel, so the request + has to run on that same replica; the caller names the rank to both + halves. + """ + if record.dp_size <= 1: + return 0 + cursor = self._replica_cursors.get(record.url, 0) + self._replica_cursors[record.url] = cursor + 1 + return cursor % record.dp_size + + def status(self) -> dict[str, Any]: + return { + role.value: { + "live": [record.url for record in self.instances(role)], + "evicted": [ + record.url + for record in self._evicted.values() + if record.role is role + ], + } + for role in InstanceRole + } + + def start_probing(self) -> None: + if self._probe_task is None and self._probe_interval > 0: + self._probe_task = asyncio.create_task(self._probe_loop()) + + async def stop_probing(self) -> None: + if self._probe_task is None: + return + self._probe_task.cancel() + with suppress(asyncio.CancelledError): + await self._probe_task + self._probe_task = None + + async def _probe_loop(self) -> None: + timeout = aiohttp.ClientTimeout(total=self._probe_timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + while True: + await asyncio.sleep(self._probe_interval) + try: + await self._probe_once(session) + except Exception: + logger.exception("EPD registry probe round failed") + + async def _probe_once(self, session: aiohttp.ClientSession) -> None: + targets = list(self._live.values()) + list(self._evicted.values()) + if not targets: + return + results = await asyncio.gather( + *(self._probe(session, record.url) for record in targets), + return_exceptions=True, + ) + now = time.monotonic() + for record, healthy in zip(targets, results): + current = self._live.get(record.url) or self._evicted.get(record.url) + if current is not record: + continue + if healthy is True: + self._on_probe_success(record) + else: + self._on_probe_failure(record, now) + self._drop_expired(now) + + async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool: + async with session.get(f"{url}/health") as resp: + return resp.status == 200 + + def _on_probe_success(self, record: InstanceRecord) -> None: + key = record.url + self._fail_counts.pop(key, None) + if key in self._evicted: + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._live[key] = record + logger.info( + "Instance %s (%s) is healthy again; routing resumed", + record.url, + record.role.value, + ) + + def _on_probe_failure(self, record: InstanceRecord, now: float) -> None: + key = record.url + if key not in self._live: + # Either already evicted, or unregistered while this probe was in + # flight. A round snapshots its targets and then awaits, so a + # removal inside that window would otherwise be undone by a result + # describing a registry that no longer exists. + return + failures = self._fail_counts.get(key, 0) + 1 + self._fail_counts[key] = failures + if failures < self._fail_threshold: + return + self._live.pop(key, None) + self._evicted[key] = record + self._evicted_since[key] = now + logger.warning( + "Instance %s (%s) failed %d consecutive probes; stopped routing " + "to it. It rejoins on its own once it responds again.", + record.url, + record.role.value, + failures, + ) + + def _drop_expired(self, now: float) -> None: + if self._evicted_ttl <= 0: + return + for key, since in list(self._evicted_since.items()): + if now - since < self._evicted_ttl: + continue + record = self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + if record is not None: + logger.warning("Instance %s stayed down; forgetting it", record.url) + + +class InstanceRegistration(BaseModel): + role: InstanceRole + url: AnyHttpUrl + ec_zmq_addrs: list[str] = Field(default_factory=list) + dp_size: int = Field(default=1, ge=1) + + @model_validator(mode="after") + def validate_consumer_addresses(self): + if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size: + raise ValueError("Provide one Mooncake control address per DP replica") + return self + + +def require_admin_key(x_api_key: str = Header(default="")) -> None: + expected = os.getenv("ADMIN_API_KEY", "") + if not expected or not secrets.compare_digest(x_api_key, expected): + raise HTTPException(403, "Invalid admin API key") + + +@dataclass +class EPDProxyConfig: + probe_interval: float = 5.0 + probe_timeout: float = 2.0 + fail_threshold: int = 3 + evicted_ttl: float = 900.0 + + +@dataclass +class _Route: + """The instances one request was assigned to. + + Attributes: + consumer: The stage that receives the embedding, if any stage does. + consumer_zmq: The receive address named to the encoders. + dp_rank: Which replica of `consumer` was named, so the request can be + pinned to it. + + """ + + encoder_urls: list[str] + prefill: InstanceRecord | None + decode: InstanceRecord + consumer: InstanceRecord | None = None + consumer_zmq: str | None = None + dp_rank: int | None = None + + +class EPDProxy: + def __init__(self, registry: InstanceRegistry): + self.registry = registry + + # ---------------------------------------------------------------- # + # Routing # + # ---------------------------------------------------------------- # + def route(self, num_items: int) -> _Route: + decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + prefill = ( + self.registry.pick(InstanceRole.PREFILL) + if decode.role is InstanceRole.DECODE + else None + ) + if decode.role is InstanceRole.DECODE and prefill is None: + raise HTTPException( + status_code=503, detail="No prefill instance is registered" + ) + encoders = self.registry.urls(InstanceRole.ENCODE) + if num_items and not encoders: + raise HTTPException( + status_code=503, detail="No encode instance is registered" + ) + route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode) + self._name_consumer(route) + return route + + def _name_consumer(self, route: _Route) -> None: + """Pin the EC consumer replica and, for push connectors, its endpoint.""" + candidate = route.prefill or route.decode + route.consumer = candidate + rank = self.registry.next_replica(candidate) + route.dp_rank = rank if candidate.dp_size > 1 else None + if candidate.ec_zmq_addrs: + route.consumer_zmq = candidate.ec_zmq_addrs[rank] + + app = FastAPI() encode_session: aiohttp.ClientSession | None = None prefill_session: aiohttp.ClientSession | None = None @@ -448,6 +788,7 @@ async def maybe_prefill( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """- Do prefill-only task if p_url exist; - Return a new body carrying kv transfer params (for nixl connector) @@ -459,7 +800,7 @@ async def maybe_prefill( if p_url: logger.info("[%s] Processing through prefill: %s", req_id, p_url) - prefill_response = await process_prefill_stage(req_data, p_url, req_id) + prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank) # for nixl connector to facilitate kv transfer... prefill_response_json = msgspec.json.decode(await prefill_response.read()) kv_transfer_params = prefill_response_json.get("kv_transfer_params", {}) @@ -473,6 +814,7 @@ async def process_prefill_stage( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """Process request through Prefill stage and return kv_transfer_params""" logger.info("[%s] Sending prefill request to: %s", req_id, p_url) @@ -494,6 +836,8 @@ async def process_prefill_stage( del prefill_request["stream_options"] headers = {"x-request-id": req_id, "Content-Type": "application/json"} + if dp_rank is not None: + headers["X-data-parallel-rank"] = str(dp_rank) try: prefill_response = await prefill_session.post( f"{p_url}/v1/chat/completions", @@ -594,9 +938,7 @@ async def on_startup() -> None: ), ) encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) - if app.state.p_urls: - # only setup if prefill instance(s) exist - prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) + prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) @@ -622,6 +964,7 @@ async def prepare_for_decode( e_urls: list[str], p_url: str, consumer_zmq: str | None, + prefill_dp_rank: int | None = None, ) -> tuple[dict, float, float]: """Encode, rewrite and prefill, returning the body to send to decode. @@ -642,7 +985,7 @@ async def prepare_for_decode( handles.update(ec_params) prepared = {**prepared, "ec_transfer_params": handles} _t2 = time.perf_counter() - prepared = await maybe_prefill(prepared, p_url, req_id) + prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank) return prepared, _t1 - _t0, _t2 - _t1 @@ -654,12 +997,13 @@ async def forward_non_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> Response: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -733,12 +1077,13 @@ async def forward_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> AsyncIterator[bytes]: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -976,10 +1321,147 @@ async def stop_profile(request: Request): return await _profile_cmd("stop", body, e_url, p_url, d_url) +def build_app(config: EPDProxyConfig | None = None) -> FastAPI: + config = config or EPDProxyConfig() + registry = InstanceRegistry( + probe_interval=config.probe_interval, + probe_timeout=config.probe_timeout, + fail_threshold=config.fail_threshold, + evicted_ttl=config.evicted_ttl, + ) + proxy = EPDProxy(registry) + + @asynccontextmanager + async def lifespan(app: FastAPI): + await on_startup() + try: + registry.start_probing() + yield + finally: + await registry.stop_probing() + await on_shutdown() + + app = FastAPI(lifespan=lifespan) + app.state.proxy = proxy + app.state.registry = registry + + @app.get("/instances") + async def list_instances(): + return registry.status() + + @app.post("/instances", dependencies=[Depends(require_admin_key)]) + async def register_instance(body: InstanceRegistration): + if body.ec_zmq_addrs and body.role not in ( + InstanceRole.PREFILL, + InstanceRole.PREFILL_DECODE, + ): + raise HTTPException(400, "Only EC consumers accept Mooncake addresses") + try: + created = registry.register( + InstanceRecord( + body.role, + str(body.url).rstrip("/"), + body.ec_zmq_addrs, + body.dp_size, + ) + ) + except ValueError as error: + raise HTTPException(409, str(error)) from error + return {"registered": created} + + @app.delete("/instances", dependencies=[Depends(require_admin_key)]) + async def unregister_instance(url: AnyHttpUrl): + return {"removed": registry.unregister(str(url).rstrip("/"))} + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request): + req_data = msgspec.json.decode(await request.body()) + req_id = request.headers.get("x-request-id", str(uuid.uuid4())) + route = proxy.route(len(extract_mm_items(req_data))) + args = ( + req_data, + req_id, + route.encoder_urls, + route.prefill.url if route.prefill else None, + route.decode.url, + route.consumer_zmq, + ) + ranks = { + "dp_rank": route.dp_rank if route.prefill is None else None, + "prefill_dp_rank": route.dp_rank if route.prefill is not None else None, + } + if req_data.get("stream", False): + return StreamingResponse( + forward_stream(*args, **ranks), media_type="text/event-stream" + ) + return await forward_non_stream(*args, **ranks) + + @app.get("/v1/models") + async def list_models(): + decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + async with decode_session.get(f"{decode.url}/v1/models") as resp: + resp.raise_for_status() + return await resp.json() + + @app.get("/health") + async def health(): + status = registry.status() + # An empty roster is not unhealthy: the proxy is meant to come up + # before anything registers with it. + return JSONResponse({"proxy": "healthy", "instances": status}) + + @app.post("/start_profile") + async def start_profile(request: Request): + return await _profile(registry, "start", await request.json()) + + @app.post("/stop_profile") + async def stop_profile(request: Request): + return await _profile(registry, "stop", await request.json()) + + return app + + +async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict: + headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} + targets = {role.value: registry.pick(role) for role in InstanceRole} + results = await asyncio.gather( + *( + _post_if_available( + decode_session, f"{record.url}/{cmd}_profile", payload, headers + ) + for record in targets.values() + if record is not None + ) + ) + reachable = [result for result in results if result is not None] + if not reachable: + raise HTTPException( + status_code=503, + detail="Profiling endpoints are disabled on every instance", + ) + live = [name for name, record in targets.items() if record is not None] + return dict(zip(live, results)) + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--dynamic-registration", + action="store_true", + help="Enable launcher-managed HTTP registration for E+PD or E+P+D.", + ) + parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL) + parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT) + parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD) + parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL) parser.add_argument( "--log-requests", action="store_true", @@ -995,12 +1477,12 @@ if __name__ == "__main__": ) parser.add_argument( "--encode-servers-urls", - required=True, + default="", help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")', ) parser.add_argument( "--prefill-servers-urls", - required=True, + default="none", help=( 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ' 'to enable E->P->D, set "disable" or "none" to enable E->PD' @@ -1008,7 +1490,7 @@ if __name__ == "__main__": ) parser.add_argument( "--decode-servers-urls", - required=True, + default="", help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")', ) parser.add_argument( @@ -1043,58 +1525,87 @@ if __name__ == "__main__": ) args = parser.parse_args() + if args.fail_threshold < 1: + parser.error("--fail-threshold must be at least 1") if args.log_requests: logging.getLogger().setLevel(logging.DEBUG) app.middleware("http")(log_requests) NO_REWRITE = args.no_rewrite DECODE_RETRIES = max(0, args.decode_retries) - app.state.e_urls = [ - u.strip() for u in args.encode_servers_urls.split(",") if u.strip() - ] - app.state.d_urls = [ - u.strip() for u in args.decode_servers_urls.split(",") if u.strip() - ] - app.state.d_ec_urls = [ - u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() - ] - if args.ec_consumer_dp_size < 1: - parser.error("--ec-consumer-dp-size must be at least 1") - app.state.ec_consumer_dp_size = args.ec_consumer_dp_size - app.state.replica_counter = itertools.count() - expected = len(app.state.d_urls) * args.ec_consumer_dp_size - if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: - parser.error( - "--ec-consumer-zmq-addrs must contain one address per consumer " - f"replica: expected {expected} " - f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), " - f"got {len(app.state.d_ec_urls)}" - ) - # handle prefill instances - if args.prefill_servers_urls.lower() in ("disable", "none", ""): - app.state.p_urls = [] - logger.info( - "Disaggregated prefill phase explicitly disabled by user. Running E + PD..." + if args.dynamic_registration: + if not os.getenv("ADMIN_API_KEY"): + parser.error("--dynamic-registration requires ADMIN_API_KEY") + if ( + args.encode_servers_urls + or args.prefill_servers_urls.lower() not in ("disable", "none", "") + or args.decode_servers_urls + or args.ec_consumer_zmq_addrs + ): + parser.error( + "With --dynamic-registration, instances register through /instances" + ) + app = build_app( + EPDProxyConfig( + probe_interval=args.probe_interval, + probe_timeout=args.probe_timeout, + fail_threshold=args.fail_threshold, + evicted_ttl=args.evicted_ttl, + ) ) + if args.log_requests: + app.middleware("http")(log_requests) else: - app.state.p_urls = [ - u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + if not args.encode_servers_urls or not args.decode_servers_urls: + parser.error( + "Static routing requires --encode-servers-urls " + "and --decode-servers-urls" + ) + app.state.e_urls = [ + u.strip() for u in args.encode_servers_urls.split(",") if u.strip() ] - logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") - try: - validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) - except ValueError as exc: - parser.error(str(exc)) + app.state.d_urls = [ + u.strip() for u in args.decode_servers_urls.split(",") if u.strip() + ] + app.state.d_ec_urls = [ + u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() + ] + if args.ec_consumer_dp_size < 1: + parser.error("--ec-consumer-dp-size must be at least 1") + app.state.ec_consumer_dp_size = args.ec_consumer_dp_size + app.state.replica_counter = itertools.count() + expected = len(app.state.d_urls) * args.ec_consumer_dp_size + if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: + parser.error( + "--ec-consumer-zmq-addrs must contain one address per consumer " + f"replica: expected {expected} " + f"({len(app.state.d_urls)} servers x " + f"{args.ec_consumer_dp_size} replicas), " + f"got {len(app.state.d_ec_urls)}" + ) + # handle prefill instances + if args.prefill_servers_urls.lower() in ("disable", "none", ""): + app.state.p_urls = [] + logger.info("Disaggregated prefill disabled. Running E + PD...") + else: + app.state.p_urls = [ + u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + ] + logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") + try: + validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) + except ValueError as exc: + parser.error(str(exc)) - logger.info("Proxy listening on %s:%s", args.host, args.port) - logger.info("Encode servers: %s", app.state.e_urls) - logger.info("Prefill instances %s", app.state.p_urls) - logger.info("Decode servers: %s", app.state.d_urls) - if app.state.ec_consumer_dp_size > 1: - logger.info( - "EC consumer replicas per server: %d (control addresses: %s)", - app.state.ec_consumer_dp_size, - app.state.d_ec_urls, - ) + logger.info("Proxy listening on %s:%s", args.host, args.port) + logger.info("Encode servers: %s", app.state.e_urls) + logger.info("Prefill instances %s", app.state.p_urls) + logger.info("Decode servers: %s", app.state.d_urls) + if app.state.ec_consumer_dp_size > 1: + logger.info( + "EC consumer replicas per server: %d (control addresses: %s)", + app.state.ec_consumer_dp_size, + app.state.d_ec_urls, + ) uvicorn.run( app, diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index a45cbe3731..0edd0ce788 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -130,10 +130,36 @@ pub struct EngineCoreOutput { #[serde(default)] pub new_sampling_mask: Option, /// Per-request speculative-decoding acceptance metrics, set on the final - /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here; - /// the Rust frontend does not yet surface it in responses. + /// output when `--per-request-spec-decode-metrics` is enabled. #[serde(default)] - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option, +} + +/// Raw per-sequence speculative-decoding accumulator. +/// +/// Python models this as a plain `@dataclass`, so it is serialized by msgspec +/// as a map (named fields). Derived values such as acceptance rates are +/// computed by the frontend, not carried on the wire. +/// +/// Original Python definition: `RequestSpecDecodeMetrics` in +/// `vllm/v1/metrics/stats.py`. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RequestSpecDecodeMetrics { + /// Configured `num_speculative_tokens`. + #[serde(default)] + pub num_spec_tokens: u64, + /// Verify-step counts indexed by accepted draft-token count. + #[serde(default)] + pub histogram: Vec, + /// Total proposed draft tokens. + #[serde(default)] + pub num_draft_tokens: u64, + /// Accepted draft count per verify step; empty unless `detailed`. + #[serde(default)] + pub per_step_accepted: Vec, + /// Proposed draft count per verify step; empty unless `detailed`. + #[serde(default)] + pub per_step_drafted: Vec, } impl EngineCoreOutput { diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index d0cab4a507..34eae25682 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -11,7 +11,9 @@ use futures::stream::FusedStream; use futures::{Stream, StreamExt as _, pin_mut}; use serde::{Deserialize, Serialize}; use vllm_engine_core_client::protocol::logprobs::Logprobs; -use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason}; +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason, +}; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; @@ -47,6 +49,8 @@ pub struct CollectedGenerateOutput { pub ec_transfer_params: Option, /// Sampling support sets aligned one-to-one with generated token positions. pub sampling_mask: Option, + /// Per-request speculative-decoding metrics from the terminal output. + pub spec_decode_metrics: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -160,6 +164,8 @@ pub struct GenerateOutput { pub ec_transfer_params: Option, /// Sampling support sets aligned one-to-one with `token_ids`. pub sampling_mask: Option, + /// Per-request speculative-decoding metrics, present on terminal outputs. + pub spec_decode_metrics: Option, } impl GenerateOutput { @@ -208,6 +214,7 @@ impl GenerateOutput { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, } } } @@ -308,6 +315,7 @@ impl Stream for GenerateOutputStream { kv_transfer_params: raw.kv_transfer_params, ec_transfer_params: raw.ec_transfer_params, sampling_mask, + spec_decode_metrics: raw.spec_decode_metrics, }; Poll::Ready(Some(Ok(output))) @@ -396,6 +404,7 @@ impl> + Send> T { kv_transfer_params: None, ec_transfer_params: None, sampling_mask, + spec_decode_metrics: None, }); } @@ -409,6 +418,7 @@ impl> + Send> T { }; collected.kv_transfer_params = output.kv_transfer_params; collected.ec_transfer_params = output.ec_transfer_params; + collected.spec_decode_metrics = output.spec_decode_metrics; if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 09b5199767..c341737a3c 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -413,6 +413,7 @@ async fn collect_output_rejects_partial_sampling_mask() { sampling_mask: Some(SamplingMask { rows: vec![vec![1, 33, 99]], }), + spec_decode_metrics: None, }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index f82059dffd..eb3338a67f 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -29,7 +29,8 @@ use vllm_llm::{ use self::convert::{ResponseOptions, prepare_generate_request}; use self::types::{ GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, - GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo, + GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo, + SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics, }; pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; @@ -160,10 +161,14 @@ async fn generate_chunk_stream( let mut prompt_tokens = None; let mut prompt_token_ids = None; let mut usage = TokenUsage::default(); + let mut spec_decode_metrics = None; while let Some(next) = stream.next().await { match next { Ok(output) => { + if let Some(metrics) = output.spec_decode_metrics { + spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics)); + } if prompt_tokens.is_none() && let Some(info) = output.prompt_info.as_ref() { @@ -223,6 +228,7 @@ async fn generate_chunk_stream( .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()), prompt_token_ids, + metrics: None, }) .await; } @@ -243,6 +249,11 @@ async fn generate_chunk_stream( usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)), prompt_token_ids: None, mm_placeholders: None, + metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics { + speculative_decoding: StreamingSpeculativeDecodingMetrics::from( + speculative_decoding, + ), + }), }) .await; } @@ -318,6 +329,9 @@ fn collect_generate( mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(), kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, + metrics: collected.spec_decode_metrics.map(|metrics| PerRequestMetrics { + speculative_decoding: SpeculativeDecodingMetrics::from(metrics), + }), }) } @@ -465,6 +479,7 @@ mod tests { use futures::{TryStreamExt as _, stream}; use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; + use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_llm::GeneratePromptInfo; use super::*; @@ -482,6 +497,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -496,6 +512,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, }), ]); @@ -578,6 +595,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, } } @@ -645,6 +663,80 @@ mod tests { assert!(chunks[1].mm_placeholders.is_none()); } + #[tokio::test] + async fn generate_chunk_stream_returns_terminal_spec_decode_metrics() { + let metrics = RequestSpecDecodeMetrics { + num_spec_tokens: 3, + histogram: vec![0, 1, 1, 0], + num_draft_tokens: 6, + ..Default::default() + }; + let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos())); + output.spec_decode_metrics = Some(metrics.clone()); + + for include_usage in [false, true] { + let chunks: Vec<_> = generate_chunk_stream( + stream::iter(vec![Ok(output.clone())]), + "raw-stream".to_string(), + ApiServerOptions::default(), + ResponseOptions { + include_usage, + ..Default::default() + }, + None, + ) + .try_collect() + .await + .expect("collect chunks"); + assert_eq!(chunks.len(), 1 + usize::from(include_usage)); + assert!(chunks[0].metrics.is_none()); + if include_usage { + assert!(chunks[1].choices.is_empty()); + assert_eq!( + chunks[1].metrics, + Some(PerRequestMetrics { + speculative_decoding: StreamingSpeculativeDecodingMetrics::from( + SpeculativeDecodingMetrics::from(metrics.clone()) + ) + }) + ); + let json = serde_json::to_value(&chunks[1]).expect("serialize chunk"); + let stats = &json["metrics"]["speculative_decoding"]; + assert_eq!(stats.as_object().expect("metrics object").len(), 7); + assert_eq!(stats["num_draft_tokens"], 6); + assert_eq!(stats["num_accepted_draft_tokens"], 3); + assert_eq!(stats["mean_acceptance_length"], 2.5); + assert_eq!( + stats["acceptance_histogram"], + serde_json::json!([0, 1, 1, 0]) + ); + assert!(stats.get("per_step_accepted").is_none()); + } + } + } + + #[test] + fn spec_decode_metrics_zero_steps_and_detailed() { + let empty = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics { + num_spec_tokens: 2, + histogram: vec![0, 0, 0], + ..Default::default() + }); + assert_eq!(empty.mean_acceptance_length, 1.0); + assert_eq!(empty.draft_acceptance_rate, 0.0); + assert_eq!(empty.per_step_accepted, None); + + let detailed = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics { + num_spec_tokens: 2, + histogram: vec![0, 1, 1], + num_draft_tokens: 4, + per_step_accepted: vec![1, 2], + per_step_drafted: vec![2, 2], + }); + assert_eq!(detailed.per_step_accepted, Some(vec![1, 2])); + assert_eq!(detailed.per_step_drafted, Some(vec![2, 2])); + } + #[test] fn collect_generate_omits_prompt_metadata_by_default() { let output = CollectedGenerateOutput { @@ -662,6 +754,7 @@ mod tests { ec_transfer_params: None, prompt_token_ids: vec![10, 20], sampling_mask: None, + spec_decode_metrics: None, }; let response = collect_generate( @@ -683,6 +776,50 @@ mod tests { assert!(response.mm_placeholders.is_none()); } + #[test] + fn collect_generate_returns_spec_decode_metrics() { + let metrics = RequestSpecDecodeMetrics { + num_spec_tokens: 3, + histogram: vec![0, 1, 1, 0], + num_draft_tokens: 6, + ..Default::default() + }; + let output = CollectedGenerateOutput { + request_id: "raw-metrics".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage::default(), + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10], + sampling_mask: None, + spec_decode_metrics: Some(metrics.clone()), + }; + + let response = collect_generate( + output, + "raw-metrics".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + None, + ) + .expect("response"); + + assert_eq!( + response.metrics, + Some(PerRequestMetrics { + speculative_decoding: SpeculativeDecodingMetrics::from(metrics) + }) + ); + let json = serde_json::to_value(&response).expect("serialize response"); + let stats = &json["metrics"]["speculative_decoding"]; + assert_eq!(stats.as_object().expect("metrics object").len(), 9); + assert!(stats["per_step_accepted"].is_null()); + assert!(stats["per_step_drafted"].is_null()); + } + #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); @@ -749,6 +886,7 @@ mod tests { ec_transfer_params: None, prompt_token_ids: vec![10, 20], sampling_mask: None, + spec_decode_metrics: None, }; let response = collect_generate( @@ -782,6 +920,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + spec_decode_metrics: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 406ec58569..86d8ec00f4 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -7,6 +7,7 @@ use llm_multimodal::MediaContentPart; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; +use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage}; @@ -84,6 +85,7 @@ pub(super) struct GenerateStreamResponse { pub usage: Option, pub prompt_token_ids: Option>, pub mm_placeholders: Option, + pub metrics: Option>, } /// Mirrors the Python vLLM `GenerateResponse` class. @@ -96,6 +98,91 @@ pub(super) struct GenerateResponse { pub mm_placeholders: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, + pub metrics: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct PerRequestMetrics { + pub speculative_decoding: T, +} + +/// Mirrors the Python vLLM `SpeculativeDecodingMetrics` class. +/// +/// Derived from the raw engine accumulator the same way as Python +/// `RequestSpecDecodeMetrics.to_dict`. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct SpeculativeDecodingMetrics { + pub mean_acceptance_length: f64, + pub draft_acceptance_rate: f64, + pub acceptance_histogram: Vec, + pub num_spec_steps: u64, + pub num_accepted_draft_tokens: u64, + pub num_draft_tokens: u64, + pub num_spec_tokens: u64, + pub per_step_accepted: Option>, + pub per_step_drafted: Option>, +} + +impl From for SpeculativeDecodingMetrics { + fn from(raw: RequestSpecDecodeMetrics) -> Self { + let num_spec_steps: u64 = raw.histogram.iter().sum(); + let num_accepted_draft_tokens: u64 = + (0u64..).zip(&raw.histogram).map(|(accepted, count)| accepted * count).sum(); + let ratio = |num: u64, den: u64| { + if den == 0 { + 0.0 + } else { + num as f64 / den as f64 + } + }; + let detailed = !raw.per_step_accepted.is_empty(); + Self { + mean_acceptance_length: if num_spec_steps == 0 { + 1.0 + } else { + 1.0 + ratio(num_accepted_draft_tokens, num_spec_steps) + }, + draft_acceptance_rate: ratio(num_accepted_draft_tokens, raw.num_draft_tokens), + acceptance_histogram: raw.histogram, + num_spec_steps, + num_accepted_draft_tokens, + num_draft_tokens: raw.num_draft_tokens, + num_spec_tokens: raw.num_spec_tokens, + per_step_accepted: detailed.then_some(raw.per_step_accepted), + per_step_drafted: detailed.then_some(raw.per_step_drafted), + } + } +} + +/// Streaming form omits detailed fields when summary metrics are requested. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(super) struct StreamingSpeculativeDecodingMetrics { + pub mean_acceptance_length: f64, + pub draft_acceptance_rate: f64, + pub acceptance_histogram: Vec, + pub num_spec_steps: u64, + pub num_accepted_draft_tokens: u64, + pub num_draft_tokens: u64, + pub num_spec_tokens: u64, + pub per_step_accepted: Option>, + pub per_step_drafted: Option>, +} + +impl From for StreamingSpeculativeDecodingMetrics { + fn from(metrics: SpeculativeDecodingMetrics) -> Self { + Self { + mean_acceptance_length: metrics.mean_acceptance_length, + draft_acceptance_rate: metrics.draft_acceptance_rate, + acceptance_histogram: metrics.acceptance_histogram, + num_spec_steps: metrics.num_spec_steps, + num_accepted_draft_tokens: metrics.num_accepted_draft_tokens, + num_draft_tokens: metrics.num_draft_tokens, + num_spec_tokens: metrics.num_spec_tokens, + per_step_accepted: metrics.per_step_accepted, + per_step_drafted: metrics.per_step_drafted, + } + } } pub(super) type MultiModalPlaceholders = HashMap>; diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 1cd6752460..362ffe4b15 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -14,6 +14,7 @@ import torch import torch.multiprocessing as mp from vllm.distributed import get_ep_group +from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, @@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port from ..utils import init_test_distributed_environment +DEVICE = current_platform.device_type + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -483,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_ep_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 64 tokens_per_rank = 16 @@ -624,7 +627,7 @@ def _two_sided_data_worker(rank, world_size): # Use DP group because MnnvlMoe workspace allocation calls get_dp_group() # internally and requires dp_size == ep_size. cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") num_gpus = torch.accelerator.device_count() hidden_size = 128 @@ -769,7 +772,7 @@ def _one_sided_data_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 256 tokens_per_rank = 32 diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index d0a417a788..cd94138654 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -15,6 +15,7 @@ import pytest from tests.utils import multi_gpu_test from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.executor import multiproc_executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor @@ -378,8 +379,8 @@ def test_multiproc_executor_properties(): def test_multiproc_executor_multi_node(): """Test MultiprocExecutor with multi-node configuration. This simulates 2 nodes with TP=4: - - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2 - - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 + - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2 + - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2 Total world_size = 4, nnodes = 2 """ # Python 3.14+ changed default multiprocessing start method to 'forkserver' @@ -392,15 +393,20 @@ def test_multiproc_executor_multi_node(): # symm_mem does not work for simulating multi instance in single node os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" + # Use the platform-specific device control env var (e.g. + # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each + # node only sees its own devices on different platform. + device_control_env_var = current_platform.device_control_env_var + def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int): """Run a single node's executor.""" executor = None try: - # Set CUDA_VISIBLE_DEVICES for this node + # Restrict visible devices for this node if node_rank == 0: - os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" + os.environ[device_control_env_var] = "0,1" else: - os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" + os.environ[device_control_env_var] = "2,3" # Create config for this node vllm_config = create_vllm_config( diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 9cc6124393..80ecc0f49c 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -52,6 +52,7 @@ def _make_generate_response( logprobs: dict | None = None, prompt_logprobs: list | None = None, kv_transfer_params: dict | None = None, + metrics: dict | None = None, ) -> dict: choice: dict = { "index": 0, @@ -64,6 +65,7 @@ def _make_generate_response( "choices": [choice], "prompt_logprobs": prompt_logprobs, "kv_transfer_params": kv_transfer_params, + "metrics": metrics, } @@ -266,6 +268,49 @@ async def test_derender_chat_kv_transfer_params_passthrough(client): assert response.json()["kv_transfer_params"] == kv +@pytest.mark.asyncio +async def test_derender_chat_metrics_passthrough(client): + gen_req = await _render_chat(client) + metrics = { + "speculative_decoding": { + "mean_acceptance_length": 2.0, + "draft_acceptance_rate": 0.5, + "acceptance_histogram": [0, 1], + "num_spec_steps": 1, + "num_accepted_draft_tokens": 1, + "num_draft_tokens": 2, + "num_spec_tokens": 1, + } + } + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": _make_generate_response( + gen_req["token_ids"][:3], metrics=metrics + ), + }, + ) + assert response.status_code == 200 + actual = response.json()["metrics"] + assert set(actual) == { + "time_to_first_token_ms", + "generation_time_ms", + "queue_time_ms", + "mean_itl_ms", + "tokens_per_second", + "speculative_decoding", + } + assert all( + actual[name] is None for name in actual if name != "speculative_decoding" + ) + assert actual["speculative_decoding"] == { + **metrics["speculative_decoding"], + "per_step_accepted": None, + "per_step_drafted": None, + } + + @pytest.mark.asyncio async def test_derender_chat_empty_token_ids(client): """Empty token_ids list returns 400.""" diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 4032d9cbf3..e7a2e3b856 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -27,6 +27,8 @@ from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, + PerRequestMetrics, + SpeculativeDecodingMetrics, ) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderStreamState, @@ -164,6 +166,7 @@ def _make_stream_chunk( finish_reason: str | None = None, request_id: str = "test-req", usage: dict | None = None, + metrics: PerRequestMetrics | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" return GenerateStreamResponse( @@ -176,6 +179,7 @@ def _make_stream_chunk( ) ], usage=UsageInfo(**usage) if usage else None, + metrics=metrics, ) @@ -512,6 +516,26 @@ class TestDerenderChatStream: assert chunk1.choices[0].delta.role == "assistant" assert chunk2.choices[0].delta.role is None + @pytest.mark.asyncio + async def test_metrics_passthrough(self, derenderer): + metrics = PerRequestMetrics( + speculative_decoding=SpeculativeDecodingMetrics( + mean_acceptance_length=2.0, + draft_acceptance_rate=0.5, + acceptance_histogram=[0, 1], + num_spec_steps=1, + num_accepted_draft_tokens=1, + num_draft_tokens=2, + num_spec_tokens=1, + ) + ) + chunk, _ = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([], metrics=metrics), + state=DerenderStreamState(), + ) + assert chunk.metrics == metrics + @pytest.mark.asyncio async def test_chunked_equals_oneshot(self, derenderer, tokenizer): """Sum of streaming content deltas == one shot decode.""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index d9ae424cc8..ba5839990c 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -28,6 +28,7 @@ from vllm.renderers import renderer_from_config from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics MODEL_NAME = "openai-community/gpt2" BASE_MODEL_PATHS = [ @@ -130,6 +131,7 @@ def _make_request_output( logprobs: list[dict[int, Any] | None] | None = None, num_cached_tokens: int | None = None, index: int = 0, + spec_decode_metrics: RequestSpecDecodeMetrics | None = None, ) -> RequestOutput: return RequestOutput( request_id=request_id, @@ -144,6 +146,7 @@ def _make_request_output( cumulative_logprob=None, logprobs=logprobs, finish_reason=finish_reason, + spec_decode_metrics=spec_decode_metrics, ) ], finished=finished, @@ -321,6 +324,145 @@ async def test_stream_basic(): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_serve_tokens_returns_spec_decode_metrics(stream: bool): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=stream, + stream_options=StreamOptions(include_usage=True) if stream else None, + ) + + response = await serving.serve_tokens(request) + if stream: + chunks = [chunk async for chunk in response] + payload = next( + chunk + for chunk in _parse_sse_chunks(chunks) + if isinstance(chunk, dict) and chunk.get("metrics") + )["metrics"]["speculative_decoding"] + else: + assert isinstance(response, GenerateResponse) + assert response.metrics is not None + assert response.metrics.speculative_decoding is not None + payload = response.metrics.speculative_decoding.model_dump() + + assert payload["num_draft_tokens"] == 3 + assert payload["num_accepted_draft_tokens"] == 2 + assert payload["num_spec_steps"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output( + include_usage, +): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", + token_ids=[], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=include_usage), + ) + + chunks = [chunk async for chunk in await serving.serve_tokens(request)] + data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"] + + assert "metrics" not in data_chunks[0] + assert len(data_chunks) == 1 + int(include_usage) + if include_usage: + assert data_chunks[1]["choices"] == [] + assert ( + data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling( + stream: bool, +): + engine = _mock_engine() + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + metrics.observe(num_draft_tokens=3, num_accepted=2) + result = _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + spec_decode_metrics=metrics, + ) + result.outputs.append( + CompletionOutput( + index=1, + text="", + token_ids=[20], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + spec_decode_metrics=metrics, + ) + ) + + async def mock_generate(*args, **kwargs): + yield result + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1, n=2), + model=MODEL_NAME, + stream=stream, + stream_options=StreamOptions(include_usage=True) if stream else None, + ) + + response = await serving.serve_tokens(request) + if stream: + chunks = [chunk async for chunk in response] + assert all( + "metrics" not in chunk + for chunk in _parse_sse_chunks(chunks) + if isinstance(chunk, dict) + ) + else: + assert isinstance(response, GenerateResponse) + assert response.metrics is None + + @pytest.mark.asyncio async def test_non_stream_returns_final_prompt_token_ids_when_requested(): engine = _mock_engine() diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8b4efaf913..1aca8593ed 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1384,24 +1384,33 @@ def test_humming_global_valid_shape_m( assert result == expected -def test_humming_permute_scratch_is_keyed_by_runtime_topk( +def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): from types import SimpleNamespace from unittest.mock import Mock import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming + import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute + import vllm.v1.worker.workspace as workspace + + manager = workspace.WorkspaceManager(torch.device("cpu")) + monkeypatch.setattr(workspace, "_manager", manager) scratch_topk6 = Mock() scratch_topk1 = Mock() scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) - monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type) + monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) moe_config.moe_parallel_config.dp_size = 2 - experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config) + experts = SimpleNamespace(moe_config=moe_config) + other_layer = SimpleNamespace(moe_config=moe_config) assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + assert ( + humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 + ) assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 @@ -1412,6 +1421,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk( assert second_call.kwargs["max_num_tokens"] == 6144 assert second_call.kwargs["topk"] == 1 + manager.lock() + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + with pytest.raises(AssertionError, match="was not allocated during warmup"): + humming.HummingExpertsBase._get_permute_scratch(experts, 2) + def test_humming_delegates_to_instance_activation(): from types import SimpleNamespace diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76d1e57f58..52fc8963ec 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -15,12 +15,14 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import ( ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import current_workspace_manager NUM_EXPERTS = [16, 64, 256] TOP_KS = [2, 6, 8] @@ -120,6 +122,7 @@ def torch_unpermute( @pytest.mark.parametrize("topk", TOP_KS) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("ep_size", EP_SIZE) +@pytest.mark.parametrize("use_scratch", [False, True]) def test_moe_permute_unpermute( n_token: int, n_hidden: int, @@ -127,6 +130,8 @@ def test_moe_permute_unpermute( n_expert: int, ep_size: int, dtype: torch.dtype, + use_scratch: bool, + workspace_init, ): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") @@ -160,6 +165,18 @@ def test_moe_permute_unpermute( expert_map=expert_map, ) + scratch = None + if use_scratch: + scratch = get_moe_permute_scratch( + max_num_tokens=n_token, + topk=topk, + num_experts=n_expert, + num_local_experts=n_local_expert, + device=hidden_states.device, + hidden_size=n_hidden, + hidden_dtype=dtype, + ) + ( permuted_hidden_states, _, @@ -173,6 +190,7 @@ def test_moe_permute_unpermute( n_expert=n_expert, n_local_expert=n_local_expert, expert_map=expert_map, + scratch=scratch, ) # check expert_first_token_offset @@ -212,20 +230,22 @@ def test_moe_permute_unpermute( torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("n_token", [1, 33, 128]) +@pytest.mark.parametrize("topk", [1, 6, 8]) +def test_moe_permute_reuses_scratch_buffers( + dtype: torch.dtype, n_token: int, topk: int, workspace_init +): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") - n_token = 64 n_hidden = 2048 n_expert = 16 - topk = 4 hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype) gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype) _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) - scratch = MoEPermuteScratch( + scratch_config = dict( max_num_tokens=n_token, topk=topk, num_experts=n_expert, @@ -234,6 +254,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): hidden_size=n_hidden, hidden_dtype=hidden_states.dtype, ) + scratch = get_moe_permute_scratch(**scratch_config) first = moe_permute( hidden_states=hidden_states, @@ -242,13 +263,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): n_expert=n_expert, scratch=scratch, ) - second = moe_permute( - hidden_states=hidden_states, - a1q_scale=None, - topk_ids=topk_ids, - n_expert=n_expert, - scratch=scratch, - ) + current_workspace_manager().lock() + assert get_moe_permute_scratch(**scratch_config) is scratch + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + second = moe_permute( + hidden_states=hidden_states, + a1q_scale=None, + topk_ids=topk_ids, + n_expert=n_expert, + scratch=get_moe_permute_scratch(**scratch_config), + ) + + for _ in range(2): + hidden_states.add_(1) + topk_ids.copy_(topk_ids.roll(1, 0)) + graph.replay() + expected = moe_permute(hidden_states, None, topk_ids, n_expert) + for actual, reference in zip(second, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) ( permuted_hidden_states_1, @@ -288,6 +322,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): ) +def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: + """Switching graph sizes must not expose metadata left by a larger batch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + device = torch.device("cuda:0") + scratch = get_moe_permute_scratch( + max_num_tokens=128, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + current_workspace_manager().lock() + runs = [] + for n_token in (128, 1, 33): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((graph, hidden, topk_ids, output)) + + for graph, hidden, topk_ids, output in runs + runs[::-1]: + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None: + """Concurrent graph replays in different ubatches/lanes cannot share scratch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + device = torch.device("cuda:0") + manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2) + monkeypatch.setattr(workspace, "_manager", manager) + ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) + scratch_config = dict( + max_num_tokens=64, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + runs = [] + for i, n_token in enumerate((1, 7, 33, 64)): + ubatch = i // 2 + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + scratch = get_moe_permute_scratch(**scratch_config) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((stream, graph, hidden, topk_ids, output, scratch)) + torch.cuda.current_stream().wait_stream(stream) + + assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 + manager.lock() + for _ in range(3): + for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): + ubatch = i // 2 + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + assert get_moe_permute_scratch(**scratch_config) is scratch + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + for stream, _, hidden, topk_ids, output, _ in runs: + torch.cuda.current_stream().wait_stream(stream) + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_without_manager(monkeypatch) -> None: + """Standalone calls get independent scratch with initialized row indices.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + monkeypatch.setattr(workspace, "_manager", None) + config = dict( + max_num_tokens=4, + topk=2, + num_experts=4, + num_local_experts=4, + device=torch.device("cuda"), + ) + first = get_moe_permute_scratch(**config) + second = get_moe_permute_scratch(**config) + + assert ( + first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() + ) + torch.testing.assert_close( + first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda") + ) + + def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py index 4c290dada6..29910390bb 100644 --- a/tests/kernels/quantization/test_marlin_gemm.py +++ b/tests/kernels/quantization/test_marlin_gemm.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_permute_bias, @@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): pytest.skip( @@ -477,6 +479,62 @@ def test_marlin_gemm( assert max_diff < 0.04 +def test_marlin_persistent_locks_cuda_graph_streams(workspace_init): + """Compiled GEMMs select stream-local locks at runtime and survive replay.""" + size_m, size_k, size_n = 32, 1024, 2048 + inputs = [rand_data((size_m, size_k)) for _ in range(2)] + weight = rand_data((size_k, size_n)) + quant_type = scalar_types.uint4b8 + w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128) + zeros = marlin_make_empty(scales.device) + streams = [torch.cuda.Stream() for _ in inputs] + graphs = [torch.cuda.CUDAGraph() for _ in inputs] + locks = [] + outputs = [] + + @torch.compile(backend="eager", fullgraph=True) + def run(x): + return torch.ops.vllm.marlin_gemm( + x, + qweight, + None, + scales, + None, + None, + zeros, + None, + quant_type.id, + size_m, + size_n, + size_k, + use_atomic_add=False, + use_fp32_reduce=True, + ) + + for stream, graph, x in zip(streams, graphs, inputs): + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + locks.append(get_marlin_workspace(x.device)) + run(x) + with torch.cuda.graph(graph, stream=stream): + outputs.append(run(x)) + torch.cuda.current_stream().wait_stream(stream) + + assert locks[0].data_ptr() != locks[1].data_ptr() + current_workspace_manager().lock() + for _ in range(3): + for stream, graph, x, lock in zip(streams, graphs, inputs, locks): + with torch.cuda.stream(stream): + assert get_marlin_workspace(x.device) is lock + x.add_(0.1) + graph.replay() + for stream in streams: + torch.cuda.current_stream().wait_stream(stream) + for x, output, lock in zip(inputs, outputs, locks): + assert compute_max_diff(output, x @ w_ref) < 0.04 + assert torch.count_nonzero(lock) == 0 + + def test_marlin_gemm_subset_input(): quant_type = scalar_types.uint4b8 group_size = 128 diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c..d224b0a2b2 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -14,6 +14,7 @@ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils.marlin_utils import ( GPTQ_MARLIN_TILE, apply_gptq_marlin_linear, + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_moe_padded_intermediate, @@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=layer.bias if use_bias else None, @@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale_inv, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=None, @@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -756,7 +757,7 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( @@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index afc0c1f836..37ba4b58c8 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot(): _assert_eq(r_ref, r_kern) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") -def test_amd_prefill_seed_honors_padded_tail_block_stride(): - """The tail shares a padded indexer allocation in production.""" +def test_prefill_seed_honors_padded_tail_block_stride(): + """The tail shares a padded indexer allocation in production. + + ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its + indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash + vs a dense 2048 B tail block), so a seed kernel that addresses blocks + densely writes into an unrelated indexer block and leaves the request's + tail block untouched. Runs on every platform; the NVIDIA kernel had this + bug while the AMD kernel did not. + """ kpool = 4 num_blocks = 6 logical_block_elems = 2 * kpool * HEAD_DIM diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index db49b8091b..b6c469181f 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -1871,17 +1871,32 @@ def test_sparse_indexer_topk_backend_resolution() -> None: has_coop = _has_cooperative_topk() has_fi = _has_flashinfer_topk() - # "auto" is exactly the pre-existing chain: cooperative -> persistent -> - # per_row. It must never select the opt-in backends by itself. + # "auto" is the chain cooperative -> persistent -> per_row, with the + # cooperative/persistent switch at cooperative_topk's cluster-wave row + # limit. It must never select the opt-in backends by itself. if has_coop: - assert resolve("auto") == "cooperative" + from vllm.model_executor.layers.indexer_topk import ( + AUTO_COOPERATIVE_MAX_ROWS, + ) + + # cooperative for every batch within its cluster-wave row limit... + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative" + ) assert resolve("auto", k=512, num_rows=8) == "cooperative" - # Past cooperative's 64-row limit -> persistent, even where - # DeepSelect would be applicable. + # ...persistent past the row limit, even where DeepSelect would be + # applicable. + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1) + == "persistent" + ) assert resolve("auto", k=512, num_rows=128) == "persistent" - # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: - # still cooperative at <= 64 rows... - assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative" + # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the + # same row-limit logic applies. + assert ( + resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS) + == "cooperative" + ) # ...and persistent past the row limit. assert resolve("auto", unaligned_logits, num_rows=128) == "persistent" # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index fde8b86cda..155957160f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -330,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer): ) -def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init): - """Marlin must reuse the workspace storage captured by CUDA graphs.""" +def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init): + """Weight reload must not create layer-owned Marlin lock storage.""" _stub_marlin_ops(monkeypatch) kernel = _make_marlin_kernel() @@ -339,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - workspace_ptr = kernel.workspace.data_ptr() + assert not hasattr(kernel, "workspace") _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - assert kernel.workspace.data_ptr() == workspace_ptr - assert torch.all(kernel.workspace == 0) + assert not hasattr(kernel, "workspace") @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) -def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): - """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold.""" +def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant): + """Weight preparation must not attach runtime workspace to model layers.""" from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -417,34 +415,57 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): load_checkpoint_format_weights() prepare(layer) - workspace_ptr = layer.workspace.data_ptr() + assert not hasattr(layer, "workspace") # Reload: fresh checkpoint-format tensors, prepare runs again load_checkpoint_format_weights() prepare(layer) - assert layer.workspace.data_ptr() == workspace_ptr - assert torch.all(layer.workspace == 0) + assert not hasattr(layer, "workspace") -def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): - """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that.""" +def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch): + """Calls reuse initialized locks; independent streams get separate storage.""" from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.utils import torch_utils + from vllm.v1.worker import workspace as workspace_module monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) device = torch.device("cpu") + manager = workspace_module.WorkspaceManager(device) + monkeypatch.setattr(workspace_module, "_manager", manager) + stream = "main" + monkeypatch.setattr(torch_utils, "current_stream", lambda: stream) - workspace = marlin_utils.marlin_make_workspace_new(device) - reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace) - assert reused is workspace + workspace = marlin_utils.get_marlin_workspace(device) + assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert workspace.dtype == torch.int32 + assert torch.count_nonzero(workspace) == 0 - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace) - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new( - device, existing=workspace.to(torch.int64) - ) + workspace.fill_(1) + assert marlin_utils.get_marlin_workspace(device) is workspace + assert torch.all(workspace == 1) + + stream = "aux" + aux_workspace = marlin_utils.get_marlin_workspace(device) + assert aux_workspace.data_ptr() != workspace.data_ptr() + assert torch.count_nonzero(aux_workspace) == 0 + + manager.lock() + assert marlin_utils.get_marlin_workspace(device) is aux_workspace + + +def test_marlin_workspace_without_manager(monkeypatch): + from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.v1.worker import workspace as workspace_module + + monkeypatch.setattr(workspace_module, "_manager", None) + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + first = marlin_utils.get_marlin_workspace(torch.device("cpu")) + second = marlin_utils.get_marlin_workspace(torch.device("cpu")) + assert first.data_ptr() != second.data_ptr() + assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert torch.count_nonzero(first) == 0 def test_model_cleanup(dist_init, default_vllm_config): diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py index 83dfdc9f6f..9f94da459c 100644 --- a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Focused unit tests for the GLM decode top-k dispatch.""" +"""Tests for the GLM kpool indexer's top-k backend wiring.""" import pytest import torch -from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk +from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -13,29 +13,31 @@ pytestmark = pytest.mark.skipif( ) -@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"]) -@torch.inference_mode() -def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int): - logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda") - assert _use_cooperative_topk(logits, 512, num_rows) +def _require_deep_gemm() -> None: + from vllm.utils.deep_gemm import has_deep_gemm + + if not has_deep_gemm(): + pytest.skip("kpool indexer requires DeepGEMM") -@pytest.mark.parametrize( - ("num_rows", "select_k", "stride_padding"), - [ - (65, 512, 0), - (1, 256, 0), - (1, 512, 1), - ], - ids=["too_many_rows", "unsupported_k", "unaligned_stride"], -) -@torch.inference_mode() -def test_cooperative_topk_falls_back_on_unsupported_batch( - num_rows: int, select_k: int, stride_padding: int -): - logits = torch.empty( - num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda" - ) - if stride_padding: - logits = logits[:, :8192] - assert not _use_cooperative_topk(logits, select_k, num_rows) +@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: + """The kpool indexer must read kernel_config.sparse_indexer_topk_backend + and hand it to the shared SparseIndexerTopk dispatcher, rather than + hard-coding its own cooperative/persistent/per_row choice.""" + _require_deep_gemm() + from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool + + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + op = SparseAttnIndexerKpool( + k_cache=None, + quant_block_size=128, + scale_fmt="ue8m0", + topk_tokens=2048, + head_dim=128, + max_pool_len=4096, + max_total_seq_len=8192, + topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"), + ) + assert op.topk_backend == backend diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 1ed84ebc52..9bda7ec85d 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str: ) +def deepseek_v41_chat_template(content: str) -> str: + return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜>" + + def step3_vl_chat_template(content: str) -> str: return ( "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n " @@ -221,6 +225,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { vllm_runner_kwargs={"gpu_memory_utilization": 0.80}, marks=[pytest.mark.core_model], ), + "deepseek_v41": VitCudagraphTestConfig( + model="deepseek-ai/DeepSeek-V4.1-Flash", + modalities=["image"], + image_prompt=deepseek_v41_chat_template( + "<｜deepseek_image｜>\n\nWhat is in this image?" + ), + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1024], + }, + vllm_runner_kwargs={ + "load_format": "dummy", + "attention_backend": "FLASHMLA_SPARSE_DSV41", + "hf_overrides": partial( + dummy_hf_overrides, + model_arch="DeepseekV41ForCausalLM", + exist_overrides={"vision_n_layers": 1}, + ), + }, + ), "step3_vl": VitCudagraphTestConfig( model="stepfun-ai/Step3-VL-10B", modalities=["image"], diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 403a01a9e2..1e1722537d 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -218,8 +218,14 @@ def test_model_tensor_schema(model_id: str): factories = model_cls._processor_factory + # Capture helpers return capture-buffer containers (e.g. + # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs. + capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"} + inputs_parse_methods = [] for attr_name in dir(model_cls): + if attr_name in capture_helpers: + continue attr = getattr(model_cls, attr_name) if hasattr(attr, "__annotations__"): return_type = attr.__annotations__.get("return", None) diff --git a/tests/models/test_deepseek_v4_rocm_wo_a.py b/tests/models/test_deepseek_v4_rocm_wo_a.py new file mode 100644 index 0000000000..69c6ed9530 --- /dev/null +++ b/tests/models/test_deepseek_v4_rocm_wo_a.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0 + + +def test_wo_a_block_scale_to_e8m0_from_float(): + scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32) + + encoded = _wo_a_block_scale_to_e8m0(scale) + + assert encoded is not None + torch.testing.assert_close( + encoded, + torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8), + ) + assert encoded.is_contiguous() + + +def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales(): + raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8) + encoded = raw.view(torch.float8_e8m0fnu) + + converted = _wo_a_block_scale_to_e8m0(encoded) + + assert converted is not None + torch.testing.assert_close(converted, raw) + + +@pytest.mark.parametrize( + "scale", + [ + torch.tensor([[0.0, 1.0]]), + torch.tensor([[-1.0, 1.0]]), + torch.tensor([[0.75, 1.0]]), + torch.tensor([[float("inf"), 1.0]]), + torch.ones(1, dtype=torch.int32), + ], +) +def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor): + assert _wo_a_block_scale_to_e8m0(scale) is None diff --git a/tests/renderers/test_truncate_after_mm_expansion.py b/tests/renderers/test_truncate_after_mm_expansion.py new file mode 100644 index 0000000000..96cdfb12c2 --- /dev/null +++ b/tests/renderers/test_truncate_after_mm_expansion.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Truncation must bound the prompt *after* multimodal placeholders expand. + +``TokenizeParams.apply_post_tokenization`` truncates the unexpanded prompt, so +one placeholder token can expand into many and push the result back over +``truncate_prompt_tokens`` (vllm-project/vllm#57038). +""" + +import numpy as np +import pytest +import torch + +from vllm.config import ModelConfig, VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.multimodal.inputs import PlaceholderRange +from vllm.renderers import renderer_from_config +from vllm.renderers.base import BaseRenderer +from vllm.renderers.hf import _ensure_prompt_embeds_placeholder_token +from vllm.renderers.params import TokenizeParams + + +def _renderer(): + """A BaseRenderer stub: only the tokenizer attribute is needed here.""" + + class _StubRenderer(BaseRenderer): + def __init__(self): + self.tokenizer = None + + def render_messages(self, *args, **kwargs): + raise NotImplementedError + + return _StubRenderer() + + +def _mm_input(num_tokens: int, offset: int, length: int, **extra): + return { + "type": "multimodal", + "prompt_token_ids": list(range(num_tokens)), + "mm_kwargs": {}, + "mm_hashes": {}, + "mm_placeholders": {"audio": [PlaceholderRange(offset=offset, length=length)]}, + **extra, + } + + +def _params(limit: int | None, side: str = "right"): + return TokenizeParams( + max_total_tokens=4096, + truncate_prompt_tokens=limit, + truncation_side=side, + ) + + +def test_expanded_prompt_is_bounded(): + """The reported case: audio fits, trailing text is cut to the limit.""" + engine_input = _mm_input(num_tokens=147, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(96)) + assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange( + offset=1, length=52 + ) + + +def test_prompt_within_limit_is_untouched(): + engine_input = _mm_input(num_tokens=64, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(64)) + + +def test_split_multimodal_span_is_rejected(): + """Cutting mid-span would desync features from token ids, so reject.""" + engine_input = _mm_input(num_tokens=147, offset=80, length=50) + + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + +def test_dropped_multimodal_span_is_rejected(): + """A span entirely past the cut would be discarded silently.""" + engine_input = _mm_input(num_tokens=147, offset=100, length=20) + + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + +def test_left_truncation_shifts_placeholder_offsets(): + engine_input = _mm_input(num_tokens=147, offset=100, length=20) + + _renderer()._truncate_expanded_prompt(engine_input, _params(96, side="left")) + + assert engine_input["prompt_token_ids"] == list(range(51, 147)) + # The span started at 100 and 51 leading tokens were dropped. + assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange( + offset=49, length=20 + ) + + +def test_text_only_input_is_untouched(): + """Only multimodal prompts grow after tokenization, so nothing else runs.""" + engine_input = { + "type": "token", + "prompt_token_ids": list(range(147)), + "prompt_token_offsets": [(i, i + 1) for i in range(147)], + } + + _renderer()._truncate_expanded_prompt(engine_input, _params(96)) + + assert engine_input["prompt_token_ids"] == list(range(147)) + assert len(engine_input["prompt_token_offsets"]) == 147 + + +def test_no_truncation_requested_is_noop(): + engine_input = _mm_input(num_tokens=147, offset=1, length=52) + + _renderer()._truncate_expanded_prompt(engine_input, _params(None)) + + assert engine_input["prompt_token_ids"] == list(range(147)) + + +def test_mixed_prompt_embeds_truncation_covers_embeds_spans(): + """Bound in _process_singleton, where prompt_embeds spans are already attached.""" + model_config = ModelConfig( + "Qwen/Qwen3-ASR-0.6B-hf", + revision="7f1569a48a89f3e3f4dc3a5c9d28bddd903bc76c", + dtype="float32", + max_model_len=4096, + mm_processor_cache_gb=0, + enable_prompt_embeds=True, + ) + renderer = renderer_from_config(VllmConfig(model_config=model_config)) + tok = renderer.tokenizer + pe_id = _ensure_prompt_embeds_placeholder_token(tok) + hidden = model_config.get_hidden_size() + audio = np.zeros(4 * 16000, dtype=np.float32) + + pe_len = 20 + tensors = [torch.zeros(pe_len, hidden, dtype=torch.float32)] + + def _build_prompt(): + text = "<|audio_start|><|audio_pad|><|audio_end|>" + ids = tok.encode(text) + [pe_id] + tok.encode(" tail" * 40) + return { + "prompt_token_ids": ids, + "multi_modal_data": {"audio": (audio, 16000)}, + "_prompt_embeds": (tensors, pe_id), + } + + # A limit past every span (audio ends at 52, prompt_embeds at 74): + # truncation should succeed and bound the prompt. + out = renderer._process_singleton( + _build_prompt(), + tok_params=_params(79), + ) + assert len(out["prompt_token_ids"]) == 79 + + # A limit inside the prompt_embeds span (54..74): must be rejected. + with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"): + renderer._process_singleton( + _build_prompt(), + tok_params=_params(64), + ) diff --git a/tests/utils.py b/tests/utils.py index 81f82d98c8..7ee0e66803 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1427,8 +1427,10 @@ def init_test_distributed_environment( get_current_vllm_config_or_none, set_current_vllm_config, ) + from vllm.platforms import current_platform distributed_init_method = f"tcp://localhost:{distributed_init_port}" + backend = current_platform.dist_backend if data_parallel_size > 1: # For DP we need to set a common DP master port @@ -1449,6 +1451,7 @@ def init_test_distributed_environment( rank=rank % tp_pp_world, distributed_init_method=distributed_init_method, local_rank=local_rank if local_rank >= 0 else rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) return @@ -1460,6 +1463,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) else: @@ -1470,6 +1474,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index cb39e54037..ed0e250b18 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -391,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular(): # The circular mapping keeps the rings isolated under interleaving. torch.testing.assert_close(circular, ground_truth) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize( + "per_req,num_actual,padded_len", + [ + ([list(range(10)), list(range(12))], 22, 22), + ([list(range(10)), list(range(12))], 22, 30), + ([[3, 4], [0], [7, 8, 9]], 6, 8), + ([[5]], 1, 1), + ], +) +def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): + """The CUDA (Triton) path must match the CPU torch reference, including + tokens between the last request boundary and num_actual_tokens (mapped to + the last request) and untouched padding beyond num_actual.""" + positions, qsl, slot_mapping, _, num_reqs = make_batch( + per_req, padded_len=padded_len + ) + # Replace the all--1 placeholder slots with sentinel values to check the + # padding range is copied through untouched. + slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 + bt = make_tail_block_table(list(range(5, 5 + num_reqs))) + + ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) + got = circular_tail_slots( + slot_mapping.cuda(), + bt.cuda(), + qsl.cuda().to(torch.int32), + positions.cuda(), + num_actual, + num_reqs, + ) + torch.testing.assert_close(got.cpu(), ref) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_triton_mapping_reads_strided_block_table(): + """The kernel must address the block table through its real stride(0), + not a dense assumption (the #57477 class of bug): a tail block table that + is a column view of a wider allocation must still read column 0.""" + per_req = [list(range(10)), list(range(12))] + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) + own_blocks = [11, 22] + wide = torch.zeros(num_reqs, 8, dtype=torch.int32) + bt = wide[:, 1:4] # non-contiguous view, stride(0) == 8 + bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) + + ref = circular_tail_slots( + slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs + ) + got = circular_tail_slots( + slot_mapping.cuda(), + bt.cuda(), + qsl.cuda().to(torch.int32), + positions.cuda(), + num_actual, + num_reqs, + ) + torch.testing.assert_close(got.cpu(), ref) + for req, blk in enumerate(own_blocks): + start, end = int(qsl[req]), int(qsl[req + 1]) + assert (got[start:end] // KPOOL == blk).all() diff --git a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py index c75591b697..ff62abab8c 100644 --- a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py +++ b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py @@ -4,6 +4,7 @@ import asyncio import json +from math import lcm from uuid import uuid4 import httpx @@ -12,6 +13,7 @@ from openai.types.chat import ChatCompletionMessageParam from prometheus_client.parser import text_string_to_metric_families from tests.utils import RemoteOpenAIServer +from vllm.utils.math_utils import round_down MODEL = "Qwen/Qwen3.5-0.8B" NUM_CONVERSATIONS = 8 @@ -34,6 +36,29 @@ def _user_turns(conversation: int) -> list[str]: ] +async def _scheduler_block_size(client: httpx.AsyncClient) -> int: + """Return the coarsest boundary an expected cache hit can land on. + + Mirrors ``scheduler_block_size`` in ``resolve_kv_cache_block_sizes``, which + the engine keeps internal: the LCM of the resolved group block sizes. A + finer ``prefix_match_unit`` only adds hits, so this stays conservative. + """ + response = await client.get("/metrics") + response.raise_for_status() + for family in text_string_to_metric_families(response.text): + for sample in family.samples: + if sample.name != "vllm:cache_config_info": + continue + sizes = [ + int(value) + for name in ("block_size", "mamba_block_size") + if (value := sample.labels.get(name, "None")) != "None" + ] + assert sizes, f"no block size in {sample.labels}" + return lcm(*sizes) + raise AssertionError("missing vllm:cache_config_info metric") + + async def _prefix_cache_counters( client: httpx.AsyncClient, ) -> tuple[float, float]: @@ -83,13 +108,13 @@ def server(request): @pytest.mark.asyncio @pytest.mark.parametrize( - "server, min_hit_rate", - # MTP re-prefills an extra block; both floors reject one more missed block. - [pytest.param(False, 0.90, id="base"), pytest.param(True, 0.75, id="mtp")], + "server, dropped_blocks", + # MTP's prefill lookahead pollutes the tail block, so EAGLE drops one. + [pytest.param(False, 0, id="base"), pytest.param(True, 1, id="mtp")], indirect=["server"], ) async def test_prefix_cache_hit_rate( - server: RemoteOpenAIServer, min_hit_rate: float + server: RemoteOpenAIServer, dropped_blocks: int ) -> None: # Isolate sessions so another conversation cannot hide a cache miss. salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)] @@ -102,6 +127,8 @@ async def test_prefix_cache_hit_rate( server.get_async_client() as client, httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics, ): + block_size = await _scheduler_block_size(metrics) + print(f"{block_size=}") async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]: history = histories[session] @@ -126,15 +153,30 @@ async def test_prefix_cache_hit_rate( assert cached is not None, context assert 0 <= cached < prompt, context assert prompt >= MIN_PROMPT_TOKENS, context - if not replay and turn == 0: - assert cached == 0, context - else: - assert cached / prompt >= min_hit_rate, ( - f"{context}, {min_hit_rate=}" - ) if replay: assert prompt == prompt_lengths[session][turn], context + # The identical request computed this whole prompt, but at + # least one token is always recomputed. + computed = prompt - 1 + elif turn: + computed = prompt_lengths[session][turn - 1] else: + assert cached == 0, context + computed = 0 + # Every token below the last reachable boundary must be reused. + floor = max( + 0, round_down(computed, block_size) - dropped_blocks * block_size + ) + print( + f"{session=} {turn=} {replay=} {prompt=} {cached=} {floor=} " + f"hit_rate={cached / prompt:.4f}" + ) + assert cached >= floor, ( + f"{context}, {computed=}, {block_size=}, " + f"hit_rate={cached / prompt:.4f}, " + f"expected_hit_rate={floor / prompt:.4f}" + ) + if not replay: prompt_lengths[session].append(prompt) history.append( { diff --git a/tests/v1/ec_connector/unit/test_epd_proxy.py b/tests/v1/ec_connector/unit/test_epd_proxy.py new file mode 100644 index 0000000000..13496a7ffb --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_proxy.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routing and registration behaviour of the EPD proxy.""" + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, Mock + +import msgspec +import pytest +from fastapi import HTTPException +from fastapi.testclient import TestClient + +from examples.disaggregated.disaggregated_encoder import ( + disagg_epd_proxy as proxy_module, +) +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + EPDProxy, + EPDProxyConfig, + InstanceRecord, + InstanceRegistry, + InstanceRole, + build_app, + extract_mm_items, +) + +ENCODE = InstanceRole.ENCODE +PREFILL = InstanceRole.PREFILL +DECODE = InstanceRole.DECODE +PD = InstanceRole.PREFILL_DECODE + +IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}} + + +@pytest.fixture +def proxy(): + registry = InstanceRegistry(probe_interval=0) + return EPDProxy(registry) + + +class TestRouting: + def test_no_decode_instance_is_service_unavailable(self, proxy): + """The proxy comes up before anything registers, so this is normal.""" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + def test_media_without_an_encoder_is_service_unavailable(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).decode.url == "http://d0:8000" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=1) + assert excinfo.value.status_code == 503 + + def test_prefill_is_optional(self, proxy): + """An E+PD deployment registers no prefill instance at all.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).prefill is None + + def test_encoder_roster_tracks_new_instances(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + for index in range(2): + proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000")) + assert proxy.route(num_items=3).encoder_urls == [ + "http://e0:8000", + "http://e1:8000", + ] + proxy.registry.unregister("http://e0:8000") + assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"] + + @pytest.mark.parametrize("unregister", [False, True]) + def test_standalone_decode_always_requires_healthy_prefill(self, unregister): + app = build_app() + proxy = app.state.proxy + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000")) + route = proxy.route(num_items=0) + assert route.decode.url == "http://d0:8000" + assert route.consumer.url == "http://p0:8000" + if unregister: + proxy.registry.unregister("http://p0:8000") + else: + for _ in range(proxy.registry._fail_threshold): + proxy.registry._on_probe_failure(route.prefill, 0) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + +class TestConsumerAddress: + """Which stage receives the embedding depends on the topology.""" + + def test_shared_storage_connectors_name_no_target(self, proxy): + """Rotate replicas without requiring a push endpoint.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2)) + route = proxy.route(num_items=0) + assert route.consumer_zmq is None + assert route.dp_rank == 0 + assert proxy.route(num_items=0).dp_rank == 1 + + def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy): + proxy.registry.register( + InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"]) + ) + assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001" + + def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy): + """In E+P+D the prefill instance consumes the embedding, not decode. + + Reading the address out of whichever list it came from -- rather than + from the record that reported one -- sends the push to the wrong + instance in this topology. + """ + proxy.registry.register( + InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"]) + ) + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001" + + +class TestProxyApi: + @pytest.fixture + def client(self, monkeypatch): + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + config = EPDProxyConfig(probe_interval=0) + with TestClient(build_app(config)) as client: + client.headers["X-API-Key"] = "test-key" + yield client + + def test_proxy_serves_before_anything_registers(self, client): + assert client.get("/health").status_code == 200 + assert client.get("/instances").json()["encode"]["live"] == [] + + def test_requests_are_refused_until_a_decoder_registers(self, client): + response = client.post( + "/v1/chat/completions", json={"model": "m", "messages": []} + ) + assert response.status_code == 503 + assert client.get("/v1/models").status_code == 503 + + def test_http_registration_keeps_consumer_replica_and_address_together( + self, client + ): + payload: dict[str, Any] = { + "role": "prefill_decode", + "url": "http://d0:8000/", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"], + } + assert client.post("/instances", json=payload).json() == {"registered": True} + assert client.post("/instances", json=payload).json() == {"registered": False} + proxy = client.app.state.proxy + for rank in [0, 1, 0]: + route = proxy.route(0) + assert route.dp_rank == rank + assert route.consumer_zmq == payload["ec_zmq_addrs"][rank] + assert client.delete("/instances", params={"url": payload["url"]}).json() == { + "removed": True + } + assert proxy.registry.pick(PD) is None + + def test_registration_rejects_invalid_topology(self, client): + cases: list[tuple[dict[str, Any], int]] = [ + ({"dp_size": 0}, 422), + ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422), + ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ] + for fields, status in cases: + response = client.post( + "/instances", + json={"url": "http://d0:8000", "role": "prefill_decode", **fields}, + ) + assert response.status_code == status + + def test_registration_requires_admin_key(self, client): + client.headers.pop("X-API-Key") + assert ( + client.post( + "/instances", json={"role": "encode", "url": "http://e0:8000"} + ).status_code + == 403 + ) + assert ( + client.delete("/instances", params={"url": "http://e0:8000"}).status_code + == 403 + ) + + @pytest.mark.parametrize("split_role", ["prefill", "decode"]) + @pytest.mark.parametrize("combined_first", [False, True]) + def test_registration_rejects_mixed_topologies( + self, client, split_role, combined_first + ): + roles = ["prefill_decode", split_role] + if not combined_first: + roles.reverse() + first = {"role": roles[0], "url": "http://first:8000"} + second = {"role": roles[1], "url": "http://second:8000"} + assert client.post("/instances", json=first).status_code == 200 + assert client.post("/instances", json=second).status_code == 409 + # Health eviction must not silently permit a topology switch. + registry = client.app.state.registry + record = registry._live[first["url"]] + for _ in range(registry._fail_threshold): + registry._on_probe_failure(record, 0) + assert client.post("/instances", json=second).status_code == 409 + assert client.delete("/instances", params={"url": first["url"]}).json() == { + "removed": True + } + assert client.post("/instances", json=second).status_code == 200 + + +def test_extract_mm_items_finds_media_across_messages(): + req = { + "messages": [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]}, + {"role": "user", "content": [IMAGE_ITEM]}, + ] + } + assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM] + + +@pytest.mark.parametrize("push", [False, True]) +def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch): + """Keep main's hash-keyed handles and per-item push IDs after the proxy move.""" + proxy.registry.register(InstanceRecord(ENCODE, "http://encoder")) + proxy.registry.register( + InstanceRecord( + PD, + "http://decode", + ec_zmq_addrs=["tcp://decode:14579"] if push else [], + ) + ) + handle = { + "metadata": {"image_grid_thw": [[1, 2, 3]]}, + "item_indices": [0, 1], + "peer_port": 5601, + } + response = Mock( + status=200, + read=AsyncMock( + return_value=msgspec.json.encode( + {"ec_transfer_params": {"engine-hash": handle}} + ) + ), + ) + session = Mock(post=AsyncMock(return_value=response)) + monkeypatch.setattr(proxy_module, "encode_session", session) + original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]} + route = proxy.route(2) + prepared, _, _ = asyncio.run( + proxy_module.prepare_for_decode( + original, "request", route.encoder_urls, None, route.consumer_zmq + ) + ) + assert "ec_transfer_params" not in original + assert prepared["ec_transfer_params"]["engine-hash"] == handle + items = prepared["ec_transfer_params"]["ec_items"] + assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"] + assert items[0]["transfer_id"] != items[1]["transfer_id"] + assert session.post.call_count == 1 + sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get( + "ec_transfer_params" + ) + if push: + assert sent == { + "consumer_zmq": "tcp://decode:14579", + "ec_items": [{"transfer_id": item["transfer_id"]} for item in items], + } + else: + assert sent is None + for index in range(2): + assert prepared["messages"][0]["content"][index]["image_embeds"] == { + "image_grid_thw": [1, 2, 3] + } diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 0c48dc84d6..8db22fe1b9 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -8,6 +8,7 @@ change to them is what these tests catch. import asyncio import importlib.util +import sys from pathlib import Path from types import SimpleNamespace from typing import Any @@ -27,6 +28,7 @@ def proxy(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module @@ -48,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch): """ served = [{"remote_block_ids": [1, 2]}, {}] - async def _stage(req_data, p_url, req_id): + async def _stage(req_data, p_url, req_id, dp_rank=None): assert "kv_transfer_params" not in req_data return _Response(served.pop(0)) @@ -222,7 +224,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch _EncoderSession([{"encoder-side-hash": handle}, {}]), ) - async def _no_prefill(req_data, p_url, req_id): + async def _no_prefill(req_data, p_url, req_id, dp_rank=None): return req_data monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill) @@ -280,6 +282,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "request_options", [ @@ -293,7 +296,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): ], ) async def test_http_roundtrip_preserves_payload_and_response_bytes( - proxy, monkeypatch, stream, prefill, request_options + proxy, monkeypatch, stream, prefill, dynamic, request_options ): """Preserve media and scheduling context across HTTP hops and retries.""" seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} @@ -307,6 +310,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( async def backend(request): stage = request.match_info["stage"] + if dynamic: + consumer = "prefill" if prefill else "decode" + expected_rank = "0" if stage == consumer else None + assert request.headers.get("X-data-parallel-rank") == expected_rank assert request.content_type == "application/json" body = await request.json() seen[stage].append(body) @@ -354,6 +361,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False) monkeypatch.setattr(proxy, "DECODE_RETRIES", 1) monkeypatch.setattr(proxy, "NO_REWRITE", False) + target_app = proxy.app + if dynamic: + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0)) item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}} body = { "model": "test", @@ -377,9 +388,37 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( await proxy.on_startup() try: async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy" + transport=httpx.ASGITransport(app=target_app), base_url="http://proxy" ) as client: + if dynamic: + registrations = [("encode", "encode")] + registrations += ( + [("prefill", "prefill"), ("decode", "decode")] + if prefill + else [("prefill_decode", "decode")] + ) + for role, stage in registrations: + consumer = "prefill" if prefill else "decode" + registered = await client.post( + "/instances", + headers={"X-API-Key": "test-key"}, + json={ + "role": role, + "url": str(server.make_url(f"/{stage}")), + "dp_size": 2 if stage == consumer else 1, + }, + ) + assert registered.status_code == 200 response = await client.post("/v1/chat/completions", json=body) + if dynamic: + removed = await client.delete( + "/instances", + headers={"X-API-Key": "test-key"}, + params={"url": str(server.make_url("/decode"))}, + ) + assert removed.json() == {"removed": True} + unavailable = await client.post("/v1/chat/completions", json=body) + assert unavailable.status_code == 503 assert response.status_code == 200 assert response.content == expected assert response.headers["content-type"].startswith( diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py index 558a80f9cf..8e1a730ca6 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -14,6 +14,7 @@ across calls so consecutive requests keep advancing through every URL. """ import importlib.util +import sys from collections import Counter from pathlib import Path @@ -27,6 +28,7 @@ def _load_proxy_module(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module diff --git a/tests/v1/ec_connector/unit/test_epd_registry.py b/tests/v1/ec_connector/unit/test_epd_registry.py new file mode 100644 index 0000000000..be5a8e9b9b --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_registry.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dynamic registration and liveness for the EPD proxy's instance registry.""" + +from unittest.mock import patch + +import pytest + +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + InstanceRecord, + InstanceRegistry, + InstanceRole, +) + +ENCODE = InstanceRole.ENCODE +DECODE = InstanceRole.DECODE + + +@pytest.fixture +def registry(): + return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0) + + +def _fill(registry, role, count, prefix="http://h"): + for index in range(count): + registry.register(InstanceRecord(role, f"{prefix}{index}:8000")) + + +async def _probe_round(registry, healthy: set[str], now: float = 0.0): + """Run one probe pass where only `healthy` URLs answer.""" + + async def fake_probe(self, session, url): + return url in healthy + + with ( + patch.object(InstanceRegistry, "_probe", fake_probe), + patch( + "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic", + return_value=now, + ), + ): + await registry._probe_once(session=None) + + +class TestRegistration: + def test_registering_is_idempotent_and_reports_novelty(self, registry): + record = InstanceRecord(ENCODE, "http://e0:8000") + assert registry.register(record) is True + assert registry.register(record) is False + assert registry.urls(ENCODE) == ["http://e0:8000"] + + def test_transfer_addresses_stay_with_the_registered_consumer(self, registry): + """Roster updates must not shift a consumer's transfer addresses.""" + registry.register(InstanceRecord(DECODE, "http://d0:8000")) + registry.register( + InstanceRecord( + DECODE, + "http://d0:8000", + ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"], + dp_size=2, + ) + ) + picked = registry.pick(DECODE) + assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"] + assert picked.dp_size == 2 + + def test_unregister_is_terminal(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + assert registry.unregister("http://e0:8000") is True + assert registry.unregister("http://e0:8000") is False + assert registry.pick(ENCODE) is None + + def test_round_robin_survives_registration(self, registry): + _fill(registry, ENCODE, 2) + assert [record.url for record in registry.pick_many(ENCODE, 2)] == [ + "http://h0:8000", + "http://h1:8000", + ] + registry.register(InstanceRecord(ENCODE, "http://h2:8000")) + assert registry.pick(ENCODE).url == "http://h2:8000" + + +class TestLiveness: + @pytest.mark.asyncio + async def test_identical_reregistration_preserves_inflight_probe_results( + self, registry + ): + """Repeated registrations must not suppress eviction or recovery.""" + url = "http://e0:8000" + registry.register(InstanceRecord(ENCODE, url)) + healthy = False + + async def probe(self, session, url): + registry.register(InstanceRecord(ENCODE, url)) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + for _ in range(registry._fail_threshold): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": [url]} + + healthy = True + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [url], "evicted": []} + + @pytest.mark.asyncio + async def test_consecutive_failures_stop_routing(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + registry.register(InstanceRecord(ENCODE, "http://e1:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy={"http://e1:8000"}) + assert registry.urls(ENCODE) == ["http://e1:8000"] + assert registry.status()["encode"]["evicted"] == ["http://e0:8000"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", [ENCODE, DECODE]) + async def test_a_recovered_instance_rejoins_without_re_registering( + self, registry, role + ): + """Restarting every encoder to recover from a blip is not acceptable.""" + registry.register(InstanceRecord(role, "http://e0:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + assert registry.urls(role) == [] + + await _probe_round(registry, healthy={"http://e0:8000"}, now=1) + assert registry.urls(role) == ["http://e0:8000"] + + @pytest.mark.asyncio + async def test_re_registration_does_not_override_failed_health_checks( + self, registry + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + registry.register(record) + assert registry.urls(ENCODE) == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("healthy", [False, True]) + async def test_inflight_probe_cannot_restore_unregistered_instance( + self, registry, healthy + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + + async def probe(self, session, url): + registry.unregister(url) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": []} diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index f21dae7e0b..68121f17e0 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -19,8 +19,10 @@ from tests.v1.simple_kv_offload.test_scheduler import ( SchedulerFixture, _alloc_and_register, _allocate_cp_gpu_blocks, + _allocate_gpu_blocks, _make_cp_request, _make_kv_cache_config, + _make_scratch_kv_cache_config, _make_vllm_config, make_request, make_scheduler_output, @@ -757,3 +759,37 @@ def test_mamba_align_skips_positional_event_metadata() -> None: get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0)) ) assert ev.block_hashes == [expected_hash] + + +def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None: + """Event metadata is resolved only for prefix-cacheable groups. + + A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not + divide the hash block size, so viewing the request hashes at its block + size is impossible; the store must skip it and still emit BlockStored + for the attention blocks. + """ + fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16)) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + events = list(sched.take_events()) + assert len(events) == num_blocks + assert all(isinstance(e, BlockStored) for e in events) + assert {e.group_idx for e in events} == {0} + assert all(len(e.token_ids) == BLOCK_SIZE for e in events) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index ffc8b20478..ed07718de6 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -41,7 +41,9 @@ from vllm.v1.core.single_type_kv_cache_manager import ( register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, + KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -125,6 +127,45 @@ def _make_kv_cache_config( ) +def _make_scratch_kv_cache_config( + num_blocks: int, scratch_block_size: int = 4 +) -> KVCacheConfig: + """FullAttention group plus a non-prefix-cacheable scratch group. + + Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request + scratch block of ``index_kpool`` tokens: it is excluded from prefix + caching, so its block size neither divides nor is divided by the + hash block size in general. + """ + register_all_kvcache_specs(vllm_config=None) + fa_config = _make_kv_cache_config(num_blocks, num_groups=1) + scratch_layers = ["layer_scratch"] + scratch_spec = KpoolTailSpec( + block_size=scratch_block_size, + num_kv_heads=2, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + sliding_window=scratch_block_size, + ) + assert not scratch_spec.prefix_cacheable + scratch_bytes = scratch_spec.page_size_bytes * num_blocks + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=fa_config.kv_cache_tensors + + [ + KVCacheTensor( + size=scratch_bytes, + layers=scratch_layers, + layer_stride=scratch_bytes, + block_stride=scratch_spec.page_size_bytes, + ) + ], + kv_cache_groups=fa_config.kv_cache_groups + + [KVCacheGroupSpec(scratch_layers, scratch_spec)], + ) + + def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig: """Minimal VllmConfig for scheduler tests (no GPU).""" model_config = ModelConfig( @@ -174,9 +215,13 @@ def make_scheduler( num_gpu_blocks: int = 16, num_groups: int = 1, lazy: bool = False, + kv_cache_config: KVCacheConfig | None = None, ) -> SchedulerFixture: """Build a SimpleCPUOffloadScheduler with small block pools.""" - kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + if kv_cache_config is None: + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + else: + num_groups = len(kv_cache_config.kv_cache_groups) vllm_config = _make_vllm_config() cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups @@ -1392,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None: ) +# --------------------------------------------------------------------------- +# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load +# --------------------------------------------------------------------------- +def test_scratch_group_excluded_from_store_and_load() -> None: + """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded. + + Its single per-request block carries no hash, so the store must offload + only the attention blocks and the load must pair only attention blocks. + """ + fix = make_scheduler( + num_cpu_blocks=8, + num_gpu_blocks=16, + kv_cache_config=_make_scratch_kv_cache_config(16), + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-scratch-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + scratch_block2 = gpu_pool.get_new_blocks(1) + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2)) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + sched_out2 = make_scheduler_output( + {req2.request_id: 1}, + new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, + ) + meta2 = sched.build_connector_meta(sched_out2) + assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2) + assert len(meta2.load_cpu_blocks) == num_blocks + + # --------------------------------------------------------------------------- # Test 9: Chunked prefill accumulates block_ids across steps # --------------------------------------------------------------------------- @@ -2844,3 +2946,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True + + +QSA_RING_BLOCK_SIZE = 4 + + +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: + register_all_kvcache_specs(vllm_config=None) + full_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring_spec = CircularBufferSpec( + block_size=QSA_RING_BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + ) + groups = [ + KVCacheGroupSpec(["layer_0"], full_spec), + KVCacheGroupSpec(["qsa_ring"], ring_spec), + ] + tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, + ) + for group in groups + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _make_qsa_scheduler( + lazy: bool = False, +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + sched = SimpleCPUOffloadScheduler( + vllm_config=_make_vllm_config(), + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + lazy_offload=lazy, + ) + gpu_pool = BlockPool( + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE + ) + sched.bind_gpu_block_pool(gpu_pool) + return sched, gpu_pool + + +def test_qsa_ring_group_is_never_stored_or_loaded() -> None: + sched, gpu_pool = _make_qsa_scheduler() + assert sched.prefix_cacheable_group_ids == (0,) + + num_blocks = 2 + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + ring_block = gpu_pool.get_new_blocks(1)[0] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} + assert ring_block.block_id not in meta.store_gpu_blocks + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-qsa-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + ring_block2 = gpu_pool.get_new_blocks(1)[0] + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} + assert ring_block2.block_id not in meta2.load_gpu_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + without_ring = KVCacheConfig( + num_blocks=16, + kv_cache_tensors=with_ring.kv_cache_tensors[:1], + kv_cache_groups=with_ring.kv_cache_groups[:1], + ) + max_batched = 64 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + with_ring, max_batched + ) + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + without_ring, max_batched + ) + assert target_with == target_without diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 7824c8000d..3d03f25e0e 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast +from unittest.mock import Mock import pytest import torch @@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None: with pytest.raises(ValueError, match="at least one"): workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) + + +def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None: + """Only the requesting slot allocates; each ubatch/lane owns its resource.""" + active_ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)}) + resources = [] + for ubatch in range(2): + active_ubatch = ubatch + for lane in range(2): + with workspace.use_workspace_lane(lane): + resource = manager.get_persistent_resource("scratch", factory) + resources.append(resource) + resource["buffer"].fill_(len(resources)) + assert manager.get_persistent_resource("scratch", factory) is resource + assert factory.call_count == len(resources) + + assert len({r["buffer"].data_ptr() for r in resources}) == 4 + for i, resource in enumerate(resources, 1): + assert torch.all(resource["buffer"] == i) + + +def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None: + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + factory = Mock(side_effect=object) + first = manager.get_persistent_resource("scratch", factory) + manager.lock() + assert manager.get_persistent_resource("scratch", factory) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("new", factory) + with workspace.use_workspace_lane(1): + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("scratch", factory) + assert factory.call_count == 1 + manager.unlock() + assert manager.get_persistent_resource("scratch", factory) is not first + assert manager.get_persistent_resource("scratch", factory) is first + + +def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None: + """Initialization happens once and transient scratch cannot overwrite locks.""" + manager = workspace.WorkspaceManager(torch.device("cpu")) + first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True) + assert torch.count_nonzero(first) == 0 + first.fill_(1) + (scratch,) = manager.get_simultaneous(((1024,), torch.uint8)) + scratch.zero_() + assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first + assert torch.all(first == 1) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (16,), torch.int32) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (8,), torch.int64) + manager.lock() + assert manager.get_persistent("locks", (8,), torch.int32) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent("new", (8,), torch.int32) diff --git a/tools/recipes/README.md b/tools/recipes/README.md index 058470167a..253bd7a726 100644 --- a/tools/recipes/README.md +++ b/tools/recipes/README.md @@ -21,6 +21,21 @@ flowchart LR The recipe is the baseline. Hardware and workload information can optionally refine the initial configuration. Sweep tuning is an optional validation step. +## Getting Started + +Use `serve_with_recipe.sh` to generate the recipe configuration and start vLLM +in one step: + +```bash +tools/recipes/serve_with_recipe.sh \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 +``` + +For `xeon6`, the script enables hardware detection automatically. Generated +`config.yml` and `env.sh` files remain in the directory where the script was +invoked. + ## 1. vLLM Recipes Only Use the converter directly when the recipe already contains the deployment diff --git a/tools/recipes/serve_with_recipe.sh b/tools/recipes/serve_with_recipe.sh new file mode 100755 index 0000000000..cad1c06f3f --- /dev/null +++ b/tools/recipes/serve_with_recipe.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +CONVERTER="${SCRIPT_DIR}/recipe_json_to_vllm_config.py" +OUTPUT_DIR="$(pwd -P)" +CONFIG_PATH="${OUTPUT_DIR}/config.yml" +ENV_PATH="${OUTPUT_DIR}/env.sh" + +args=("$@") +hardware="" + +for ((i = 0; i < ${#args[@]}; i++)); do + case "${args[$i]}" in + --hardware) + if ((i + 1 < ${#args[@]})); then + hardware="${args[$((i + 1))]}" + fi + ;; + --hardware=*) + hardware="${args[$i]#*=}" + ;; + esac +done + +if [[ "${hardware,,}" == "xeon6" ]]; then + args+=(--detect-hardware) +fi + +echo "==> Generating vLLM configuration" +python3 "${CONVERTER}" "${args[@]}" + +echo +echo "==> Loading recipe environment" +# shellcheck disable=SC1090 +source "${ENV_PATH}" + +# Released vLLM images keep recipe-relative assets such as examples/*.jinja +# under /vllm-workspace. Launch from there when available while keeping the +# generated config and environment files in the caller's output directory. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +echo +echo "==> Starting vLLM" +exec vllm serve --config "${CONFIG_PATH}" diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index d3c5f3aac1..5da7366403 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -179,6 +179,7 @@ class ServingDerender(BaseServing): usage=usage, prompt_logprobs=gen.prompt_logprobs, kv_transfer_params=gen.kv_transfer_params, + metrics=gen.metrics, ) async def derender_completion_response( @@ -245,6 +246,9 @@ class ServingDerender(BaseServing): choices=choices, usage=usage, kv_transfer_params=kv_params, + # Metrics describe one prompt. Parallel samples are already + # suppressed by /generate; multi-prompt responses cannot be merged. + metrics=first.metrics if len(request.generate_responses) == 1 else None, ) async def derender_chat_stream_response( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index ae0fd62b6e..088ec4d293 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -11,7 +11,11 @@ from pydantic import ( ) from vllm.config import ModelConfig -from vllm.entrypoints.generate.base.protocol import StreamOptions, validate_cache_salt +from vllm.entrypoints.generate.base.protocol import ( + PerRequestMetrics, + StreamOptions, + validate_cache_salt, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProbs, ChatCompletionRequest, @@ -360,6 +364,7 @@ class GenerateStreamResponse(BaseModel): usage: UsageInfo | None = Field(default=None) prompt_token_ids: list[int] | None = None mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None + metrics: PerRequestMetrics | None = None class GenerateResponse(BaseModel): @@ -378,6 +383,7 @@ class GenerateResponse(BaseModel): prompt_logprobs: list[dict[int, Logprob] | None] | None = None prompt_token_ids: list[int] | None = None mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None + metrics: PerRequestMetrics | None = None kv_transfer_params: dict[str, Any] | None = Field( default=None, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 7813c7facc..59fda2dbf8 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -12,9 +12,13 @@ from fastapi import Request from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker -from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.generate.base.protocol import ( + PerRequestMetrics, + RequestResponseMetadata, +) from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, + build_spec_decoding_metrics, clamp_prompt_logprobs, ) from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -373,6 +377,11 @@ class ServingTokens(GenerateBaseServing): request_metadata.final_usage_info = usage + per_request_metrics = None + if request.sampling_params.n == 1: + spec_stats = build_spec_decoding_metrics(final_res) + if spec_stats is not None: + per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats) response = GenerateResponse( request_id=request_id, created=created_time, @@ -384,6 +393,7 @@ class ServingTokens(GenerateBaseServing): final_res.prompt_token_ids if request.return_token_ids else None ), mm_placeholders=request._response_mm_placeholders, + metrics=per_request_metrics, kv_transfer_params=final_res.kv_transfer_params, ec_transfer_params=final_res.ec_transfer_params, ) @@ -423,6 +433,7 @@ class ServingTokens(GenerateBaseServing): prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params + last_res: RequestOutput | None = None include_usage, include_continuous_usage = should_include_usage( request.stream_options, False @@ -430,6 +441,7 @@ class ServingTokens(GenerateBaseServing): try: async for res in result_generator: + last_res = res if first_iteration: if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) @@ -496,10 +508,10 @@ class ServingTokens(GenerateBaseServing): total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent prompt metadata (Rust skips None fields). + # Omit fields that are absent from token-bearing chunks. exclude = { name - for name in ("prompt_token_ids", "mm_placeholders") + for name in ("prompt_token_ids", "mm_placeholders", "metrics") if getattr(chunk, name) is None } yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" @@ -517,10 +529,18 @@ class ServingTokens(GenerateBaseServing): ) if include_usage: + per_request_metrics = None + if sampling_params.n == 1: + spec_stats = build_spec_decoding_metrics(last_res) + if spec_stats is not None: + per_request_metrics = PerRequestMetrics( + speculative_decoding=spec_stats + ) final_chunk = GenerateStreamResponse( request_id=request_id, choices=[], usage=final_usage_info, + metrics=per_request_metrics, ) yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 387c33c110..60687fc860 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -10,7 +10,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( apply_gptq_marlin_linear, marlin_act_int8_process_scales, marlin_make_empty, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -91,11 +90,6 @@ class MarlinLinearKernel(MPLinearKernel): size_k, size_n = c.partition_weight_shape padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) - # Allocate marlin workspace, reusing existing storage on reload. - self.workspace = marlin_make_workspace_new( - device, existing=getattr(self, "workspace", None) - ) - # Default name since marlin requires empty parameter for zp, # TODO: remove this requirement from marlin (allow optional tensors) if self.w_zp_name is None: @@ -197,7 +191,7 @@ class MarlinLinearKernel(MPLinearKernel): weight=w_q, weight_scale=w_s, weight_zp=w_zp, # type: ignore - workspace=self.workspace, + workspace=None, wtype=c.weight_type, input_size_per_partition=c.partition_weight_shape[0], output_size_per_partition=c.partition_weight_shape[1], diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py index 30ff0d58d4..108f835e64 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py @@ -58,7 +58,7 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=None, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py index bec54cd942..6ebdbd6474 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -46,7 +46,7 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/nvfp4/marlin.py b/vllm/model_executor/kernels/linear/nvfp4/marlin.py index a05d6823c8..1ab5336876 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/nvfp4/marlin.py @@ -50,7 +50,7 @@ class MarlinNvFp4LinearKernel(NvFp4LinearKernel): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 3ca89d3cea..1df6ad5541 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -93,7 +93,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): input=x, weight=layer.weight, weight_scale=weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, input_dtype=self.marlin_input_dtype, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 52b6d4e753..18a2fb5cc6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -176,7 +177,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): max_num_tokens=max_num_tokens, num_dispatchers=num_dispatchers, ) - self._permute_scratch: dict[int, MoEPermuteScratch] = {} def init_humming_moe(self): from vllm.utils.humming import get_heuristics_config @@ -259,24 +259,20 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): if not moe_permute_unpermute_supported(): return None - scratch = self._permute_scratch.get(topk) - if scratch is None: - max_expanded_rows = ( - self.moe_config.max_num_tokens - * self.moe_config.dp_size - * self.moe_config.experts_per_token - ) - scratch = MoEPermuteScratch( - max_num_tokens=math.ceil(max_expanded_rows / topk), - topk=topk, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - hidden_size=self.moe_config.hidden_dim, - hidden_dtype=self.moe_config.in_dtype, - ) - self._permute_scratch[topk] = scratch - return scratch + max_expanded_rows = ( + self.moe_config.max_num_tokens + * self.moe_config.dp_size + * self.moe_config.experts_per_token + ) + return get_moe_permute_scratch( + max_num_tokens=math.ceil(max_expanded_rows / topk), + topk=topk, + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + hidden_size=self.moe_config.hidden_dim, + hidden_dtype=self.moe_config.in_dtype, + ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): ctx = get_forward_context() diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index a6a70ca952..cfd78de28d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -35,7 +35,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( from vllm.model_executor.layers.fused_moe.utils import _resize_cache from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, - marlin_make_workspace_new, + get_marlin_workspace, marlin_moe_intermediate_size, marlin_quant_input, ) @@ -95,7 +95,7 @@ def _fused_marlin_moe( N = marlin_moe_intermediate_size(w1, w2) w13_num_shards = 2 if activation.is_gated else 1 if workspace is None: - workspace = marlin_make_workspace_new(hidden_states.device, 4) + workspace = get_marlin_workspace(hidden_states.device) if intermediate_cache13 is None: intermediate_cache13 = torch.empty( diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index 7a53a960e8..0aa9a38071 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -5,6 +5,11 @@ from dataclasses import dataclass, field import torch +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + @dataclass class MoEPermuteScratch: @@ -102,6 +107,48 @@ class MoEPermuteScratch: return topk_ids_int32 +def get_moe_permute_scratch( + *, + max_num_tokens: int, + topk: int, + num_experts: int, + num_local_experts: int, + device: torch.device, + hidden_size: int | None = None, + hidden_dtype: torch.dtype | None = None, +) -> MoEPermuteScratch: + """Share scratch across sequential layers in the current ubatch and lane. + + Without a workspace manager, allocate new scratch for each call. + """ + + def create_scratch() -> MoEPermuteScratch: + return MoEPermuteScratch( + max_num_tokens=max_num_tokens, + topk=topk, + num_experts=num_experts, + num_local_experts=num_local_experts, + device=device, + hidden_size=hidden_size, + hidden_dtype=hidden_dtype, + ) + + if not is_workspace_manager_initialized(): + return create_scratch() + + key = ( + MoEPermuteScratch, + max_num_tokens, + topk, + num_experts, + num_local_experts, + device, + hidden_size, + hidden_dtype, + ) + return current_workspace_manager().get_persistent_resource(key, create_scratch) + + def moe_permute( hidden_states: torch.Tensor, a1q_scale: torch.Tensor | None, diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index 2ddf710703..c3422445b2 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -14,6 +14,12 @@ from vllm.v1.worker.workspace import current_workspace_manager RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 +# cooperative_topk's hard row limit: one cluster wave covers at most this many +# rows. "auto" uses cooperative_topk for every batch within the limit and falls +# back to persistent_topk past it; explicit ``cooperative`` requests are +# validated against the same bound. +AUTO_COOPERATIVE_MAX_ROWS = 64 + # --------------------------------------------------------------------------- # DeepSelect (vllm._deepselect_C) # --------------------------------------------------------------------------- @@ -209,8 +215,13 @@ class SparseIndexerTopk(torch.nn.Module): def _resolve_auto( self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: - """The pre-existing priority chain: cooperative -> persistent -> - per_row. deep_select/flashinfer/torch are opt-in only.""" + """The priority chain: cooperative -> persistent -> per_row. + deep_select/flashinfer/torch are opt-in only. + + cooperative_topk is preferred whenever it is applicable, i.e. within + its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to + persistent_topk. + """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" if self._is_cuda and topk_tokens in (512, 1024, 2048): @@ -228,8 +239,10 @@ class SparseIndexerTopk(torch.nn.Module): failures.append( f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" ) - if num_rows > 64: - failures.append(f"num_rows must be <= 64, got {num_rows}") + if num_rows > AUTO_COOPERATIVE_MAX_ROWS: + failures.append( + f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}" + ) if logits.stride(0) % 4 != 0: failures.append( f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}" @@ -293,13 +306,18 @@ class SparseIndexerTopk(torch.nn.Module): (topk_workspace,) = current_workspace_manager().get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) + # persistent_topk's last argument is the host-side "max seq_len + # across rows" bound (it gates the sampled_topk path and clamps + # per-row lengths), not the padded row width: logits is sized to + # max_model_len, so passing its width would incorrectly enable the + # long-row kernels on short-context batches. torch.ops._C.persistent_topk( logits, seq_lens, topk_indices, topk_workspace, topk_tokens, - logits.shape[1], + max_seq_len, ) elif backend == "flashinfer": # Deferred: importing flashinfer initializes CUDA at import time. diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index ccf1c60f0f..1ca07ad0cf 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -167,6 +167,8 @@ def _log_gdn_backend_decision( head_k_dim, ) return + elif current_platform.is_xpu(): + return chosen = { "flashinfer": "FlashInfer", @@ -524,6 +526,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): self.gdn_decode_kernel = "triton" elif current_platform.is_cpu(): self.gdn_decode_kernel = "CPU" + elif current_platform.is_xpu(): + self.gdn_decode_kernel = "XPU" self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 71b858e919..92132f3ffb 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -48,7 +48,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_marlin_supports_layer, check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, verify_marlin_supported, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -642,9 +641,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_qzeros", w2_qzeros) set_weight_attrs(w2_qzeros, extra_weight_attrs) - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index e4478e897d..ffbc6fba22 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -16,7 +16,6 @@ from vllm.model_executor.kernels.linear import ( ) from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, - FusedMoEExpertsModular, FusedMoEMethodBase, FusedMoEQuantConfig, FusedMoeWeightScaleSupported, @@ -46,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import ( from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_repeat_scales_on_all_ranks, verify_marlin_supported, ) @@ -601,12 +599,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_bias", w2_bias) set_weight_attrs(w2_bias, extra_weight_attrs) - if self.experts_cls is not None and issubclass( - self.experts_cls, FusedMoEExpertsModular - ): - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: def replace_or_register(name: str, val: torch.Tensor | None): if val is None: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index bcbb687e3b..feddb77cff 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -11,7 +11,6 @@ from compressed_tensors.quantization import ( from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoEExpertsModular, RoutedExperts, SharedExperts, ) @@ -36,7 +35,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_config, get_marlin_input_dtype, - marlin_make_workspace_new, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -466,18 +464,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - # Marlin workspace — only needed for Marlin-family backends, not emulation. - if ( - self.experts_cls is not None - and issubclass(self.experts_cls, FusedMoEExpertsModular) - and self.wna16_backend != WNA16MoEBackend.EMULATION - ): - layer.workspace = marlin_make_workspace_new( - layer.w13_weight_packed.device, - 4, - existing=getattr(layer, "workspace", None), - ) - # Alias packed weights to w13_weight/w2_weight for the modular kernel interface layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 82201dff3b..1ee52a1780 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -8,7 +8,6 @@ import numpy import torch import vllm.envs as envs -from vllm import _custom_ops as ops from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -23,6 +22,7 @@ from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.math_utils import round_up from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import direct_register_custom_op from .quant_utils import pack_cols, unpack_cols @@ -32,6 +32,7 @@ GPTQ_MARLIN_TILE = 16 GPTQ_MARLIN_MIN_THREAD_N = 64 GPTQ_MARLIN_MIN_THREAD_K = 128 GPTQ_MARLIN_MAX_PARALLEL = 16 +MARLIN_MAX_BLOCKS_PER_SM = 4 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] @@ -407,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso def marlin_make_workspace_new( device: torch.device, max_blocks_per_sm: int = 1, - existing: torch.Tensor | None = None, ) -> torch.Tensor: # In the new marlin kernel, we use the num of threadblocks as workspace # size. The num of threadblocks is sms_count * max_blocks_per_sm. sms = num_compute_units(device.index) - size = sms * max_blocks_per_sm - # On weight reload, reuse the existing storage so the workspace address - # captured by CUDA graphs stays valid. - if existing is not None: - if ( - existing.device != device - or existing.dtype != torch.int - or existing.numel() != size - ): - raise ValueError( - f"Existing Marlin workspace is incompatible " - f"(device={existing.device}, dtype={existing.dtype}, " - f"numel={existing.numel()}; expected device={device}, " - f"dtype={torch.int}, numel={size}). Reload must reuse the " - f"workspace storage captured by CUDA graphs." - ) - return existing.zero_() - return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False) + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device) + + +def get_marlin_workspace(device: torch.device) -> torch.Tensor: + """Get zero-initialized locks for the current stream, ubatch, and lane. + + Marlin resets its locks after each call, so ordered calls can reuse them. + """ + from vllm.utils.torch_utils import current_stream + from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, + ) + + if not is_workspace_manager_initialized(): + return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM) + + return current_workspace_manager().get_persistent( + ("marlin_locks", current_stream()), + (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,), + torch.int, + zero_init=True, + ) + + +def _marlin_gemm( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + if workspace is None: + workspace = get_marlin_workspace(a.device) + return torch.ops._C.marlin_gemm( + a, + None, + b_q_weight, + b_bias, + b_scales, + a_scales, + global_scale, + b_zeros, + workspace, + b_q_type_id, + size_m, + size_n, + size_k, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ) + + +def _marlin_gemm_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype + return torch.empty((size_m, size_n), device=a.device, dtype=dtype) + + +direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake) def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool: @@ -686,7 +754,7 @@ def apply_gptq_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_zp: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, wtype: ScalarType, output_size_per_partition: int, input_size_per_partition: int, @@ -723,9 +791,8 @@ def apply_gptq_marlin_linear( reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( reshaped_x, - None, weight, bias, weight_scale, @@ -733,7 +800,7 @@ def apply_gptq_marlin_linear( None, weight_zp, workspace, - wtype, + wtype.id, size_m=reshaped_x.shape[0], size_n=padded_n, size_k=padded_k, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index a12d8b87e9..7fb5621f2f 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -161,7 +160,7 @@ def apply_fp4_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_global_scale: torch.Tensor | None, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -196,9 +195,8 @@ def apply_fp4_marlin_linear( inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -206,7 +204,7 @@ def apply_fp4_marlin_linear( global_scale=weight_global_scale, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float4_e2m1f, + b_q_type_id=scalar_types.float4_e2m1f.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -237,13 +235,6 @@ def prepare_fp4_layer_for_marlin( assert layer.weight.shape == (part_size_n, part_size_k // 2) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = layer.weight.view(torch.int32).T.contiguous() @@ -389,15 +380,9 @@ def prepare_nvfp4_moe_layer_for_marlin( return x return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) - device = w13.device param_dtype = layer.params_dtype is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: @@ -481,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin( k = layer.moe_config.hidden_dim n = layer.moe_config.intermediate_size_per_partition - # WORKSPACE - device = layer.w13_weight.device param_dtype = layer.params_dtype - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 # WEIGHT diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 4a5f5ee59a..257c5bf445 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -9,7 +9,6 @@ from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_moe_padded_intermediate, marlin_pad_dim, marlin_pad_qweight, @@ -50,7 +49,7 @@ def apply_fp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None, @@ -80,9 +79,8 @@ def apply_fp8_marlin_linear( # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) raise RuntimeError("Marlin W8A8 is not supported.") - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -90,7 +88,7 @@ def apply_fp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin( else: assert layer.weight.shape == (part_size_n, part_size_k) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first) @@ -241,12 +232,7 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight_scale: torch.Tensor, w2_weight_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Shuffle weights and scales into marlin format. - - Note that this function has the side effect of adding a `workspace` - attribute to the layer. This `workspace` does not need to be - registered as a Parameter as it is not used during weight reloading. - """ + """Shuffle weights and scales into marlin format.""" logger.warning_once( "Your GPU does not have native support for FP8 computation but " "FP8 quantization is being used. Weight-only FP8 compression will " @@ -272,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n) w2_weight = _moe_pad_last(w2_weight, n, padded_n) - # WORKSPACE - device = layer.w13_weight.device - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor: @@ -396,7 +376,7 @@ def apply_mxfp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -416,9 +396,8 @@ def apply_mxfp8_marlin_linear( dtype=input.dtype, ) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=reshaped_x, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -426,7 +405,7 @@ def apply_mxfp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -451,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: group_size = 32 # MX standard block size padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT - repack FP8 weights to Marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) qweight = qweight.T.contiguous() @@ -512,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin( """Repack MXFP8 MoE weights and scales into Marlin kernel format. Args: - layer: MoE layer (used to read params_dtype and attach workspace). + layer: MoE layer. w13: [E, 2*N, K] float8_e4m3fn weights. w2: [E, K, N] float8_e4m3fn weights. w13_scale: [E, 2*N, K//32] uint8 e8m0 scales. @@ -537,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin( n = padded_n w13_n = w13.shape[1] - device = w13.device param_dtype = torch.get_default_dtype() - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: if "w13" in name: size_n, size_k = w13_n, k diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 0c1cc8c55e..10e0d8a2c2 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -543,7 +543,7 @@ def _compiled_sample_step( if tp_size > 1: soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name) soft_embeds = soft_embeds * normalizer - sc_embeds[decode_slots] = soft_embeds * sc_keep + sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype) # Overwrite canvas with argmax for newly converged denoise requests newly_converged = (converged & is_denoise).unsqueeze(1) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index be665e7e06..dfbed25655 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -156,14 +157,16 @@ class MiMoV2MoE(nn.Module): dtype = getattr(config, "moe_router_dtype", "float32") self.gate_dtype = str_dtype_to_torch_dtype(dtype) - self.gate = nn.Linear( + self.gate = GateLinear( config.hidden_size, config.n_routed_experts, bias=False, - dtype=self.gate_dtype, + params_dtype=self.gate_dtype, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", ) self.gate.e_score_correction_bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=self.gate_dtype) + torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype) ) self.experts = FusedMoEFactory( @@ -182,7 +185,7 @@ class MiMoV2MoE(nn.Module): num_expert_group=config.n_group, topk_group=config.topk_group, scoring_func="sigmoid", - router_logits_dtype=self.gate_dtype, + router_logits_dtype=self.gate.out_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -198,7 +201,7 @@ class MiMoV2MoE(nn.Module): gate_input = hidden_states.to(self.gate_dtype) else: gate_input = hidden_states - router_logits = self.gate(gate_input) + router_logits, _ = self.gate(gate_input) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 91a3177b7f..f55cf44f37 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -52,6 +52,7 @@ from vllm.transformers_utils.processors.mimo_v2_omni import ( from .interfaces import ( MultiModalEmbeddings, + SupportsEagle3, SupportsMultiModal, SupportsPP, SupportsQuant, @@ -1209,7 +1210,9 @@ class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingIn info=MiMoV2OmniProcessingInfo, dummy_inputs=MiMoV2OmniDummyInputsBuilder, ) -class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant): +class MiMoV2OmniForCausalLM( + nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3 +): # To ensure correct weight loading and mapping. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 5504304c73..1636e3852c 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -168,6 +168,7 @@ class DFlashQwen3Attention(nn.Module): rms_norm_eps: float = 1e-06, attention_bias: bool = False, add_swa_attention_sink_bias: bool = False, + v_scale: float | None = None, sliding_window: int | None = None, causal: bool = False, is_neox_style: bool = True, @@ -193,6 +194,7 @@ class DFlashQwen3Attention(nn.Module): self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 + self.v_scale = v_scale self.qkv_proj = QKVParallelLinear( hidden_size, @@ -264,6 +266,8 @@ class DFlashQwen3Attention(nn.Module): q, k = self.rotary_emb(positions, q, k) + if self.v_scale is not None: + v = v * self.v_scale attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -312,6 +316,7 @@ class DFlashQwen3DecoderLayer(nn.Module): rms_norm_eps=config.rms_norm_eps, attention_bias=getattr(config, "attention_bias", False), add_swa_attention_sink_bias=add_swa_attention_sink_bias, + v_scale=dflash_config.get("attention_value_scale"), sliding_window=sliding_window, causal=causal, is_neox_style=is_neox_style, @@ -633,6 +638,10 @@ class DFlashQwen3Model(nn.Module): if context_slot_mapping is None: return + v_scale = getattr(self.layers[0].self_attn, "v_scale", None) + if v_scale is not None: + all_v.mul_(v_scale) + # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) per_layer = isinstance(context_slot_mapping, (list, tuple)) diff --git a/vllm/model_executor/models/voyage.py b/vllm/model_executor/models/voyage.py index d131ffff4b..9f4885a0f9 100644 --- a/vllm/model_executor/models/voyage.py +++ b/vllm/model_executor/models/voyage.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn from vllm.config import VllmConfig +from vllm.model_executor.models.interfaces import SupportsLoRA from vllm.model_executor.models.qwen3 import Qwen3Model from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix from vllm.sequence import IntermediateTensors @@ -15,7 +16,7 @@ from vllm.sequence import IntermediateTensors WeightItem = tuple[str, torch.Tensor] -class VoyageQwen3BidirectionalEmbedModel(nn.Module): +class VoyageQwen3BidirectionalEmbedModel(nn.Module, SupportsLoRA): """Qwen3Model + Voyage embedding head + bidirectional attention. Checkpoint conventions (HF): @@ -30,6 +31,11 @@ class VoyageQwen3BidirectionalEmbedModel(nn.Module): - No "model." prefix """ + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3cf26f8767..a70a23e3ab 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -7,6 +7,7 @@ from typing import cast import torch +from vllm import envs from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -42,6 +43,49 @@ from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) +def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None: + """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes. + + E8M0 is an unsigned exponent-only scale format with bias 127. A finite + encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``; + ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used + by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP + Microscaling Formats (MX) Specification, section 5.4.1: + https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + + Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``, + or decode it to a floating-point power of two. Floating-point inputs are + accepted only when the original E8M0 byte can be recovered losslessly. + This function does not quantize or round arbitrary scales. + """ + if scale.dtype == torch.float8_e8m0fnu: + # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a + # numeric conversion instead of preserving the encoded exponent byte. + return scale.view(torch.uint8).contiguous() + if scale.dtype == torch.uint8: + # The checkpoint loader already exposed the E8M0 wire representation. + return scale.contiguous() + if not scale.dtype.is_floating_point: + return None + + scale_f32 = scale.detach().float() + if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()): + return None + + # With no mantissa, E8M0 can represent only exact powers of two. Rebuild + # the value before encoding so this adapter never silently quantizes a + # general floating-point checkpoint scale. + exponent = torch.round(torch.log2(scale_f32)) + if not torch.equal(torch.exp2(exponent), scale_f32): + return None + + encoded = exponent.to(torch.int32) + 127 + # 0xFF is NaN in OCP E8M0, not a finite exponent. + if int(encoded.min()) < 0 or int(encoded.max()) > 254: + return None + return encoded.to(torch.uint8).contiguous() + + def _trust_dsv4_extra_cache_nan_free( kv_cache_dtype: str, has_kv_transfer: bool, @@ -555,6 +599,10 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None + self._wo_a_fp8_weight: torch.Tensor | None = None + self._wo_a_e8m0_scale: torch.Tensor | None = None + self._wo_a_cos_cache: torch.Tensor | None = None + self._wo_a_sin_cache: torch.Tensor | None = None self._fused_compressor_weight: torch.Tensor | None self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None @@ -815,6 +863,93 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) self._wo_b_scale = _prep(self.wo_b) + if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM: + self._prepare_fp8_wo_a() + + def _prepare_fp8_wo_a(self) -> None: + try: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale as mxscale_op, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant as inverse_quant_op, + ) + except ImportError: + logger.warning_once( + "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; " + "falling back to BF16 WO_A." + ) + return + del mxscale_op, inverse_quant_op + + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_fp8_block_weight_scale, + is_fp8, + ) + + weight = getattr(self.wo_a, "weight", None) + scale = get_fp8_block_weight_scale(self.wo_a) + if scale is None: + # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the + # historical ``_inv`` suffix. + scale = getattr(self.wo_a, "weight_scale", None) + if weight is None or scale is None: + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; " + "the layer exposes no weight/weight scale. Falling back to " + "BF16 WO_A." + ) + return + if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype): + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D " + "block scale, got weight %s%s and scale %s. Falling back to " + "BF16 WO_A.", + weight.dtype, + tuple(weight.shape), + tuple(scale.shape), + ) + return + + groups = self.n_local_groups + out_per_group = self.o_lora_rank + out_features, in_features = weight.shape + if ( + out_features != groups * out_per_group + or out_per_group % 128 != 0 + or in_features % 128 != 0 + or scale.shape != (out_features // 128, in_features // 128) + ): + logger.warning_once( + "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of " + "%d outputs, got weight %s and scale %s. Falling back to BF16 " + "WO_A.", + groups, + out_per_group, + tuple(weight.shape), + tuple(scale.shape), + ) + return + + e8m0_scale = _wo_a_block_scale_to_e8m0(scale) + if e8m0_scale is None: + logger.warning_once( + "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a " + "block scale as OCP E8M0. Falling back to BF16 WO_A.", + scale.dtype, + ) + return + + self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features) + self._wo_a_e8m0_scale = e8m0_scale.view( + groups, out_per_group // 128, in_features // 128 + ) + cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None) + if cache is None: + cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16) + cos_cache, sin_cache = cache.chunk(2, dim=-1) + self._wo_a_cos_cache = cos_cache.contiguous() + self._wo_a_sin_cache = sin_cache.contiguous() def prepare_compressor_gemm_fusion(self) -> bool: if self._fused_compressor_weight is not None: @@ -973,17 +1108,44 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. - z = rocm_inv_rope_einsum( - self.rotary_emb, - o, - positions, - self.rope_head_dim, - self.n_local_groups, - self.o_lora_rank, - self.wo_a, - ) - zf = z.flatten(1) + if self._wo_a_fp8_weight is not None: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant, + ) + + assert self._wo_a_cos_cache is not None + assert self._wo_a_sin_cache is not None + o_fp8, o_scale = inverse_rope_group_quant( + o.view(o.shape[0], self.n_local_heads, self.head_dim), + positions.to(torch.int64), + self._wo_a_cos_cache, + self._wo_a_sin_cache, + num_groups=self.n_local_groups, + quant_group_size=128, + ) + assert self._wo_a_e8m0_scale is not None + zf = batched_gemm_a8w8_mxscale( + o_fp8, + self._wo_a_fp8_weight, + o_scale, + self._wo_a_e8m0_scale, + dtype=o.dtype, + ).flatten(1) + else: + # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. + z = rocm_inv_rope_einsum( + self.rotary_emb, + o, + positions, + self.rope_head_dim, + self.n_local_groups, + self.o_lora_rank, + self.wo_a, + ) + zf = z.flatten(1) if self._wo_b_scale is not None and zf.dim() == 2: return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) return self.wo_b(zf) diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py index 4c633da6c5..bfc02768c0 100644 --- a/vllm/models/deepseek_v4/common/vision.py +++ b/vllm/models/deepseek_v4/common/vision.py @@ -40,8 +40,7 @@ from vllm.model_executor.models.vision import ( ) -@lru_cache(8) -def get_vision_cos_sin( +def _compute_vision_cos_sin( n_h: int, n_w: int, dim: int, theta: float ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) @@ -52,6 +51,13 @@ def get_vision_cos_sin( return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) +@lru_cache(8) +def get_vision_cos_sin( + n_h: int, n_w: int, dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor]: + return _compute_vision_cos_sin(n_h, n_w, dim, theta) + + def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: dtype = x.dtype x1, x2 = x.float().chunk(2, dim=-1) @@ -123,15 +129,23 @@ class DeepseekV4VisionAttention(nn.Module): ) def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: n = x.size(0) qkv, _ = self.wqkv(x) q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) k = apply_rotary(k, cos, sin).unsqueeze(0) - # One image per call: a dense batch, no varlen packing metadata. - o = self.attn(q, k, v.unsqueeze(0)) + # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens + # set): packed multi-image batch for encoder CUDA graph replay. + o = self.attn( + q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen + ) out, _ = self.wo(o.reshape(n, -1)) return out @@ -173,9 +187,14 @@ class DeepseekV4VisionBlock(nn.Module): self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), cos, sin) + x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen) return x + self.mlp(self.norm2(x)) @@ -206,6 +225,115 @@ class DeepseekV4ViT(nn.Module): x = block(x, cos, sin) return self.norm(x) + def forward_packed( + self, + patches: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + ) -> torch.Tensor: + """Varlen path for encoder CUDA graphs: multiple images packed along + rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed + on device. Pure-tensor (no host reads, no H2D), safe to capture.""" + x = self.patch_embed(patches) + for block in self.blocks: + x = block(x, cos, sin, cu_seqlens, max_seqlen) + return self.norm(x) + + +def build_packed_vit_metadata( + grids: list[list[int]] | list[tuple[int, int]], + *, + rope_dim: int, + rope_theta: float, + device: torch.device, + max_seqlen_override: int | None = None, + cached: bool = True, +) -> dict[str, torch.Tensor]: + """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``. + + Args: + grids: ``[n_vit_h, n_vit_w]`` per image, in packing order. + rope_dim: RoPE embedding dimension of the ViT. + rope_theta: RoPE base frequency of the ViT. + device: Device for the returned tensors (except ``max_seqlen``, + which stays on the host). + max_seqlen_override: Worst-case value baked in at CUDA graph capture + (the attention wrapper reads ``max_seqlen`` on the host, so the + capture-time value becomes a graph constant). + cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time + dummy grids pass ``False`` to avoid evicting real entries. + + """ + cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin + cos_list: list[torch.Tensor] = [] + sin_list: list[torch.Tensor] = [] + cu_seqlens = [0] + max_seqlen = 0 + for n_h, n_w in grids: + cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta) + cos_list.append(cos) + sin_list.append(sin) + n = n_h * n_w + cu_seqlens.append(cu_seqlens[-1] + n) + max_seqlen = max(max_seqlen, n) + if cos_list: + cos = torch.cat(cos_list).to(device) + sin = torch.cat(sin_list).to(device) + else: + cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + if max_seqlen_override is not None: + max_seqlen = max_seqlen_override + return { + "vit_cos": cos, + "vit_sin": sin, + "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device), + # Read on the host by the attention wrapper; keep on CPU. + "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32), + } + + +def build_packed_merge_metadata( + grids: list[list[int]] | list[tuple[int, int]], + downsample_ratio: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``. + + For each aligner output row, the ``r x r`` source patch rows in the + packed ViT output; positions past the image edge (the eager path pads + them with zeros) get mask 0 and a clamped in-range index. + """ + r = downsample_ratio + idx_list: list[torch.Tensor] = [] + mask_list: list[torch.Tensor] = [] + offset = 0 + for n_h, n_w in grids: + rows = -(-n_h // r) + cols = -(-n_w // r) + bi = torch.arange(rows).repeat_interleave(cols) + bj = torch.arange(cols).repeat(rows) + di = torch.arange(r).repeat_interleave(r) + dj = torch.arange(r).repeat(r) + hi = bi[:, None] * r + di[None, :] + wj = bj[:, None] * r + dj[None, :] + valid = (hi < n_h) & (wj < n_w) + idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1) + idx_list.append(idx) + mask_list.append(valid) + offset += n_h * n_w + if idx_list: + merge_idx = torch.cat(idx_list).to(device) + merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype) + else: + merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device) + merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device) + return {"merge_idx": merge_idx, "merge_mask": merge_mask} + class DeepseekV4Aligner(nn.Module): """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" @@ -240,6 +368,27 @@ class DeepseekV4Aligner(nn.Module): out, _ = self.w2(F.gelu(hidden)) return out + def forward_packed( + self, + x: torch.Tensor, + merge_idx: torch.Tensor, + merge_mask: torch.Tensor, + ) -> torch.Tensor: + """Packed multi-image path for encoder CUDA graphs. + + ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``; + gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge + exactly (unfold rows are channel-major, so the gathered patch rows + are transposed before flattening). + """ + r2 = self.downsample_ratio**2 + m = merge_idx.shape[0] + gathered = (x[merge_idx] * merge_mask).view(m, r2, -1) + gathered = gathered.permute(0, 2, 1).reshape(m, -1) + hidden, _ = self.w1(gathered) + out, _ = self.w2(F.gelu(hidden)) + return out + def run_dp_sharded_vision_tower( vision_model: DeepseekV4ViT, diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 4d838e121e..b7a67e85bd 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -13,7 +13,10 @@ from vllm.distributed import ( ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.attention import ( + DeepseekV4Attention, + _replace_layer_index, +) from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, @@ -391,6 +394,10 @@ def _copy_ragged_to_graph_buffers( return ragged_out, indptr_out +# (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode. +_TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + @dataclass class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None @@ -500,6 +507,23 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): self._fused_compressor_weight: torch.Tensor | None self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None + # Decode ragged topk metadata is a pure function of the indices its + # index source published, so every consumer of a source rebuilds the + # same thing. The source owns one cache per compress ratio, refreshed + # when it runs; consumers below it read through. + self._topk_ragged_cache: dict[int, _TopkRagged] = {} + self._index_source_prefix: str | None = None + if self.compress_ratio > 0: + assert self.index_source_layer_id is not None + self._index_source_prefix = _replace_layer_index( + self.prefix, self.index_source_layer_id + ) + if self._index_source_prefix not in self._static_forward_context: + raise NotImplementedError( + f"Index source {self._index_source_prefix} not found on " + "this rank; PP splits inside a v4.1 index-sharing group " + "are not supported." + ) @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -733,6 +757,42 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): output=output[:num_decode_tokens], ) + def _decode_topk_ragged( + self, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + num_decodes: int, + num_decode_tokens: int, + ) -> _TopkRagged: + """Ragged form of the topk indices this layer's index source published. + + The packing depends only on the shared ``topk_indices_buffer`` and + per-step metadata, apart from the layer's own compress ratio, so the + source memoizes one result per ratio for all the layers below it. + """ + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + assert self.topk_indices_buffer is not None + assert self._index_source_prefix is not None + + source = self._static_forward_context[self._index_source_prefix] + if source is self: + # Fresh indices as of this layer; drop what the last step cached. + source._topk_ragged_cache = {} + cached = source._topk_ragged_cache.get(self.compress_ratio) + if cached is not None: + return cached + + built = compute_global_topk_ragged_indices_and_indptr( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + attn_metadata.block_size // self.compress_ratio, + swa_metadata.is_valid_token[:num_decode_tokens], + ) + source._topk_ragged_cache[self.compress_ratio] = built + return built + def _forward_decode( self, q: torch.Tensor, @@ -749,22 +809,15 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): topk_ragged_indices = None topk_ragged_indptr = None if not swa_only: - # Local indices filled by the index-source layer's indexer. - assert attn_metadata is not None - assert swa_metadata.is_valid_token is not None - assert self.topk_indices_buffer is not None - block_size = attn_metadata.block_size // self.compress_ratio - is_valid = swa_metadata.is_valid_token[:num_decode_tokens] ( topk_ragged_indices, topk_ragged_indptr, topk_lens, - ) = compute_global_topk_ragged_indices_and_indptr( - self.topk_indices_buffer[:num_decode_tokens], - swa_metadata.token_to_req_indices, - attn_metadata.block_table[:num_decodes], - block_size, - is_valid, + ) = self._decode_topk_ragged( + swa_metadata=swa_metadata, + attn_metadata=attn_metadata, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, ) rocm_sparse_attn_decode( diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 3227b50a4e..3c57df2ad2 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py new file mode 100644 index 0000000000..e2a718b745 --- /dev/null +++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower. + +Mixin implementing the ``SupportsEncoderCudaGraph`` protocol +(``vllm/model_executor/models/interfaces.py``) on top of the shared +``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole +batch of images into one varlen run: ViT blocks attend per image via +``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask +against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/ +IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is +batched over all items of the packed batch. + +Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the +manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP +sharding and eager fallback. +""" + +from collections.abc import Hashable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.models.deepseek_v4.common.vision import ( + build_packed_merge_metadata, + build_packed_vit_metadata, +) +from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, +) + +from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.config.multimodal import MultiModalConfig + from vllm.models.deepseek_v4.common.vision import ( + DeepseekV4Aligner, + DeepseekV4ViT, + ) + from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config + + +class DeepseekV4VLEncoderCudaGraphMixin: + """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``. + + Expects ``self.vision``/``self.aligner``/``self.config``/ + ``self.multimodal_config`` and the ``image_start``/``image_newline``/ + ``image_end`` parameters from the host model. + """ + + # Provided by the host model; declared here so mixin methods type-check. + vision: "DeepseekV4ViT" + aligner: "DeepseekV4Aligner" + config: "DeepseekV41Config" + multimodal_config: "MultiModalConfig | None" + image_start: torch.nn.Parameter + image_end: torch.nn.Parameter + image_newline: torch.nn.Parameter + + # -- helpers shared with the eager multimodal path -- + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + ) -> torch.Tensor: + # Aligner rows in reading order, one per IMAGE slot. + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def _build_image_span( + self, image_embeds: torch.Tensor, types: torch.Tensor + ) -> torch.Tensor: + """Full image span: aligner rows at IMAGE slots, the learned + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" + types = types.to(image_embeds.device) + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) + dtype = image_embeds.dtype + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + span[types == IMAGE] = image_embeds + return span + + @staticmethod + def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: + grid = mm_kwargs[key] + if isinstance(grid, torch.Tensor): + grid = grid.tolist() + return [[int(v) for v in row] for row in grid] + + _encoder_cg_pad_totals: dict[int, int] = {} + + # -- SupportsEncoderCudaGraph protocol -- + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + mm_cfg = self.multimodal_config + if ( + mm_cfg is not None + and mm_cfg.mm_encoder_attn_dtype == "fp8" + and mm_cfg.mm_encoder_fp8_scale_path is None + ): + raise ValueError( + "cudagraph_mm_encoder is incompatible with dynamic FP8 " + "scaling for the vision encoder (the amax history is " + "host-side state that cannot be captured). Provide a static " + "scale file via --mm-encoder-fp8-scale-path or disable " + "cudagraph_mm_encoder." + ) + + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + backend = self.vision.blocks[0].attn.attn.attn_backend + # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls + # .tolist() on CUDA cu_seqlens; neither can be captured. + if backend in ( + AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.TORCH_SDPA, + ): + raise ValueError( + f"cudagraph_mm_encoder is not supported with the {backend.name} " + "ViT attention backend for DeepSeek-V4.1. Set " + "--mm-encoder-attn-backend FLASH_ATTN or disable " + "cudagraph_mm_encoder." + ) + + pad_totals = self._encoder_cg_pad_totals + + def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: + # Varlen attention requires cu_seqlens[-1] to equal the number of + # rows actually passed in. The captured buffers are sized for the + # full token budget, so a smaller real batch is completed with one + # trailing padding sequence; declaring fewer rows than the buffer + # holds is undefined behaviour and returns NaN on FlashAttn. + total = pad_totals.get(dst.data_ptr()) + if total is None: + raise RuntimeError( + "cu_seqlens replay buffer was not registered at capture " + "time; the manager must replay into the capture-time " + "buffers." + ) + n = min(src.shape[0], dst.shape[0]) + dst[:n].copy_(src[:n]) + dst[n:] = total + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "patches", + "vit_cos", + "vit_sin", + "cu_seqlens", + "max_seqlen", + "merge_idx", + "merge_mask", + ], + out_hidden_size=self.config.hidden_size, + padding_logics={"cu_seqlens": pad_cu_seqlens}, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: "VllmConfig", + ) -> tuple[int, int]: + # Min: comfortably above the smallest possible image span (4 tokens + # for a 1x1 grid); also sets max_batch_size = 64 via the manager's + # auto-inference (max_batch_size <= min_budget invariant). + min_budget = 64 + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ) -> list[EncoderItemSpec]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + return [ + EncoderItemSpec( + input_size=h * w, + output_tokens=lh * (lw + 1) + 2, + ) + for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True) + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + patches = mm_kwargs["patches"] + types = mm_kwargs["types"] + + if len(indices) == 0: + return { + "patches": patches[:0], + "vit_grid": torch.zeros((0, 2), dtype=torch.int64), + "llm_grid": torch.zeros((0, 2), dtype=torch.int64), + "types": types[:0], + } + + cum_patches = [0] + for h, w in vit_grid: + cum_patches.append(cum_patches[-1] + h * w) + cum_spans = [0] + for lh, lw in llm_grid: + cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2) + + return { + "patches": torch.cat( + [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ), + "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64), + "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64), + "types": torch.cat( + [types[cum_spans[i] : cum_spans[i + 1]] for i in indices] + ), + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, + ) -> EncoderCudaGraphCaptureInputs: + if ( + self.multimodal_config is not None + and self.multimodal_config.get_limit_per_prompt("image") == 0 + ): + raise RuntimeError( + "cudagraph_mm_encoder requires the vision tower, but it was " + "stubbed out because the image limit is 0. Disable " + "cudagraph_mm_encoder." + ) + + config = self.config + r = config.vision_downsample_ratio + # Span rows per dummy item (ceiling so the buffers tile the budget). + # Each dummy image is a single aligner row of per_item_out blocks: + # r x (r * per_item_out) patches. One real image using the whole + # budget needs at most r^2 * token_budget patches, which this fits. + per_item_out = -(-token_budget // max_batch_size) + grids = [[r, r * per_item_out]] * max_batch_size + total_patches = r * r * per_item_out * max_batch_size + + p = config.vision_patch_size + patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype) + + metadata = build_packed_vit_metadata( + grids, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=device, + max_seqlen_override=total_patches, + cached=False, + ) + # Spare cu_seqlens slots let replay append a padding sequence that + # covers the rows a smaller real batch does not fill. + real_cu = metadata.pop("cu_seqlens") + cu_seqlens = torch.full( + (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device + ) + cu_seqlens[: real_cu.numel()] = real_cu + + merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) + + self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches + + return EncoderCudaGraphCaptureInputs( + values={ + "patches": patches, + **metadata, + "cu_seqlens": cu_seqlens, + **merge, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ) -> EncoderCudaGraphReplayBuffers: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + patches = mm_kwargs["patches"] + dtype = self.aligner.w1.weight.dtype + if patches.dtype != dtype: + patches = patches.to(dtype) + + # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and + # pad_cu_seqlens appends the padding sequence covering the tail rows. + metadata = build_packed_vit_metadata( + vit_grid, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=patches.device, + ) + merge = build_packed_merge_metadata( + vit_grid, + self.config.vision_downsample_ratio, + device=patches.device, + dtype=dtype, + ) + return EncoderCudaGraphReplayBuffers( + values={"patches": patches, **metadata, **merge} + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + vit_out = self.vision.forward_packed( + inputs["patches"], + inputs["vit_cos"], + inputs["vit_sin"], + inputs["cu_seqlens"], + inputs["max_seqlen"], + ) + return self.aligner.forward_packed( + vit_out, inputs["merge_idx"], inputs["merge_mask"] + ) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype) + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + outs: list[torch.Tensor] = [] + offset = 0 + for n_vit_h, n_vit_w in vit_grid: + n_vit = n_vit_h * n_vit_w + outs.append( + self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w) + ) + offset += n_vit + if not outs: + return patches.new_zeros((0, self.config.hidden_size)) + return torch.cat(outs, dim=0) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + if not indices: + return + assert batch_mm_kwargs is not None + aligner_out = outputs["default"] + r = self.config.vision_downsample_ratio + vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid") + types = batch_mm_kwargs["types"].to(aligner_out.device) + + n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid) + span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid] + + # Batched span assembly: one masked fill per role across all items. + dtype = aligner_out.dtype + span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size) + span[types == IMAGE] = aligner_out[:n_rows] + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + + # Freshly allocated, so later replays cannot clobber the results. + offset = 0 + for idx, span_len in zip(indices, span_lens): + dest[idx] = span[offset : offset + span_len] + offset += span_len diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index bfbcaa89a1..901985bdd7 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs, diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 75ccfa8c2a..416d3b5c48 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -383,6 +383,8 @@ def _kpool_tail_seed_kernel( tslot_ptr, tail_ptr, n_tokens, + TAIL_BLOCK_ELEMS: tl.constexpr, + KPOOL_HEAD: tl.constexpr, HEAD_DIM: tl.constexpr, KPOOL: tl.constexpr, BLOCK_D: tl.constexpr, @@ -393,6 +395,11 @@ def _kpool_tail_seed_kernel( ahead belongs to a different tail block (or is past the batch / padding, slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``. + + The tail cache aliases the indexer cache with the indexer's (padded) block + stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` / + ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense + ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array. """ i = tl.program_id(0) t = tl.load(tslot_ptr + i).to(tl.int64) @@ -409,11 +416,11 @@ def _kpool_tail_seed_kernel( return offs = tl.arange(0, BLOCK_D) m = offs < HEAD_DIM - base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM + base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m) s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m) tl.store(tail_ptr + base + offs, k, mask=m) - tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m) + tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m) def kpool_seed_tail_cache( @@ -426,6 +433,8 @@ def kpool_seed_tail_cache( ) -> None: """Seed the paged tail cache from a prefill batch (see the kernel).""" assert tail_kv_cache.dtype == torch.bfloat16 + assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2 + assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim assert key.dtype == torch.bfloat16 n = tslot.shape[0] if n == 0: @@ -436,6 +445,8 @@ def kpool_seed_tail_cache( tslot, tail_kv_cache, n, + TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0), + KPOOL_HEAD=tail_kv_cache.stride(1), HEAD_DIM=head_dim, KPOOL=kpool, BLOCK_D=triton.next_power_of_2(head_dim), diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index f7204949fc..c80e442015 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -10,6 +10,7 @@ from vllm.config import get_current_vllm_config_or_none from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -39,21 +40,6 @@ if current_platform.is_cuda_alike(): logger = init_logger(__name__) -def _use_cooperative_topk( - logits: torch.Tensor, - select_k: int, - num_rows: int, -) -> bool: - return ( - current_platform.is_cuda() - and select_k in (512, 1024, 2048) - and num_rows <= 64 - and logits.stride(0) % 4 == 0 - and current_platform.has_device_capability(90) - and not current_platform.is_device_capability_family(120) - ) - - # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -134,6 +120,7 @@ def sparse_attn_indexer_kpool( # path and when the tail cache is disabled. tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -584,35 +571,18 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and select_k in (512, 1024, 2048): - workspace_manager = current_workspace_manager() - (topk_workspace,) = workspace_manager.get_simultaneous( - ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) - topk_op = ( - torch.ops._C.cooperative_topk - if _use_cooperative_topk(logits, select_k, num_rows) - else torch.ops._C.persistent_topk - ) - topk_op( - logits, - seq_lens, - topk_dst, - topk_workspace, - select_k, - attn_metadata_narrowed.max_seq_len, - ) - else: - torch.ops._C.top_k_per_row_decode( - logits, - next_n, - seq_lens, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) + # Shared dispatcher with the DSA sparse indexer, so the kpool select + # honors kernel_config.sparse_indexer_topk_backend and picks between + # cooperative/persistent by batch size instead of always taking the + # same kernel. + get_indexer_topk(topk_backend)( + logits, + seq_lens, + next_n, + topk_dst, + select_k, + attn_metadata_narrowed.max_seq_len, + ) # Resolve to token-level indices in the output buffer. if index_kpool > 1: @@ -692,6 +662,11 @@ class SparseAttnIndexerKpool(CustomOp): "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." ) _cfg = get_current_vllm_config_or_none() + self.topk_backend = ( + _cfg.kernel_config.sparse_indexer_topk_backend + if _cfg is not None + else "auto" + ) _parallel = _cfg.parallel_config if _cfg is not None else None if ( _parallel is not None @@ -766,4 +741,5 @@ class SparseAttnIndexerKpool(CustomOp): positions, self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, + self.topk_backend, ) diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index f4916faf6b..9b769686a7 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -7,11 +7,13 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from concurrent.futures import Future, ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import replace from functools import cached_property from typing import TYPE_CHECKING, Any, Generic, overload from typing_extensions import TypeVar +from vllm.exceptions import VLLMValidationError from vllm.inputs import ( EmbedsInput, EmbedsPrompt, @@ -962,27 +964,103 @@ class BaseRenderer(ABC, Generic[_T]): return engine_input + def _truncate_expanded_prompt( + self, + engine_input: TokensInput | MultiModalInput, + params: TokenizeParams | None, + ) -> None: + """Re-apply `truncate_prompt_tokens` once placeholders are expanded. + + `TokenizeParams.apply_post_tokenization` truncates the prompt before + multimodal expansion, so a single placeholder token can expand into + many and push the result back over the requested bound. + + Multimodal spans are never cut: a request whose spans would not all + survive intact is rejected, matching the post-expansion rejection in + `vllm.v1.engine.input_processor._validate_model_input`. + + Only multimodal inputs can grow after `apply_post_tokenization`, so + everything else returns immediately. + """ + if engine_input["type"] != "multimodal": + return + + if params is None or params.truncate_prompt_tokens is None: + return + + prompt_token_ids = engine_input.get("prompt_token_ids") + if prompt_token_ids is None: + return + + num_tokens = len(prompt_token_ids) + truncation = params._truncation_slice(self.tokenizer, num_tokens) + if truncation is None: + return + + start, stop, _ = truncation.indices(num_tokens) + + mm_placeholders = engine_input["mm_placeholders"] + for modality, positions in mm_placeholders.items(): + for position in positions: + end = position.offset + position.length + if position.offset >= start and end <= stop: + continue + raise VLLMValidationError( + f"truncate_prompt_tokens={params.truncate_prompt_tokens} " + f"would split or drop a(n) {modality} item, which spans " + f"tokens {position.offset} to {end} of the " + f"{num_tokens}-token prompt. Multimodal placeholders " + f"expand after truncation is applied, so the limit must " + f"leave every multimodal item intact. Please raise " + f"truncate_prompt_tokens or shorten the multimodal input.", + parameter="truncate_prompt_tokens", + value=params.truncate_prompt_tokens, + ) + + if start: + engine_input["mm_placeholders"] = { + modality: [ + replace(position, offset=position.offset - start) + for position in positions + ] + for modality, positions in mm_placeholders.items() + } + + engine_input["prompt_token_ids"] = prompt_token_ids[truncation] + def _process_singleton( self, prompt: SingletonTokPrompt, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> SingletonInput: if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] - return self._process_tokens(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] + engine_input = self._process_tokens( # type: ignore[arg-type] + prompt, # type: ignore[arg-type] + skip_mm_cache=skip_mm_cache, + ) + self._truncate_expanded_prompt(engine_input, tok_params) + return engine_input async def _process_singleton_async( self, prompt: SingletonTokPrompt, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> SingletonInput: if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] - return await self._process_tokens_async(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] + engine_input = await self._process_tokens_async( # type: ignore[arg-type] + prompt, # type: ignore[arg-type] + skip_mm_cache=skip_mm_cache, + ) + self._truncate_expanded_prompt(engine_input, tok_params) + return engine_input def _get_skip_decoder_start_token(self) -> bool: """Whether the multimodal processor supplies a complete decoder prefix.""" @@ -1048,12 +1126,15 @@ class BaseRenderer(ABC, Generic[_T]): arrival_time: float, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: engine_input = self._process_enc_dec(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] else: - engine_input = self._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + engine_input = self._process_singleton( + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) engine_input["arrival_time"] = arrival_time @@ -1065,6 +1146,7 @@ class BaseRenderer(ABC, Generic[_T]): arrival_time: float, *, skip_mm_cache: bool = False, + tok_params: TokenizeParams | None = None, ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: @@ -1074,7 +1156,7 @@ class BaseRenderer(ABC, Generic[_T]): ) else: engine_input = await self._process_singleton_async( - prompt, skip_mm_cache=skip_mm_cache + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) engine_input["arrival_time"] = arrival_time @@ -1101,7 +1183,9 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) return [ - self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + self.process_for_engine( + prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) for prompt in tok_prompts ] @@ -1126,7 +1210,7 @@ class BaseRenderer(ABC, Generic[_T]): return await asyncio.gather( *( self.process_for_engine_async( - p, arrival_time, skip_mm_cache=skip_mm_cache + p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) for p in tok_prompts ) @@ -1164,7 +1248,9 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) eng_prompts = [ - self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + self.process_for_engine( + prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) for prompt in tok_prompts ] @@ -1204,7 +1290,7 @@ class BaseRenderer(ABC, Generic[_T]): eng_prompts = await asyncio.gather( *( self.process_for_engine_async( - p, arrival_time, skip_mm_cache=skip_mm_cache + p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params ) for p in tok_prompts ) diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 1c35802d0d..36605c62b9 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1194,22 +1194,20 @@ class HfRenderer(BaseRenderer[HfTokenizer]): coordinate space, no offset shifting needed afterwards. """ prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) - if prompt_embeds_info is not None: - tensors, placeholder_token_id = prompt_embeds_info - mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) - cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( - list(prompt["prompt_token_ids"]), mm_updates - ) + if prompt_embeds_info is None: + return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) + tensors, placeholder_token_id = prompt_embeds_info + mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) + cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( + list(prompt["prompt_token_ids"]), mm_updates + ) engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) - - if prompt_embeds_info is not None: - tensors, _ = prompt_embeds_info - self._apply_prompt_embeds_to_engine_input( - cast(MultiModalInput, engine_input), - tensors, - mm_updates, - ) + self._apply_prompt_embeds_to_engine_input( + cast(MultiModalInput, engine_input), + tensors, + mm_updates, + ) return engine_input @@ -1222,24 +1220,24 @@ class HfRenderer(BaseRenderer[HfTokenizer]): ) -> TokensInput | MultiModalInput: """Async equivalent of `_process_tokens`.""" prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) - if prompt_embeds_info is not None: - tensors, placeholder_token_id = prompt_embeds_info - mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) - cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( - list(prompt["prompt_token_ids"]), mm_updates + if prompt_embeds_info is None: + return await super()._process_tokens_async( + prompt, skip_mm_cache=skip_mm_cache ) + tensors, placeholder_token_id = prompt_embeds_info + mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) + cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( + list(prompt["prompt_token_ids"]), mm_updates + ) engine_input = await super()._process_tokens_async( prompt, skip_mm_cache=skip_mm_cache ) - - if prompt_embeds_info is not None: - tensors, _ = prompt_embeds_info - self._apply_prompt_embeds_to_engine_input( - cast(MultiModalInput, engine_input), - tensors, - mm_updates, - ) + self._apply_prompt_embeds_to_engine_input( + cast(MultiModalInput, engine_input), + tensors, + mm_updates, + ) return engine_input diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index a89bc843b8..9dfa0c412f 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -398,6 +398,7 @@ class OnlineDerenderer: model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state @@ -595,6 +596,7 @@ class OnlineDerenderer: model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state @@ -741,6 +743,7 @@ class OnlineDerenderer: model=model, choices=stream_choices, usage=usage, + metrics=generate_chunk.metrics, ) return chunk, updated_state diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b6864897bf..24be5b8ec9 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -512,11 +512,10 @@ def _get_attn_isa( ) if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0: return "amx" + elif supports_arm: + return "neon" elif block_size % 32 == 0: - if supports_arm: - # support ARM NEON FMLA and BFMMLA (bf16) for block size 32 - return "neon" - elif supports_riscv and _riscv_supports_rvv(): + if supports_riscv and _riscv_supports_rvv(): return "rvv" elif supports_vxe: return "vxe" diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 430bd10c99..8af81c8d29 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -634,6 +634,38 @@ class DeepseekV32IndexerMetadata: prefill: DeepseekV32IndexerPrefillMetadata | None = None +@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) +def _kpool_tail_slot_mapping_kernel( + slot_mapping_ptr, + block_table_ptr, + block_table_stride, + query_start_loc_ptr, + positions_ptr, + out_ptr, + num_reqs, + num_actual_tokens, + num_tokens, + kpool, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + if pid < num_reqs: + start = tl.load(query_start_loc_ptr + pid) + # Tokens past the last request's boundary (if any) also map to it. + end = tl.load(query_start_loc_ptr + pid + 1) + end = tl.where(pid == num_reqs - 1, num_actual_tokens, end) + own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64) + for i in range(start, end, BLOCK): + offs = i + tl.arange(0, BLOCK) + mask = offs < end + pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64) + tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask) + else: + offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK) + mask = offs < num_tokens + tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask) + + def compute_kpool_tail_slot_mapping( slot_mapping: torch.Tensor, block_table: torch.Tensor, @@ -646,10 +678,33 @@ def compute_kpool_tail_slot_mapping( ) -> torch.Tensor: """Map every token to its request's one circular tail block.""" if out is None: - out = slot_mapping.clone() + out = torch.empty_like(slot_mapping) else: assert out.shape == slot_mapping.shape - out.copy_(slot_mapping) + if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0: + block = 256 + num_tokens = slot_mapping.shape[0] + num_actual_tokens = min(num_actual_tokens, num_tokens) + grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),) + _kpool_tail_slot_mapping_kernel[grid]( + slot_mapping, + block_table, + block_table.stride(0), + query_start_loc, + positions, + out, + num_reqs, + num_actual_tokens, + num_tokens, + kpool, + BLOCK=block, + num_warps=4, + ) + return out + # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs. + # Production always takes the Triton path above — spec-decode tokens arrive + # flattened token-major, so slot_mapping is 1D there too. + out.copy_(slot_mapping) if num_actual_tokens == 0: return out tokens = torch.arange(num_actual_tokens, device=slot_mapping.device) diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 85dde88f2d..96bfb35a67 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -165,6 +165,9 @@ class SimpleCPUOffloadScheduler: kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks + self.prefix_cacheable_group_ids = ( + self.cpu_kv_cache_config.prefix_cacheable_group_ids + ) self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU # Find the full attention kv group for prefix cache matching. self.fa_gidx = -1 @@ -297,7 +300,7 @@ class SimpleCPUOffloadScheduler: """GPU blocks to keep available (free/offloaded) per step in lazy mode.""" WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks target = 0 - for g in kv_cache_config.kv_cache_groups: + for g in kv_cache_config.prefix_cacheable_groups: spec = g.kv_cache_spec # Only full attention is sharded across DCP ranks; replicated specs # (mamba, sliding window, chunked-local) keep their own block size. @@ -442,6 +445,9 @@ class SimpleCPUOffloadScheduler: # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): + if g not in self.prefix_cacheable_group_ids: + cpu_hit_blocks.append([]) + continue g_block_size = self.group_block_sizes[g] n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -828,6 +834,8 @@ class SimpleCPUOffloadScheduler: num_free = self.cpu_block_pool.get_num_free_blocks() for g, group_gpu_ids in enumerate(block_ids_by_group): + if g not in self.prefix_cacheable_group_ids: + continue if len(gpu_block_ids) >= num_free: break group_manager = self.cpu_coordinator.single_type_managers[g] diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4690c7eae7..ddab7799d4 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -662,6 +662,7 @@ class ModelCudaGraphManager(CudaGraphManager): model_inputs = { "input_ids": input_buffers.input_ids[:num_tokens], "positions": input_buffers.positions[:num_tokens], + "intermediate_tensors": None, **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e76c447ea5..8c736394cc 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1291,7 +1291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Get the number of draft tokens for each request. draft_tokens = scheduler_output.scheduled_spec_decode_tokens num_draft_tokens_per_req = None - if not draft_tokens: + if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1: # No draft token scheduled (common case). total_num_draft_tokens = 0 total_num_logits = num_reqs @@ -1322,9 +1322,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): np.cumsum(num_logits, out=cu_num_logits_np[1:]) cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) - adaptive_verification = ( - self.adaptive_verification if num_draft_tokens_per_req is not None else None - ) + # The general branch also serves a no-draft batch with k != 1, and + # compact_batch's budget is only primed when drafts are scheduled. + adaptive_verification = self.adaptive_verification if draft_tokens else None num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across @@ -1354,7 +1354,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): adaptive_verification.reallocate_drafts(req_ids, idx_mapping) ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens - if draft_tokens: + if num_draft_tokens_per_req is not None: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len ) @@ -1534,7 +1534,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): sample_hidden_states = hidden_states[input_batch.logits_indices] logits = self.model.compute_logits(sample_hidden_states) - if grammar_output is not None: + # A diffusion prefill has no logit rows even when a bitmask row + # arrived for it. + if grammar_output is not None and logits.shape[0] > 0: # Apply grammar bitmask to the logits in-place. assert self.structured_outputs_worker is not None self.structured_outputs_worker.apply_grammar_bitmask( diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index cddb760acd..a86110f2f8 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -3,11 +3,12 @@ import inspect import os -from collections.abc import Iterator +from collections.abc import Callable, Hashable, Iterator from contextlib import contextmanager from contextvars import ContextVar from itertools import accumulate from math import prod +from typing import TypeVar, cast import torch @@ -17,6 +18,7 @@ from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) +T = TypeVar("T") def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: @@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]: class WorkspaceManager: """Manager for workspace allocation. - Manages one workspace buffer per active ``(ubatch, lane)`` slot. - Can be locked to prevent further growth during execution. + Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``. + Can be locked to prevent growth or new resources during execution. """ def __init__( @@ -66,6 +68,9 @@ class WorkspaceManager: self._current_workspaces: list[torch.Tensor | None] = [None] * ( self._num_ubatches * self._num_lanes ) + self._persistent_resources: list[dict[Hashable, object]] = [ + {} for _ in self._current_workspaces + ] self._locked: bool = False @staticmethod @@ -76,10 +81,10 @@ class WorkspaceManager: return workspace.numel() * workspace.element_size() def lock(self) -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After locking, any attempt to allocate a larger workspace will raise - an assertion error. This ensures workspace size is fixed during execution. + Larger scratch requests and new persistent resources raise an assertion + error. Existing resources retain their storage throughout execution. """ self._locked = True if envs.VLLM_DEBUG_WORKSPACE: @@ -113,6 +118,63 @@ class WorkspaceManager: """Check if workspace is locked.""" return self._locked + def _get_workspace_id(self) -> int: + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + return dbo_current_ubatch_id() * self._num_lanes + lane + + def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: + """Return the resource cached for ``key`` in the current ubatch and lane. + + The first request calls ``factory`` and stores its result for the lifetime + of this manager. Later requests return the same object without calling + the factory again. Each ubatch/lane pair has a separate cache. Once the + manager is locked, requesting a key missing from that cache raises an + AssertionError. + + Keys must identify the resource type and configuration. Callers must not + resize or replace tensor storage held by a cached resource. Uses of the + same resource must not overlap; include the CUDA stream in the key when + streams can execute concurrently. + + With torch.compile, call this inside a runtime custom op so resource + selection happens during execution. + """ + resources = self._persistent_resources[self._get_workspace_id()] + if key not in resources: + if self._locked: + raise AssertionError( + f"Workspace is locked but {key!r} was not allocated during " + "warmup for the current ubatch and lane." + ) + resources[key] = factory() + return cast(T, resources[key]) + + def get_persistent( + self, + key: Hashable, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + zero_init: bool = False, + ) -> torch.Tensor: + """Get a fixed-shape tensor, optionally zeroed on first allocation only.""" + factory = torch.zeros if zero_init else torch.empty + workspace = self.get_persistent_resource( + (torch.Tensor, key), + lambda: factory(shape, dtype=dtype, device=self._device), + ) + if workspace.shape != shape or workspace.dtype != dtype: + raise ValueError( + f"Persistent workspace {key!r} has shape={workspace.shape}, " + f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}." + ) + return workspace + def get_simultaneous( self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] ) -> list[torch.Tensor]: @@ -151,14 +213,7 @@ class WorkspaceManager: The current workspace tensor. """ - ubatch_id = dbo_current_ubatch_id() - lane = _workspace_lane.get() - if lane >= self._num_lanes: - raise RuntimeError( - f"Workspace lane {lane} is not configured; manager has " - f"{self._num_lanes} lane(s)." - ) - workspace_id = ubatch_id * self._num_lanes + lane + workspace_id = self._get_workspace_id() current_workspace = self._current_workspaces[workspace_id] current_size = self._workspace_size_bytes(current_workspace) @@ -218,8 +273,8 @@ class WorkspaceManager: get_caller_info(), current_size / _MB, required_bytes / _MB, - ubatch_id, - lane, + workspace_id // self._num_lanes, + workspace_id % self._num_lanes, ) return current_workspace @@ -277,12 +332,11 @@ def init_workspace_manager( def lock_workspace() -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After calling this function, any attempt to allocate a workspace larger - than the current size will raise an AssertionError. This ensures that - workspace size is fixed during execution and prevents unexpected memory - allocations in the hot path. + Larger scratch requests and new persistent resources raise an AssertionError. + All required resources must be initialized in each execution slot during + warmup, so their storage remains fixed during execution. Example: # During initialization