diff --git a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu new file mode 100644 index 000000000000..7b6188c74eca --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu @@ -0,0 +1,372 @@ +/* + * Copyright (c) 2026, 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/deepseekV4BlockTable.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +constexpr int32_t kBadPageIndex = -1; +constexpr int32_t kThreadsPerBlock = 256; +constexpr int32_t kVecThreadsPerBlock = 128; +constexpr int32_t kRowKernelMinBlocks = 256; +constexpr int32_t kVecRowsPerBlock = 8; + +__device__ __forceinline__ int32_t computeBasePageIndex(int32_t const* __restrict__ blockOffsets, + int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, + int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, int32_t numPools, + int32_t copyIdxCapacity, int32_t numAttnTypes, int32_t maxBlocksPerSeq, int32_t layerId, int32_t attnTypeId, + int32_t tableId, int32_t blockId) +{ + int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; + int64_t const poolId64 = poolIds[layerAttnOffset]; + bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; + if (!isValidPool) + { + return kBadPageIndex; + } + + int32_t const mappedTableId = copyIdx[tableId]; + if (mappedTableId < 0 || mappedTableId >= copyIdxCapacity) + { + return kBadPageIndex; + } + + auto const poolId = static_cast(poolId64); + int64_t const blockOffsetsIndex + = (((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq) + blockId; + int32_t const base = blockOffsets[blockOffsetsIndex]; + if (base == kBadPageIndex) + { + return kBadPageIndex; + } + + return base * scales[layerAttnOffset] + layerOffsets[layerAttnOffset]; +} + +__device__ __forceinline__ int32_t applyScaleAndOffset(int32_t base, int32_t scale, int32_t layerOffset) +{ + return base == kBadPageIndex ? kBadPageIndex : base * scale + layerOffset; +} + +__device__ __forceinline__ void fillBadSlidingBlockTableRow(int32_t* outputRow, int32_t maxBlocksPerSeq, bool useVec4) +{ + if (useVec4) + { + int4 const bad = {kBadPageIndex, kBadPageIndex, kBadPageIndex, kBadPageIndex}; + auto* outputVec = reinterpret_cast(outputRow); + int32_t const vecsPerRow = maxBlocksPerSeq / 4; + for (int32_t vecId = threadIdx.x; vecId < vecsPerRow; vecId += blockDim.x) + { + outputVec[vecId] = bad; + } + return; + } + + for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) + { + outputRow[blockId] = kBadPageIndex; + } +} + +__global__ void computeSlidingBlockTablesRowsTiledKernel(int32_t const* __restrict__ blockOffsets, + int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, + int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, int32_t* __restrict__ output, + int32_t numPools, int32_t copyIdxCapacity, int32_t numLayerAttn, int32_t numTables, int32_t maxBlocksPerSeq) +{ + bool const useVec4 = maxBlocksPerSeq % 4 == 0; + int32_t const vecsPerRow = maxBlocksPerSeq / 4; + int32_t const firstTableId = static_cast(blockIdx.x) * kVecRowsPerBlock; + int32_t const layerAttnOffset = static_cast(blockIdx.y); + if (layerAttnOffset >= numLayerAttn) + { + return; + } + + int64_t const poolId64 = poolIds[layerAttnOffset]; + bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; + if (!isValidPool) + { +#pragma unroll + for (int32_t localRow = 0; localRow < kVecRowsPerBlock; ++localRow) + { + int32_t const tableId = firstTableId + localRow; + if (tableId >= numTables) + { + continue; + } + + int64_t const outputOffset + = (static_cast(layerAttnOffset) * numTables + tableId) * maxBlocksPerSeq; + fillBadSlidingBlockTableRow(output + outputOffset, maxBlocksPerSeq, useVec4); + } + return; + } + + auto const poolId = static_cast(poolId64); + int32_t const scale = scales[layerAttnOffset]; + int32_t const layerOffset = layerOffsets[layerAttnOffset]; + +#pragma unroll + for (int32_t localRow = 0; localRow < kVecRowsPerBlock; ++localRow) + { + int32_t const tableId = firstTableId + localRow; + if (tableId >= numTables) + { + continue; + } + + int64_t const outputOffset = (static_cast(layerAttnOffset) * numTables + tableId) * maxBlocksPerSeq; + auto* outputRow = output + outputOffset; + int32_t const mappedTableId = copyIdx[tableId]; + bool const isValidTable = mappedTableId >= 0 && mappedTableId < copyIdxCapacity; + if (!isValidTable) + { + fillBadSlidingBlockTableRow(outputRow, maxBlocksPerSeq, useVec4); + continue; + } + + int64_t const blockOffsetsOffset + = ((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq; + auto const* blockOffsetsRow = blockOffsets + blockOffsetsOffset; + if (useVec4) + { + auto const* blockOffsetsVec = reinterpret_cast(blockOffsetsRow); + auto* outputVec = reinterpret_cast(outputRow); + for (int32_t vecId = threadIdx.x; vecId < vecsPerRow; vecId += blockDim.x) + { + int4 const base = blockOffsetsVec[vecId]; + int4 const value = {applyScaleAndOffset(base.x, scale, layerOffset), + applyScaleAndOffset(base.y, scale, layerOffset), applyScaleAndOffset(base.z, scale, layerOffset), + applyScaleAndOffset(base.w, scale, layerOffset)}; + outputVec[vecId] = value; + } + continue; + } + + for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) + { + int32_t const base = blockOffsetsRow[blockId]; + outputRow[blockId] = applyScaleAndOffset(base, scale, layerOffset); + } + } +} + +__global__ void computeSlidingBlockTablesWithScratchKernel(int32_t const* __restrict__ blockOffsets, + int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, + int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, + int32_t const* __restrict__ scratchPages, int32_t const* __restrict__ scratchBegs, + int32_t const* __restrict__ scratchEnds, int32_t const* __restrict__ scratchSlots, + int32_t const* __restrict__ numContexts, int32_t* __restrict__ output, int64_t totalElements, int32_t numPools, + int32_t copyIdxCapacity, int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, + int32_t maxScratchSlots) +{ + int64_t const linearIdx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linearIdx >= totalElements) + { + return; + } + + int64_t remaining = linearIdx; + int32_t const blockId = static_cast(remaining % maxBlocksPerSeq); + remaining /= maxBlocksPerSeq; + int32_t const tableId = static_cast(remaining % numTables); + remaining /= numTables; + int32_t const attnTypeId = static_cast(remaining % numAttnTypes); + int32_t const layerId = static_cast(remaining / numAttnTypes); + + int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; + int32_t const basePageIndex = computeBasePageIndex(blockOffsets, copyIdx, poolIds, validPool, scales, layerOffsets, + numPools, copyIdxCapacity, numAttnTypes, maxBlocksPerSeq, layerId, attnTypeId, tableId, blockId); + + int64_t const poolId64 = poolIds[layerAttnOffset]; + bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; + int32_t const activeContexts = numContexts[0]; + bool const canUseScratch = isValidPool && tableId < scratchCapacity && tableId < activeContexts; + + if (!canUseScratch) + { + output[linearIdx] = basePageIndex; + return; + } + + auto const poolId = static_cast(poolId64); + int64_t const scratchRangeOffset = static_cast(poolId) * scratchCapacity + tableId; + int32_t const scratchBeg = scratchBegs[scratchRangeOffset]; + int32_t const scratchEnd = scratchEnds[scratchRangeOffset]; + bool const inScratchRange = blockId >= scratchBeg && blockId < scratchEnd; + if (!inScratchRange) + { + output[linearIdx] = basePageIndex; + return; + } + + int32_t const scale = scales[layerAttnOffset]; + int32_t const rangeIndex = blockId - scratchBeg; + int32_t const totalOffset = rangeIndex * scratchPages[layerAttnOffset]; + int32_t slotIdx = totalOffset / scale; + if (slotIdx >= maxScratchSlots) + { + slotIdx = maxScratchSlots - 1; + } + + int64_t const slotOffset = scratchRangeOffset * maxScratchSlots + slotIdx; + int32_t const slotId = scratchSlots[slotOffset]; + int32_t const offset = totalOffset % scale; + output[linearIdx] = slotId * scale + ((offset + layerOffsets[layerAttnOffset]) % scale); +} + +__global__ void computeSlidingBlockTablesWithScratchRowsKernel(int32_t const* __restrict__ blockOffsets, + int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, + int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, + int32_t const* __restrict__ scratchPages, int32_t const* __restrict__ scratchBegs, + int32_t const* __restrict__ scratchEnds, int32_t const* __restrict__ scratchSlots, + int32_t const* __restrict__ numContexts, int32_t* __restrict__ output, int32_t numPools, int32_t copyIdxCapacity, + int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots) +{ + int32_t const rowIdx = static_cast(blockIdx.x); + int32_t const tableId = rowIdx % numTables; + int32_t const layerAttnIdx = rowIdx / numTables; + int32_t const attnTypeId = layerAttnIdx % numAttnTypes; + int32_t const layerId = layerAttnIdx / numAttnTypes; + int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; + int32_t const outputOffset = rowIdx * maxBlocksPerSeq; + + int64_t const poolId64 = poolIds[layerAttnOffset]; + bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; + if (!isValidPool) + { + for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) + { + output[outputOffset + blockId] = kBadPageIndex; + } + return; + } + + auto const poolId = static_cast(poolId64); + int32_t const scale = scales[layerAttnOffset]; + int32_t const layerOffset = layerOffsets[layerAttnOffset]; + int32_t const activeContexts = numContexts[0]; + bool const canUseScratch = tableId < scratchCapacity && tableId < activeContexts; + int64_t const scratchRangeOffset = static_cast(poolId) * scratchCapacity + tableId; + int32_t const scratchBeg = canUseScratch ? scratchBegs[scratchRangeOffset] : 0; + int32_t const scratchEnd = canUseScratch ? scratchEnds[scratchRangeOffset] : 0; + int32_t const scratchPageCount = scratchPages[layerAttnOffset]; + + int32_t const mappedTableId = copyIdx[tableId]; + bool const isValidTable = mappedTableId >= 0 && mappedTableId < copyIdxCapacity; + int64_t const blockOffsetsOffset + = ((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq; + + for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) + { + bool const inScratchRange = canUseScratch && blockId >= scratchBeg && blockId < scratchEnd; + if (inScratchRange) + { + int32_t const rangeIndex = blockId - scratchBeg; + int32_t const totalOffset = rangeIndex * scratchPageCount; + int32_t slotIdx = totalOffset / scale; + if (slotIdx >= maxScratchSlots) + { + slotIdx = maxScratchSlots - 1; + } + + int64_t const slotOffset = scratchRangeOffset * maxScratchSlots + slotIdx; + int32_t const slotId = scratchSlots[slotOffset]; + int32_t const offset = totalOffset % scale; + output[outputOffset + blockId] = slotId * scale + ((offset + layerOffset) % scale); + continue; + } + + if (!isValidTable) + { + output[outputOffset + blockId] = kBadPageIndex; + continue; + } + + int32_t const base = blockOffsets[blockOffsetsOffset + blockId]; + output[outputOffset + blockId] = base == kBadPageIndex ? kBadPageIndex : base * scale + layerOffset; + } +} + +} // namespace + +void invokeDeepseekV4ComputeSlidingBlockTables(int32_t const* blockOffsets, int32_t const* copyIdx, + int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, int32_t* output, + int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, int32_t numAttnTypes, int32_t numTables, + int32_t maxBlocksPerSeq, cudaStream_t stream) +{ + int64_t const totalElements = static_cast(numLayers) * numAttnTypes * numTables * maxBlocksPerSeq; + if (totalElements == 0) + { + return; + } + + int32_t const numLayerAttn = numLayers * numAttnTypes; + int32_t const itemsPerRow = maxBlocksPerSeq % 4 == 0 ? maxBlocksPerSeq / 4 : maxBlocksPerSeq; + int32_t threadsPerBlock = itemsPerRow >= kVecThreadsPerBlock ? kVecThreadsPerBlock : itemsPerRow; + if (threadsPerBlock < 64) + { + threadsPerBlock = 64; + } + + dim3 const block(static_cast(threadsPerBlock)); + dim3 const grid(static_cast((numTables + kVecRowsPerBlock - 1) / kVecRowsPerBlock), + static_cast(numLayerAttn)); + computeSlidingBlockTablesRowsTiledKernel<<>>(blockOffsets, copyIdx, poolIds, validPool, + scales, layerOffsets, output, numPools, copyIdxCapacity, numLayerAttn, numTables, maxBlocksPerSeq); +} + +void invokeDeepseekV4ComputeSlidingBlockTablesWithScratch(int32_t const* blockOffsets, int32_t const* copyIdx, + int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, + int32_t const* scratchPages, int32_t const* scratchBegs, int32_t const* scratchEnds, int32_t const* scratchSlots, + int32_t const* numContexts, int32_t* output, int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, + int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots, + cudaStream_t stream) +{ + int64_t const totalElements = static_cast(numLayers) * numAttnTypes * numTables * maxBlocksPerSeq; + if (totalElements == 0) + { + return; + } + + if (maxBlocksPerSeq >= kRowKernelMinBlocks) + { + int32_t const numRows = numLayers * numAttnTypes * numTables; + dim3 const block(kThreadsPerBlock); + dim3 const grid(static_cast(numRows)); + computeSlidingBlockTablesWithScratchRowsKernel<<>>(blockOffsets, copyIdx, poolIds, + validPool, scales, layerOffsets, scratchPages, scratchBegs, scratchEnds, scratchSlots, numContexts, output, + numPools, copyIdxCapacity, numAttnTypes, numTables, maxBlocksPerSeq, scratchCapacity, maxScratchSlots); + return; + } + + dim3 const block(kThreadsPerBlock); + dim3 const grid(static_cast((totalElements + kThreadsPerBlock - 1) / kThreadsPerBlock)); + computeSlidingBlockTablesWithScratchKernel<<>>(blockOffsets, copyIdx, poolIds, validPool, + scales, layerOffsets, scratchPages, scratchBegs, scratchEnds, scratchSlots, numContexts, output, totalElements, + numPools, copyIdxCapacity, numAttnTypes, numTables, maxBlocksPerSeq, scratchCapacity, maxScratchSlots); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h new file mode 100644 index 000000000000..f57b757faaf3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) 2026, 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void invokeDeepseekV4ComputeSlidingBlockTables(int32_t const* blockOffsets, int32_t const* copyIdx, + int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, int32_t* output, + int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, int32_t numAttnTypes, int32_t numTables, + int32_t maxBlocksPerSeq, cudaStream_t stream); + +void invokeDeepseekV4ComputeSlidingBlockTablesWithScratch(int32_t const* blockOffsets, int32_t const* copyIdx, + int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, + int32_t const* scratchPages, int32_t const* scratchBegs, int32_t const* scratchEnds, int32_t const* scratchSlots, + int32_t const* numContexts, int32_t* output, int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, + int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index cfc72d9788ce..25694680218d 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -66,6 +66,7 @@ add_library( fp8Quantize.cpp dsv3FusedAGemmOp.cpp deepseekV4QNormOp.cpp + deepseekV4BlockTableOp.cpp inverseRopeFp8QuantOp.cpp fusedQKNormRopeOp.cpp fusedDiTQKNormRopeOp.cpp diff --git a/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp b/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp new file mode 100644 index 000000000000..8eb20a88bc1c --- /dev/null +++ b/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp @@ -0,0 +1,201 @@ +/* + * Copyright (c) 2026, 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/deepseekV4BlockTable.h" + +#include +#include +#include +#include +#include + +namespace th = torch; +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +void checkInt32Tensor(th::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.scalar_type() == th::kInt32, name, " must be int32"); +} + +void checkCudaContiguousTensor(th::Tensor const& tensor, char const* name, int device) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.get_device() == device, name, " must be on the same CUDA device as output"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +void checkLayerAttnShape(th::Tensor const& tensor, char const* name, int64_t numLayers, int64_t numAttnTypes) +{ + TORCH_CHECK(tensor.dim() == 2, name, " must be 2D [num_layers, num_attn_types]"); + TORCH_CHECK(tensor.size(0) == numLayers && tensor.size(1) == numAttnTypes, name, + " must match pool_ids shape [num_layers, num_attn_types]"); +} + +int32_t checkedInt32Size(int64_t value, char const* name) +{ + TORCH_CHECK(value <= std::numeric_limits::max(), name, " exceeds int32 range"); + return static_cast(value); +} + +void checkCommonInputs(th::Tensor const& blockOffsets, th::Tensor const& copyIdx, th::Tensor const& poolIds, + th::Tensor const& validPool, th::Tensor const& scales, th::Tensor const& layerOffsets, th::Tensor const& output) +{ + TORCH_CHECK(output.is_cuda(), "output must be a CUDA tensor"); + int const device = output.get_device(); + TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); + checkCudaContiguousTensor(blockOffsets, "block_offsets", device); + checkCudaContiguousTensor(copyIdx, "copy_idx", device); + checkCudaContiguousTensor(poolIds, "pool_ids", device); + checkCudaContiguousTensor(validPool, "valid_pool", device); + checkCudaContiguousTensor(scales, "scales", device); + checkCudaContiguousTensor(layerOffsets, "layer_offsets", device); + + checkInt32Tensor(blockOffsets, "block_offsets"); + checkInt32Tensor(copyIdx, "copy_idx"); + TORCH_CHECK(poolIds.scalar_type() == th::kInt64, "pool_ids must be int64"); + TORCH_CHECK(validPool.scalar_type() == th::kBool, "valid_pool must be bool"); + checkInt32Tensor(scales, "scales"); + checkInt32Tensor(layerOffsets, "layer_offsets"); + checkInt32Tensor(output, "output"); + + TORCH_CHECK(blockOffsets.dim() == 4, "block_offsets must be 4D [num_pools, table_capacity, 2, max_blocks]"); + TORCH_CHECK(blockOffsets.size(2) == 2, "block_offsets dim 2 must be 2"); + TORCH_CHECK(copyIdx.dim() == 1, "copy_idx must be 1D"); + TORCH_CHECK(poolIds.dim() == 2, "pool_ids must be 2D [num_layers, num_attn_types]"); + TORCH_CHECK(output.dim() == 4, "output must be 4D [num_layers, num_attn_types, num_tables, max_blocks]"); + + int64_t const numLayers = poolIds.size(0); + int64_t const numAttnTypes = poolIds.size(1); + int64_t const numTables = copyIdx.size(0); + int64_t const maxBlocksPerSeq = blockOffsets.size(3); + + checkLayerAttnShape(validPool, "valid_pool", numLayers, numAttnTypes); + checkLayerAttnShape(scales, "scales", numLayers, numAttnTypes); + checkLayerAttnShape(layerOffsets, "layer_offsets", numLayers, numAttnTypes); + + TORCH_CHECK(output.size(0) == numLayers && output.size(1) == numAttnTypes && output.size(2) == numTables + && output.size(3) == maxBlocksPerSeq, + "output shape must be [pool_ids.size(0), pool_ids.size(1), copy_idx.size(0), block_offsets.size(3)]"); +} + +} // namespace + +void deepseekV4ComputeSlidingBlockTables(th::Tensor const& blockOffsets, th::Tensor const& copyIdx, + th::Tensor const& poolIds, th::Tensor const& validPool, th::Tensor const& scales, th::Tensor const& layerOffsets, + th::Tensor const& output) +{ + checkCommonInputs(blockOffsets, copyIdx, poolIds, validPool, scales, layerOffsets, output); + c10::cuda::CUDAGuard const deviceGuard(output.device()); + + int32_t const numPools = checkedInt32Size(blockOffsets.size(0), "num_pools"); + int32_t const copyIdxCapacity = checkedInt32Size(blockOffsets.size(1), "copy_idx_capacity"); + int32_t const numLayers = checkedInt32Size(poolIds.size(0), "num_layers"); + int32_t const numAttnTypes = checkedInt32Size(poolIds.size(1), "num_attn_types"); + int32_t const numTables = checkedInt32Size(copyIdx.size(0), "num_tables"); + int32_t const maxBlocksPerSeq = checkedInt32Size(blockOffsets.size(3), "max_blocks_per_seq"); + + auto stream = at::cuda::getCurrentCUDAStream(output.get_device()); + tk::invokeDeepseekV4ComputeSlidingBlockTables(blockOffsets.data_ptr(), copyIdx.data_ptr(), + poolIds.data_ptr(), validPool.data_ptr(), scales.data_ptr(), + layerOffsets.data_ptr(), output.data_ptr(), numPools, copyIdxCapacity, numLayers, + numAttnTypes, numTables, maxBlocksPerSeq, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void deepseekV4ComputeSlidingBlockTablesWithScratch(th::Tensor const& blockOffsets, th::Tensor const& copyIdx, + th::Tensor const& poolIds, th::Tensor const& validPool, th::Tensor const& scales, th::Tensor const& layerOffsets, + th::Tensor const& scratchPages, th::Tensor const& scratchBegs, th::Tensor const& scratchEnds, + th::Tensor const& scratchSlots, th::Tensor const& numContexts, th::Tensor const& output) +{ + checkCommonInputs(blockOffsets, copyIdx, poolIds, validPool, scales, layerOffsets, output); + int const device = output.get_device(); + checkCudaContiguousTensor(scratchPages, "scratch_pages", device); + checkCudaContiguousTensor(scratchBegs, "scratch_begs", device); + checkCudaContiguousTensor(scratchEnds, "scratch_ends", device); + checkCudaContiguousTensor(scratchSlots, "scratch_slots", device); + checkCudaContiguousTensor(numContexts, "num_contexts", device); + + checkInt32Tensor(scratchPages, "scratch_pages"); + checkInt32Tensor(scratchBegs, "scratch_begs"); + checkInt32Tensor(scratchEnds, "scratch_ends"); + checkInt32Tensor(scratchSlots, "scratch_slots"); + checkInt32Tensor(numContexts, "num_contexts"); + + int64_t const numLayers = poolIds.size(0); + int64_t const numAttnTypes = poolIds.size(1); + checkLayerAttnShape(scratchPages, "scratch_pages", numLayers, numAttnTypes); + TORCH_CHECK(scratchBegs.dim() == 2, "scratch_begs must be 2D [num_pools, scratch_capacity]"); + TORCH_CHECK(scratchEnds.dim() == 2, "scratch_ends must be 2D [num_pools, scratch_capacity]"); + TORCH_CHECK(scratchSlots.dim() == 3, "scratch_slots must be 3D [num_pools, scratch_capacity, max_scratch_slots]"); + TORCH_CHECK(scratchBegs.size(0) == blockOffsets.size(0), "scratch_begs.size(0) must match num_pools"); + TORCH_CHECK(scratchEnds.size(0) == scratchBegs.size(0) && scratchEnds.size(1) == scratchBegs.size(1), + "scratch_ends shape must match scratch_begs"); + TORCH_CHECK(scratchSlots.size(0) == scratchBegs.size(0) && scratchSlots.size(1) == scratchBegs.size(1), + "scratch_slots first two dimensions must match scratch_begs"); + TORCH_CHECK(scratchBegs.size(1) <= output.size(2), "scratch_capacity must not exceed num_tables"); + TORCH_CHECK(numContexts.dim() == 0 && numContexts.numel() == 1, "num_contexts must be a scalar tensor"); + TORCH_CHECK(scratchSlots.size(2) > 0 || scratchBegs.size(1) == 0, + "max_scratch_slots must be positive when scratch_capacity is nonzero"); + + c10::cuda::CUDAGuard const deviceGuard(output.device()); + int32_t const numPools = checkedInt32Size(blockOffsets.size(0), "num_pools"); + int32_t const copyIdxCapacity = checkedInt32Size(blockOffsets.size(1), "copy_idx_capacity"); + int32_t const numLayers32 = checkedInt32Size(numLayers, "num_layers"); + int32_t const numAttnTypes32 = checkedInt32Size(numAttnTypes, "num_attn_types"); + int32_t const numTables = checkedInt32Size(copyIdx.size(0), "num_tables"); + int32_t const maxBlocksPerSeq = checkedInt32Size(blockOffsets.size(3), "max_blocks_per_seq"); + int32_t const scratchCapacity = checkedInt32Size(scratchBegs.size(1), "scratch_capacity"); + int32_t const maxScratchSlots = checkedInt32Size(scratchSlots.size(2), "max_scratch_slots"); + + auto stream = at::cuda::getCurrentCUDAStream(output.get_device()); + tk::invokeDeepseekV4ComputeSlidingBlockTablesWithScratch(blockOffsets.data_ptr(), + copyIdx.data_ptr(), poolIds.data_ptr(), validPool.data_ptr(), + scales.data_ptr(), layerOffsets.data_ptr(), scratchPages.data_ptr(), + scratchBegs.data_ptr(), scratchEnds.data_ptr(), scratchSlots.data_ptr(), + numContexts.data_ptr(), output.data_ptr(), numPools, copyIdxCapacity, numLayers32, + numAttnTypes32, numTables, maxBlocksPerSeq, scratchCapacity, maxScratchSlots, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "deepseek_v4_compute_sliding_block_tables(Tensor block_offsets, Tensor copy_idx, Tensor pool_ids, " + "Tensor valid_pool, Tensor scales, Tensor layer_offsets, Tensor(a!) output) -> ()"); + m.def( + "deepseek_v4_compute_sliding_block_tables_with_scratch(Tensor block_offsets, Tensor copy_idx, " + "Tensor pool_ids, Tensor valid_pool, Tensor scales, Tensor layer_offsets, Tensor scratch_pages, " + "Tensor scratch_begs, Tensor scratch_ends, Tensor scratch_slots, Tensor num_contexts, " + "Tensor(a!) output) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("deepseek_v4_compute_sliding_block_tables", &tensorrt_llm::torch_ext::deepseekV4ComputeSlidingBlockTables); + m.impl("deepseek_v4_compute_sliding_block_tables_with_scratch", + &tensorrt_llm::torch_ext::deepseekV4ComputeSlidingBlockTablesWithScratch); +} diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index b61b040e375a..4ee0977f04d0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -20,7 +20,6 @@ from tensorrt_llm._torch.pyexecutor import llm_request from tensorrt_llm._torch.pyexecutor.resource_manager import GPU_LEVEL, KVCacheManagerV2 -from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import ( TensorWrapper, convert_to_torch_tensor, @@ -167,77 +166,6 @@ def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode: return PageIndexMode.SHARED -@maybe_compile(options={"max-autotune": True}) -def _compute_sliding_block_tables_compiled( - block_offsets: torch.Tensor, - copy_idx: torch.Tensor, - pool_ids: torch.Tensor, - valid_pool: torch.Tensor, - scales: torch.Tensor, - layer_offsets: torch.Tensor, - output: torch.Tensor, -) -> None: - base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] - scaled_base = torch.where( - (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), - BAD_PAGE_INDEX, - base * scales[:, :, None, None] + layer_offsets[:, :, None, None], - ) - output.copy_(scaled_base) - - -@maybe_compile(options={"max-autotune": True}) -def _compute_sliding_block_tables_with_scratch_compiled( - block_offsets: torch.Tensor, - copy_idx: torch.Tensor, - pool_ids: torch.Tensor, - valid_pool: torch.Tensor, - scales: torch.Tensor, - layer_offsets: torch.Tensor, - block_positions: torch.Tensor, - scratch_pages: torch.Tensor, - scratch_begs: torch.Tensor, - scratch_ends: torch.Tensor, - scratch_slots: torch.Tensor, - num_contexts: torch.Tensor, - output: torch.Tensor, -) -> None: - base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] - scaled_base = torch.where( - (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), - BAD_PAGE_INDEX, - base * scales[:, :, None, None] + layer_offsets[:, :, None, None], - ) - output.copy_(scaled_base) - - context_positions = torch.arange( - scratch_begs.shape[1], - dtype=torch.int32, - device=scratch_begs.device, - ) - active_context = context_positions < num_contexts - mask = ( - (block_positions >= scratch_begs[:, :, None]) - & (block_positions < scratch_ends[:, :, None]) - & active_context[None, :, None] - ) - range_index = torch.where(mask, block_positions - scratch_begs[:, :, None], 0) - total_offset = range_index[pool_ids] * scratch_pages[:, :, None, None] - slot_idx = (total_offset // scales[:, :, None, None]).clamp( - max=scratch_slots.shape[-1] - 1, - ) - slot_id = scratch_slots[pool_ids].gather(-1, slot_idx.long()) - offset = total_offset % scales[:, :, None, None] - scratch_index = ( - slot_id * scales[:, :, None, None] - + (offset + layer_offsets[:, :, None, None]) % scales[:, :, None, None] - ) - scratch_capacity = scratch_begs.shape[1] - scratch_rows = scaled_base[:, :, :scratch_capacity, :] - mask = mask[pool_ids] & valid_pool[:, :, None, None] - output[:, :, :scratch_capacity, :].copy_(torch.where(mask, scratch_index, scratch_rows)) - - class DeepseekV4CacheManager(KVCacheManagerV2): # This tensor is for compatibility with AttentionOp, it only contains swa attention. # kv_cache_pool_pointers contains one virtual attention-op pool per local @@ -1360,14 +1288,13 @@ def compute_sliding_block_tables( self._host_scratch_slots_staging, ) self._device_num_contexts.fill_(num_contexts) - _compute_sliding_block_tables_with_scratch_compiled( + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables_with_scratch( self._device_kv_cache_block_offsets_input, device_copy_idx, self._device_layer_attn_pool_ids, self._device_valid_sliding_pool, self._device_layer_attn_scales, self._device_layer_offsets, - self._device_block_positions, self._device_scratch_pages, scratch_begs, scratch_ends, @@ -1376,7 +1303,7 @@ def compute_sliding_block_tables( self._precomputed_sliding_block_tables, ) else: - _compute_sliding_block_tables_compiled( + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables( self._device_kv_cache_block_offsets_input, device_copy_idx, self._device_layer_attn_pool_ids, diff --git a/tests/microbenchmarks/dsv4_block_table_perf.py b/tests/microbenchmarks/dsv4_block_table_perf.py new file mode 100644 index 000000000000..599a1d97941e --- /dev/null +++ b/tests/microbenchmarks/dsv4_block_table_perf.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +from collections.abc import Callable + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 + +BAD_PAGE_INDEX = -1 + + +def _torch_sliding_block_tables( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + valid_pool: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, + output: torch.Tensor, +) -> None: + base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] + scaled_base = torch.where( + (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), + BAD_PAGE_INDEX, + base * scales[:, :, None, None] + layer_offsets[:, :, None, None], + ) + output.copy_(scaled_base) + + +def _torch_sliding_block_tables_with_scratch( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + valid_pool: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, + block_positions: torch.Tensor, + scratch_pages: torch.Tensor, + scratch_begs: torch.Tensor, + scratch_ends: torch.Tensor, + scratch_slots: torch.Tensor, + num_contexts: torch.Tensor, + output: torch.Tensor, +) -> None: + base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] + scaled_base = torch.where( + (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), + BAD_PAGE_INDEX, + base * scales[:, :, None, None] + layer_offsets[:, :, None, None], + ) + output.copy_(scaled_base) + + context_positions = torch.arange( + scratch_begs.shape[1], + dtype=torch.int32, + device=scratch_begs.device, + ) + active_context = context_positions < num_contexts + mask = ( + (block_positions >= scratch_begs[:, :, None]) + & (block_positions < scratch_ends[:, :, None]) + & active_context[None, :, None] + ) + range_index = torch.where(mask, block_positions - scratch_begs[:, :, None], 0) + total_offset = range_index[pool_ids] * scratch_pages[:, :, None, None] + slot_idx = (total_offset // scales[:, :, None, None]).clamp(max=scratch_slots.shape[-1] - 1) + slot_id = scratch_slots[pool_ids].gather(-1, slot_idx.long()) + offset = total_offset % scales[:, :, None, None] + scratch_index = ( + slot_id * scales[:, :, None, None] + + (offset + layer_offsets[:, :, None, None]) % scales[:, :, None, None] + ) + scratch_capacity = scratch_begs.shape[1] + scratch_rows = scaled_base[:, :, :scratch_capacity, :] + mask = mask[pool_ids] & valid_pool[:, :, None, None] + output[:, :, :scratch_capacity, :].copy_(torch.where(mask, scratch_index, scratch_rows)) + + +def _custom_sliding_block_tables(*args) -> None: + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables(*args) + + +def _custom_sliding_block_tables_with_scratch( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + valid_pool: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, + block_positions: torch.Tensor, + scratch_pages: torch.Tensor, + scratch_begs: torch.Tensor, + scratch_ends: torch.Tensor, + scratch_slots: torch.Tensor, + num_contexts: torch.Tensor, + output: torch.Tensor, +) -> None: + del block_positions + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables_with_scratch( + block_offsets, + copy_idx, + pool_ids, + valid_pool, + scales, + layer_offsets, + scratch_pages, + scratch_begs, + scratch_ends, + scratch_slots, + num_contexts, + output, + ) + + +def _make_inputs(args: argparse.Namespace) -> dict[str, torch.Tensor]: + torch.manual_seed(args.seed) + device = "cuda" + block_offsets = torch.randint( + 0, + 1_000_000, + (args.num_pools, args.table_capacity, 2, args.max_blocks), + dtype=torch.int32, + device=device, + ) + block_offsets.masked_fill_( + torch.rand(block_offsets.shape, device=device) < args.bad_ratio, BAD_PAGE_INDEX + ) + copy_idx = ( + torch.randperm(args.table_capacity, device=device)[: args.num_tables] + .to(torch.int32) + .contiguous() + ) + pool_ids = torch.randint( + 0, + args.num_pools, + (args.num_layers, args.num_attn_types), + dtype=torch.int64, + device=device, + ) + valid_pool = torch.ones((args.num_layers, args.num_attn_types), dtype=torch.bool, device=device) + if args.num_layers > 0 and args.num_attn_types > 0: + pool_ids[0, -1] = -1 + valid_pool[0, -1] = False + valid_pool[-1, 0] = False + scales = torch.randint(1, args.max_scale + 1, pool_ids.shape, dtype=torch.int32, device=device) + layer_offsets = torch.randint( + 0, args.max_scale, pool_ids.shape, dtype=torch.int32, device=device + ) + block_positions = torch.arange(args.max_blocks, dtype=torch.int32, device=device) + + scratch_begs = torch.randint( + 0, + max(1, args.max_blocks // 2), + (args.num_pools, args.num_tables), + dtype=torch.int32, + device=device, + ) + scratch_widths = torch.randint( + 1, + max(2, args.max_blocks // 4), + (args.num_pools, args.num_tables), + dtype=torch.int32, + device=device, + ) + scratch_ends = torch.minimum( + scratch_begs + scratch_widths, + torch.tensor(args.max_blocks, dtype=torch.int32, device=device), + ) + scratch_slots = torch.randint( + 0, + 1_000_000, + (args.num_pools, args.num_tables, args.max_scratch_slots), + dtype=torch.int32, + device=device, + ) + scratch_pages = torch.randint( + 1, args.max_scratch_pages + 1, pool_ids.shape, dtype=torch.int32, device=device + ) + num_contexts = torch.tensor(args.num_contexts, dtype=torch.int32, device=device) + output_shape = (args.num_layers, args.num_attn_types, args.num_tables, args.max_blocks) + return { + "block_offsets": block_offsets, + "copy_idx": copy_idx, + "pool_ids": pool_ids, + "valid_pool": valid_pool, + "scales": scales, + "layer_offsets": layer_offsets, + "block_positions": block_positions, + "scratch_pages": scratch_pages, + "scratch_begs": scratch_begs, + "scratch_ends": scratch_ends, + "scratch_slots": scratch_slots, + "num_contexts": num_contexts, + "reference_output": torch.empty(output_shape, dtype=torch.int32, device=device), + "custom_output": torch.empty(output_shape, dtype=torch.int32, device=device), + } + + +def _time_cuda(fn: Callable, fn_args: tuple, warmup: int, iters: int) -> float: + for _ in range(warmup): + fn(*fn_args) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn(*fn_args) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _maybe_compile(fn: Callable, reference: str) -> Callable: + if reference == "eager": + return fn + return torch.compile(fn, options={"max-autotune": True}) + + +def _run_case( + name: str, reference_fn: Callable, custom_fn: Callable, fn_args: tuple, args: argparse.Namespace +) -> None: + reference_fn = _maybe_compile(reference_fn, args.reference) + reference_output = fn_args[-1] + custom_output = args.inputs["custom_output"] + custom_args = (*fn_args[:-1], custom_output) + + reference_fn(*fn_args) + custom_fn(*custom_args) + torch.cuda.synchronize() + if not torch.equal(custom_output, reference_output): + mismatch = (custom_output != reference_output).nonzero() + first = mismatch[0].tolist() if mismatch.numel() > 0 else [] + raise AssertionError(f"{name}: custom output differs from reference at {first}") + + reference_ms = _time_cuda(reference_fn, fn_args, args.warmup, args.iters) + custom_ms = _time_cuda(custom_fn, custom_args, args.warmup, args.iters) + print( + f"{name}: reference={reference_ms:.4f} ms custom={custom_ms:.4f} ms " + f"speedup={reference_ms / custom_ms:.3f}x" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--reference", choices=["compile", "eager"], default="compile") + parser.add_argument("--case", choices=["all", "no-scratch", "scratch"], default="all") + parser.add_argument("--num-pools", type=int, default=8) + parser.add_argument("--table-capacity", type=int, default=256) + parser.add_argument("--num-layers", type=int, default=61) + parser.add_argument("--num-attn-types", type=int, default=4) + parser.add_argument("--num-tables", type=int, default=256) + parser.add_argument("--max-blocks", type=int, default=1024) + parser.add_argument("--max-scale", type=int, default=128) + parser.add_argument("--max-scratch-pages", type=int, default=16) + parser.add_argument("--max-scratch-slots", type=int, default=64) + parser.add_argument("--num-contexts", type=int, default=256) + parser.add_argument("--bad-ratio", type=float, default=0.05) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.num_contexts > args.num_tables: + raise ValueError("--num-contexts must be <= --num-tables") + + args.inputs = _make_inputs(args) + common_args = ( + args.inputs["block_offsets"], + args.inputs["copy_idx"], + args.inputs["pool_ids"], + args.inputs["valid_pool"], + args.inputs["scales"], + args.inputs["layer_offsets"], + ) + scratch_args = ( + args.inputs["block_positions"], + args.inputs["scratch_pages"], + args.inputs["scratch_begs"], + args.inputs["scratch_ends"], + args.inputs["scratch_slots"], + args.inputs["num_contexts"], + ) + + shape = tuple(args.inputs["reference_output"].shape) + print(f"shape={shape} reference={args.reference} warmup={args.warmup} iters={args.iters}") + if args.case in ("all", "no-scratch"): + _run_case( + "no-scratch", + _torch_sliding_block_tables, + _custom_sliding_block_tables, + (*common_args, args.inputs["reference_output"]), + args, + ) + if args.case in ("all", "scratch"): + _run_case( + "scratch", + _torch_sliding_block_tables_with_scratch, + _custom_sliding_block_tables_with_scratch, + (*common_args, *scratch_args, args.inputs["reference_output"]), + args, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py b/tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py new file mode 100644 index 000000000000..f070a4a291e1 --- /dev/null +++ b/tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 + +BAD_PAGE_INDEX = -1 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + + +def _make_common_inputs( + *, + num_pools: int = 5, + table_capacity: int = 9, + num_layers: int = 4, + num_attn_types: int = 4, + num_tables: int = 7, + max_blocks: int = 17, +): + torch.manual_seed(0) + device = "cuda" + block_offsets = torch.randint( + 0, + 4096, + (num_pools, table_capacity, 2, max_blocks), + dtype=torch.int32, + device=device, + ) + block_offsets.masked_fill_( + torch.rand(block_offsets.shape, device=device) < 0.17, BAD_PAGE_INDEX + ) + copy_idx = ( + torch.randperm(table_capacity, device=device)[:num_tables].to(torch.int32).contiguous() + ) + pool_ids = torch.randint( + 0, num_pools, (num_layers, num_attn_types), dtype=torch.int64, device=device + ) + valid_pool = torch.ones((num_layers, num_attn_types), dtype=torch.bool, device=device) + pool_ids[0, -1] = -1 + valid_pool[0, -1] = False + valid_pool[-1, 0] = False + scales = torch.randint(1, 8, (num_layers, num_attn_types), dtype=torch.int32, device=device) + layer_offsets = torch.randint( + 0, 32, (num_layers, num_attn_types), dtype=torch.int32, device=device + ) + return block_offsets, copy_idx, pool_ids, valid_pool, scales, layer_offsets + + +def _reference_sliding_block_tables( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + valid_pool: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, +) -> torch.Tensor: + base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] + scaled_base = torch.where( + (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), + BAD_PAGE_INDEX, + base * scales[:, :, None, None] + layer_offsets[:, :, None, None], + ) + output = torch.empty_like(scaled_base) + output.copy_(scaled_base) + return output + + +def _reference_sliding_block_tables_with_scratch( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + valid_pool: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, + scratch_pages: torch.Tensor, + scratch_begs: torch.Tensor, + scratch_ends: torch.Tensor, + scratch_slots: torch.Tensor, + num_contexts: torch.Tensor, +) -> torch.Tensor: + base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] + scaled_base = torch.where( + (base == BAD_PAGE_INDEX) | ~(valid_pool[:, :, None, None]), + BAD_PAGE_INDEX, + base * scales[:, :, None, None] + layer_offsets[:, :, None, None], + ) + output = torch.empty_like(scaled_base) + output.copy_(scaled_base) + + block_positions = torch.arange( + block_offsets.shape[-1], dtype=torch.int32, device=block_offsets.device + ) + context_positions = torch.arange( + scratch_begs.shape[1], + dtype=torch.int32, + device=scratch_begs.device, + ) + active_context = context_positions < num_contexts + mask = ( + (block_positions >= scratch_begs[:, :, None]) + & (block_positions < scratch_ends[:, :, None]) + & active_context[None, :, None] + ) + range_index = torch.where(mask, block_positions - scratch_begs[:, :, None], 0) + total_offset = range_index[pool_ids] * scratch_pages[:, :, None, None] + slot_idx = (total_offset // scales[:, :, None, None]).clamp(max=scratch_slots.shape[-1] - 1) + slot_id = scratch_slots[pool_ids].gather(-1, slot_idx.long()) + offset = total_offset % scales[:, :, None, None] + scratch_index = ( + slot_id * scales[:, :, None, None] + + (offset + layer_offsets[:, :, None, None]) % scales[:, :, None, None] + ) + scratch_capacity = scratch_begs.shape[1] + scratch_rows = output[:, :, :scratch_capacity, :] + mask = mask[pool_ids] & valid_pool[:, :, None, None] + output[:, :, :scratch_capacity, :].copy_(torch.where(mask, scratch_index, scratch_rows)) + return output + + +@pytest.mark.parametrize("max_blocks", [17, 256]) +def test_deepseek_v4_compute_sliding_block_tables_matches_reference(max_blocks): + block_offsets, copy_idx, pool_ids, valid_pool, scales, layer_offsets = _make_common_inputs( + max_blocks=max_blocks + ) + expected = _reference_sliding_block_tables( + block_offsets, + copy_idx, + pool_ids, + valid_pool, + scales, + layer_offsets, + ) + output = torch.empty_like(expected) + + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables( + block_offsets, + copy_idx, + pool_ids, + valid_pool, + scales, + layer_offsets, + output, + ) + + assert torch.equal(output, expected) + + +@pytest.mark.parametrize("max_blocks", [17, 256]) +def test_deepseek_v4_compute_sliding_block_tables_with_scratch_matches_reference(max_blocks): + block_offsets, copy_idx, pool_ids, valid_pool, scales, layer_offsets = _make_common_inputs( + max_blocks=max_blocks + ) + num_pools = block_offsets.shape[0] + scratch_capacity = copy_idx.shape[0] + max_scratch_slots = 5 + device = block_offsets.device + torch.manual_seed(1) + scratch_begs = torch.randint( + 0, + block_offsets.shape[-1] // 2, + (num_pools, scratch_capacity), + dtype=torch.int32, + device=device, + ) + scratch_widths = torch.randint( + 1, + max(2, block_offsets.shape[-1] // 3), + (num_pools, scratch_capacity), + dtype=torch.int32, + device=device, + ) + scratch_ends = torch.minimum( + scratch_begs + scratch_widths, + torch.tensor(block_offsets.shape[-1], dtype=torch.int32, device=device), + ) + scratch_slots = torch.randint( + 0, + 4096, + (num_pools, scratch_capacity, max_scratch_slots), + dtype=torch.int32, + device=device, + ) + scratch_pages = torch.randint(1, 8, scales.shape, dtype=torch.int32, device=device) + num_contexts = torch.tensor(scratch_capacity - 1, dtype=torch.int32, device=device) + + expected = _reference_sliding_block_tables_with_scratch( + block_offsets, + copy_idx, + pool_ids, + valid_pool, + scales, + layer_offsets, + scratch_pages, + scratch_begs, + scratch_ends, + scratch_slots, + num_contexts, + ) + output = torch.empty_like(expected) + + torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables_with_scratch( + block_offsets, + copy_idx, + pool_ids, + valid_pool, + scales, + layer_offsets, + scratch_pages, + scratch_begs, + scratch_ends, + scratch_slots, + num_contexts, + output, + ) + + assert torch.equal(output, expected)