mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-13 22:18:36 +08:00
210 lines
9.2 KiB
C++
210 lines
9.2 KiB
C++
/*
|
|
* Copyright (c) 2020-2023, 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.
|
|
*/
|
|
#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/tensorMapUtils.h"
|
|
#include "tensorrt_llm/common/config.h"
|
|
#include "tensorrt_llm/kernels/kvCacheUtils.h"
|
|
|
|
#include <cstdint>
|
|
#include <type_traits>
|
|
|
|
TRTLLM_NAMESPACE_BEGIN
|
|
|
|
namespace kernels
|
|
{
|
|
|
|
namespace
|
|
{
|
|
|
|
using tensorrt_llm::common::CUDADriverWrapper;
|
|
|
|
uint32_t getElemBytes(CUtensorMapDataType_enum dataType)
|
|
{
|
|
switch (dataType)
|
|
{
|
|
case CU_TENSOR_MAP_DATA_TYPE_UINT8: return 1;
|
|
case CU_TENSOR_MAP_DATA_TYPE_UINT16: return 2;
|
|
case CU_TENSOR_MAP_DATA_TYPE_UINT32: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_INT32: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_UINT64: return 8;
|
|
case CU_TENSOR_MAP_DATA_TYPE_INT64: return 8;
|
|
case CU_TENSOR_MAP_DATA_TYPE_FLOAT16: return 2;
|
|
case CU_TENSOR_MAP_DATA_TYPE_FLOAT32: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_FLOAT64: return 8;
|
|
case CU_TENSOR_MAP_DATA_TYPE_BFLOAT16: return 2;
|
|
case CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_TFLOAT32: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ: return 4;
|
|
case CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B: return 8;
|
|
case CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B: return 16;
|
|
case CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B: return 16;
|
|
}
|
|
throw std::runtime_error("unsupported data type");
|
|
}
|
|
|
|
CUtensorMapSwizzle getSwizzleMode(uint32_t partBytes)
|
|
{
|
|
switch (partBytes)
|
|
{
|
|
case 128: return CU_TENSOR_MAP_SWIZZLE_128B;
|
|
case 64: return CU_TENSOR_MAP_SWIZZLE_64B;
|
|
default: TLLM_THROW("unsupported cache head size");
|
|
}
|
|
};
|
|
|
|
CUtensorMapDataType_enum getDataTypeFromXqaParams(XQAParams const& xqaParams)
|
|
{
|
|
if (xqaParams.kv_cache_data_type == DATA_TYPE_BF16)
|
|
{
|
|
return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
|
|
}
|
|
else if (xqaParams.kv_cache_data_type == DATA_TYPE_FP16)
|
|
{
|
|
return CU_TENSOR_MAP_DATA_TYPE_FLOAT16;
|
|
}
|
|
TLLM_CHECK(xqaParams.kv_cache_data_type == DATA_TYPE_E4M3 || xqaParams.kv_cache_data_type == DATA_TYPE_E5M2
|
|
|| xqaParams.kv_cache_data_type == DATA_TYPE_INT8);
|
|
return CU_TENSOR_MAP_DATA_TYPE_UINT8;
|
|
}
|
|
|
|
CUtensorMap makeTensorMapForQ(std::shared_ptr<CUDADriverWrapper> const& driver, void const* addr,
|
|
CUtensorMapDataType_enum dataType, uint32_t headElems, uint32_t totalNbHeads, uint32_t partElems, uint32_t boxHeads)
|
|
{
|
|
CUtensorMap tensorMap{};
|
|
uint64_t const globalDims[] = {headElems, totalNbHeads};
|
|
uint32_t elemBytes = getElemBytes(dataType);
|
|
uint32_t const headBytes = elemBytes * headElems;
|
|
uint64_t const globalStrides[] = {headBytes};
|
|
uint32_t const boxDims[] = {partElems, boxHeads};
|
|
uint32_t const elemStrides[] = {1, 1};
|
|
auto const swizzle = getSwizzleMode(elemBytes * partElems);
|
|
|
|
TLLM_CU_CHECK(driver->cuTensorMapEncodeTiled(&tensorMap, dataType, 2, const_cast<void*>(addr), globalDims,
|
|
globalStrides, boxDims, elemStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_NONE,
|
|
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
|
return tensorMap;
|
|
}
|
|
|
|
CUtensorMap makeTensorMapForPagedKVCache(std::shared_ptr<CUDADriverWrapper> const& driver, void const* addr,
|
|
CUtensorMapDataType_enum dataType, uint32_t headElems, uint32_t nbKHeads, uint32_t tokensPerPage,
|
|
uint32_t partElems, uint32_t nbTokensPerTile = 64)
|
|
{
|
|
CUtensorMap tensorMap{};
|
|
uint32_t elemBytes = getElemBytes(dataType);
|
|
uint64_t const globalDims[] = {headElems, tokensPerPage, nbKHeads, 1U << 31};
|
|
uint32_t const headBytes = elemBytes * headElems;
|
|
uint64_t const globalStrides[] = {headBytes, headBytes * tokensPerPage, headBytes * tokensPerPage * nbKHeads};
|
|
uint32_t const boxDims[] = {partElems, std::min(tokensPerPage, nbTokensPerTile), 1, 1};
|
|
uint32_t const elemStrides[] = {1, 1, 1, 1};
|
|
|
|
auto const swizzle = getSwizzleMode(elemBytes * partElems);
|
|
|
|
TLLM_CU_CHECK(driver->cuTensorMapEncodeTiled(&tensorMap, dataType, 4, const_cast<void*>(addr), globalDims,
|
|
globalStrides, boxDims, elemStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_NONE,
|
|
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
|
return tensorMap;
|
|
}
|
|
|
|
CUtensorMap makeTensorMapForContiguousKVCache(std::shared_ptr<CUDADriverWrapper> const& driver, void const* addr,
|
|
CUtensorMapDataType_enum dataType, uint32_t headElems, uint32_t nbKHeads, uint32_t maxCacheLen, uint32_t beamWidth,
|
|
uint32_t batchSize, uint32_t partElems, uint32_t nbTokensPerTile = 64)
|
|
{
|
|
CUtensorMap tensorMap{};
|
|
uint64_t const globalDims[] = {headElems, maxCacheLen, nbKHeads, 2 * beamWidth * batchSize};
|
|
uint32_t elemBytes = getElemBytes(dataType);
|
|
uint32_t const headBytes = elemBytes * headElems;
|
|
uint64_t const globalStrides[] = {headBytes, headBytes * maxCacheLen, headBytes * maxCacheLen * nbKHeads};
|
|
uint32_t const boxDims[] = {partElems, nbTokensPerTile, 1, 1};
|
|
uint32_t const elemStrides[] = {1, 1, 1, 1};
|
|
|
|
auto const swizzle = getSwizzleMode(elemBytes * partElems);
|
|
|
|
TLLM_CU_CHECK(driver->cuTensorMapEncodeTiled(&tensorMap, dataType, 4, const_cast<void*>(addr), globalDims,
|
|
globalStrides, boxDims, elemStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_NONE,
|
|
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
|
return tensorMap;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
template <typename KVCacheBuffer>
|
|
CUtensorMap makeTensorMapForHopperXqaKVCache(
|
|
std::shared_ptr<CUDADriverWrapper> const& driver, XQAParams const& xqaParams, KVCacheBuffer const& kv_cache_buffer)
|
|
{
|
|
if constexpr (std::is_same_v<KVCacheBuffer, KVBlockArray>)
|
|
{
|
|
uint32_t const headElems = xqaParams.head_size;
|
|
CUtensorMapDataType_enum const dataType = getDataTypeFromXqaParams(xqaParams);
|
|
uint32_t const elemBytes = getElemBytes(dataType);
|
|
TLLM_CHECK(headElems <= 256);
|
|
uint32_t const paddedHeadElems = headElems <= 64 ? 64 : (headElems <= 128 ? 128 : 256);
|
|
uint32_t const partElems = std::min(elemBytes * paddedHeadElems, 128U) / elemBytes;
|
|
return makeTensorMapForPagedKVCache(driver, kv_cache_buffer.mPrimaryPoolPtr, dataType, xqaParams.head_size,
|
|
xqaParams.num_kv_heads, xqaParams.tokens_per_block, partElems);
|
|
}
|
|
else
|
|
{
|
|
static_assert(std::is_same_v<KVCacheBuffer, KVLinearBuffer>);
|
|
uint32_t const headElems = xqaParams.head_size;
|
|
CUtensorMapDataType_enum const dataType = getDataTypeFromXqaParams(xqaParams);
|
|
uint32_t const elemBytes = getElemBytes(dataType);
|
|
TLLM_CHECK(headElems <= 256);
|
|
uint32_t const paddedHeadElems = headElems <= 64 ? 64 : (headElems <= 128 ? 128 : 256);
|
|
uint32_t const partElems = std::min(elemBytes * paddedHeadElems, 128U) / elemBytes;
|
|
return makeTensorMapForContiguousKVCache(driver, kv_cache_buffer.data, dataType, xqaParams.head_size,
|
|
xqaParams.num_kv_heads, xqaParams.max_attention_window_size, xqaParams.beam_width, xqaParams.batch_size,
|
|
partElems);
|
|
}
|
|
}
|
|
|
|
template CUtensorMap makeTensorMapForHopperXqaKVCache(
|
|
std::shared_ptr<CUDADriverWrapper> const&, XQAParams const&, KVBlockArray const&);
|
|
template CUtensorMap makeTensorMapForHopperXqaKVCache(
|
|
std::shared_ptr<CUDADriverWrapper> const&, XQAParams const&, KVLinearBuffer const&);
|
|
|
|
template <typename KVCacheBuffer>
|
|
CUtensorMap makeTensorMapForXqaMlaKVCache(std::shared_ptr<tensorrt_llm::common::CUDADriverWrapper> const& driver,
|
|
XQAParams const& xqaParams, KVCacheBuffer const& kv_cache_buffer, bool forK)
|
|
{
|
|
CUtensorMapDataType_enum const dataType = getDataTypeFromXqaParams(xqaParams);
|
|
uint32_t const partElems = (forK ? 64 : 128);
|
|
if constexpr (std::is_same_v<KVCacheBuffer, KVBlockArray>)
|
|
{
|
|
return makeTensorMapForPagedKVCache(driver, kv_cache_buffer.mPrimaryPoolPtr, dataType, xqaParams.head_size,
|
|
xqaParams.num_kv_heads, xqaParams.tokens_per_block, partElems);
|
|
}
|
|
else
|
|
{
|
|
// The current implementation will waste memory. Don't use this for now.
|
|
TLLM_THROW("unsupported kv cache buffer type");
|
|
}
|
|
}
|
|
|
|
template CUtensorMap makeTensorMapForXqaMlaKVCache(
|
|
std::shared_ptr<CUDADriverWrapper> const&, XQAParams const&, KVBlockArray const&, bool);
|
|
template CUtensorMap makeTensorMapForXqaMlaKVCache(
|
|
std::shared_ptr<CUDADriverWrapper> const&, XQAParams const&, KVLinearBuffer const&, bool);
|
|
|
|
CUtensorMap makeTensorMapForXqaMlaQ(
|
|
std::shared_ptr<tensorrt_llm::common::CUDADriverWrapper> const& driver, XQAParams const& xqaParams, void const* q)
|
|
{
|
|
uint32_t const partElems = 64;
|
|
return makeTensorMapForQ(driver, q, getDataTypeFromXqaParams(xqaParams), xqaParams.head_size,
|
|
xqaParams.num_q_heads * xqaParams.total_num_input_tokens, partElems, xqaParams.num_q_heads);
|
|
}
|
|
} // namespace kernels
|
|
|
|
TRTLLM_NAMESPACE_END
|