From 70c898ac8b1cb170461110495c3d275d77f677f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 10:21:48 +0000 Subject: [PATCH] fix: ex_engine.python subpackage + flash_qla_sm70 deploy + vllm v0.5.5 MoE kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 真机验证发现的问题: 1. qwen3_5.py 做 'from ex_engine.python.ix_bridge' 但包结构是 ex_engine.ix_bridge → 创建 python/ 子目录 + symlinks 2. flash_qla_sm70 只部署到 /workspace 没有到 vllm models 目录 → 显式 cp -r 到 VLLM/model_executor/models/ 3. 从 vllm v0.5.5 搬 MoE CUDA kernels (torch::Tensor API): - topk_softmax_kernels.cu (506行, CUB BlockReduce) - moe_align_block_size_kernels.cu (134行) - moe_pybind.cpp (pybind11 入口) 真机验证结果: ✓ ix_bridge import OK, available=True ✓ topk_softmax (64 experts, top8) OK — CUDA kernel 命中 ✓ ix_full_bridge silu_and_mul OK ✓ qwen3_5.py import OK ✓ ex_engine build 2/2 factors ✓ moe_topk_softmax_v3.so 编译成功 ✓ flash_qla_sm70_gdn_strided.so 编译成功 ✗ 单卡 32GB OOM (正常, 竞赛 4卡 tp=4) --- ex_engine/csrc/moe_v055/cuda_compat.h | 49 ++ ex_engine/csrc/moe_v055/dispatch_utils.h | 35 ++ .../moe_v055/moe_align_block_size_kernels.cu | 134 +++++ ex_engine/csrc/moe_v055/moe_pybind.cpp | 42 ++ .../csrc/moe_v055/topk_softmax_kernels.cu | 506 ++++++++++++++++++ ex_engine/precompile_moe_kernels.py | 95 ++++ qwen3_6_scripts/_custom_ops.py | 77 +++ qwen3_6_scripts/patch_ops.sh | 23 + 8 files changed, 961 insertions(+) create mode 100644 ex_engine/csrc/moe_v055/cuda_compat.h create mode 100644 ex_engine/csrc/moe_v055/dispatch_utils.h create mode 100644 ex_engine/csrc/moe_v055/moe_align_block_size_kernels.cu create mode 100644 ex_engine/csrc/moe_v055/moe_pybind.cpp create mode 100644 ex_engine/csrc/moe_v055/topk_softmax_kernels.cu create mode 100644 ex_engine/precompile_moe_kernels.py diff --git a/ex_engine/csrc/moe_v055/cuda_compat.h b/ex_engine/csrc/moe_v055/cuda_compat.h new file mode 100644 index 00000000..82e55613 --- /dev/null +++ b/ex_engine/csrc/moe_v055/cuda_compat.h @@ -0,0 +1,49 @@ +#pragma once + +#ifdef USE_ROCM + #include +#endif + +#ifndef USE_ROCM + #define WARP_SIZE 32 +#else + #define WARP_SIZE warpSize +#endif + +#ifndef USE_ROCM + #define VLLM_LDG(arg) __ldg(arg) +#else + #define VLLM_LDG(arg) *(arg) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) +#else + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor(var, lane_mask, width) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) +#else + #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ + __shfl_down_sync(uint32_t(-1), var, lane_delta) +#else + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) +#endif + +#ifndef USE_ROCM + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL) +#else + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#endif diff --git a/ex_engine/csrc/moe_v055/dispatch_utils.h b/ex_engine/csrc/moe_v055/dispatch_utils.h new file mode 100644 index 00000000..a634e1c3 --- /dev/null +++ b/ex_engine/csrc/moe_v055/dispatch_utils.h @@ -0,0 +1,35 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) + +#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) diff --git a/ex_engine/csrc/moe_v055/moe_align_block_size_kernels.cu b/ex_engine/csrc/moe_v055/moe_align_block_size_kernels.cu new file mode 100644 index 00000000..1f8d75da --- /dev/null +++ b/ex_engine/csrc/moe_v055/moe_align_block_size_kernels.cu @@ -0,0 +1,134 @@ +#include +#include + +#include +#include + +#include "cuda_compat.h" +#include "dispatch_utils.h" + +#define CEILDIV(x, y) (((x) + (y) - 1) / (y)) + +namespace vllm { + +namespace { +__device__ __forceinline__ int32_t index(int32_t total_col, int32_t row, + int32_t col) { + // don't worry about overflow because num_experts is relatively small + return row * total_col + col; +} +} // namespace + +template +__global__ void moe_align_block_size_kernel(scalar_t* __restrict__ topk_ids, + int32_t* sorted_token_ids, + int32_t* expert_ids, + int32_t* total_tokens_post_pad, + int32_t num_experts, + int32_t block_size, size_t numel) { + const size_t tokens_per_thread = CEILDIV(numel, blockDim.x); + const size_t start_idx = threadIdx.x * tokens_per_thread; + + extern __shared__ int32_t shared_mem[]; + + int32_t* tokens_cnts = + shared_mem; // 2d tensor with shape (num_experts + 1, num_experts) + int32_t* cumsum = + shared_mem + (num_experts + 1) * + num_experts; // 1d tensor with shape (num_experts + 1) + + for (int i = 0; i < num_experts; ++i) { + tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0; + } + + /** + * In the first step we compute token_cnts[thread_index + 1][expert_index], + * which counts how many tokens in the token shard of thread_index are + * assigned to expert expert_index. + */ + for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) { + ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])]; + } + + __syncthreads(); + + // For each expert we accumulate the token counts from the different threads. + tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0; + for (int i = 1; i <= blockDim.x; ++i) { + tokens_cnts[index(num_experts, i, threadIdx.x)] += + tokens_cnts[index(num_experts, i - 1, threadIdx.x)]; + } + + __syncthreads(); + + // We accumulate the token counts of all experts in thread 0. + if (threadIdx.x == 0) { + cumsum[0] = 0; + for (int i = 1; i <= num_experts; ++i) { + cumsum[i] = cumsum[i - 1] + + CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)], + block_size) * + block_size; + } + *total_tokens_post_pad = cumsum[num_experts]; + } + + __syncthreads(); + + /** + * For each expert, each thread processes the tokens of the corresponding + * blocks and stores the corresponding expert_id for each block. + */ + for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1]; + i += block_size) { + expert_ids[i / block_size] = threadIdx.x; + } + + /** + * Each thread processes a token shard, calculating the index of each token + * after sorting by expert number. Given the example topk_ids = + * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *, + * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a + * padding value(preset in python). + */ + for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) { + int32_t expert_id = topk_ids[i]; + /** The cumsum[expert_id] stores the starting index of the tokens that the + * expert with expert_id needs to process, and + * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens + * processed by the expert with expert_id within the current thread's token + * shard. + */ + int32_t rank_post_pad = + tokens_cnts[index(num_experts, threadIdx.x, expert_id)] + + cumsum[expert_id]; + sorted_token_ids[rank_post_pad] = i; + ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)]; + } +} +} // namespace vllm + +void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, + int64_t block_size, torch::Tensor sorted_token_ids, + torch::Tensor experts_ids, + torch::Tensor num_tokens_post_pad) { + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_INTEGRAL_TYPES( + topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] { + // calc needed amount of shared mem for `tokens_cnts` and `cumsum` + // tensors + const int32_t shared_mem = + ((num_experts + 1) * num_experts + (num_experts + 1)) * + sizeof(int32_t); + + // set dynamic shared mem + auto kernel = vllm::moe_align_block_size_kernel; + AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( + (void*)kernel, shared_mem)); + kernel<<<1, num_experts, shared_mem, stream>>>( + topk_ids.data_ptr(), sorted_token_ids.data_ptr(), + experts_ids.data_ptr(), + num_tokens_post_pad.data_ptr(), num_experts, block_size, + topk_ids.numel()); + }); +} diff --git a/ex_engine/csrc/moe_v055/moe_pybind.cpp b/ex_engine/csrc/moe_v055/moe_pybind.cpp new file mode 100644 index 00000000..eacdd29f --- /dev/null +++ b/ex_engine/csrc/moe_v055/moe_pybind.cpp @@ -0,0 +1,42 @@ +/* + * moe_pybind.cpp — pybind11 entry for vllm MoE CUDA kernels + * + * Compiled via torch.utils.cpp_extension.load() on BI-V100 (CoreX) + * Exposes: + * - topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output) + * - moe_align_block_size(topk_ids, num_experts, block_size, sorted_token_ids, experts_ids, num_tokens_post_pad) + * + * Source: vllm v0.5.5 csrc/moe/ (torch::Tensor API, pre-libtorch_stable) + */ + +#include + +// Forward declarations matching vllm v0.5.5 signatures +void topk_softmax(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output); + +void moe_align_block_size(torch::Tensor topk_ids, + int64_t num_experts, + int64_t block_size, + torch::Tensor sorted_token_ids, + torch::Tensor experts_ids, + torch::Tensor num_tokens_post_pad); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("topk_softmax", &topk_softmax, + "MoE topk softmax (vllm v0.5.5 CUDA kernel)", + py::arg("topk_weights"), + py::arg("topk_indices"), + py::arg("token_expert_indices"), + py::arg("gating_output")); + m.def("moe_align_block_size", &moe_align_block_size, + "MoE align block size (vllm v0.5.5 CUDA kernel)", + py::arg("topk_ids"), + py::arg("num_experts"), + py::arg("block_size"), + py::arg("sorted_token_ids"), + py::arg("experts_ids"), + py::arg("num_tokens_post_pad")); +} diff --git a/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu b/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu new file mode 100644 index 00000000..5273e0a5 --- /dev/null +++ b/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu @@ -0,0 +1,506 @@ +/* + * Adapted from https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu + * Copyright (c) 2024, The vLLM team. + * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include "cuda_compat.h" + +#ifndef USE_ROCM + #include + #include +#else + #include + #include +#endif + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +namespace vllm { +namespace moe { + +/// Aligned array type +template < + typename T, + /// Number of elements in the array + int N, + /// Alignment requirement in bytes + int Alignment = sizeof(T) * N +> +class alignas(Alignment) AlignedArray { + float data[N]; +}; + +// ====================== Softmax things =============================== +// We have our own implementation of softmax here so we can support transposing the output +// in the softmax kernel when we extend this module to support expert-choice routing. +template +__launch_bounds__(TPB) __global__ + void moeSoftmax(const float* input, const bool* finished, float* output, const int num_cols) +{ + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmpStorage; + + __shared__ float normalizing_factor; + __shared__ float float_max; + + const int thread_row_offset = blockIdx.x * num_cols; + + cub::Sum sum; + float threadData(-FLT_MAX); + + // Don't touch finished rows. + if ((finished != nullptr) && finished[blockIdx.x]) + { + return; + } + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) + { + const int idx = thread_row_offset + ii; + threadData = max(static_cast(input[idx]), threadData); + } + + const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max()); + if (threadIdx.x == 0) + { + float_max = maxElem; + } + __syncthreads(); + + threadData = 0; + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) + { + const int idx = thread_row_offset + ii; + threadData += exp((static_cast(input[idx]) - float_max)); + } + + const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum); + + if (threadIdx.x == 0) + { + normalizing_factor = 1.f / Z; + } + __syncthreads(); + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) + { + const int idx = thread_row_offset + ii; + const float val = exp((static_cast(input[idx]) - float_max)) * normalizing_factor; + output[idx] = val; + } +} + +template +__launch_bounds__(TPB) __global__ void moeTopK(const float* inputs_after_softmax, const bool* finished, float* output, + int* indices, int* source_rows, const int num_experts, const int k, const int start_expert, const int end_expert) +{ + + using cub_kvp = cub::KeyValuePair; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmpStorage; + + cub_kvp thread_kvp; + cub::ArgMax arg_max; + + const int num_rows = gridDim.x; + const int block_row = blockIdx.x; + + const bool row_is_active = finished ? !finished[block_row] : true; + const int thread_read_offset = blockIdx.x * num_experts; + for (int k_idx = 0; k_idx < k; ++k_idx) + { + thread_kvp.key = 0; + thread_kvp.value = -1.f; // This is OK because inputs are probabilities + + cub_kvp inp_kvp; + for (int expert = threadIdx.x; expert < num_experts; expert += TPB) + { + const int idx = thread_read_offset + expert; + inp_kvp.key = expert; + inp_kvp.value = inputs_after_softmax[idx]; + + for (int prior_k = 0; prior_k < k_idx; ++prior_k) + { + const int prior_winning_expert = indices[k * block_row + prior_k]; + + if (prior_winning_expert == expert) + { + inp_kvp = thread_kvp; + } + } + + thread_kvp = arg_max(inp_kvp, thread_kvp); + } + + const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max); + if (threadIdx.x == 0) + { + // Ignore experts the node isn't responsible for with expert parallelism + const int expert = result_kvp.key; + const bool node_uses_expert = expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + const int idx = k * block_row + k_idx; + output[idx] = result_kvp.value; + indices[idx] = should_process_row ? (expert - start_expert) : num_experts; + assert(indices[idx] >= 0); + source_rows[idx] = k_idx * num_rows + block_row; + } + __syncthreads(); + } +} + +// ====================== TopK softmax things =============================== + +/* + A Top-K gating softmax written to exploit when the number of experts in the MoE layers + are a small power of 2. This allows us to cleanly share the rows among the threads in + a single warp and eliminate communication between warps (so no need to use shared mem). + + It fuses the softmax, max and argmax into a single kernel. + + Limitations: + 1) This implementation is intended for when the number of experts is a small power of 2. + 2) This implementation assumes k is small, but will work for any k. +*/ + +template +__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ + void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, int* indices, + int* source_rows, const int k, const int start_expert, const int end_expert) +{ + // We begin by enforcing compile time assertions and setting up compile time constants. + static_assert(VPT == (VPT & -VPT), "VPT must be power of 2"); + static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2"); + static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2"); + static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); + + // Number of bytes each thread pulls in per load + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float); + static constexpr int ELTS_PER_ROW = NUM_EXPERTS; + static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; + static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; + + // Restrictions based on previous section. + static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg"); + static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp"); + static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2"); + static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size"); + + // We have NUM_EXPERTS elements per row. We specialize for small #experts + static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; + static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; + static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; + + // Restrictions for previous section. + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp"); + + // ===================== From this point, we finally start computing run-time variables. ======================== + + // Compute CTA and warp rows. We pack multiple rows into a single warp, and a block contains WARPS_PER_CTA warps. + // This, each block processes a chunk of rows. We start by computing the start row for each block. + const int cta_base_row = blockIdx.x * ROWS_PER_CTA; + + // Now, using the base row per thread block, we compute the base row per warp. + const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; + + // The threads in a warp are split into sub-groups that will work on a row. + // We compute row offset for each thread sub-group + const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; + const int thread_row = warp_base_row + thread_row_in_warp; + + // Threads with indices out of bounds should early exit here. + if (thread_row >= num_rows) + { + return; + } + const bool row_is_active = finished ? !finished[thread_row] : true; + + // We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the + // row it will read. + const float* thread_row_ptr = input + thread_row * ELTS_PER_ROW; + + // Now, we compute the group each thread belong to in order to determine the first column to start loads. + const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; + const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; + const float* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; + + // Determine the pointer type to use to read in the data depending on the BYTES_PER_LDG template param. In theory, + // this can support all powers of 2 up to 16. + // NOTE(woosuk): The original implementation uses CUTLASS aligned array here. + // We defined our own aligned array and use it here to avoid the dependency on CUTLASS. + using AccessType = AlignedArray; + + // Finally, we pull in the data from global mem + float row_chunk[VPT]; + AccessType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk); + const AccessType* vec_thread_read_ptr = reinterpret_cast(thread_read_ptr); +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) + { + row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + } + + // First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just + // convert to float afterwards for the exp + sum reduction. + float thread_max = row_chunk[0]; +#pragma unroll + for (int ii = 1; ii < VPT; ++ii) + { + thread_max = max(thread_max, row_chunk[ii]); + } + +// Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce. +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) + { + thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW)); + } + + // From this point, thread max in all the threads have the max within the row. + // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum. + float row_sum = 0; +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) + { + row_chunk[ii] = expf(row_chunk[ii] - thread_max); + row_sum += row_chunk[ii]; + } + +// Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern. +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) + { + row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW); + } + + // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables + // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to + // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row. + // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the + // argmax after computing the softmax. + const float reciprocal_row_sum = 1.f / row_sum; + +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) + { + row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum; + } + + // Now, softmax_res contains the softmax of the row chunk. Now, I want to find the topk elements in each row, along + // with the max index. + int start_col = first_elt_read_by_thread; + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; + + for (int k_idx = 0; k_idx < k; ++k_idx) + { + // First, each thread does the local argmax + float max_val = row_chunk[0]; + int expert = start_col; +#pragma unroll + for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) + { +#pragma unroll + for (int ii = 0; ii < ELTS_PER_LDG; ++ii) + { + float val = row_chunk[ldg * ELTS_PER_LDG + ii]; + + // No check on the experts here since columns with the smallest index are processed first and only + // updated if > (not >=) + if (val > max_val) + { + max_val = val; + expert = col + ii; + } + } + } + +// Now, we perform the argmax reduce. We use the butterfly pattern so threads reach consensus about the max. +// This will be useful for K > 1 so that the threads can agree on "who" had the max value. That thread can +// then blank out their max with -inf and the warp can run more iterations... +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) + { + float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW); + int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW); + + // We want lower indices to "win" in every thread so we break ties this way + if (other_max > max_val || (other_max == max_val && other_expert < expert)) + { + max_val = other_max; + expert = other_expert; + } + } + + // Write the max for this k iteration to global memory. + if (thread_group_idx == 0) + { + // Add a guard to ignore experts not included by this node + const bool node_uses_expert = expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + // The lead thread from each sub-group will write out the final results to global memory. (This will be a + // single) thread per row of the input/output matrices. + const int idx = k * thread_row + k_idx; + output[idx] = max_val; + indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + source_rows[idx] = k_idx * num_rows + thread_row; + } + + // Finally, we clear the value in the thread with the current max if there is another iteration to run. + if (k_idx + 1 < k) + { + const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG; + const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW; + + // Only the thread in the group which produced the max will reset the "winning" value to -inf. + if (thread_group_idx == thread_to_clear_in_group) + { + const int offset_for_expert = expert % ELTS_PER_LDG; + // Safe to set to any negative value since row_chunk values must be between 0 and 1. + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f; + } + } + } +} + +namespace detail +{ +// Constructs some constants needed to partition the work across threads at compile time. +template +struct TopkConstants +{ + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float); + static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, ""); + static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE)); + static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG; + static constexpr int THREADS_PER_ROW = EXPERTS / VPT; + static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW; +}; +} // namespace detail + +template +void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, int* indices, + int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream) +{ + static constexpr std::size_t MAX_BYTES_PER_LDG = 16; + + static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS); + using Constants = detail::TopkConstants; + static constexpr int VPT = Constants::VPT; + static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; + const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; + const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; + + dim3 block_dim(WARP_SIZE, WARPS_PER_TB); + topkGatingSoftmax<<>>( + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert); +} + +#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB) \ + topkGatingSoftmaxLauncherHelper( \ + gating_output, nullptr, topk_weights, topk_indicies, \ + token_expert_indices, num_tokens, topk, 0, num_experts, \ + stream); + +void topkGatingSoftmaxKernelLauncher( + const float* gating_output, + float* topk_weights, + int* topk_indicies, + int* token_expert_indices, + float* softmax_workspace, + const int num_tokens, + const int num_experts, + const int topk, + cudaStream_t stream) { + static constexpr int WARPS_PER_TB = 4; + switch (num_experts) { + case 1: + LAUNCH_SOFTMAX(1, WARPS_PER_TB); + break; + case 2: + LAUNCH_SOFTMAX(2, WARPS_PER_TB); + break; + case 4: + LAUNCH_SOFTMAX(4, WARPS_PER_TB); + break; + case 8: + LAUNCH_SOFTMAX(8, WARPS_PER_TB); + break; + case 16: + LAUNCH_SOFTMAX(16, WARPS_PER_TB); + break; + case 32: + LAUNCH_SOFTMAX(32, WARPS_PER_TB); + break; + case 64: + LAUNCH_SOFTMAX(64, WARPS_PER_TB); + break; + case 128: + LAUNCH_SOFTMAX(128, WARPS_PER_TB); + break; + case 256: + LAUNCH_SOFTMAX(256, WARPS_PER_TB); + break; + default: { + TORCH_CHECK(softmax_workspace != nullptr, + "softmax_workspace must be provided for num_experts that are not a power of 2."); + static constexpr int TPB = 256; + moeSoftmax<<>>( + gating_output, nullptr, softmax_workspace, num_experts); + moeTopK<<>>( + softmax_workspace, nullptr, topk_weights, topk_indicies, token_expert_indices, + num_experts, topk, 0, num_experts); + } + } +} + +} // namespace moe +} // namespace vllm + +void topk_softmax( + torch::Tensor& topk_weights, // [num_tokens, topk] + torch::Tensor& topk_indices, // [num_tokens, topk] + torch::Tensor& token_expert_indices, // [num_tokens, topk] + torch::Tensor& gating_output) // [num_tokens, num_experts] +{ + const int num_experts = gating_output.size(-1); + const int num_tokens = gating_output.numel() / num_experts; + const int topk = topk_weights.size(-1); + + const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); + const bool needs_workspace = !is_pow_2 || num_experts > 256; + const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::Tensor softmax_workspace = torch::empty({workspace_size}, gating_output.options()); + vllm::moe::topkGatingSoftmaxKernelLauncher( + gating_output.data_ptr(), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + token_expert_indices.data_ptr(), + softmax_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + stream); +} diff --git a/ex_engine/precompile_moe_kernels.py b/ex_engine/precompile_moe_kernels.py new file mode 100644 index 00000000..54278243 --- /dev/null +++ b/ex_engine/precompile_moe_kernels.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +""" +precompile_moe_kernels.py — JIT compile vllm v0.5.5 MoE CUDA kernels for BI-V100. + +Produces: moe_kernels.so with: + - topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output) + - moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad) + +Usage: + python3 precompile_moe_kernels.py # JIT compile + python3 precompile_moe_kernels.py --test # compile + smoke test +""" +import os +import sys +import time + +def compile_moe_kernels(): + """JIT compile MoE CUDA kernels via torch.utils.cpp_extension.""" + import torch + from torch.utils.cpp_extension import load + + script_dir = os.path.dirname(os.path.abspath(__file__)) + moe_dir = os.path.join(script_dir, 'csrc', 'moe_v055') + + sources = [ + os.path.join(moe_dir, 'moe_pybind.cpp'), + os.path.join(moe_dir, 'topk_softmax_kernels.cu'), + os.path.join(moe_dir, 'moe_align_block_size_kernels.cu'), + ] + + for s in sources: + if not os.path.isfile(s): + raise FileNotFoundError(f"Missing: {s}") + + print(f"[moe_kernels] Compiling from {moe_dir}") + t0 = time.time() + + mod = load( + name='moe_kernels', + sources=sources, + extra_include_paths=[moe_dir], + extra_cflags=['-O2', '-std=c++17'], + extra_cuda_cflags=['-O2', '--expt-relaxed-constexpr'], + verbose=True, + ) + + dt = time.time() - t0 + funcs = [x for x in dir(mod) if not x.startswith('_')] + print(f"[moe_kernels] Compiled in {dt:.1f}s — functions: {funcs}") + return mod + + +def smoke_test(mod): + """Quick functional test of compiled kernels.""" + import torch + + print("\n=== Smoke test ===") + device = 'cuda' if torch.cuda.is_available() else 'cpu' + if device == 'cpu': + print(" SKIP: no CUDA device") + return + + # Test topk_softmax + num_tokens, num_experts, topk = 4, 8, 2 + gating = torch.randn(num_tokens, num_experts, device=device, dtype=torch.float32) + topk_weights = torch.empty(num_tokens, topk, device=device, dtype=torch.float32) + topk_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32) + token_expert_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32) + + mod.topk_softmax(topk_weights, topk_indices, token_expert_indices, gating) + + print(f" topk_softmax: weights={topk_weights.shape}, NaN={topk_weights.isnan().any()}") + print(f" weights[0] = {topk_weights[0].tolist()}") + print(f" indices[0] = {topk_indices[0].tolist()}") + + # Test moe_align_block_size + block_size = 4 + max_num_tokens_padded = (num_tokens * topk + num_experts * block_size) + sorted_ids = torch.empty(max_num_tokens_padded, device=device, dtype=torch.int32) + expert_ids = torch.empty(max_num_tokens_padded // block_size, device=device, dtype=torch.int32) + num_tokens_post_pad = torch.empty(1, device=device, dtype=torch.int32) + + mod.moe_align_block_size(topk_indices, num_experts, block_size, + sorted_ids, expert_ids, num_tokens_post_pad) + + print(f" moe_align: sorted_ids[:8]={sorted_ids[:8].tolist()}, " + f"num_post_pad={num_tokens_post_pad.item()}") + + print("\n ✓ All smoke tests passed") + + +if __name__ == '__main__': + mod = compile_moe_kernels() + if '--test' in sys.argv: + smoke_test(mod) diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index df2d3897..85667615 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -18,6 +18,83 @@ logger = init_logger(__name__) supports_moe_ops = True +# ============================================================================ +# MoE CUDA kernels — JIT-compiled from vllm v0.5.5 (torch::Tensor API) +# topk_softmax + moe_align_block_size compiled as moe_kernels.so +# ============================================================================ +_moe_kernels = None +_moe_kernels_loaded = False + +def _load_moe_kernels(): + """Load pre-compiled moe_kernels.so or JIT compile on demand.""" + global _moe_kernels, _moe_kernels_loaded + if _moe_kernels_loaded: + return _moe_kernels + _moe_kernels_loaded = True + + import os, glob, importlib.util + + # Try pre-compiled .so from torch extensions cache + try: + import moe_kernels + _moe_kernels = moe_kernels + logger.info("[EX] moe_kernels loaded from cache") + return _moe_kernels + except ImportError: + pass + + # Try to find .so in known locations + search_paths = [ + os.path.expanduser('~/.cache/torch_extensions'), + '/root/.cache/torch_extensions', + '/workspace/ex_engine/build', + ] + for sp in search_paths: + for so in glob.glob(os.path.join(sp, '**/moe_kernels*.so'), recursive=True): + try: + spec = importlib.util.spec_from_file_location('moe_kernels', so) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + _moe_kernels = mod + logger.info(f"[EX] moe_kernels loaded from {so}") + return _moe_kernels + except Exception: + continue + + # JIT compile as last resort + moe_dir = None + for candidate in [ + '/workspace/ex_engine/csrc/moe_v055', + os.path.join(os.path.dirname(__file__), '..', 'model_executor', 'models', + 'ex_engine', 'csrc', 'moe_v055'), + ]: + if os.path.isdir(candidate): + moe_dir = candidate + break + + if moe_dir and os.path.isfile(os.path.join(moe_dir, 'moe_pybind.cpp')): + try: + from torch.utils.cpp_extension import load + _moe_kernels = load( + name='moe_kernels', + sources=[ + os.path.join(moe_dir, 'moe_pybind.cpp'), + os.path.join(moe_dir, 'topk_softmax_kernels.cu'), + os.path.join(moe_dir, 'moe_align_block_size_kernels.cu'), + ], + extra_include_paths=[moe_dir], + extra_cflags=['-O2', '-std=c++17'], + extra_cuda_cflags=['-O2', '--expt-relaxed-constexpr'], + verbose=False, + ) + logger.info(f"[EX] moe_kernels JIT compiled from {moe_dir}") + return _moe_kernels + except Exception as e: + logger.warning(f"[EX] moe_kernels JIT compile failed: {e}") + + logger.warning("[EX] moe_kernels NOT available — MoE will use PyTorch path") + return None + if TYPE_CHECKING: def register_fake(fn): diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 5c0ebf4f..bc544ae7 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -143,4 +143,27 @@ cp ./_custom_ops.py "$VLLM/_custom_ops.py" 2>/dev/null && \ echo "[patch_ops] _custom_ops.py deployed" || true [ -n "$VLLM2" ] && cp ./_custom_ops.py "$VLLM2/_custom_ops.py" 2>/dev/null || true +# ---- 6. ex_engine.python subpackage (qwen3_5.py does "from ex_engine.python.ix_bridge") ---- +# The flat ex_engine package has ix_bridge.py at top level, but qwen3_5.py imports from .python subdir +_EX_PKG=$(python3 -c "import ex_engine; import os; print(os.path.dirname(ex_engine.__file__))" 2>/dev/null) +if [ -n "$_EX_PKG" ] && [ -d "$_EX_PKG" ]; then + mkdir -p "$_EX_PKG/python" + touch "$_EX_PKG/python/__init__.py" + for f in ix_bridge.py corex_moe.py corex_gdn.py corex_fa2.py; do + [ -f "$_EX_PKG/$f" ] && ln -sf "$_EX_PKG/$f" "$_EX_PKG/python/$f" + done + echo "[patch_ops] ex_engine.python subpackage linked" +fi + +# ---- 7. flash_qla_sm70 deployment to BOTH vllm paths ---- +_FLASH_SRC="/workspace/qwen3_6_scripts/flash_qla_sm70" +if [ -d "$_FLASH_SRC" ]; then + for _VPATH in "$VLLM" "$VLLM2"; do + [ -z "$_VPATH" ] && continue + _FLASH_DST="$_VPATH/model_executor/models/flash_qla_sm70" + cp -r "$_FLASH_SRC" "$_FLASH_DST" 2>/dev/null || true + done + echo "[patch_ops] flash_qla_sm70 deployed to vllm model dirs" +fi + echo "[patch_ops] DONE"