mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
52 lines
2.1 KiB
C++
52 lines
2.1 KiB
C++
/*
|
|
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* 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.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "tensorrt_llm/common/config.h"
|
|
#include <cuda_runtime.h>
|
|
#include <stdint.h>
|
|
|
|
TRTLLM_NAMESPACE_BEGIN
|
|
|
|
namespace kernels
|
|
{
|
|
|
|
/**
|
|
* @brief Aligns token distribution across experts to be compatible with block size for matrix multiplication.
|
|
*
|
|
* This kernel sorts tokens by expert assignment and pads the distribution to match block size requirements.
|
|
* Inspired by vLLM's moe_align_kernel and ported to TensorRT-LLM.
|
|
*
|
|
* @param topk_ids Input tensor with expert IDs per token [total_tokens, top_k]
|
|
* @param topk_ids_dtype_size Size of the dtype (e.g., sizeof(int32_t) or sizeof(int64_t))
|
|
* @param sorted_token_ids Output tensor for sorted token indices
|
|
* @param expert_ids Output tensor for expert IDs per block
|
|
* @param num_tokens_post_pad Output tensor for total tokens after padding (single int32)
|
|
* @param num_experts Total number of experts
|
|
* @param block_size Block size for matrix multiplication alignment
|
|
* @param numel Total number of elements in topk_ids (topk_ids.numel())
|
|
* @param max_num_tokens_padded Maximum number of tokens after padding (sorted_token_ids.size(0))
|
|
* @param stream CUDA stream for kernel execution
|
|
*/
|
|
void invokeMoeAlignBlockSize(void const* topk_ids, int32_t topk_ids_dtype_size, int32_t* sorted_token_ids,
|
|
int32_t* expert_ids, int32_t* num_tokens_post_pad, int32_t num_experts, int32_t block_size, int32_t numel,
|
|
int32_t max_num_tokens_padded, cudaStream_t stream);
|
|
|
|
} // namespace kernels
|
|
|
|
TRTLLM_NAMESPACE_END
|