diff --git a/src/infinicore/ops/cat/cat.cc b/src/infinicore/ops/cat/cat.cc index a4e0a8ae0..3eafeae9e 100644 --- a/src/infinicore/ops/cat/cat.cc +++ b/src/infinicore/ops/cat/cat.cc @@ -6,6 +6,17 @@ namespace infinicore::op { namespace { +bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { + // Keep this path limited to CUDA-like backends we have validated or can + // trace to the nvidia rearrange implementation. Other devices may be + // correct through copy_from, but their performance impact is unverified. + return dim == ndim - 1 + && (device_type == Device::Type::NVIDIA + || device_type == Device::Type::HYGON + || device_type == Device::Type::ILUVATAR + || device_type == Device::Type::ALI); +} + class CatInfo { CatInfo() = default; @@ -156,6 +167,25 @@ void cat_(Tensor out, std::vector tensors, int dim) { } assert(dim_shape == out->shape()[dim]); + if (use_slice_copy_cat(out->device().getType(), dim, ndim)) { + // The generic path recursively issues one memcpy for every outer + // index. Concatenating MLA tensors on the last dimension can therefore + // enqueue hundreds of tiny D2D copy calls per layer. A strided output + // slice is semantically identical and copy_from lowers it to a single + // rearrange kernel per input tensor on CUDA-like backends. + size_t offset = 0; + for (auto &tensor : tensors) { + if (tensor->ndim() == 1) { + continue; + } + const size_t length = tensor->shape()[dim]; + auto output_slice = out->narrow({{static_cast(dim), offset, length}}); + output_slice->copy_from(tensor); + offset += length; + } + return; + } + // Get info CatInfo info = CatInfo::create(out, tensors, dim); std::vector tensors_ptr(tensors.size()); diff --git a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu b/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu index 9c63f97d2..e938d42f1 100644 --- a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu +++ b/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu @@ -6,6 +6,10 @@ #include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "deepseek_moe_nvidia.cuh" +#include +#include +#include + namespace op::deepseek_moe::nvidia { struct Descriptor::Opaque { @@ -18,10 +22,57 @@ Descriptor::~Descriptor() { namespace { +constexpr size_t ROUTE_BATCHED_GEMM_MAX_TOKENS = 16384; +constexpr size_t EXPERT_GROUPED_GEMM_MIN_TOKENS = 2048; +constexpr size_t FUSED_DECODE_MAX_TOKENS = 32; +constexpr int FUSED_KERNEL_THREADS = 64; +constexpr size_t CHUNKED_GROUPED_GEMM_TOKENS = 4096; + constexpr size_t align_up(size_t value, size_t alignment) { return (value + alignment - 1) / alignment * alignment; } +inline char *advance_workspace(char *&ptr, size_t bytes, size_t alignment = 256) { + ptr = reinterpret_cast(align_up(reinterpret_cast(ptr), alignment)); + char *out = ptr; + ptr += bytes; + return out; +} + +template +size_t route_batched_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { + const size_t routes = info.ntokens * info.topk; + size_t bytes = base_offset; + bytes = align_up(bytes, alignof(void *)); + bytes += routes * sizeof(void *) * 9; + bytes = align_up(bytes, 256); + bytes += routes * info.intermediate_size * sizeof(T) * 3; + bytes = align_up(bytes, 256); + bytes += routes * info.hidden_size * sizeof(T); + return align_up(bytes, 256); +} + +template +size_t expert_grouped_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { + const size_t routes = info.ntokens * info.topk; + size_t bytes = base_offset; + for (int i = 0; i < 3; ++i) { + bytes = align_up(bytes, 256); + bytes += (info.num_experts + 1) * sizeof(int); + } + for (int i = 0; i < 2; ++i) { + bytes = align_up(bytes, 256); + bytes += routes * sizeof(int); + } + bytes = align_up(bytes, 256); + bytes += routes * info.hidden_size * sizeof(T); + for (int i = 0; i < 2; ++i) { + bytes = align_up(bytes, 256); + bytes += routes * info.intermediate_size * sizeof(T); + } + return align_up(bytes, 256); +} + template __device__ float to_float(T value) { return static_cast(value); @@ -52,6 +103,19 @@ __device__ __nv_bfloat16 from_float<__nv_bfloat16>(float value) { return __float2bfloat16_rn(value); } +__device__ __forceinline__ float dot_bf16x8(const uint4 &a, const uint4 &b) { + const auto *a2 = reinterpret_cast(&a); + const auto *b2 = reinterpret_cast(&b); + float sum = 0.0f; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const float2 av = __bfloat1622float2(a2[pair]); + const float2 bv = __bfloat1622float2(b2[pair]); + sum += av.x * bv.x + av.y * bv.y; + } + return sum; +} + template __global__ void gate_up_kernel( T *intermediate, @@ -82,14 +146,46 @@ __global__ void gate_up_kernel( float gate_sum = 0.0f; float up_sum = 0.0f; - for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { - const float xv = to_float(x[h]); - gate_sum += xv * to_float(gate[h]); - up_sum += xv * to_float(up[h]); + if constexpr (std::is_same_v) { + if ((hidden_size & 7) == 0) { + const auto *x8 = reinterpret_cast(x); + const auto *gate8 = reinterpret_cast(gate); + const auto *up8 = reinterpret_cast(up); + const size_t hidden_chunks = hidden_size / 8; + for (size_t h8 = threadIdx.x; h8 < hidden_chunks; h8 += blockDim.x) { + const uint4 xv = x8[h8]; + gate_sum += dot_bf16x8(xv, gate8[h8]); + up_sum += dot_bf16x8(xv, up8[h8]); + } + } else if ((hidden_size & 1) == 0) { + const auto *x2 = reinterpret_cast(x); + const auto *gate2 = reinterpret_cast(gate); + const auto *up2 = reinterpret_cast(up); + const size_t hidden_pairs = hidden_size / 2; + for (size_t h2 = threadIdx.x; h2 < hidden_pairs; h2 += blockDim.x) { + const float2 xv = __bfloat1622float2(x2[h2]); + const float2 gv = __bfloat1622float2(gate2[h2]); + const float2 uv = __bfloat1622float2(up2[h2]); + gate_sum += xv.x * gv.x + xv.y * gv.y; + up_sum += xv.x * uv.x + xv.y * uv.y; + } + } else { + for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { + const float xv = to_float(x[h]); + gate_sum += xv * to_float(gate[h]); + up_sum += xv * to_float(up[h]); + } + } + } else { + for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { + const float xv = to_float(x[h]); + gate_sum += xv * to_float(gate[h]); + up_sum += xv * to_float(up[h]); + } } - __shared__ float gate_shared[256]; - __shared__ float up_shared[256]; + __shared__ float gate_shared[FUSED_KERNEL_THREADS]; + __shared__ float up_shared[FUSED_KERNEL_THREADS]; gate_shared[threadIdx.x] = gate_sum; up_shared[threadIdx.x] = up_sum; __syncthreads(); @@ -131,18 +227,67 @@ __global__ void down_kernel( float acc = 0.0f; const size_t route_base = token * topk; const size_t count = topk * intermediate_size; - for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { - const size_t k = idx / intermediate_size; - const size_t j = idx - k * intermediate_size; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; - acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); + if constexpr (std::is_same_v) { + if ((intermediate_size & 7) == 0) { + const size_t chunks_per_expert = intermediate_size / 8; + const size_t chunk_count = topk * chunks_per_expert; + for (size_t chunk_idx = threadIdx.x; chunk_idx < chunk_count; chunk_idx += blockDim.x) { + const size_t k = chunk_idx / chunks_per_expert; + const size_t j8 = chunk_idx - k * chunks_per_expert; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const auto *intermediate8 = reinterpret_cast( + intermediate + route * intermediate_size); + const auto *down8 = reinterpret_cast( + reinterpret_cast(down_weights[expert]) + h * intermediate_size); + acc += dot_bf16x8(intermediate8[j8], down8[j8]); + } + } + } else if ((intermediate_size & 1) == 0) { + const size_t pairs_per_expert = intermediate_size / 2; + const size_t pair_count = topk * pairs_per_expert; + for (size_t pair_idx = threadIdx.x; pair_idx < pair_count; pair_idx += blockDim.x) { + const size_t k = pair_idx / pairs_per_expert; + const size_t j2 = pair_idx - k * pairs_per_expert; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const auto *intermediate2 = reinterpret_cast( + intermediate + route * intermediate_size); + const auto *down2 = reinterpret_cast( + reinterpret_cast(down_weights[expert]) + h * intermediate_size); + const float2 iv = __bfloat1622float2(intermediate2[j2]); + const float2 dv = __bfloat1622float2(down2[j2]); + acc += iv.x * dv.x + iv.y * dv.y; + } + } + } else { + for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { + const size_t k = idx / intermediate_size; + const size_t j = idx - k * intermediate_size; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; + acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); + } + } + } + } else { + for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { + const size_t k = idx / intermediate_size; + const size_t j = idx - k * intermediate_size; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; + acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); + } } } - __shared__ float shared[256]; + __shared__ float shared[FUSED_KERNEL_THREADS]; shared[threadIdx.x] = acc; __syncthreads(); for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { @@ -156,6 +301,505 @@ __global__ void down_kernel( } } +__global__ void count_experts_kernel( + const int *topk_indices, + int *counts, + int routes, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert >= 0 && expert < num_experts) { + atomicAdd(counts + expert, 1); + } +} + +__global__ void prefix_counts_kernel( + const int *counts, + int *offsets, + int num_experts) { + extern __shared__ int scan[]; + int tid = threadIdx.x; + if (tid < num_experts) { + scan[tid] = counts[tid]; + } else { + scan[tid] = 0; + } + __syncthreads(); + + for (int stride = 1; stride < blockDim.x; stride <<= 1) { + int value = 0; + if (tid >= stride) { + value = scan[tid - stride]; + } + __syncthreads(); + scan[tid] += value; + __syncthreads(); + } + + if (tid == 0) { + offsets[0] = 0; + } + if (tid < num_experts) { + offsets[tid + 1] = scan[tid]; + } +} + +template +__global__ void pack_grouped_hidden_kernel( + const T *hidden, + const int *topk_indices, + const int *offsets, + int *positions, + int *row_to_route, + int *route_to_row, + T *packed_hidden, + int routes, + int hidden_size, + int topk, + int num_experts) { + int route = blockIdx.x; + int tid = threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + return; + } + __shared__ int row; + if (tid == 0) { + int local = atomicAdd(positions + expert, 1); + row = offsets[expert] + local; + row_to_route[row] = route; + route_to_row[route] = row; + } + __syncthreads(); + int token = route / topk; + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = hidden[static_cast(token) * hidden_size + h]; + } +} + +template +__global__ void swiglu_weight_grouped_kernel( + T *activated, + const T *gate, + const T *up, + const int *row_to_route, + const float *topk_weights, + int routes, + int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = routes * intermediate_size; + if (idx >= total) { + return; + } + int row = idx / intermediate_size; + int route = row_to_route[row]; + float g = to_float(gate[idx]); + float u = to_float(up[idx]); + float silu = g / (1.0f + __expf(-g)); + activated[idx] = from_float(silu * u * topk_weights[route]); +} + +template +__global__ void sum_grouped_out_kernel( + T *out, + const T *packed_out, + const int *route_to_row, + int ntokens, + int hidden_size, + int topk) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = ntokens * hidden_size; + if (idx >= total) { + return; + } + int token = idx / hidden_size; + int h = idx - token * hidden_size; + float acc = 0.0f; + for (int k = 0; k < topk; ++k) { + int route = token * topk + k; + int row = route_to_row[route]; + acc += to_float(packed_out[static_cast(row) * hidden_size + h]); + } + out[idx] = from_float(acc); +} + +template +__global__ void setup_route_batched_ptrs_kernel( + const void **a_array, + const void **b_array, + void **c_array, + const T *hidden, + const int *topk_indices, + const void *const *weights, + T *gemm_out, + int routes, + int hidden_size, + int topk, + int out_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + int token = route / topk; + a_array[route] = weights[expert]; + b_array[route] = hidden + static_cast(token) * hidden_size; + c_array[route] = gemm_out + static_cast(route) * out_size; +} + +template +__global__ void setup_down_batched_ptrs_kernel( + const void **a_array, + const void **b_array, + void **c_array, + const T *activated, + const int *topk_indices, + const void *const *down_weights, + T *route_out, + int routes, + int hidden_size, + int topk, + int intermediate_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + a_array[route] = down_weights[expert]; + b_array[route] = activated + static_cast(route) * intermediate_size; + c_array[route] = route_out + static_cast(route) * hidden_size; +} + +template +__global__ void setup_all_batched_ptrs_kernel( + const void **gate_a_array, + const void **gate_b_array, + void **gate_c_array, + const void **up_a_array, + const void **up_b_array, + void **up_c_array, + const void **down_a_array, + const void **down_b_array, + void **down_c_array, + const T *hidden, + const int *topk_indices, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + T *gate_buf, + T *up_buf, + T *activated, + T *route_out, + int routes, + int hidden_size, + int topk, + int intermediate_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + int token = route / topk; + const T *token_hidden = hidden + static_cast(token) * hidden_size; + + gate_a_array[route] = gate_weights[expert]; + gate_b_array[route] = token_hidden; + gate_c_array[route] = gate_buf + static_cast(route) * intermediate_size; + + up_a_array[route] = up_weights[expert]; + up_b_array[route] = token_hidden; + up_c_array[route] = up_buf + static_cast(route) * intermediate_size; + + down_a_array[route] = down_weights[expert]; + down_b_array[route] = activated + static_cast(route) * intermediate_size; + down_c_array[route] = route_out + static_cast(route) * hidden_size; +} + +template +__global__ void swiglu_weight_kernel( + T *activated, + const T *gate, + const T *up, + const float *topk_weights, + int routes, + int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = routes * intermediate_size; + if (idx >= total) { + return; + } + int route = idx / intermediate_size; + float g = to_float(gate[idx]); + float u = to_float(up[idx]); + float silu = g / (1.0f + __expf(-g)); + activated[idx] = from_float(silu * u * topk_weights[route]); +} + +template +__global__ void sum_route_out_kernel( + T *out, + const T *route_out, + int ntokens, + int hidden_size, + int topk) { + int linear = blockIdx.x * blockDim.x + threadIdx.x; + int total = ntokens * hidden_size; + if (linear >= total) { + return; + } + int token = linear / hidden_size; + int h = linear - token * hidden_size; + float acc = 0.0f; + for (int k = 0; k < topk; ++k) { + int route = token * topk + k; + acc += to_float(route_out[static_cast(route) * hidden_size + h]); + } + out[linear] = from_float(acc); +} + +template +infiniStatus_t run_batched_gemm( + const std::shared_ptr &internal, + cudaStream_t stream, + const void **a_array, + const void **b_array, + void **c_array, + int m, + int n, + int k, + int lda, + int ldb, + int ldc, + int batch_count, + cudaDataType dtype) { + float alpha = 1.0f; + float beta = 0.0f; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + cudaDataType compute_type = CUDA_R_32F; +#else + cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; +#endif + CHECK_STATUS(internal->useCublas( + stream, + [&](cublasHandle_t handle) { + CHECK_CUBLAS(cublasGemmBatchedEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + m, + n, + k, + &alpha, + reinterpret_cast(a_array), + dtype, + lda, + reinterpret_cast(b_array), + dtype, + ldb, + &beta, + reinterpret_cast(c_array), + dtype, + ldc, + batch_count, + compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + return INFINI_STATUS_SUCCESS; + })); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t run_gemm( + const std::shared_ptr &internal, + cudaStream_t stream, + const void *a, + const void *b, + void *c, + int m, + int n, + int k, + int lda, + int ldb, + int ldc, + cudaDataType dtype) { + float alpha = 1.0f; + float beta = 0.0f; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + cudaDataType compute_type = CUDA_R_32F; +#else + cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; +#endif + CHECK_STATUS(internal->useCublas( + stream, + [&](cublasHandle_t handle) { + CHECK_CUBLAS(cublasGemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + m, + n, + k, + &alpha, + a, + dtype, + lda, + b, + dtype, + ldb, + &beta, + c, + dtype, + ldc, + compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + return INFINI_STATUS_SUCCESS; + })); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t run_chunked_grouped_gemm( + char *base, + size_t workspace_size, + size_t intermediate_offset, + const DeepseekMoeInfo &info, + T *out, + const T *hidden, + const int *topk_indices, + const float *topk_weights, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + cudaStream_t stream, + bool weight_ptrs_on_device, + const std::shared_ptr &internal) { + + const int hidden_size = static_cast(info.hidden_size); + const int topk = static_cast(info.topk); + const int intermediate_size = static_cast(info.intermediate_size); + const int num_experts = static_cast(info.num_experts); + const cudaDataType blas_dtype = info.dtype == INFINI_DTYPE_F16 ? CUDA_R_16F : CUDA_R_16BF; + + std::vector host_counts(info.num_experts + 1); + std::vector host_offsets(info.num_experts + 1); + std::vector host_gate(info.num_experts); + std::vector host_up(info.num_experts); + std::vector host_down(info.num_experts); + if (weight_ptrs_on_device) { + CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + } else { + for (size_t e = 0; e < info.num_experts; ++e) { + host_gate[e] = gate_weights[e]; + host_up[e] = up_weights[e]; + host_down[e] = down_weights[e]; + } + } + + for (size_t token_begin = 0; token_begin < info.ntokens; token_begin += CHUNKED_GROUPED_GEMM_TOKENS) { + const int chunk_tokens = static_cast(std::min( + CHUNKED_GROUPED_GEMM_TOKENS, info.ntokens - token_begin)); + const int routes = chunk_tokens * topk; + + char *gptr = base + intermediate_offset; + auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); + auto *gate_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + if (static_cast(gptr - base) > workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + const T *hidden_chunk = hidden + token_begin * info.hidden_size; + const int *indices_chunk = topk_indices + token_begin * info.topk; + const float *weights_chunk = topk_weights + token_begin * info.topk; + T *out_chunk = out + token_begin * info.hidden_size; + + CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); + CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); + count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( + indices_chunk, counts, routes, num_experts); + prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( + counts, offsets, num_experts); + pack_grouped_hidden_kernel<<>>( + hidden_chunk, indices_chunk, offsets, positions, + row_to_route, route_to_row, packed_hidden, + routes, hidden_size, topk, num_experts); + + CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + for (int e = 0; e < num_experts; ++e) { + const int count = host_counts[e]; + if (count == 0) { + continue; + } + const int offset = host_offsets[e]; + CHECK_STATUS(run_gemm( + internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size, + gate_buf + static_cast(offset) * intermediate_size, + intermediate_size, count, hidden_size, + hidden_size, hidden_size, intermediate_size, blas_dtype)); + CHECK_STATUS(run_gemm( + internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size, + up_buf + static_cast(offset) * intermediate_size, + intermediate_size, count, hidden_size, + hidden_size, hidden_size, intermediate_size, blas_dtype)); + } + + swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + gate_buf, gate_buf, up_buf, row_to_route, weights_chunk, routes, intermediate_size); + + for (int e = 0; e < num_experts; ++e) { + const int count = host_counts[e]; + if (count == 0) { + continue; + } + const int offset = host_offsets[e]; + CHECK_STATUS(run_gemm( + internal, stream, host_down[e], gate_buf + static_cast(offset) * intermediate_size, + packed_hidden + static_cast(offset) * hidden_size, + hidden_size, count, intermediate_size, + intermediate_size, intermediate_size, hidden_size, blas_dtype)); + } + + sum_grouped_out_kernel<<<(static_cast(chunk_tokens) * info.hidden_size + 255) / 256, 256, 0, stream>>>( + out_chunk, packed_hidden, route_to_row, chunk_tokens, hidden_size, topk); + } + + // Keep the large-prefill path memory-bounded across transformer layers. + // All earlier chunks are synchronized by the next chunk's routing copy; + // this final sync drains the last chunk before its workspace can be reused. + CHECK_CUDA(cudaStreamSynchronize(stream)); + + return INFINI_STATUS_SUCCESS; +} + template infiniStatus_t launch_typed( void *workspace, @@ -169,12 +813,17 @@ infiniStatus_t launch_typed( const void *const *up_weights, const void *const *down_weights, cudaStream_t stream, - bool weight_ptrs_on_device) { + bool weight_ptrs_on_device, + const std::shared_ptr &internal) { const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); const size_t ptr_workspace = ptr_bytes * 3; const size_t intermediate_offset = align_up(ptr_workspace, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * sizeof(T); + size_t intermediate_tokens = info.ntokens; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + intermediate_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); + } + const size_t intermediate_bytes = intermediate_tokens * info.topk * info.intermediate_size * sizeof(T); if (workspace_size < intermediate_offset + intermediate_bytes) { return INFINI_STATUS_INSUFFICIENT_WORKSPACE; } @@ -201,7 +850,186 @@ infiniStatus_t launch_typed( down_ptrs = down_workspace; } - constexpr int threads = 256; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + return run_chunked_grouped_gemm( + base, workspace_size, intermediate_offset, info, + reinterpret_cast(out), reinterpret_cast(hidden), + reinterpret_cast(topk_indices), reinterpret_cast(topk_weights), + gate_weights, up_weights, down_weights, + stream, weight_ptrs_on_device, internal); + } + + if (info.ntokens > FUSED_DECODE_MAX_TOKENS + && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { + const int routes = static_cast(info.ntokens * info.topk); + const int hidden_size_i = static_cast(info.hidden_size); + const int topk_i = static_cast(info.topk); + const int intermediate_size_i = static_cast(info.intermediate_size); + const int num_experts_i = static_cast(info.num_experts); + const cudaDataType blas_dtype = info.dtype == INFINI_DTYPE_F16 ? CUDA_R_16F : CUDA_R_16BF; + + if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS && info.num_experts <= 256) { + char *gptr = base + intermediate_offset; + auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); + auto *gate_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + if (static_cast(gptr - base) <= workspace_size) { + CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); + CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); + count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(topk_indices), counts, routes, num_experts_i); + prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( + counts, offsets, num_experts_i); + pack_grouped_hidden_kernel<<>>( + reinterpret_cast(hidden), reinterpret_cast(topk_indices), + offsets, positions, row_to_route, route_to_row, packed_hidden, + routes, hidden_size_i, topk_i, num_experts_i); + + std::vector host_counts(info.num_experts + 1); + std::vector host_offsets(info.num_experts + 1); + std::vector host_gate(info.num_experts); + std::vector host_up(info.num_experts); + std::vector host_down(info.num_experts); + CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + if (weight_ptrs_on_device) { + CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + } else { + for (size_t e = 0; e < info.num_experts; ++e) { + host_gate[e] = gate_weights[e]; + host_up[e] = up_weights[e]; + host_down[e] = down_weights[e]; + } + } + CHECK_CUDA(cudaStreamSynchronize(stream)); + + for (int e = 0; e < num_experts_i; ++e) { + int count = host_counts[e]; + if (count == 0) { + continue; + } + int offset = host_offsets[e]; + CHECK_STATUS(run_gemm( + internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size_i, + gate_buf_g + static_cast(offset) * intermediate_size_i, + intermediate_size_i, count, hidden_size_i, + hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); + CHECK_STATUS(run_gemm( + internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size_i, + up_buf_g + static_cast(offset) * intermediate_size_i, + intermediate_size_i, count, hidden_size_i, + hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); + } + + swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + gate_buf_g, gate_buf_g, up_buf_g, row_to_route, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + for (int e = 0; e < num_experts_i; ++e) { + int count = host_counts[e]; + if (count == 0) { + continue; + } + int offset = host_offsets[e]; + CHECK_STATUS(run_gemm( + internal, stream, host_down[e], gate_buf_g + static_cast(offset) * intermediate_size_i, + packed_hidden + static_cast(offset) * hidden_size_i, + hidden_size_i, count, intermediate_size_i, + intermediate_size_i, intermediate_size_i, hidden_size_i, blas_dtype)); + } + + sum_grouped_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(out), packed_hidden, route_to_row, static_cast(info.ntokens), hidden_size_i, topk_i); + return INFINI_STATUS_SUCCESS; + } + } + + char *bptr = base + intermediate_offset; + auto **gate_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **gate_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **gate_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto *gate_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *activated = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *route_out = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.hidden_size * sizeof(T))); + if (static_cast(bptr - base) <= workspace_size) { + const int blocks = (routes + 255) / 256; + if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS) { + setup_all_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + up_a_array, up_b_array, up_c_array, + down_a_array, down_b_array, down_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), + gate_ptrs, up_ptrs, down_ptrs, gate_buf, up_buf, activated, route_out, + routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + CHECK_STATUS(run_batched_gemm( + internal, stream, up_a_array, up_b_array, up_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + CHECK_STATUS(run_batched_gemm( + internal, stream, down_a_array, down_b_array, down_c_array, + hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, + routes, blas_dtype)); + } else { + setup_route_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), gate_ptrs, + gate_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + setup_route_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), up_ptrs, + up_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + setup_down_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + activated, reinterpret_cast(topk_indices), down_ptrs, + route_out, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, + routes, blas_dtype)); + } + + sum_route_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(out), route_out, static_cast(info.ntokens), hidden_size_i, topk_i); + return INFINI_STATUS_SUCCESS; + } + } + + constexpr int threads = FUSED_KERNEL_THREADS; const dim3 gate_blocks(static_cast(info.ntokens * info.topk * info.intermediate_size)); gate_up_kernel<<>>( intermediate, @@ -250,8 +1078,33 @@ infiniStatus_t Descriptor::create( const size_t dtype_size = info.dtype == INFINI_DTYPE_F16 ? sizeof(half) : sizeof(__nv_bfloat16); const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); const size_t intermediate_offset = align_up(ptr_bytes * 3, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * dtype_size; - const size_t workspace_size = intermediate_offset + intermediate_bytes; + size_t workspace_tokens = info.ntokens; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + workspace_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); + } + const size_t intermediate_bytes = workspace_tokens * info.topk * info.intermediate_size * dtype_size; + const size_t old_workspace_size = intermediate_offset + intermediate_bytes; + size_t batched_workspace_size = old_workspace_size; + size_t grouped_workspace_size = old_workspace_size; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + auto chunk_info = info; + chunk_info.ntokens = workspace_tokens; + if (info.dtype == INFINI_DTYPE_F16) { + grouped_workspace_size = expert_grouped_workspace_size(chunk_info, intermediate_offset); + } else { + grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(chunk_info, intermediate_offset); + } + } else if (info.ntokens > FUSED_DECODE_MAX_TOKENS + && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { + if (info.dtype == INFINI_DTYPE_F16) { + batched_workspace_size = route_batched_workspace_size(info, intermediate_offset); + grouped_workspace_size = expert_grouped_workspace_size(info, intermediate_offset); + } else { + batched_workspace_size = route_batched_workspace_size<__nv_bfloat16>(info, intermediate_offset); + grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(info, intermediate_offset); + } + } + const size_t workspace_size = std::max(old_workspace_size, std::max(batched_workspace_size, grouped_workspace_size)); *desc_ptr = new Descriptor( new Opaque{reinterpret_cast(handle)->internal()}, @@ -276,10 +1129,10 @@ infiniStatus_t Descriptor::calculate( auto stream = reinterpret_cast(stream_); if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); + return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false, _opaque->internal); } if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); + return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false, _opaque->internal); } return INFINI_STATUS_BAD_TENSOR_DTYPE; } @@ -301,10 +1154,10 @@ infiniStatus_t Descriptor::calculateWithDevicePtrs( auto up_weights = reinterpret_cast(up_weight_ptrs); auto down_weights = reinterpret_cast(down_weight_ptrs); if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); + return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true, _opaque->internal); } if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); + return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true, _opaque->internal); } return INFINI_STATUS_BAD_TENSOR_DTYPE; } diff --git a/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu b/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu index 03501ed19..3c449a168 100644 --- a/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu +++ b/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu @@ -494,7 +494,11 @@ infiniStatus_t launch_decode_mla_hd576_v512_impl( if (const char *env = std::getenv("INFINIOP_FLASH_DECODE_SPLITKV")) { use_split = !(std::strcmp(env, "0") == 0 || std::strcmp(env, "false") == 0); } +#if defined(ENABLE_HYGON_API) + int num_splits = 8; +#else int num_splits = 4; +#endif if (const char *env = std::getenv("INFINIOP_FLASH_NUM_SPLITS")) { const int v = std::atoi(env); if (v > 0) {