diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h new file mode 100644 index 000000000000..60422dea5910 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2022-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 "tensorrt_llm/common/cudaUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mapping_fp8, + int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens, + int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, + int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3, + cudaStream_t stream = 0); + +} + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu new file mode 100644 index 000000000000..aaf792e33c1d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu @@ -0,0 +1,97 @@ +/* + * Copyright (c) 2022-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 "convertReqIndexToGlobal.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Each thread handles one element at (token_id, col). +// Grid: (num_tokens, ceil(numTopkTokens / blockDim.x)) +__global__ void convertReqIndexToGlobalKernel(int32_t const* __restrict__ reqId, int32_t const* __restrict__ blockTable, + int32_t const* __restrict__ tokenIndices, int32_t* __restrict__ output, int32_t numTopkTokens, + int32_t maxNumBlocksPerReq, int32_t blockSize, int32_t strideFactor, int32_t layerId, int64_t btStride0, + int64_t btStride1, int64_t tiStride0, int64_t tiStride1, int64_t outStride0, int64_t outStride1) +{ + int32_t const tokenId = blockIdx.x; + int32_t const col = blockIdx.y * blockDim.x + threadIdx.x; + + if (col >= numTopkTokens) + { + return; + } + + // Load request id for this token + int32_t const req = reqId[tokenId]; + + // Load token index + int32_t const tok = tokenIndices[tokenId * tiStride0 + col * tiStride1]; + + // Invalid token → output -1 + if (tok < 0) + { + output[tokenId * outStride0 + col * outStride1] = -1; + return; + } + + // Compute block id and in-block offset + int32_t const blockId = tok / blockSize; + int32_t const inblockOff = tok % blockSize + layerId * blockSize; + + // Guard block_table access + if (blockId >= maxNumBlocksPerReq) + { + output[tokenId * outStride0 + col * outStride1] = -1; + return; + } + + int32_t const base = blockTable[req * btStride0 + blockId * btStride1]; + + // Padding entry in block table + if (base < 0) + { + output[tokenId * outStride0 + col * outStride1] = -1; + return; + } + + output[tokenId * outStride0 + col * outStride1] = base * strideFactor + inblockOff; +} + +void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices, + int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize, + int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1, + int64_t outStride0, int64_t outStride1, cudaStream_t stream) +{ + if (numTokens == 0 || numTopkTokens == 0) + { + return; + } + + constexpr int32_t kThreadsPerBlock = 256; + int32_t const tilesPerRow = (numTopkTokens + kThreadsPerBlock - 1) / kThreadsPerBlock; + dim3 const grid(numTokens, tilesPerRow); + dim3 const block(kThreadsPerBlock); + + convertReqIndexToGlobalKernel<<>>(reqId, blockTable, tokenIndices, output, numTopkTokens, + maxNumBlocksPerReq, blockSize, strideFactor, layerId, btStride0, btStride1, tiStride0, tiStride1, outStride0, + outStride1); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h new file mode 100644 index 000000000000..74aa7783ee61 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2022-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 "tensorrt_llm/common/cudaUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices, + int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize, + int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1, + int64_t outStride0, int64_t outStride1, cudaStream_t stream = 0); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu new file mode 100644 index 000000000000..dcb5a212b319 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu @@ -0,0 +1,163 @@ +/* + * Copyright (c) 2022-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 "IndexerKCacheGather.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/cudaUtils.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +/** + * Given a flat element index and tensor shape [d0, d1, d2, d3] with strides [s0, s1, s2, s3], + * find the actual memory offset within the given k cache pool using the strides. + */ +__device__ __forceinline__ int64_t flatIndexToMemoryOffset( + int64_t flat_idx, int32_t d0, int32_t d1, int32_t d2, int32_t d3, int64_t s0, int64_t s1, int64_t s2, int64_t s3) +{ + // Unravel from innermost to outermost dimension + int32_t i3 = flat_idx % d3; + flat_idx /= d3; + + int32_t i2 = flat_idx % d2; + flat_idx /= d2; + + int32_t i1 = flat_idx % d1; + flat_idx /= d1; + + int32_t i0 = flat_idx; + + // Compute memory offset using strides + return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3; +} + +} // anonymous namespace + +/** + * CUDA kernel to gather both FP8 K values and scales from the indexer k cache pool. + * This is the inverse of indexerKCacheScatterUnifiedKernel. + * + * @param k_cache Indexer k cache pool with shape [num_blocks, block_size, 1, per_token_size] + * (can be non-contiguous) + * @param slot_mapping_fp8 Flat element index for FP8 data start position [total_kv_len] + * @param slot_mapping_scale Flat element index for scale data start position [total_kv_len] + * @param out_fp8 Output FP8 data [num_tokens, head_dim] contiguous + * @param out_scale Output scale data [num_tokens, scale_size] contiguous + * @param k_token_start Start offset into slot_mapping arrays + * @param num_tokens Number of tokens to gather + * @param head_dim Head dimension (must be 128) + * @param scale_size Scale size in bytes (must be 4) + * @param cache_stride_0 Stride for k_cache dimension 0 (in bytes) + * @param cache_stride_1 Stride for k_cache dimension 1 (in bytes) + * @param cache_stride_2 Stride for k_cache dimension 2 (in bytes) + * @param cache_stride_3 Stride for k_cache dimension 3 (in bytes) + * @param cache_dim_0 Size of k_cache dimension 0 + * @param cache_dim_1 Size of k_cache dimension 1 + * @param cache_dim_2 Size of k_cache dimension 2 + * @param cache_dim_3 Size of k_cache dimension 3 + */ +__global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_cache, + int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale, + uint8_t* __restrict__ out_fp8, uint8_t* __restrict__ out_scale, int32_t k_token_start, int32_t num_tokens, + int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, + int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3) +{ + // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction + constexpr int VEC_SIZE = 4; + + // Token index from block.x + int32_t token_idx = blockIdx.x; + + if (token_idx >= num_tokens) + { + return; + } + + // Index into slot_mapping with k_token_start offset + int32_t slot_idx = k_token_start + token_idx; + + int64_t flat_idx_fp8_base = slot_mapping_fp8[slot_idx]; + int64_t flat_idx_scale_base = slot_mapping_scale[slot_idx]; + + if (flat_idx_fp8_base < 0 || flat_idx_scale_base < 0) + { + return; + } + + int32_t head_dim_idx = threadIdx.x * VEC_SIZE; + int64_t flat_idx = flat_idx_fp8_base + head_dim_idx; + + // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous) + int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, + cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); + int64_t dst_offset = token_idx * head_dim + head_dim_idx; + + // 4 bytes read from non-contiguous cache, write to contiguous output + *reinterpret_cast(&out_fp8[dst_offset]) = *reinterpret_cast(&k_cache[src_offset]); + + // Only thread 0 reads the single 4 bytes scale value + if (threadIdx.x == 0) + { + int64_t src_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2, + cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); + int64_t dst_offset_scale = token_idx * scale_size; // scale_size = 4 + + // 4 bytes read for scale + *reinterpret_cast(&out_scale[dst_offset_scale]) + = *reinterpret_cast(&k_cache[src_offset_scale]); + } +} + +void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mapping_fp8, + int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens, + int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, + int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3, + cudaStream_t stream) +{ + if (num_tokens == 0) + { + return; + } + + // Assertions for DeepSeek-V3.2 configuration + constexpr int32_t QUANT_BLOCK_SIZE = 128; + TLLM_CHECK_WITH_INFO( + head_dim == QUANT_BLOCK_SIZE, "head_dim must equal 128 for DeepSeek-V3 indexer cache (got %d)", head_dim); + TLLM_CHECK_WITH_INFO( + scale_size == 4, "scale_size must equal 4 bytes (1 float32 scale per token, got %d)", scale_size); + + // For head_dim=128, we use 32 threads to handle 128 bytes per token and extra 4 bytes for scale + constexpr int32_t THREADS_PER_BLOCK = 32; + + dim3 block(THREADS_PER_BLOCK); + dim3 grid(num_tokens); + + indexerKCacheGatherUnifiedKernel<<>>(k_cache, slot_mapping_fp8, slot_mapping_scale, out_fp8, + out_scale, k_token_start, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2, + cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3); + + // Check for kernel launch errors + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 28f468960b90..0df032e3bd2f 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -89,6 +89,7 @@ add_library( fp4BlockScaleMoe.cpp noAuxTcOp.cpp fusedCatFp8Op.cpp + IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp ncclCommunicatorOp.cpp @@ -111,6 +112,7 @@ add_library( tinygemm2.cpp dsv3RopeOp.cpp fusedGemmAllreduceOp.cpp + convertReqIndexToGlobalOp.cpp trtllmGenQKVProcessOp.cpp) set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_libraries( diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp new file mode 100644 index 000000000000..27fbc33bd1be --- /dev/null +++ b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp @@ -0,0 +1,121 @@ +/* + * Copyright (c) 2022-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/common/opUtils.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +#include "tensorrt_llm/kernels/IndexerKCacheGather.h" + +namespace th = torch; +namespace tl = tensorrt_llm; +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +std::tuple indexer_k_cache_gather_op(th::Tensor const& k_cache, + th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t k_token_start, int64_t num_tokens) +{ + constexpr int32_t HEAD_DIM = 128; + constexpr int32_t SCALE_SIZE = 4; + + auto device = k_cache.device(); + + // Early return for empty gather + if (num_tokens == 0) + { + auto k_fp8 = th::empty({0, HEAD_DIM}, th::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(device)); + auto k_scale = th::empty({0, 1}, th::TensorOptions().dtype(torch::kFloat32).device(device)); + return std::make_tuple(std::move(k_fp8), std::move(k_scale)); + } + + // Validate all tensors are CUDA tensors + TORCH_CHECK(k_cache.is_cuda() && slot_mapping_fp8.is_cuda() && slot_mapping_scale.is_cuda(), + "All tensors must be CUDA tensors"); + + // Validate tensor dimensions + TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be a 1D Tensor"); + TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be a 1D Tensor"); + + // Enforce k_cache is 4D tensor + TORCH_CHECK(k_cache.dim() == 4, + "k_cache must be a 4D Tensor [num_blocks, block_size, 1, per_token_size], got %d dimensions", + static_cast(k_cache.dim())); + + // Validate tensor dtypes + TORCH_CHECK(slot_mapping_fp8.scalar_type() == torch::kInt64, "slot_mapping_fp8 must be int64"); + TORCH_CHECK(slot_mapping_scale.scalar_type() == torch::kInt64, "slot_mapping_scale must be int64"); + + // Validate tensors are contiguous (except k_cache which may be non-contiguous) + // k_cache can be non-contiguous - we handle this via strides + TORCH_CHECK(slot_mapping_fp8.is_contiguous(), "slot_mapping_fp8 must be contiguous"); + TORCH_CHECK(slot_mapping_scale.is_contiguous(), "slot_mapping_scale must be contiguous"); + + // Validate slot_mapping has enough elements + TORCH_CHECK(slot_mapping_fp8.size(0) >= k_token_start + num_tokens, + "slot_mapping_fp8 too short for k_token_start + num_tokens"); + TORCH_CHECK(slot_mapping_scale.size(0) >= k_token_start + num_tokens, + "slot_mapping_scale too short for k_token_start + num_tokens"); + + int32_t cache_dim_0 = static_cast(k_cache.size(0)); // num_blocks + int32_t cache_dim_1 = static_cast(k_cache.size(1)); // block_size + int32_t cache_dim_2 = static_cast(k_cache.size(2)); // num_kv_heads + int32_t cache_dim_3 = static_cast(k_cache.size(3)); // per_token_size + + // Validation for indexer k cache pool for DeepSeek-V3.2 constraints + TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2); + + int64_t cache_stride_0 = static_cast(k_cache.stride(0)); + int64_t cache_stride_1 = static_cast(k_cache.stride(1)); + int64_t cache_stride_2 = static_cast(k_cache.stride(2)); + int64_t cache_stride_3 = static_cast(k_cache.stride(3)); + + // Allocate output buffers + auto num_tokens_i32 = static_cast(num_tokens); + auto out_fp8 = th::empty({num_tokens, HEAD_DIM}, th::TensorOptions().dtype(torch::kUInt8).device(device)); + auto out_scale = th::empty({num_tokens, SCALE_SIZE}, th::TensorOptions().dtype(torch::kUInt8).device(device)); + + auto stream = at::cuda::getCurrentCUDAStream(k_cache.get_device()); + + tk::invokeIndexerKCacheGather(k_cache.data_ptr(), slot_mapping_fp8.data_ptr(), + slot_mapping_scale.data_ptr(), out_fp8.data_ptr(), out_scale.data_ptr(), + static_cast(k_token_start), num_tokens_i32, HEAD_DIM, SCALE_SIZE, cache_dim_0, cache_dim_1, + cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3, stream); + + // View-cast to final dtypes (no copy) + auto k_fp8 = out_fp8.view(torch::kFloat8_e4m3fn); + auto k_scale = out_scale.view(torch::kFloat32).view({num_tokens, 1}); + + return std::make_tuple(std::move(k_fp8), std::move(k_scale)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "indexer_k_cache_gather_op(Tensor k_cache, Tensor slot_mapping_fp8, " + "Tensor slot_mapping_scale, int k_token_start, int num_tokens) -> (Tensor, Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("indexer_k_cache_gather_op", &tensorrt_llm::torch_ext::indexer_k_cache_gather_op); +} diff --git a/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp b/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp new file mode 100644 index 000000000000..45b862289de5 --- /dev/null +++ b/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp @@ -0,0 +1,84 @@ +/* + * Copyright (c) 2022-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/common/opUtils.h" +#include "tensorrt_llm/kernels/convertReqIndexToGlobal.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +namespace th = torch; +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +th::Tensor convertReqIndexToGlobal(th::Tensor const& reqId, th::Tensor const& blockTable, + th::Tensor const& tokenIndices, int64_t blockSize, int64_t numTopkTokens, int64_t strideFactor, int64_t layerId) +{ + TORCH_CHECK(reqId.is_cuda() && blockTable.is_cuda() && tokenIndices.is_cuda(), "All tensors must be CUDA tensors"); + TORCH_CHECK(reqId.scalar_type() == th::kInt32, "req_id must be int32"); + TORCH_CHECK(blockTable.scalar_type() == th::kInt32, "block_table must be int32"); + TORCH_CHECK(tokenIndices.scalar_type() == th::kInt32, "token_indices must be int32"); + + TORCH_CHECK(reqId.dim() == 1, "req_id must be 1D"); + TORCH_CHECK(blockTable.dim() == 2, "block_table must be 2D"); + TORCH_CHECK(tokenIndices.dim() == 2, "token_indices must be 2D"); + + // Ensure contiguous + auto reqIdC = reqId.contiguous(); + auto blockTableC = blockTable.contiguous(); + auto tokenIndicesC = tokenIndices.contiguous(); + + int32_t const numTokens = static_cast(reqIdC.size(0)); + int32_t const maxNumBlocksPerReq = static_cast(blockTableC.size(1)); + + // Allocate output + auto out = th::empty_like(tokenIndicesC); + + // Extract strides + int64_t const btStride0 = blockTableC.stride(0); + int64_t const btStride1 = blockTableC.stride(1); + int64_t const tiStride0 = tokenIndicesC.stride(0); + int64_t const tiStride1 = tokenIndicesC.stride(1); + int64_t const outStride0 = out.stride(0); + int64_t const outStride1 = out.stride(1); + + auto stream = at::cuda::getCurrentCUDAStream(reqIdC.get_device()).stream(); + + tk::invokeConvertReqIndexToGlobal(reqIdC.data_ptr(), blockTableC.data_ptr(), + tokenIndicesC.data_ptr(), out.data_ptr(), numTokens, static_cast(numTopkTokens), + maxNumBlocksPerReq, static_cast(blockSize), static_cast(strideFactor), + static_cast(layerId), btStride0, btStride1, tiStride0, tiStride1, outStride0, outStride1, stream); + + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "convert_req_index_to_global(Tensor req_id, Tensor block_table, Tensor token_indices, int block_size, int " + "num_topk_tokens, int stride_factor, int layer_id) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("convert_req_index_to_global", &tensorrt_llm::torch_ext::convertReqIndexToGlobal); +} diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 9cc38da0f1ec..b3a5fbeb3136 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -32,9 +32,6 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig -from .kernel import (triton_convert_req_index_to_global_index, - triton_gather_k_cache) - ModelConfig = tensorrt_llm.bindings.ModelConfig if TYPE_CHECKING: @@ -119,59 +116,33 @@ def transform_local_topk_and_prepare_pool_view( layer_idx: int, is_generation: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Convert local topk indices to global pool indices and prepare KV pool. - Auto-detects stride and handles both contiguous/strided layouts. + """Convert local topk indices to global pool indices and prepare KV pool. - Args: - topk_indices: [num_tokens, NUM_TOPK] - attn_metadata: Metadata with block_table and request mappings - kv_cache_manager: KV cache manager - layer_idx: Layer index - is_generation: Generation vs context phase - - Returns: - (global_indices, kv_pool): - - global_indices: [num_tokens, NUM_TOPK] - - kv_pool: [total_tokens, 1, head_dim] + Uses cached values from attn_metadata._ensure_pool_view_cached() + to avoid redundant Python/CUDA overhead across layers. """ assert topk_indices.dtype == torch.int32 - # Get all layer KV cache pool: [num_blocks, num_layers, kv_factor, blockSize] - kv_cache_manager = attn_metadata.kv_cache_manager - all_layer_kv_pool = kv_cache_manager.get_unique_primary_pool( - ) # [num_blocks, num_layers, kv_factor, blockSize] - num_blocks, num_layers, _, _ = all_layer_kv_pool.shape - tokens_per_block = kv_cache_manager.tokens_per_block - head_dim = kv_cache_manager.head_dim - assert all_layer_kv_pool.is_contiguous( - ), "all_layer_kv_pool should be contiguous" - all_layer_kv_pool = all_layer_kv_pool.squeeze(2).view(-1, 1, head_dim) - stride_factor = num_layers * tokens_per_block - - # Get block_table and request indices for this phase + attn_metadata._ensure_pool_view_cached() + if is_generation: - block_table = attn_metadata.block_table[ - attn_metadata.num_contexts:attn_metadata.num_seqs] - req_idx = attn_metadata.req_idx_per_token[ - attn_metadata.num_ctx_tokens:attn_metadata.num_tokens] - req_idx = req_idx - attn_metadata.num_contexts + block_table = attn_metadata._cached_block_table_gen + req_idx = attn_metadata._cached_req_idx_gen else: - block_table = attn_metadata.block_table[:attn_metadata.num_contexts] - req_idx = attn_metadata.req_idx_per_token[:attn_metadata.num_ctx_tokens] + block_table = attn_metadata._cached_block_table_ctx + req_idx = attn_metadata._cached_req_idx_ctx - # Convert to global indices - global_indices = triton_convert_req_index_to_global_index( + global_indices = torch.ops.trtllm.convert_req_index_to_global( req_idx, block_table, topk_indices, - BLOCK_SIZE=tokens_per_block, - NUM_TOPK_TOKENS=topk_indices.shape[1], - stride_factor=stride_factor, - layer_id=layer_idx, + attn_metadata._cached_tokens_per_block, + topk_indices.shape[1], + attn_metadata._cached_stride_factor, + layer_idx, ) - return global_indices, all_layer_kv_pool + return global_indices, attn_metadata._cached_pool_view def split_prefill_chunks( @@ -346,6 +317,20 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): def __init__(self, *args, **kwargs): self.num_sms = tensorrt_llm.deep_gemm.get_num_sms() + # Cached step-invariant values for transform_local_topk_and_prepare_pool_view. + # These are recomputed once per step in _ensure_pool_view_cached() and + # reused across all layers to avoid redundant Python/CUDA overhead. + # Initialized here as plain instance attributes (not class-level + # annotations) to stay invisible to dataclass/torch.compile introspection. + self._pool_cache_valid = False + self._cached_kv_mgr_id = 0 + self._cached_pool_view = None + self._cached_stride_factor = 0 + self._cached_tokens_per_block = 0 + self._cached_block_table_ctx = None + self._cached_block_table_gen = None + self._cached_req_idx_ctx = None + self._cached_req_idx_gen = None super().__init__(*args, **kwargs) if self.sparse_attention_config.indexer_max_chunk_size is not None: self.indexer_max_chunk_size = self.sparse_attention_config.indexer_max_chunk_size @@ -582,40 +567,83 @@ def update_spec_dec_param( capture_graph = self.is_cuda_graph self.create_expanded_buffers(capture_graph=capture_graph) + def _invalidate_pool_view_cache(self): + """Invalidate the cached pool view and related step-invariant values. + + Must be called at the start of each forward step (in prepare()) so that + _ensure_pool_view_cached() recomputes them for the new batch. + """ + self._pool_cache_valid = False + + def _ensure_pool_view_cached(self): + """Compute and cache values used by + transform_local_topk_and_prepare_pool_view(). + + These values (pool view, stride factor, block table slices, request + index slices) are constant across all layers sharing the same KV pool + and batch dimensions within a forward pass. Caching them avoids + redundant Python/CUDA overhead per layer. + + Safety: _invalidate_pool_view_cache() is called unconditionally at the + start of every step (prepare() and on_update_kv_lens()), so the boolean + flag is always cleared before the first per-layer call within a step. + """ + if self._pool_cache_valid and self._cached_kv_mgr_id == id( + self.kv_cache_manager): + return + + pool = self.kv_cache_manager.get_unique_primary_pool() + kv_cache_manager = self.kv_cache_manager + num_blocks, num_layers, _, _ = pool.shape + self._cached_tokens_per_block = kv_cache_manager.tokens_per_block + head_dim = kv_cache_manager.head_dim + self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim) + self._cached_stride_factor = (num_layers * + self._cached_tokens_per_block) + self._cached_block_table_ctx = self.block_table[:self.num_contexts] + self._cached_block_table_gen = self.block_table[self.num_contexts:self. + num_seqs] + self._cached_req_idx_ctx = self.req_idx_per_token[:self.num_ctx_tokens] + self._cached_req_idx_gen = ( + self.req_idx_per_token[self.num_ctx_tokens:self.num_tokens] - + self.num_contexts) + self._cached_kv_mgr_id = id(kv_cache_manager) + self._pool_cache_valid = True + + @maybe_compile(dynamic=True) + def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens): + device = kv_lens.device + past_kv_lens = kv_lens - seq_lens + # get position ids + seq_ends = torch.cumsum(seq_lens, dim=0) + seq_starts = seq_ends - seq_lens + per_seq_offsets = past_kv_lens - seq_starts # Shape: [batch_size] + global_indices = torch.arange(num_tokens, device=device) + batch_indices = torch.searchsorted(seq_ends, + global_indices, + side='right') + repeated_offsets = per_seq_offsets[batch_indices] + position_ids = global_indices + repeated_offsets + # get the dense topk indices with causal mask + range_row = torch.arange(self.sparse_mla_topk, device=device) + mask = range_row <= position_ids.unsqueeze(1) + return torch.where(mask, range_row, -1) + def prepare_dense_topk_indices(self, kv_lens, device=False): # device=False means use CPU - @maybe_compile(dynamic=True) - def _get_dense_topk_indices(seq_lens, kv_lens, num_tokens): - device = kv_lens.device - past_kv_lens = kv_lens - seq_lens - # get position ids - seq_ends = torch.cumsum(seq_lens, dim=0) - seq_starts = seq_ends - seq_lens - per_seq_offsets = past_kv_lens - seq_starts # Shape: [batch_size] - global_indices = torch.arange(num_tokens, device=device) - batch_indices = torch.searchsorted(seq_ends, - global_indices, - side='right') - repeated_offsets = per_seq_offsets[batch_indices] - position_ids = global_indices + repeated_offsets - # get the dense topk indices with causal mask - range_row = torch.arange(self.sparse_mla_topk, device=device) - mask = range_row <= position_ids.unsqueeze(1) - return torch.where(mask, range_row, -1) - if self.num_contexts > 0 and self.skip_indexer_for_ctx_reqs: ctx_range = slice(self.num_ctx_tokens) if device: self.topk_indices_buffer[ctx_range, :].copy_( - _get_dense_topk_indices( + self._get_dense_topk_indices( self.seq_lens_cuda[:self.num_contexts], kv_lens[:self.num_contexts], self.num_ctx_tokens), non_blocking=True) else: self.host_topk_indices_buffer[ - ctx_range, :] = _get_dense_topk_indices( + ctx_range, :] = self._get_dense_topk_indices( self.seq_lens[:self.num_contexts], kv_lens[:self.num_contexts], self.num_ctx_tokens) self.topk_indices_buffer[ctx_range, :].copy_( @@ -626,14 +654,14 @@ def _get_dense_topk_indices(seq_lens, kv_lens, num_tokens): gen_range = slice(self.num_ctx_tokens, self.num_tokens) if device: self.topk_indices_buffer[gen_range, :].copy_( - _get_dense_topk_indices( + self._get_dense_topk_indices( self.seq_lens_cuda[self.num_contexts:self.num_seqs], kv_lens[self.num_contexts:self.num_seqs], self.num_tokens - self.num_ctx_tokens), non_blocking=True) else: self.host_topk_indices_buffer[ - gen_range, :] = _get_dense_topk_indices( + gen_range, :] = self._get_dense_topk_indices( self.seq_lens[self.num_contexts:self.num_seqs], kv_lens[self.num_contexts:self.num_seqs], self.num_tokens - self.num_ctx_tokens) @@ -671,6 +699,7 @@ def _get_pool_block_indices(self) -> torch.Tensor: def prepare(self): super().prepare() + self._invalidate_pool_view_cache() # Get kv lengths assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" @@ -856,6 +885,13 @@ def on_update_kv_lens(self): # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer # slot_mapping_* buffers also depend on these effective cached lengths. If we do not # refresh slot mappings here, indexer K-cache updates can be written with stale offsets. + + # _preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass + # caches so they are recomputed (and captured) on every _forward_step". Invalidate the + # pool_view cache here so it is recomputed on the next + # transform_local_topk_and_prepare_pool_view() call. + self._invalidate_pool_view_cache() + if self.kv_cache_manager is not None and self.num_tokens > 0: seq_lens = self.seq_lens_cuda[:self.num_seqs] # Runtime cached lengths after overlap/spec-dec correction. @@ -1429,23 +1465,15 @@ def sparse_attn_indexer( tp_rank = metadata.mapping.tp_rank tp_size = metadata.mapping.tp_size - # Use the 2D pool data directly (contiguous) instead of the - # 4D view, because the 4D view may have strides that - # prevent flattening via .view(-1). - layer_offset = metadata.kv_cache_manager.layer_offsets[ - self.layer_idx] - gather_k_cache_pool = metadata.kv_cache_manager.indexer_k_cache_pool_per_layer[ - layer_offset] + k_cache_4d = metadata.kv_cache_manager.get_indexer_k_cache_buffers( + self.layer_idx) for chunk in metadata.indexer_prefill_chunks: - chunk_k_fp8, chunk_k_scale = triton_gather_k_cache( - gather_k_cache_pool, - metadata.slot_mapping_fp8_fullkv, - metadata.slot_mapping_scale_fullkv, - chunk.k_token_start, - chunk.k_token_end, - self.head_dim, - ) + num_k_tokens = chunk.k_token_end - chunk.k_token_start + chunk_k_fp8, chunk_k_scale = torch.ops.trtllm.indexer_k_cache_gather_op( + k_cache_4d, metadata.slot_mapping_fp8_fullkv, + metadata.slot_mapping_scale_fullkv, chunk.k_token_start, + num_k_tokens) chunk_num_token = chunk.token_end - chunk.token_start apply_q_split = q_split_eligible and chunk_num_token >= q_split_threshold diff --git a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py index 9daec9b0add0..6412610064fa 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py @@ -1799,12 +1799,12 @@ def _convert_req_index_to_global_index_kernel_with_stride_factor( block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] - # shapes (compile-time where possible) - max_num_blocks_per_req: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) - stride_factor: tl.constexpr, # for strided memory layout adjustment - layer_id: tl.constexpr, # for layer interleaving layout + # shapes + max_num_blocks_per_req, + BLOCK_SIZE, + BLOCK_N: tl.constexpr, # tile width along columns (used in tl.arange) + stride_factor, # for strided memory layout adjustment + layer_id, # for layer interleaving layout bt_stride0, bt_stride1, ti_stride0, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 21b7de23c373..ff2ca7b65017 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1105,3 +1105,17 @@ def _(seq_q_offsets, seq_kv_offsets, padding_offsets, tokens_info, rotary_embedding_scale, rotary_embedding_base, rotary_embedding_dim, rotary_scaling_type, rotary_embedding_max_positions): return True + + @torch.library.register_fake("trtllm::convert_req_index_to_global") + def _(req_id: torch.Tensor, block_table: torch.Tensor, + token_indices: torch.Tensor, block_size: int, num_topk_tokens: int, + stride_factor: int, layer_id: int) -> torch.Tensor: + return torch.empty_like(token_indices) + + @torch.library.register_fake("trtllm::indexer_k_cache_gather_op") + def _(k_cache: torch.Tensor, slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, k_token_start: int, + num_tokens: int) -> Tuple[torch.Tensor, torch.Tensor]: + k_fp8 = k_cache.new_empty([num_tokens, 128], dtype=torch.float8_e4m3fn) + k_scale = k_cache.new_empty([num_tokens, 1], dtype=torch.float32) + return k_fp8, k_scale diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 1ac648cadf52..cda368e972cc 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1132,6 +1132,13 @@ def update_draft_tokens(self, next_draft_tokens, new_draft_token, position_ids = inputs["position_ids"][gather_ids] + 1 return hidden_states, position_ids + @torch.compile(options={"max-autotune": True}) + def prepare_position_ids_and_last_tokens(self, position_ids, seq_lens_cuda): + position_ids = position_ids.squeeze(0) + last_tokens_idx = torch.cumsum(seq_lens_cuda, dim=0, + dtype=torch.long) - 1 + return position_ids, last_tokens_idx + def forward( self, input_ids, @@ -1167,16 +1174,8 @@ def forward( # Save the old attn_metadata and spec_metadata self._prepare_attn_metadata_for_spec_dec(attn_metadata) - # Prepare inputs for the 1st MTP layer - @torch.compile(options={"max-autotune": True}) - def prepare_position_ids_and_last_tokens(position_ids, attn_metadata): - position_ids = position_ids.squeeze(0) - last_tokens_idx = torch.cumsum( - attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 - return position_ids, last_tokens_idx - - position_ids, last_tokens_idx = prepare_position_ids_and_last_tokens( - position_ids, attn_metadata) + position_ids, last_tokens_idx = self.prepare_position_ids_and_last_tokens( + position_ids, attn_metadata.seq_lens_cuda) inputs = self.prepare_drafter_inputs(input_ids=input_ids, position_ids=position_ids, last_tokens_idx=last_tokens_idx, @@ -1301,12 +1300,9 @@ def prepare_position_ids_and_last_tokens(position_ids, attn_metadata): # as draft model only infer 1 token for the subsequent inference. attn_metadata.use_spec_decoding = False elif hasattr(attn_metadata, 'kv_lens_cuda'): + # update kv_lens_cuda + attn_metadata.kv_lens_cuda[:batch_size] += 1 - @torch.compile(options={"max-autotune": True}) - def update_kv_lens(kv_lens_cuda, batch_size): - kv_lens_cuda[:batch_size] += 1 - - update_kv_lens(attn_metadata.kv_lens_cuda, batch_size) # update metadata # some attention metadata needs to be updated when changing kv_lens attn_metadata.update_for_spec_dec() diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py new file mode 100644 index 000000000000..fd1a00b34c62 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for DSA C++ custom ops: + +- ``torch.ops.trtllm.indexer_k_cache_gather_op`` +- ``torch.ops.trtllm.convert_req_index_to_global`` +""" + +import pytest +import torch + +# Import tensorrt_llm to load C++ custom operators +import tensorrt_llm # noqa: F401 + +# --------------------------------------------------------------------------- +# Constants matching the C++ kernel (DeepSeek-V3.2 indexer config) +# --------------------------------------------------------------------------- +HEAD_DIM = 128 +SCALE_BYTES = 4 +BYTES_PER_TOKEN = HEAD_DIM + SCALE_BYTES + + +# =================================================================== +# Test 1: indexer_k_cache_gather_op +# =================================================================== + + +def _reference_indexer_k_cache_gather( + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + k_token_start: int, + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Python reference for indexer_k_cache_gather_op. + + The C++ op treats the 4D k_cache as a flat byte buffer (via strides) + and gathers HEAD_DIM bytes for FP8 data and SCALE_BYTES bytes for scales + at positions given by the slot mappings. + """ + device = k_cache.device + + if num_tokens == 0: + return ( + torch.empty(0, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device), + torch.empty(0, 1, dtype=torch.float32, device=device), + ) + + # Flatten the cache to a 1D byte view using a contiguous copy + k_cache_flat = k_cache.contiguous().reshape(-1) + + fp8_bases = slot_mapping_fp8[k_token_start : k_token_start + num_tokens] + byte_offsets_fp8 = torch.arange(HEAD_DIM, device=device, dtype=torch.int64) + gather_fp8 = fp8_bases.unsqueeze(1) + byte_offsets_fp8.unsqueeze(0) + out_fp8 = k_cache_flat[gather_fp8] + + scale_bases = slot_mapping_scale[k_token_start : k_token_start + num_tokens] + byte_offsets_scale = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64) + gather_scale = scale_bases.unsqueeze(1) + byte_offsets_scale.unsqueeze(0) + out_scale = k_cache_flat[gather_scale] + + k_fp8 = out_fp8.view(torch.float8_e4m3fn) + k_scale = out_scale.view(torch.float32).view(num_tokens, 1) + return k_fp8, k_scale + + +def _create_4d_cache_and_mappings( + total_kv_len: int, + num_blocks: int, + block_size: int, + per_token_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build a random 4D k_cache with valid, non-overlapping slot mappings. + + k_cache shape: [num_blocks, block_size, 1, per_token_size] + Each token occupies ``per_token_size`` bytes at a unique (block, slot) + position. Slot mappings point into the *contiguous flat* byte layout. + """ + k_cache = torch.randint( + 0, + 256, + (num_blocks, block_size, 1, per_token_size), + dtype=torch.uint8, + device=device, + ) + + total_slots = num_blocks * block_size + assert total_kv_len <= total_slots, ( + f"total_kv_len ({total_kv_len}) exceeds total slots ({total_slots})" + ) + + # Pick random unique (block, slot) positions + perm = torch.randperm(total_slots, device=device)[:total_kv_len] + # Flat byte offset of each token in the contiguous view + # Each token starts at perm[i] * per_token_size (within the [total_slots, 1, per_token_size] view) + # But the cache is [num_blocks, block_size, 1, per_token_size], so flat offset is: + # perm[i] * (1 * per_token_size) since dim2=1 + flat_starts = perm.to(torch.int64) * per_token_size + + slot_mapping_fp8 = flat_starts + slot_mapping_scale = flat_starts + HEAD_DIM + + return k_cache, slot_mapping_fp8, slot_mapping_scale + + +@pytest.mark.parametrize( + "total_kv_len,num_blocks,block_size,k_token_start,num_tokens", + [ + # Gather all tokens + (64, 8, 8, 0, 64), + # Sub-range from the middle + (128, 16, 8, 32, 64), + # Single token + (10, 2, 8, 3, 1), + # Larger test + (512, 64, 8, 100, 300), + # Non-power-of-2 token count + (100, 16, 8, 10, 37), + # Very small + (3, 1, 4, 0, 3), + # Larger block_size + (64, 4, 16, 0, 64), + ], +) +def test_indexer_k_cache_gather_contiguous( + total_kv_len, + num_blocks, + block_size, + k_token_start, + num_tokens, +): + """Test C++ gather op on a contiguous 4D cache against Python reference.""" + device = torch.device("cuda") + per_token_size = BYTES_PER_TOKEN # 132 + + k_cache, slot_fp8, slot_scale = _create_4d_cache_and_mappings( + total_kv_len, num_blocks, block_size, per_token_size, device + ) + + # C++ op + cpp_fp8, cpp_scale = torch.ops.trtllm.indexer_k_cache_gather_op( + k_cache, slot_fp8, slot_scale, k_token_start, num_tokens + ) + + # Reference + ref_fp8, ref_scale = _reference_indexer_k_cache_gather( + k_cache, slot_fp8, slot_scale, k_token_start, num_tokens + ) + + assert cpp_fp8.shape == ref_fp8.shape + assert cpp_scale.shape == ref_scale.shape + assert torch.equal(cpp_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)), "FP8 data mismatch" + assert torch.equal(cpp_scale.view(torch.uint8), ref_scale.view(torch.uint8)), ( + "Scale data mismatch" + ) + + +@pytest.mark.parametrize( + "total_kv_len,num_blocks,block_size,k_token_start,num_tokens", + [ + (64, 8, 8, 0, 64), + (128, 16, 8, 32, 64), + (512, 64, 8, 100, 300), + ], +) +def test_indexer_k_cache_gather_noncontiguous( + total_kv_len, + num_blocks, + block_size, + k_token_start, + num_tokens, +): + """Test C++ gather op on a non-contiguous 4D cache. + + The real KV cache pool is typically a large buffer viewed with + non-contiguous strides (layer interleaving). Simulate this by + allocating a larger buffer and taking a strided view. + """ + device = torch.device("cuda") + per_token_size = BYTES_PER_TOKEN + num_layers = 3 # simulate multi-layer interleaving + target_layer = 1 + + # Allocate the full interleaved pool: + # [num_blocks, block_size, num_layers, per_token_size] + full_pool = torch.randint( + 0, + 256, + (num_blocks, block_size, num_layers, per_token_size), + dtype=torch.uint8, + device=device, + ) + + # Select a single layer → shape [num_blocks, block_size, 1, per_token_size] + # This slice is non-contiguous because stride(1) spans across all layers. + k_cache_nc = full_pool[:, :, target_layer : target_layer + 1, :] + assert not k_cache_nc.is_contiguous(), "Expected non-contiguous cache" + + # Build slot mappings based on the *contiguous* copy (what the reference sees) + k_cache_contig = k_cache_nc.contiguous() + total_slots = num_blocks * block_size + perm = torch.randperm(total_slots, device=device)[:total_kv_len] + flat_starts = perm.to(torch.int64) * per_token_size + slot_fp8 = flat_starts + slot_scale = flat_starts + HEAD_DIM + + # C++ op (handles non-contiguous strides internally) + cpp_fp8, cpp_scale = torch.ops.trtllm.indexer_k_cache_gather_op( + k_cache_nc, slot_fp8, slot_scale, k_token_start, num_tokens + ) + + # Reference uses contiguous copy + ref_fp8, ref_scale = _reference_indexer_k_cache_gather( + k_cache_contig, slot_fp8, slot_scale, k_token_start, num_tokens + ) + + assert cpp_fp8.shape == ref_fp8.shape + assert cpp_scale.shape == ref_scale.shape + assert torch.equal(cpp_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)), ( + "FP8 data mismatch (non-contiguous)" + ) + assert torch.equal(cpp_scale.view(torch.uint8), ref_scale.view(torch.uint8)), ( + "Scale data mismatch (non-contiguous)" + ) + + +def test_indexer_k_cache_gather_empty(): + """Zero-length gather should return correctly shaped empty tensors.""" + device = torch.device("cuda") + k_cache = torch.randint(0, 256, (4, 8, 1, BYTES_PER_TOKEN), dtype=torch.uint8, device=device) + slot_fp8 = torch.zeros(10, dtype=torch.int64, device=device) + slot_scale = torch.zeros(10, dtype=torch.int64, device=device) + + k_fp8, k_scale = torch.ops.trtllm.indexer_k_cache_gather_op( + k_cache, slot_fp8, slot_scale, k_token_start=5, num_tokens=0 + ) + + assert k_fp8.shape == (0, HEAD_DIM) + assert k_scale.shape == (0, 1) + assert k_fp8.dtype == torch.float8_e4m3fn + assert k_scale.dtype == torch.float32 + + +# =================================================================== +# Test 2: convert_req_index_to_global +# =================================================================== + + +def _reference_convert_req_index_to_global( + req_id: torch.Tensor, + block_table: torch.Tensor, + token_indices: torch.Tensor, + block_size: int, + stride_factor: int, + layer_id: int, +) -> torch.Tensor: + """Python reference for convert_req_index_to_global. + + For each (token_id, indice_id): + tok = token_indices[token_id, indice_id] + if tok == -1: + out = -1 + else: + block_id = tok // block_size + inblock_off = tok % block_size + layer_id * block_size + req = req_id[token_id] + base = block_table[req, block_id] + if block_id >= max_blocks_per_req or base < 0: + out = -1 + else: + out = base * stride_factor + inblock_off + """ + num_tokens, num_topk = token_indices.shape + max_blocks_per_req = block_table.shape[1] + + out = torch.empty_like(token_indices) + # Move to CPU for the reference loop + req_id_cpu = req_id.cpu() + block_table_cpu = block_table.cpu() + token_indices_cpu = token_indices.cpu() + out_cpu = out.cpu() + + for i in range(num_tokens): + req = req_id_cpu[i].item() + for j in range(num_topk): + tok = token_indices_cpu[i, j].item() + if tok < 0: + out_cpu[i, j] = -1 + continue + block_id = tok // block_size + inblock_off = tok % block_size + layer_id * block_size + if block_id >= max_blocks_per_req: + out_cpu[i, j] = -1 + continue + base = block_table_cpu[req, block_id].item() + if base < 0: + out_cpu[i, j] = -1 + continue + out_cpu[i, j] = base * stride_factor + inblock_off + + return out_cpu.to(token_indices.device) + + +@pytest.mark.parametrize( + "num_tokens,num_requests,num_topk,block_size,stride_factor,layer_id", + [ + # Basic case: contiguous pool (stride_factor == block_size), layer 0 + (4, 2, 128, 64, 64, 0), + # Layer interleaving: stride_factor > block_size + (4, 2, 128, 64, 192, 1), + # Single token + (1, 1, 128, 64, 64, 0), + # Larger batch + (32, 8, 256, 64, 64, 0), + # Small block_size + (8, 4, 128, 16, 16, 0), + # layer_id > 0 with contiguous pool + (4, 2, 128, 64, 64, 2), + # Large topk + (4, 2, 2048, 64, 192, 1), + ], +) +def test_convert_req_index_to_global( + num_tokens, + num_requests, + num_topk, + block_size, + stride_factor, + layer_id, +): + """Test C++ op against Python reference.""" + device = torch.device("cuda") + torch.manual_seed(42) + + max_kv_len = 4096 + max_blocks_per_req = (max_kv_len + block_size - 1) // block_size + + # req_id: which request each token belongs to + req_id = torch.randint(0, num_requests, (num_tokens,), dtype=torch.int32, device=device) + + # block_table: maps (request, block_id) → physical block index + # Use positive values; some entries can be -1 (padding) + block_table = torch.randint( + 0, 1000, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device + ) + # Add some padding (-1) in later blocks + block_table[:, max_blocks_per_req // 2 :] = -1 + + # token_indices: request-local token positions, some -1 (invalid) + max_valid_token = (max_blocks_per_req // 2) * block_size - 1 + token_indices = torch.randint( + 0, max(max_valid_token, 1), (num_tokens, num_topk), dtype=torch.int32, device=device + ) + # Sprinkle -1s for invalid tokens (~20%) + invalid_mask = torch.rand(num_tokens, num_topk, device=device) < 0.2 + token_indices[invalid_mask] = -1 + + # C++ op + cpp_out = torch.ops.trtllm.convert_req_index_to_global( + req_id, block_table, token_indices, block_size, num_topk, stride_factor, layer_id + ) + + # Reference + ref_out = _reference_convert_req_index_to_global( + req_id, block_table, token_indices, block_size, stride_factor, layer_id + ) + + assert cpp_out.shape == ref_out.shape + assert torch.equal(cpp_out, ref_out), ( + f"Mismatch found.\n" + f" First diff at: {(cpp_out != ref_out).nonzero(as_tuple=False)[0].tolist()}\n" + f" cpp: {cpp_out[(cpp_out != ref_out)][:5]}\n" + f" ref: {ref_out[(cpp_out != ref_out)][:5]}" + ) + + +def test_convert_req_index_to_global_all_invalid(): + """All token_indices are -1 → entire output should be -1.""" + device = torch.device("cuda") + num_tokens, num_topk, block_size = 4, 128, 64 + + req_id = torch.zeros(num_tokens, dtype=torch.int32, device=device) + block_table = torch.ones(1, 16, dtype=torch.int32, device=device) * 100 + token_indices = torch.full((num_tokens, num_topk), -1, dtype=torch.int32, device=device) + + out = torch.ops.trtllm.convert_req_index_to_global( + req_id, block_table, token_indices, block_size, num_topk, block_size, 0 + ) + + assert (out == -1).all(), "All outputs should be -1 for all-invalid input" + + +def test_convert_req_index_to_global_block_table_padding(): + """block_table entries of -1 should produce -1 in output.""" + device = torch.device("cuda") + num_tokens, num_topk, block_size = 2, 64, 16 + + req_id = torch.zeros(num_tokens, dtype=torch.int32, device=device) + # All block_table entries are -1 (no valid blocks) + block_table = torch.full((1, 8), -1, dtype=torch.int32, device=device) + # Valid-looking token indices that would map to these blocks + token_indices = torch.randint(0, 128, (num_tokens, num_topk), dtype=torch.int32, device=device) + + out = torch.ops.trtllm.convert_req_index_to_global( + req_id, block_table, token_indices, block_size, num_topk, block_size, 0 + ) + + assert (out == -1).all(), "All outputs should be -1 when block_table is all padding" diff --git a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py b/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py index b21b481d62a0..ed0ff4a10357 100644 --- a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py +++ b/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py @@ -149,31 +149,3 @@ def test_triton_gather_k_cache_empty(): assert k_scale.shape == (0, 1) assert k_fp8.dtype == torch.float8_e4m3fn assert k_scale.dtype == torch.float32 - - -if __name__ == "__main__": - print("\n" + "=" * 80) - print("Testing Triton Gather K Cache Kernel") - print("=" * 80) - - print("\n--- Empty tokens ---") - test_triton_gather_k_cache_empty() - - print("\n--- Gather all tokens ---") - test_triton_gather_k_cache(64, 0, 64, 128, 16) - - print("\n--- Sub-range from middle ---") - test_triton_gather_k_cache(128, 32, 96, 128, 32) - - print("\n--- Single token ---") - test_triton_gather_k_cache(10, 3, 4, 128, 4) - - print("\n--- Larger test ---") - test_triton_gather_k_cache(512, 100, 400, 128, 64) - - print("\n--- Non-power-of-2 tokens ---") - test_triton_gather_k_cache(100, 10, 47, 128, 32) - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80)