From 00e20e76f775b88f47469ae9fcb0f1ecd7580bb9 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 18 May 2026 14:14:21 -0400 Subject: [PATCH] [Refactor] Remove dead cuda kernels (#42767) Signed-off-by: yewentao256 --- CMakeLists.txt | 1 - csrc/attention/vertical_slash_index.cu | 401 ------------------------- csrc/moe/torch_bindings.cpp | 10 - csrc/ops.h | 26 -- csrc/torch_bindings.cpp | 24 -- vllm/_custom_ops.py | 144 +-------- 6 files changed, 1 insertion(+), 605 deletions(-) delete mode 100644 csrc/attention/vertical_slash_index.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index b63d6e7445..35171b473d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -311,7 +311,6 @@ set(VLLM_EXT_SRC "csrc/attention/paged_attention_v1.cu" "csrc/attention/paged_attention_v2.cu" "csrc/attention/merge_attn_states.cu" - "csrc/attention/vertical_slash_index.cu" "csrc/pos_encoding_kernels.cu" "csrc/activation_kernels.cu" "csrc/layernorm_kernels.cu" diff --git a/csrc/attention/vertical_slash_index.cu b/csrc/attention/vertical_slash_index.cu deleted file mode 100644 index c1b45b143f..0000000000 --- a/csrc/attention/vertical_slash_index.cu +++ /dev/null @@ -1,401 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include - -#include - -#include - -__device__ int64_t save_blocks(int* block_offset, int64_t range_start, - int64_t range_end, int64_t block_size, - int64_t input_block_count, int64_t kv_seqlen) { - if (range_start >= kv_seqlen) { - return input_block_count; - } - if (range_end > kv_seqlen) { - range_end = kv_seqlen; - } - int64_t current_block_count = input_block_count; - for (int idx = range_start; idx < range_end; idx += block_size) { - block_offset[current_block_count++] = idx; - } - return current_block_count; -} - -__global__ void convert_vertical_slash_indexes_kernel( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, - int64_t NNZ_V, int64_t NNZ_S, - bool causal // True for intra, False for succ -) { - const int batch_idx = blockIdx.y; - const int head_idx = blockIdx.x; - const int group_idx = blockIdx.z; - - int64_t q_seqlen = q_seqlens[batch_idx]; - int64_t kv_seqlen = kv_seqlens[batch_idx]; - int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x; - int64_t start_m = block_idx_m * BLOCK_SIZE_M; - if (start_m >= q_seqlen) { - return; - } - int64_t end_m = start_m + BLOCK_SIZE_M; - vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V; - slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S; - int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m; - block_count += row_offset; - block_offset += row_offset * NNZ_S; - column_count += row_offset; - column_index += row_offset * NNZ_V; - - bool has_slash = true; - int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0; - int64_t s = 0, v = 0; - int64_t v_idx = vertical_indexes[v++]; - int64_t s_idx = slash_indexes[s++]; - if (causal) { - while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false; - s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M); - } else { - while (s_idx >= end_m + kv_seqlen && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + kv_seqlen) has_slash = false; - s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M); - } - - int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx; - if (!has_slash) { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - } - - bool slash_finished = false; - while (1) { - if (v_idx < range_end) { - if (v_idx < range_start) { - column_index[tmp_col_cnt++] = v_idx; - } - if (v < NNZ_V) { - v_idx = vertical_indexes[v++]; - } else { - if (causal) - v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen); - else - v_idx = end_m + BLOCK_SIZE_N + kv_seqlen; - } - } else { - if ((s < NNZ_S && causal) || - (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) { - if (causal) - s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], - BLOCK_SIZE_M); - else - s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M); - } else { - if (v == NNZ_V || (v_idx > range_start && causal)) { - // add the last vertical if no more slash - if (v == NNZ_V && !causal && v_idx < kv_seqlen) { - column_index[tmp_col_cnt++] = v_idx; - } - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - break; - } else { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - // if slash_finished but there are vertical left, save current - // blocks - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - slash_finished = true; - } - } - if (!slash_finished) { - if (s_idx > range_end + BLOCK_SIZE_M) { - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = s_idx - BLOCK_SIZE_M; - range_end = s_idx; - } else if (s_idx > range_end) { - range_end += BLOCK_SIZE_M; - } - } - } - } - - block_count[0] = tmp_blk_cnt; - column_count[0] = tmp_col_cnt; -} - -void convert_vertical_slash_indexes_64x64( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) { - const int N_THREADS = 64; - const dim3 dimBlock(N_THREADS); - const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS); - convert_vertical_slash_indexes_kernel<<>>( - q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count, - block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, - BLOCK_SIZE_N, NNZ_V, NNZ_S, causal); -} - -/** - * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490. - * - * This function builds the index of each row of blocks from vertical indices - * and slash indices. The vertical indices are treated as points, while the - * slash indices are converted as ranges. The output consists of the merged - * ranges and separate column indices, where the ranges are represented by - * block indices. - * - * The implementation is referenced from the original MInference repo: - * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu. - */ -void convert_vertical_slash_indexes( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int64_t context_size, int64_t block_size_M, int64_t block_size_N, - bool causal) { - cudaSetDevice(q_seqlens.get_device()); - - int batch_size = slash_indexes.size(0); - int num_heads = slash_indexes.size(1); - int nnz_slash = slash_indexes.size(2); - int nnz_vertical = vertical_indexes.size(2); - int num_rows = (context_size + block_size_M - 1) / block_size_M; - - convert_vertical_slash_indexes_64x64( - q_seqlens.data_ptr(), kv_seqlens.data_ptr(), - vertical_indexes.data_ptr(), slash_indexes.data_ptr(), - block_count.data_ptr(), block_offset.data_ptr(), - column_count.data_ptr(), column_index.data_ptr(), batch_size, - num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash, - causal); -} - -__global__ void convert_vertical_slash_indexes_kernel_mergehead( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - const int* per_head_vertical_topkv, const int* per_head_slash_topkv, - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, - int64_t NNZ_V, int64_t NNZ_S, - bool causal // True for intra, False for succ -) { - const int batch_idx = blockIdx.y; - const int head_idx = blockIdx.x; - const int group_idx = blockIdx.z; - - int64_t q_seqlen = q_seqlens[batch_idx]; - int64_t kv_seqlen = kv_seqlens[batch_idx]; - int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x; - int64_t start_m = block_idx_m * BLOCK_SIZE_M; - if (start_m >= q_seqlen) { - return; - } - int64_t end_m = start_m + BLOCK_SIZE_M; - vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V; - slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S; - int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m; - block_count += row_offset; - block_offset += row_offset * NNZ_S; - column_count += row_offset; - column_index += row_offset * NNZ_V; - - // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S - // above is buffer size, use to compute offset) - NNZ_S = per_head_slash_topkv[head_idx]; - NNZ_V = per_head_vertical_topkv[head_idx]; - - bool has_slash = true; - int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0; - int64_t s = 0, v = 0; - int64_t v_idx = vertical_indexes[v++]; - int64_t s_idx = slash_indexes[s++]; - if (causal) { - while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false; - s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M); - } else { - while (s_idx >= end_m + kv_seqlen && s < NNZ_S) { - s_idx = slash_indexes[s++]; - } - if (s_idx > end_m + kv_seqlen) has_slash = false; - s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M); - } - - int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx; - if (!has_slash) { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - } - - bool slash_finished = false; - while (1) { - if (v_idx < range_end) { - if (v_idx < range_start) { - column_index[tmp_col_cnt++] = v_idx; - } - if (v < NNZ_V) { - v_idx = vertical_indexes[v++]; - } else { - if (causal) - v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen); - else - v_idx = end_m + BLOCK_SIZE_N + kv_seqlen; - } - } else { - if ((s < NNZ_S && causal) || - (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) { - if (causal) - s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], - BLOCK_SIZE_M); - else - s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M); - } else { - if (v == NNZ_V || (v_idx > range_start && causal)) { - // add the last vertical if no more slash - if (v == NNZ_V && !causal && v_idx < kv_seqlen) { - column_index[tmp_col_cnt++] = v_idx; - } - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - break; - } else { - if (causal) { - range_start = (kv_seqlen - q_seqlen) + end_m; - range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N; - } else { - // if slash_finished but there are vertical left, save current - // blocks - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = kv_seqlen; - range_end = kv_seqlen + BLOCK_SIZE_N; - } - slash_finished = true; - } - } - if (!slash_finished) { - if (s_idx > range_end + BLOCK_SIZE_M) { - tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, - BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen); - range_start = s_idx - BLOCK_SIZE_M; - range_end = s_idx; - } else if (s_idx > range_end) { - range_end += BLOCK_SIZE_M; - } - } - } - } - - block_count[0] = tmp_blk_cnt; - column_count[0] = tmp_col_cnt; -} - -void convert_vertical_slash_indexes_64x64_mergehead( - const int* q_seqlens, // [BATCH, ] - const int* kv_seqlens, // [BATCH, ] - const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int* per_head_vertical_topkv, int* per_head_slash_topkv, - int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S] - int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)] - int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V] - int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, - int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) { - const int N_THREADS = 64; - const dim3 dimBlock(N_THREADS); - const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS); - convert_vertical_slash_indexes_kernel_mergehead<<>>( - q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, - per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset, - column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N, - NNZ_V, NNZ_S, causal); -} - -/** - * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490. - * - * Like the above convert_vertical_slash_indexes, but with - * pre-computed vertical and slash counts. - */ -void convert_vertical_slash_indexes_mergehead( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - torch::Tensor vertical_indices_count, // [N_HEADS, ] - torch::Tensor slash_indices_count, // [N_HEADS, ] - int64_t context_size, int64_t block_size_M, int64_t block_size_N, - bool causal) { - cudaSetDevice(q_seqlens.get_device()); - - int batch_size = slash_indexes.size(0); - int num_heads = slash_indexes.size(1); - int nnz_slash = slash_indexes.size(2); - int nnz_vertical = vertical_indexes.size(2); - int num_rows = (context_size + block_size_M - 1) / block_size_M; - - convert_vertical_slash_indexes_64x64_mergehead( - q_seqlens.data_ptr(), kv_seqlens.data_ptr(), - vertical_indexes.data_ptr(), slash_indexes.data_ptr(), - vertical_indices_count.data_ptr(), - slash_indices_count.data_ptr(), block_count.data_ptr(), - block_offset.data_ptr(), column_count.data_ptr(), - column_index.data_ptr(), batch_size, num_heads, num_rows, - block_size_M, block_size_N, nnz_vertical, nnz_slash, causal); -} diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 913fa07bbc..b8b91e3c38 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -92,16 +92,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { "bool use_fp32_reduce, bool is_zp_float," "int thread_k, int thread_n, int blocks_per_sm) -> Tensor"); - m.def( - "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, " - "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! " - "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, " - "int b_q_type, SymInt size_m, " - "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int " - "topk, " - "int moe_block_size, bool replicate_input, bool apply_weights)" - " -> Tensor"); - m.def( "moe_permute(Tensor input, Tensor topk_ids," "Tensor token_expert_indices, Tensor? expert_map, int n_expert," diff --git a/csrc/ops.h b/csrc/ops.h index 2a0819618e..c0689534be 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -60,32 +60,6 @@ void merge_attn_states( const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, const std::optional prefill_tokens_with_context, const std::optional& output_scale = std::nullopt); -#ifndef USE_ROCM -void convert_vertical_slash_indexes( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - int64_t context_size, int64_t block_size_M, int64_t block_size_N, - bool causal); - -void convert_vertical_slash_indexes_mergehead( - torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S] - torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS] - torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V] - torch::Tensor q_seqlens, // [BATCH, ] - torch::Tensor kv_seqlens, // [BATCH, ] - torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V] - torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S] - torch::Tensor vertical_indices_count, // [N_HEADS, ] - torch::Tensor slash_indices_count, int64_t context_size, - int64_t block_size_M, int64_t block_size_N, bool causal); -#endif void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, double epsilon); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index f0f3641d29..dde4b3028b 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -76,30 +76,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { " int!? prefill_tokens_with_context," " Tensor? output_scale=None) -> ()"); ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states); -#ifndef USE_ROCM - ops.def( - "convert_vertical_slash_indexes(" - " Tensor! block_count, Tensor! block_offset, " - " Tensor! column_count, Tensor! column_index, " - " Tensor q_seqlens, Tensor q_seqlens, " - " Tensor vertical_indexes, Tensor slash_indexes, " - " int context_size, int block_size_M, int block_size_N, " - " bool causal) -> ()"); - ops.impl("convert_vertical_slash_indexes", torch::kCUDA, - &convert_vertical_slash_indexes); - - ops.def( - "convert_vertical_slash_indexes_mergehead(" - " Tensor! block_count, Tensor! block_offset, " - " Tensor! column_count, Tensor! column_index, " - " Tensor q_seqlens, Tensor q_seqlens, " - " Tensor vertical_indexes, Tensor slash_indexes, " - " Tensor vertical_indices_count, Tensor slash_indices_count, " - " int context_size, int block_size_M, int block_size_N, " - " bool causal) -> ()"); - ops.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA, - &convert_vertical_slash_indexes_mergehead); -#endif // Activation ops // Activation function used in SwiGLU. diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2346cb40a0..6299cf15d6 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -280,123 +280,6 @@ def merge_attn_states( ) -def convert_vertical_slash_indexes( - q_seqlens: torch.Tensor, # [BATCH, ] - kv_seqlens: torch.Tensor, # [BATCH, ] - vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] - slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] - context_size: int, - block_size_M: int, - block_size_N: int, - causal: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size = slash_indexes.size(0) - num_heads = slash_indexes.size(1) - nnz_slash = slash_indexes.size(2) - nnz_vertical = vertical_indexes.size(2) - num_rows = (context_size + block_size_M - 1) // block_size_M - - block_count = torch.zeros( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - block_offset = torch.zeros( - batch_size, - num_heads, - num_rows, - nnz_slash, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - column_count = torch.zeros( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - column_index = torch.zeros( - batch_size, - num_heads, - num_rows, - nnz_vertical, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - - torch.ops._C.convert_vertical_slash_indexes( - block_count, - block_offset, - column_count, - column_index, - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - context_size, - block_size_M, - block_size_N, - causal, - ) - return block_count, block_offset, column_count, column_index - - -def convert_vertical_slash_indexes_mergehead( - q_seqlens: torch.Tensor, # [BATCH, ] - kv_seqlens: torch.Tensor, # [BATCH, ] - vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] - slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] - # [N_HEADS] : different head use different number of indices - vertical_indices_count: torch.Tensor, - slash_indices_count: torch.Tensor, - context_size: int, - block_size_M: int, - block_size_N: int, - causal: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size = slash_indexes.size(0) - num_heads = slash_indexes.size(1) - nnz_slash = slash_indexes.size(2) - nnz_vertical = vertical_indexes.size(2) - num_rows = (context_size + block_size_M - 1) // block_size_M - - block_count = torch.empty( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - block_offset = torch.empty( - batch_size, - num_heads, - num_rows, - nnz_slash, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - column_count = torch.empty( - batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device - ) - column_index = torch.empty( - batch_size, - num_heads, - num_rows, - nnz_vertical, - dtype=q_seqlens.dtype, - device=q_seqlens.device, - ) - - torch.ops._C.convert_vertical_slash_indexes_mergehead( - block_count, - block_offset, - column_count, - column_index, - q_seqlens, - kv_seqlens, - vertical_indexes, - slash_indexes, - vertical_indices_count, - slash_indices_count, - context_size, - block_size_M, - block_size_N, - causal, - ) - return block_count, block_offset, column_count, column_index - - # pos encoding ops def rotary_embedding( positions: torch.Tensor, @@ -2658,32 +2541,7 @@ def moe_wna16_marlin_gemm( ) -if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe"): - - @register_fake("_moe_C::marlin_gemm_moe") - def marlin_gemm_moe_fake( - a: torch.Tensor, - b_q_weights: torch.Tensor, - sorted_ids: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - b_scales: torch.Tensor, - b_zero_points: torch.Tensor, - g_idx: torch.Tensor, - perm: torch.Tensor, - workspace: torch.Tensor, - b_q_type: ScalarType, - size_m: torch.SymInt, - size_n: torch.SymInt, - size_k: torch.SymInt, - is_k_full: bool, - num_experts: int, - topk: int, - moe_block_size: int, - replicate_input: bool, - apply_weights: bool, - ) -> torch.Tensor: - return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device) +if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"): @register_fake("_moe_C::moe_wna16_marlin_gemm") def moe_wna16_marlin_gemm_fake(