diff --git a/cpp/tensorrt_llm/kernels/attentionMetadataKernels.cu b/cpp/tensorrt_llm/kernels/attentionMetadataKernels.cu new file mode 100644 index 000000000000..4807800d971c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/attentionMetadataKernels.cu @@ -0,0 +1,207 @@ +/* + * 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/attentionMetadataKernels.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int32_t kThreadsPerBlock = 256; +constexpr int32_t kVecThreadsPerBlock = 128; +// Mirrors kv_cache_manager_v2::kBadPageIndex; kept local so this file does not +// pull the batch_manager headers into device code. +constexpr int32_t kBadPageIndex = -1; + +// Phase 1: padded exclusive scan of seq_lens into cu_seq_lens. +// batchSize is the scheduler batch (a few hundred), so a single-block +// shared-memory scan avoids a separate cumsum launch. +template +__global__ void computeCuSeqLensKernel( + int32_t const* __restrict__ seqLens, int32_t* __restrict__ cuSeqLens, int32_t batchSize) +{ + __shared__ int32_t buffers[2][kMaxBatch]; + + int32_t const stride = static_cast(blockDim.x); + int32_t const rounded = ((batchSize + stride - 1) / stride) * stride; + + for (int32_t i = static_cast(threadIdx.x); i < rounded; i += stride) + { + if (i < batchSize) + { + buffers[0][i] = seqLens[i]; + } + } + __syncthreads(); + + int32_t src = 0; + for (int32_t offset = 1; offset < batchSize; offset <<= 1) + { + int32_t const dst = src ^ 1; + for (int32_t i = static_cast(threadIdx.x); i < rounded; i += stride) + { + if (i < batchSize) + { + buffers[dst][i] = buffers[src][i] + (i >= offset ? buffers[src][i - offset] : 0); + } + } + __syncthreads(); + src = dst; + } + + if (threadIdx.x == 0) + { + cuSeqLens[0] = 0; + } + for (int32_t i = static_cast(threadIdx.x); i < batchSize; i += stride) + { + cuSeqLens[i + 1] = buffers[src][i]; + } +} + +// Phase 2: per-token request index and absolute position. +// Replaces the CPU repeat_interleave + pinned H2D memcpy on the prepare() path +// and the arange + searchsorted + two gathers on the update path. +__global__ void computeTokenPositionsKernel(int32_t const* __restrict__ cuSeqLens, + int32_t const* __restrict__ cachedTokens, int32_t* __restrict__ reqIdxPerToken, + int32_t* __restrict__ tokenPositions, int32_t batchSize, int32_t numTokens) +{ + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < numTokens; t += gridDim.x * blockDim.x) + { + // searchsorted(cu_seq_lens[1:], t, right=True): largest j with cu[j] <= t. + int32_t lo = 0; + int32_t hi = batchSize; + while (lo < hi) + { + int32_t const mid = lo + ((hi - lo) >> 1); + if (cuSeqLens[mid + 1] <= t) + { + lo = mid + 1; + } + else + { + hi = mid; + } + } + int32_t const reqIdx = min(lo, batchSize - 1); + reqIdxPerToken[t] = reqIdx; + if (tokenPositions != nullptr) + { + tokenPositions[t] = cachedTokens[reqIdx] + (t - cuSeqLens[reqIdx]); + } + } +} + +// --------------------------------------------------------------------------- +// One shared-page block table: gather block_offsets[poolId, copyIdx, 0, :] and +// map it with where(base == kBadPageIndex, kBadPageIndex, base * scale). +// +// Keeping it on the GPU removes a host gather of a few hundred KB plus the +// subsequent host->device staging copy from the decode critical path; callers +// that build several tables per iteration pay that cost once per table. +// --------------------------------------------------------------------------- +__global__ void computeSharedBlockTableKernel(int32_t const* __restrict__ blockOffsets, + int32_t const* __restrict__ copyIdx, int32_t* __restrict__ output, int32_t poolId, int32_t scale, + int32_t copyIdxCapacity, int32_t numTables, int32_t maxBlocksPerSeq) +{ + int32_t const tableId = static_cast(blockIdx.y); + if (tableId >= numTables) + { + return; + } + + int64_t const outputOffset = static_cast(tableId) * maxBlocksPerSeq; + int32_t const mappedTableId = copyIdx[tableId]; + bool const validTable = mappedTableId >= 0 && mappedTableId < copyIdxCapacity; + + // blockOffsets layout is [numPools, copyIdxCapacity, 2, maxBlocksPerSeq]; + // the CPU path reads index 0 of the K/V dimension. + int64_t const baseOffset = ((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq; + + for (int32_t blockId + = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); + blockId < maxBlocksPerSeq; blockId += static_cast(gridDim.x) * static_cast(blockDim.x)) + { + int32_t value = kBadPageIndex; + if (validTable) + { + int32_t const base = blockOffsets[baseOffset + blockId]; + value = base == kBadPageIndex ? kBadPageIndex : base * scale; + } + output[outputOffset + blockId] = value; + } +} +} // namespace + +void invokeComputeTokenPositions(int32_t const* seqLens, int32_t const* cachedTokens, int32_t* cuSeqLens, + int32_t* reqIdxPerToken, int32_t* tokenPositions, int32_t batchSize, int32_t numTokens, bool computeCuSeqLens, + cudaStream_t stream) +{ + if (batchSize <= 0) + { + return; + } + if (computeCuSeqLens) + { + dim3 const grid(1); + dim3 const block(static_cast(kThreadsPerBlock)); + if (batchSize <= 512) + { + computeCuSeqLensKernel<512><<>>(seqLens, cuSeqLens, batchSize); + } + else if (batchSize <= 2048) + { + computeCuSeqLensKernel<2048><<>>(seqLens, cuSeqLens, batchSize); + } + else + { + computeCuSeqLensKernel + <<>>(seqLens, cuSeqLens, batchSize); + } + } + if (numTokens > 0) + { + int32_t const blocks = std::min((numTokens + kThreadsPerBlock - 1) / kThreadsPerBlock, 2048); + computeTokenPositionsKernel<<>>( + cuSeqLens, cachedTokens, reqIdxPerToken, tokenPositions, batchSize, numTokens); + } +} + +void invokeComputeSharedBlockTable(int32_t const* blockOffsets, int32_t const* copyIdx, int32_t* output, int32_t poolId, + int32_t scale, int32_t copyIdxCapacity, int32_t numTables, int32_t maxBlocksPerSeq, cudaStream_t stream) +{ + if (numTables <= 0 || maxBlocksPerSeq <= 0) + { + return; + } + + int32_t const threadsPerBlock = kVecThreadsPerBlock; + int32_t const blocksPerRow = std::min((maxBlocksPerSeq + threadsPerBlock - 1) / threadsPerBlock, 64); + dim3 const block(static_cast(threadsPerBlock)); + dim3 const grid(static_cast(blocksPerRow), static_cast(numTables)); + computeSharedBlockTableKernel<<>>( + blockOffsets, copyIdx, output, poolId, scale, copyIdxCapacity, numTables, maxBlocksPerSeq); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/attentionMetadataKernels.h b/cpp/tensorrt_llm/kernels/attentionMetadataKernels.h new file mode 100644 index 000000000000..a3c9d579bd33 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/attentionMetadataKernels.h @@ -0,0 +1,67 @@ +/* + * 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 +{ + +// Backend-agnostic helpers for the per-iteration attention-metadata rebuild. +// Nothing here depends on a particular sparse-attention algorithm: they are the +// device-side forms of tensor patterns that several backends currently build +// with element-wise ATen chains on the host critical path. + +// Upper bound on the scheduler batch handled by the single-block scan below. +constexpr int32_t kMaxTokenPositionScanBatch = 4096; + +// Computes cu_seq_lens (optional), req_idx_per_token, and token_positions. +// +// Device-side form of: +// cu_seq_lens = pad(cumsum(seq_lens), (1, 0)) +// req_idx_per_token = repeat_interleave(arange(batch_size), seq_lens) +// token_positions = cached_tokens[req_idx] + (t - cu_seq_lens[req_idx]) +// where the last line is the searchsorted(cu_seq_lens[1:], t, right=True) gather. +// +// `tokenPositions` may be null when only the request index is needed; +// `cachedTokens` is then unused. When `computeCuSeqLens` is false, `cuSeqLens` +// is read as an already-populated input and `batchSize` is not bounded by +// kMaxTokenPositionScanBatch. +void invokeComputeTokenPositions(int32_t const* seqLens, int32_t const* cachedTokens, int32_t* cuSeqLens, + int32_t* reqIdxPerToken, int32_t* tokenPositions, int32_t batchSize, int32_t numTokens, bool computeCuSeqLens, + cudaStream_t stream); + +// Builds one shared-page block table from the host block-offset buffer. +// +// Device-side form of: +// base = block_offsets[pool_id, copy_idx, 0, :] +// out = where(base == kBadPageIndex, kBadPageIndex, base * scale) +// +// `blockOffsets` is laid out [numPools, copyIdxCapacity, 2, maxBlocksPerSeq]. +// Rows past `numTables` are left untouched, so padded CUDA-graph slots keep +// whatever the caller put there. +void invokeComputeSharedBlockTable(int32_t const* blockOffsets, int32_t const* copyIdx, int32_t* output, int32_t poolId, + int32_t scale, int32_t copyIdxCapacity, int32_t numTables, int32_t maxBlocksPerSeq, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu index 7b6188c74eca..51b20bb1d509 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu @@ -16,6 +16,7 @@ #include "tensorrt_llm/kernels/deepseekV4BlockTable.h" +#include #include TRTLLM_NAMESPACE_BEGIN diff --git a/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cu b/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cu new file mode 100644 index 000000000000..bd02907f7bbf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cu @@ -0,0 +1,243 @@ +/* + * 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/deepseekV4CompressedMeta.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int32_t kThreadsPerBlock = 256; + +// Binary search for the request owning a compact output index: the largest j +// with cu[j] <= value, i.e. searchsorted(cu[1:], value, right=True). +__device__ __forceinline__ int32_t upperBoundRequest(int32_t const* __restrict__ cu, int32_t numEntries, int32_t value) +{ + int32_t lo = 0; + int32_t hi = numEntries; + while (lo < hi) + { + int32_t const mid = lo + ((hi - lo) >> 1); + if (cu[mid + 1] <= value) + { + lo = mid + 1; + } + else + { + hi = mid; + } + } + return lo; +} + +// ── Per-ratio compressed/past/new KV lens + padded exclusive scan of new_comp. +// +// One block per ratio. `batchSize` is the scheduler batch (a few hundred), so a +// single-block shared-memory scan keeps the step at one launch total instead of +// four ATen calls per ratio. +template +__global__ void computePerRatioKvLensKernel(int32_t const* __restrict__ kvLens, + int32_t const* __restrict__ cachedTokens, PerRatioKvLensParams params, int32_t batchSize) +{ + // Ping-pong buffers so a scan pass never reads a slot another thread is + // concurrently writing. + __shared__ int32_t buffers[2][kMaxBatch]; + + int32_t const ratioId = static_cast(blockIdx.x); + int32_t const ratio = params.ratios[ratioId]; + int32_t* __restrict__ compressedOut = params.compressedKvLens[ratioId]; + int32_t* __restrict__ pastOut = params.pastKvLens[ratioId]; + int32_t* __restrict__ newCompOut = params.newCompKvLens[ratioId]; + int32_t* __restrict__ cuOut = params.cuNewCompKv[ratioId]; + + // Thread-uniform trip count: batchSize and blockDim.x are both uniform, so + // every thread reaches every __syncthreads() below. + int32_t const stride = static_cast(blockDim.x); + int32_t const rounded = ((batchSize + stride - 1) / stride) * stride; + + for (int32_t i = static_cast(threadIdx.x); i < rounded; i += stride) + { + if (i < batchSize) + { + int32_t const compressedKv = kvLens[i] / ratio; + int32_t const pastKv = cachedTokens[i] / ratio; + compressedOut[i] = compressedKv; + pastOut[i] = pastKv; + int32_t const newComp = compressedKv - pastKv; + newCompOut[i] = newComp; + buffers[0][i] = newComp; + } + } + __syncthreads(); + + // Inclusive Hillis-Steele scan over new_comp. + int32_t src = 0; + for (int32_t offset = 1; offset < batchSize; offset <<= 1) + { + int32_t const dst = src ^ 1; + for (int32_t i = static_cast(threadIdx.x); i < rounded; i += stride) + { + if (i < batchSize) + { + buffers[dst][i] = buffers[src][i] + (i >= offset ? buffers[src][i - offset] : 0); + } + } + __syncthreads(); + src = dst; + } + + // Shifted by one on write-out: the padded exclusive scan that the python + // code built with pad(cumsum(x), (1, 0)). + if (threadIdx.x == 0) + { + cuOut[0] = 0; + } + for (int32_t i = static_cast(threadIdx.x); i < batchSize; i += stride) + { + cuOut[i + 1] = buffers[src][i]; + } +} + +// ── Compressed mask: for each compact token, real token or decode padding? +__global__ void computeCompressedMaskKernel(CompressedMaskParams params, int32_t batchSize) +{ + int32_t const ratioId = static_cast(blockIdx.y); + int32_t const total = params.totalTokens[ratioId]; + int32_t const* __restrict__ newComp = params.newCompKvLens[ratioId]; + int32_t const* __restrict__ cu = params.cuNewCompKv[ratioId]; + bool* __restrict__ out = params.mask[ratioId]; + + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < total; t += gridDim.x * blockDim.x) + { + int32_t seqIdx = upperBoundRequest(cu, batchSize, t); + seqIdx = min(seqIdx, batchSize - 1); + out[t] = (t - cu[seqIdx]) < newComp[seqIdx]; + } +} + +// ── Context compressed position ids. +__global__ void computeCtxCompressedPositionIdsKernel(CompressedPositionIdsParams params, int32_t numContexts) +{ + int32_t const ratioId = static_cast(blockIdx.y); + int32_t const total = params.counts[ratioId]; + int32_t const ratio = params.ratios[ratioId]; + int32_t const* __restrict__ pastKv = params.pastKvLens[ratioId]; + int32_t const* __restrict__ cu = params.cuNewCompKv[ratioId]; + int32_t* __restrict__ out = params.positionIds[ratioId]; + + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < total; t += gridDim.x * blockDim.x) + { + int32_t const reqIdx = upperBoundRequest(cu, numContexts, t); + out[t] = (pastKv[reqIdx] + (t - cu[reqIdx])) * ratio; + } +} + +// ── Generation compressed position ids, in compact compressor output order. +__global__ void computeGenCompressedPositionIdsKernel( + CompressedPositionIdsParams params, int32_t numContexts, int32_t batchSize) +{ + int32_t const ratioId = static_cast(blockIdx.y); + int32_t const genComp = params.counts[ratioId]; + int32_t const ratio = params.ratios[ratioId]; + int32_t const outputOffset = params.offsets[ratioId]; + int32_t const* __restrict__ pastKv = params.pastKvLens[ratioId]; + int32_t const* __restrict__ cu = params.cuNewCompKv[ratioId]; + int32_t* __restrict__ out = params.positionIds[ratioId]; + + for (int32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < genComp; i += gridDim.x * blockDim.x) + { + int32_t const outputIdx = i + outputOffset; + int32_t reqIdx = upperBoundRequest(cu, batchSize, outputIdx); + reqIdx = min(max(reqIdx, numContexts), batchSize - 1); + out[outputIdx] = (pastKv[reqIdx] + (outputIdx - cu[reqIdx])) * ratio; + } +} + +int32_t gridX(int32_t work, int32_t cap) +{ + return std::min((work + kThreadsPerBlock - 1) / kThreadsPerBlock, cap); +} +} // namespace + +void invokeDeepseekV4ComputePerRatioKvLens(int32_t const* kvLens, int32_t const* cachedTokens, + PerRatioKvLensParams const& params, int32_t numRatios, int32_t batchSize, cudaStream_t stream) +{ + if (numRatios <= 0 || batchSize <= 0) + { + return; + } + dim3 const grid(static_cast(numRatios)); + dim3 const block(static_cast(kThreadsPerBlock)); + // Templated on a shared-memory batch bound; pick the smallest that fits. + if (batchSize <= 512) + { + computePerRatioKvLensKernel<512><<>>(kvLens, cachedTokens, params, batchSize); + } + else if (batchSize <= 2048) + { + computePerRatioKvLensKernel<2048><<>>(kvLens, cachedTokens, params, batchSize); + } + else + { + computePerRatioKvLensKernel<<>>(kvLens, cachedTokens, params, batchSize); + } +} + +void invokeDeepseekV4ComputeCompressedMask(CompressedMaskParams const& params, int32_t maxTotalTokens, + int32_t numRatios, int32_t batchSize, cudaStream_t stream) +{ + if (numRatios <= 0 || batchSize <= 0 || maxTotalTokens <= 0) + { + return; + } + dim3 const grid(static_cast(gridX(maxTotalTokens, 1024)), static_cast(numRatios)); + dim3 const block(static_cast(kThreadsPerBlock)); + computeCompressedMaskKernel<<>>(params, batchSize); +} + +void invokeDeepseekV4ComputeCtxCompressedPositionIds(CompressedPositionIdsParams const& params, int32_t maxCount, + int32_t numRatios, int32_t numContexts, cudaStream_t stream) +{ + if (numRatios <= 0 || numContexts <= 0 || maxCount <= 0) + { + return; + } + dim3 const grid(static_cast(gridX(maxCount, 2048)), static_cast(numRatios)); + dim3 const block(static_cast(kThreadsPerBlock)); + computeCtxCompressedPositionIdsKernel<<>>(params, numContexts); +} + +void invokeDeepseekV4ComputeGenCompressedPositionIds(CompressedPositionIdsParams const& params, int32_t maxCount, + int32_t numRatios, int32_t numContexts, int32_t batchSize, cudaStream_t stream) +{ + if (numRatios <= 0 || maxCount <= 0 || batchSize <= 0) + { + return; + } + dim3 const grid(static_cast(gridX(maxCount, 1024)), static_cast(numRatios)); + dim3 const block(static_cast(kThreadsPerBlock)); + computeGenCompressedPositionIdsKernel<<>>(params, numContexts, batchSize); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h b/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h new file mode 100644 index 000000000000..ce3492cbb1d0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h @@ -0,0 +1,83 @@ +/* + * 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 +{ + +// DeepSeek-V4 restricts compression ratios to {1, 4, 128}; one slot spare. +constexpr int32_t kMaxCompressRatios = 4; +// Upper bound on the scheduler batch handled by the single-block shared-memory +// scan (2 * 4096 * 4B = 32KB static shared memory, within the 48KB budget). +constexpr int32_t kMaxScanBatch = 4096; + +// Per-ratio buffers are persistent across iterations, and there are at most 4 +// ratios, so they travel by value in the kernel launch packet. This keeps the +// whole step host-copy-free -- passing device pointer arrays would need an H2D +// memcpy per call and negate the point of the port. +struct PerRatioKvLensParams +{ + int32_t* compressedKvLens[kMaxCompressRatios]; + int32_t* pastKvLens[kMaxCompressRatios]; + int32_t* newCompKvLens[kMaxCompressRatios]; + int32_t* cuNewCompKv[kMaxCompressRatios]; + int32_t ratios[kMaxCompressRatios]; +}; + +struct CompressedMaskParams +{ + int32_t const* newCompKvLens[kMaxCompressRatios]; + int32_t const* cuNewCompKv[kMaxCompressRatios]; + bool* mask[kMaxCompressRatios]; + int32_t totalTokens[kMaxCompressRatios]; +}; + +// Shared by the context and generation position-id kernels. `counts` is the +// per-ratio element count; `offsets` is the compact output offset (generation +// only, zero for context). +struct CompressedPositionIdsParams +{ + int32_t const* pastKvLens[kMaxCompressRatios]; + int32_t const* cuNewCompKv[kMaxCompressRatios]; + int32_t* positionIds[kMaxCompressRatios]; + int32_t ratios[kMaxCompressRatios]; + int32_t counts[kMaxCompressRatios]; + int32_t offsets[kMaxCompressRatios]; +}; + +void invokeDeepseekV4ComputePerRatioKvLens(int32_t const* kvLens, int32_t const* cachedTokens, + PerRatioKvLensParams const& params, int32_t numRatios, int32_t batchSize, cudaStream_t stream); + +void invokeDeepseekV4ComputeCompressedMask(CompressedMaskParams const& params, int32_t maxTotalTokens, + int32_t numRatios, int32_t batchSize, cudaStream_t stream); + +void invokeDeepseekV4ComputeCtxCompressedPositionIds(CompressedPositionIdsParams const& params, int32_t maxCount, + int32_t numRatios, int32_t numContexts, cudaStream_t stream); + +void invokeDeepseekV4ComputeGenCompressedPositionIds(CompressedPositionIdsParams const& params, int32_t maxCount, + int32_t numRatios, int32_t numContexts, int32_t batchSize, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4Indices.cu b/cpp/tensorrt_llm/kernels/deepseekV4Indices.cu new file mode 100644 index 000000000000..3645da730920 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4Indices.cu @@ -0,0 +1,110 @@ +/* + * 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/deepseekV4Indices.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int32_t kInvalidIndex = -1; +constexpr int32_t kThreadsPerBlock = 256; + +// One block per token row; threads stride along the window / compressed columns. +// Replaces the @maybe_compile(max-autotune) python graph which materialised +// [num_tokens, window_size] and [num_tokens, max_compressed_indices] temporaries +// plus one [num_tokens] tensor per compression ratio. +__global__ void computeDeepseekV4IndicesKernel(int32_t const* __restrict__ tokenPositions, + int32_t* __restrict__ swaLocalIndices, int32_t* __restrict__ compressedLocalIndices, + int32_t* __restrict__ topkLensRatio1, int32_t* __restrict__ topkLensRatio4, int32_t* __restrict__ topkLensRatio128, + int32_t numTokens, int32_t windowSize, int32_t maxCompressedIndices, int32_t sparseMlaTopk, int32_t swaStride, + int32_t compressedStride) +{ + int32_t const tokenId = static_cast(blockIdx.x); + if (tokenId >= numTokens) + { + return; + } + + int32_t const position = tokenPositions[tokenId]; + int32_t const kvLen = position + 1; + + // ── SWA local indices: start = clamp(position - window + 1, min=0); entries + // beyond `position` are invalid. + int32_t const swaStart = max(position - windowSize + 1, 0); + int32_t* swaRow = swaLocalIndices + static_cast(tokenId) * swaStride; + for (int32_t col = static_cast(threadIdx.x); col < windowSize; col += static_cast(blockDim.x)) + { + int32_t const idx = swaStart + col; + swaRow[col] = idx > position ? kInvalidIndex : idx; + } + + // ── Compressed local indices (ratio 128): first `kvLen / 128` columns are + // their own index, the rest are invalid. + int32_t const numValid = kvLen / 128; + int32_t* compRow = compressedLocalIndices + static_cast(tokenId) * compressedStride; + for (int32_t col = static_cast(threadIdx.x); col < maxCompressedIndices; + col += static_cast(blockDim.x)) + { + compRow[col] = col < numValid ? col : kInvalidIndex; + } + + // ── sparse_mla_topk_lens per compression ratio (one scalar per token). + if (threadIdx.x == 0) + { + if (topkLensRatio1 != nullptr) + { + topkLensRatio1[tokenId] = min(kvLen, windowSize); + } + if (topkLensRatio4 != nullptr) + { + topkLensRatio4[tokenId] = windowSize + min(kvLen / 4, sparseMlaTopk); + } + if (topkLensRatio128 != nullptr) + { + topkLensRatio128[tokenId] = windowSize + kvLen / 128; + } + } +} +} // namespace + +void invokeDeepseekV4ComputeIndices(int32_t const* tokenPositions, int32_t* swaLocalIndices, + int32_t* compressedLocalIndices, int32_t* topkLensRatio1, int32_t* topkLensRatio4, int32_t* topkLensRatio128, + int32_t numTokens, int32_t windowSize, int32_t maxCompressedIndices, int32_t sparseMlaTopk, int32_t swaStride, + int32_t compressedStride, cudaStream_t stream) +{ + if (numTokens <= 0) + { + return; + } + int32_t const widest = max(windowSize, maxCompressedIndices); + int32_t threads = widest >= kThreadsPerBlock ? kThreadsPerBlock : max(widest, 32); + dim3 const block(static_cast(threads)); + dim3 const grid(static_cast(numTokens)); + computeDeepseekV4IndicesKernel<<>>(tokenPositions, swaLocalIndices, compressedLocalIndices, + topkLensRatio1, topkLensRatio4, topkLensRatio128, numTokens, windowSize, maxCompressedIndices, sparseMlaTopk, + swaStride, compressedStride); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4Indices.h b/cpp/tensorrt_llm/kernels/deepseekV4Indices.h new file mode 100644 index 000000000000..347b11b7b1c0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4Indices.h @@ -0,0 +1,39 @@ +/* + * 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 +{ + +// Builds the DeepSeek-V4 SWA / compressed local index tables and the per-ratio +// sparse-MLA topk lengths from `tokenPositions`. Any of the three topkLens +// pointers may be null when that compression ratio is not configured. +void invokeDeepseekV4ComputeIndices(int32_t const* tokenPositions, int32_t* swaLocalIndices, + int32_t* compressedLocalIndices, int32_t* topkLensRatio1, int32_t* topkLensRatio4, int32_t* topkLensRatio128, + int32_t numTokens, int32_t windowSize, int32_t maxCompressedIndices, int32_t sparseMlaTopk, int32_t swaStride, + int32_t compressedStride, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp b/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp index 8eb20a88bc1c..ec80e08afcb2 100644 --- a/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp +++ b/cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp @@ -14,13 +14,19 @@ * limitations under the License. */ +#include "tensorrt_llm/kernels/attentionMetadataKernels.h" #include "tensorrt_llm/kernels/deepseekV4BlockTable.h" +#include "tensorrt_llm/kernels/deepseekV4CompressedMeta.h" +#include "tensorrt_llm/kernels/deepseekV4Indices.h" #include +#include #include #include #include +#include #include +#include namespace th = torch; namespace tk = tensorrt_llm::kernels; @@ -177,6 +183,313 @@ void deepseekV4ComputeSlidingBlockTablesWithScratch(th::Tensor const& blockOffse C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void computeSharedBlockTable( + th::Tensor const& blockOffsets, th::Tensor const& copyIdx, int64_t poolId, int64_t scale, th::Tensor const& output) +{ + int const device = output.get_device(); + checkCudaContiguousTensor(blockOffsets, "block_offsets", device); + checkCudaContiguousTensor(copyIdx, "copy_idx", device); + checkCudaContiguousTensor(output, "output", device); + checkInt32Tensor(blockOffsets, "block_offsets"); + checkInt32Tensor(copyIdx, "copy_idx"); + checkInt32Tensor(output, "output"); + TORCH_CHECK(blockOffsets.dim() == 4, "block_offsets must be 4D [num_pools, copy_idx_capacity, 2, max_blocks]"); + TORCH_CHECK(copyIdx.dim() == 1, "copy_idx must be 1D [num_tables]"); + TORCH_CHECK(output.dim() == 2, "output must be 2D [num_tables, max_blocks_per_seq]"); + TORCH_CHECK(output.size(1) == blockOffsets.size(3), "output max_blocks must match block_offsets"); + TORCH_CHECK(poolId >= 0 && poolId < blockOffsets.size(0), "pool_id out of range"); + + c10::cuda::CUDAGuard const deviceGuard(output.device()); + int32_t const copyIdxCapacity = checkedInt32Size(blockOffsets.size(1), "copy_idx_capacity"); + int32_t const maxBlocksPerSeq = checkedInt32Size(blockOffsets.size(3), "max_blocks_per_seq"); + // The caller may pass a fixed-capacity copy_idx buffer; honour the smaller of + // the two so a padded staging tensor cannot write past the output rows. + int32_t const numTables + = std::min(checkedInt32Size(copyIdx.size(0), "num_tables"), checkedInt32Size(output.size(0), "output_rows")); + + auto stream = at::cuda::getCurrentCUDAStream(device); + tk::invokeComputeSharedBlockTable(blockOffsets.data_ptr(), copyIdx.data_ptr(), + output.data_ptr(), static_cast(poolId), static_cast(scale), copyIdxCapacity, + numTables, maxBlocksPerSeq, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void deepseekV4ComputeIndices(th::Tensor const& tokenPositions, int64_t windowSize, int64_t maxCompressedIndices, + int64_t sparseMlaTopk, th::Tensor const& swaLocalIndices, th::Tensor const& compressedLocalIndices, + std::optional const& topkLensRatio1, std::optional const& topkLensRatio4, + std::optional const& topkLensRatio128) +{ + int const device = swaLocalIndices.get_device(); + checkCudaContiguousTensor(tokenPositions, "token_positions", device); + checkCudaContiguousTensor(swaLocalIndices, "swa_local_indices", device); + checkCudaContiguousTensor(compressedLocalIndices, "compressed_local_indices", device); + checkInt32Tensor(tokenPositions, "token_positions"); + checkInt32Tensor(swaLocalIndices, "swa_local_indices"); + checkInt32Tensor(compressedLocalIndices, "compressed_local_indices"); + TORCH_CHECK(tokenPositions.dim() == 1, "token_positions must be 1D [num_tokens]"); + TORCH_CHECK(swaLocalIndices.dim() == 2, "swa_local_indices must be 2D [rows, window_size]"); + TORCH_CHECK(compressedLocalIndices.dim() == 2, "compressed_local_indices must be 2D"); + TORCH_CHECK( + windowSize > 0 && windowSize <= swaLocalIndices.size(1), "window_size must fit swa_local_indices columns"); + TORCH_CHECK(maxCompressedIndices > 0 && maxCompressedIndices <= compressedLocalIndices.size(1), + "max_compressed_indices must fit compressed_local_indices columns"); + + int32_t const numTokens = checkedInt32Size(tokenPositions.size(0), "num_tokens"); + TORCH_CHECK(numTokens <= swaLocalIndices.size(0) && numTokens <= compressedLocalIndices.size(0), + "output buffers must have at least num_tokens rows"); + + auto ptr = [&](std::optional const& t) -> int32_t* + { + if (!t.has_value()) + { + return nullptr; + } + checkCudaContiguousTensor(*t, "sparse_mla_topk_lens", device); + checkInt32Tensor(*t, "sparse_mla_topk_lens"); + TORCH_CHECK(t->numel() >= numTokens, "sparse_mla_topk_lens must hold num_tokens entries"); + return t->data_ptr(); + }; + + c10::cuda::CUDAGuard const deviceGuard(swaLocalIndices.device()); + auto stream = at::cuda::getCurrentCUDAStream(device); + tk::invokeDeepseekV4ComputeIndices(tokenPositions.data_ptr(), swaLocalIndices.data_ptr(), + compressedLocalIndices.data_ptr(), ptr(topkLensRatio1), ptr(topkLensRatio4), ptr(topkLensRatio128), + numTokens, static_cast(windowSize), static_cast(maxCompressedIndices), + static_cast(sparseMlaTopk), checkedInt32Size(swaLocalIndices.stride(0), "swa_stride"), + checkedInt32Size(compressedLocalIndices.stride(0), "compressed_stride"), stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +namespace +{ +// Fills the per-ratio pointer/scalar slots of a launch-packet struct from the +// python-side lists, validating each tensor. Returns the ratio count. +template +int32_t forEachRatio(int64_t expectedCount, FillFn&& fill) +{ + TORCH_CHECK(expectedCount > 0 && expectedCount <= tk::kMaxCompressRatios, + "number of compression ratios must be in [1, ", tk::kMaxCompressRatios, "], got ", expectedCount); + for (int64_t i = 0; i < expectedCount; ++i) + { + fill(static_cast(i)); + } + return static_cast(expectedCount); +} + +void checkPerRatioList( + std::vector const& list, int64_t numRatios, char const* name, int64_t minElems, int device) +{ + TORCH_CHECK(static_cast(list.size()) == numRatios, name, + " must have one entry per compression ratio (expected ", numRatios, ", got ", list.size(), ")"); + for (auto const& t : list) + { + checkCudaContiguousTensor(t, name, device); + TORCH_CHECK(t.numel() >= minElems, name, " must hold at least ", minElems, " elements, got ", t.numel()); + } +} +} // namespace + +void deepseekV4ComputePerRatioKvLens(th::Tensor const& kvLens, th::Tensor const& cachedTokens, + std::vector const& ratios, std::vector const& compressedKvLens, + std::vector const& pastKvLens, std::vector const& newCompKvLens, + std::vector const& cuNewCompKv) +{ + int const device = kvLens.get_device(); + checkCudaContiguousTensor(kvLens, "kv_lens", device); + checkCudaContiguousTensor(cachedTokens, "cached_tokens", device); + checkInt32Tensor(kvLens, "kv_lens"); + checkInt32Tensor(cachedTokens, "cached_tokens"); + int32_t const batchSize = checkedInt32Size(kvLens.size(0), "batch_size"); + TORCH_CHECK(cachedTokens.size(0) >= batchSize, "cached_tokens shorter than kv_lens"); + TORCH_CHECK(batchSize <= tk::kMaxScanBatch, "batch_size ", batchSize, " exceeds the single-block scan bound ", + tk::kMaxScanBatch); + + int64_t const numRatios = static_cast(ratios.size()); + checkPerRatioList(compressedKvLens, numRatios, "compressed_kv_lens", batchSize, device); + checkPerRatioList(pastKvLens, numRatios, "past_kv_lens", batchSize, device); + checkPerRatioList(newCompKvLens, numRatios, "new_comp_kv_lens", batchSize, device); + checkPerRatioList(cuNewCompKv, numRatios, "cu_new_comp_kv", batchSize + 1, device); + + tk::PerRatioKvLensParams params{}; + int32_t const numRatios32 = forEachRatio(numRatios, + [&](int32_t i) + { + TORCH_CHECK(ratios[i] > 0, "compression ratio must be positive"); + params.ratios[i] = static_cast(ratios[i]); + params.compressedKvLens[i] = compressedKvLens[i].data_ptr(); + params.pastKvLens[i] = pastKvLens[i].data_ptr(); + params.newCompKvLens[i] = newCompKvLens[i].data_ptr(); + params.cuNewCompKv[i] = cuNewCompKv[i].data_ptr(); + }); + + c10::cuda::CUDAGuard const deviceGuard(kvLens.device()); + tk::invokeDeepseekV4ComputePerRatioKvLens(kvLens.data_ptr(), cachedTokens.data_ptr(), params, + numRatios32, batchSize, at::cuda::getCurrentCUDAStream(device)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void deepseekV4ComputeCompressedMask(std::vector const& newCompKvLens, + std::vector const& cuNewCompKv, std::vector const& mask, + std::vector const& totalTokens, int64_t batchSizeIn) +{ + TORCH_CHECK(!mask.empty(), "compressed_mask list must not be empty"); + int const device = mask[0].get_device(); + int32_t const batchSize = checkedInt32Size(batchSizeIn, "batch_size"); + TORCH_CHECK(batchSize > 0, "batch_size must be positive"); + + int64_t const numRatios = static_cast(totalTokens.size()); + checkPerRatioList(newCompKvLens, numRatios, "new_comp_kv_lens", batchSize, device); + checkPerRatioList(cuNewCompKv, numRatios, "cu_new_comp_kv", batchSize + 1, device); + TORCH_CHECK( + static_cast(mask.size()) == numRatios, "compressed_mask must have one entry per compression ratio"); + + tk::CompressedMaskParams params{}; + int32_t maxTotal = 0; + int32_t const numRatios32 = forEachRatio(numRatios, + [&](int32_t i) + { + int32_t const total = checkedInt32Size(totalTokens[i], "total_compressed_tokens"); + TORCH_CHECK(total >= 0, "total_compressed_tokens must be non-negative"); + checkCudaContiguousTensor(mask[i], "compressed_mask", device); + TORCH_CHECK(mask[i].scalar_type() == th::kBool, "compressed_mask must be bool"); + TORCH_CHECK( + mask[i].numel() >= total, "compressed_mask buffer too small: need ", total, ", have ", mask[i].numel()); + params.newCompKvLens[i] = newCompKvLens[i].data_ptr(); + params.cuNewCompKv[i] = cuNewCompKv[i].data_ptr(); + params.mask[i] = mask[i].data_ptr(); + params.totalTokens[i] = total; + maxTotal = std::max(maxTotal, total); + }); + + c10::cuda::CUDAGuard const deviceGuard(mask[0].device()); + tk::invokeDeepseekV4ComputeCompressedMask( + params, maxTotal, numRatios32, batchSize, at::cuda::getCurrentCUDAStream(device)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +// Shared by the ctx and gen position-id entry points; `isGen` selects the +// kernel and whether `offsets` participates. +namespace +{ +void computeCompressedPositionIdsImpl(std::vector const& pastKvLens, + std::vector const& cuNewCompKv, std::vector const& positionIds, + std::vector const& ratios, std::vector const& counts, std::vector const& offsets, + int64_t numContextsIn, int64_t batchSizeIn, bool isGen) +{ + TORCH_CHECK(!positionIds.empty(), "compressed_position_ids list must not be empty"); + int const device = positionIds[0].get_device(); + int32_t const numContexts = checkedInt32Size(numContextsIn, "num_contexts"); + int32_t const batchSize = checkedInt32Size(batchSizeIn, "batch_size"); + TORCH_CHECK( + batchSize > 0 && numContexts >= 0 && numContexts <= batchSize, "num_contexts must be within [0, batch_size]"); + + int64_t const numRatios = static_cast(ratios.size()); + TORCH_CHECK(static_cast(counts.size()) == numRatios && static_cast(offsets.size()) == numRatios, + "counts and offsets must have one entry per compression ratio"); + checkPerRatioList(pastKvLens, numRatios, "past_kv_lens", batchSize, device); + checkPerRatioList(cuNewCompKv, numRatios, "cu_new_comp_kv", batchSize + 1, device); + TORCH_CHECK(static_cast(positionIds.size()) == numRatios, + "compressed_position_ids must have one entry per compression ratio"); + + tk::CompressedPositionIdsParams params{}; + int32_t maxCount = 0; + int32_t const numRatios32 = forEachRatio(numRatios, + [&](int32_t i) + { + int32_t const count = checkedInt32Size(counts[i], "count"); + int32_t const offset = checkedInt32Size(offsets[i], "offset"); + TORCH_CHECK(count >= 0 && offset >= 0, "count and offset must be non-negative"); + checkCudaContiguousTensor(positionIds[i], "compressed_position_ids", device); + checkInt32Tensor(positionIds[i], "compressed_position_ids"); + TORCH_CHECK(positionIds[i].numel() >= static_cast(offset) + count, + "compressed_position_ids buffer too small: need ", offset + count, ", have ", positionIds[i].numel()); + TORCH_CHECK(ratios[i] > 0, "compression ratio must be positive"); + params.pastKvLens[i] = pastKvLens[i].data_ptr(); + params.cuNewCompKv[i] = cuNewCompKv[i].data_ptr(); + params.positionIds[i] = positionIds[i].data_ptr(); + params.ratios[i] = static_cast(ratios[i]); + params.counts[i] = count; + params.offsets[i] = offset; + maxCount = std::max(maxCount, count); + }); + + c10::cuda::CUDAGuard const deviceGuard(positionIds[0].device()); + auto stream = at::cuda::getCurrentCUDAStream(device); + if (isGen) + { + tk::invokeDeepseekV4ComputeGenCompressedPositionIds( + params, maxCount, numRatios32, numContexts, batchSize, stream); + } + else + { + tk::invokeDeepseekV4ComputeCtxCompressedPositionIds(params, maxCount, numRatios32, numContexts, stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +void deepseekV4ComputeCtxCompressedPositionIds(std::vector const& pastKvLens, + std::vector const& cuNewCompKv, std::vector const& positionIds, + std::vector const& ratios, std::vector const& counts, int64_t numContexts) +{ + std::vector const zeroOffsets(ratios.size(), 0); + computeCompressedPositionIdsImpl(pastKvLens, cuNewCompKv, positionIds, ratios, counts, zeroOffsets, numContexts, + /*batchSize=*/numContexts, /*isGen=*/false); +} + +void deepseekV4ComputeGenCompressedPositionIds(std::vector const& pastKvLens, + std::vector const& cuNewCompKv, std::vector const& positionIds, + std::vector const& ratios, std::vector const& counts, std::vector const& offsets, + int64_t numContexts, int64_t batchSize) +{ + computeCompressedPositionIdsImpl( + pastKvLens, cuNewCompKv, positionIds, ratios, counts, offsets, numContexts, batchSize, /*isGen=*/true); +} + +void computeTokenPositions(th::Tensor const& seqLens, std::optional const& cachedTokens, + th::Tensor const& cuSeqLens, th::Tensor const& reqIdxPerToken, std::optional const& tokenPositions, + int64_t numTokensIn, bool computeCuSeqLens) +{ + int const device = cuSeqLens.get_device(); + checkCudaContiguousTensor(seqLens, "seq_lens", device); + checkCudaContiguousTensor(cuSeqLens, "cu_seq_lens", device); + checkCudaContiguousTensor(reqIdxPerToken, "req_idx_per_token", device); + checkInt32Tensor(seqLens, "seq_lens"); + checkInt32Tensor(cuSeqLens, "cu_seq_lens"); + checkInt32Tensor(reqIdxPerToken, "req_idx_per_token"); + + int32_t const batchSize = checkedInt32Size(seqLens.size(0), "batch_size"); + int32_t const numTokens = checkedInt32Size(numTokensIn, "num_tokens"); + TORCH_CHECK(batchSize > 0, "batch_size must be positive"); + TORCH_CHECK(numTokens >= 0, "num_tokens must be non-negative"); + TORCH_CHECK(!computeCuSeqLens || batchSize <= tk::kMaxTokenPositionScanBatch, "batch_size ", batchSize, + " exceeds the single-block scan bound ", tk::kMaxTokenPositionScanBatch); + TORCH_CHECK( + cuSeqLens.numel() >= static_cast(batchSize) + 1, "cu_seq_lens must hold batch_size + 1 entries"); + TORCH_CHECK(reqIdxPerToken.numel() >= numTokens, "req_idx_per_token buffer too small"); + + int32_t* positionsPtr = nullptr; + int32_t const* cachedPtr = nullptr; + if (tokenPositions.has_value()) + { + TORCH_CHECK(cachedTokens.has_value(), "cached_tokens is required when token_positions is requested"); + checkCudaContiguousTensor(*tokenPositions, "token_positions", device); + checkCudaContiguousTensor(*cachedTokens, "cached_tokens", device); + checkInt32Tensor(*tokenPositions, "token_positions"); + checkInt32Tensor(*cachedTokens, "cached_tokens"); + TORCH_CHECK(tokenPositions->numel() >= numTokens, "token_positions buffer too small"); + TORCH_CHECK(cachedTokens->size(0) >= batchSize, "cached_tokens shorter than seq_lens"); + positionsPtr = tokenPositions->data_ptr(); + cachedPtr = cachedTokens->data_ptr(); + } + + c10::cuda::CUDAGuard const deviceGuard(cuSeqLens.device()); + tk::invokeComputeTokenPositions(seqLens.data_ptr(), cachedPtr, cuSeqLens.data_ptr(), + reqIdxPerToken.data_ptr(), positionsPtr, batchSize, numTokens, computeCuSeqLens, + at::cuda::getCurrentCUDAStream(device)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -186,6 +499,32 @@ 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( + "compute_shared_block_table(Tensor block_offsets, Tensor copy_idx, int pool_id, int scale, " + "Tensor! output) -> ()"); + m.def( + "deepseek_v4_compute_per_ratio_kv_lens(Tensor kv_lens, Tensor cached_tokens, int[] ratios, " + "Tensor(a!)[] compressed_kv_lens, Tensor(b!)[] past_kv_lens, Tensor(c!)[] new_comp_kv_lens, " + "Tensor(d!)[] cu_new_comp_kv) -> ()"); + m.def( + "deepseek_v4_compute_compressed_mask(Tensor[] new_comp_kv_lens, Tensor[] cu_new_comp_kv, " + "Tensor(a!)[] compressed_mask, int[] total_tokens, int batch_size) -> ()"); + m.def( + "deepseek_v4_compute_ctx_compressed_position_ids(Tensor[] past_kv_lens, " + "Tensor[] cu_new_comp_kv, Tensor(a!)[] compressed_position_ids, int[] ratios, " + "int[] counts, int num_contexts) -> ()"); + m.def( + "deepseek_v4_compute_gen_compressed_position_ids(Tensor[] past_kv_lens, " + "Tensor[] cu_new_comp_kv, Tensor(a!)[] compressed_position_ids, int[] ratios, " + "int[] counts, int[] offsets, int num_contexts, int batch_size) -> ()"); + m.def( + "compute_token_positions(Tensor seq_lens, Tensor? cached_tokens, " + "Tensor(a!) cu_seq_lens, Tensor(b!) req_idx_per_token, Tensor(c!)? token_positions, " + "int num_tokens, bool compute_cu_seq_lens) -> ()"); + m.def( + "deepseek_v4_compute_indices(Tensor token_positions, int window_size, int max_compressed_indices, " + "int sparse_mla_topk, Tensor(a!) swa_local_indices, Tensor(b!) compressed_local_indices, " + "Tensor(c!)? topk_lens_ratio1, Tensor(d!)? topk_lens_ratio4, Tensor(e!)? topk_lens_ratio128) -> ()"); 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, " @@ -196,6 +535,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("deepseek_v4_compute_sliding_block_tables", &tensorrt_llm::torch_ext::deepseekV4ComputeSlidingBlockTables); + m.impl("compute_shared_block_table", &tensorrt_llm::torch_ext::computeSharedBlockTable); + m.impl("deepseek_v4_compute_per_ratio_kv_lens", &tensorrt_llm::torch_ext::deepseekV4ComputePerRatioKvLens); + m.impl("deepseek_v4_compute_compressed_mask", &tensorrt_llm::torch_ext::deepseekV4ComputeCompressedMask); + m.impl("deepseek_v4_compute_ctx_compressed_position_ids", + &tensorrt_llm::torch_ext::deepseekV4ComputeCtxCompressedPositionIds); + m.impl("deepseek_v4_compute_gen_compressed_position_ids", + &tensorrt_llm::torch_ext::deepseekV4ComputeGenCompressedPositionIds); + m.impl("compute_token_positions", &tensorrt_llm::torch_ext::computeTokenPositions); + m.impl("deepseek_v4_compute_indices", &tensorrt_llm::torch_ext::deepseekV4ComputeIndices); 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 933c1dc443b9..ff64adb368ac 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 @@ -1095,6 +1095,37 @@ def _compute_shared_block_table( base = self.host_kv_cache_block_offsets[pool_id, copy_idx, 0, :] return torch.where(base == BAD_PAGE_INDEX, BAD_PAGE_INDEX, base * scale) + def _compute_shared_block_table_device( + self, pool_id: int, scale: int, copy_idx: torch.Tensor, out: torch.Tensor + ) -> None: + """Device-side equivalent of `_compute_shared_block_table`. + + Writes `where(base == BAD_PAGE_INDEX, BAD_PAGE_INDEX, base * scale)` for + `host_kv_cache_block_offsets[pool_id, copy_idx, 0, :]` straight into `out` + (shape [num_seqs, max_blocks_per_seq]) without any host-side gather. + """ + device_copy_idx = self._copy_idx_to_device(copy_idx) + # host_kv_cache_block_offsets is [num_pools, capacity, 2, max_blocks] and can be + # several MB; it is uploaded once per prepare() and shared by every consumer + # (sliding + one call per compression ratio + the indexer table). + self._ensure_device_block_offsets() + torch.ops.trtllm.compute_shared_block_table( + self._device_kv_cache_block_offsets_input, + device_copy_idx[: copy_idx.size(0)], + int(pool_id), + int(scale), + out, + ) + + def _ensure_device_block_offsets(self) -> None: + """Upload host_kv_cache_block_offsets to the device at most once per iteration.""" + if getattr(self, "_block_offsets_uploaded", False): + return + self._device_kv_cache_block_offsets_input.copy_( + self.host_kv_cache_block_offsets, non_blocking=True + ) + self._block_offsets_uploaded = True + def _copy_idx_to_device(self, copy_idx: torch.Tensor) -> torch.Tensor: num_tables = copy_idx.size(0) device_copy_idx = self._device_copy_idx_staging[:num_tables] @@ -1145,6 +1176,22 @@ def _copy_scratch_metadata_to_device( self._device_scratch_slots_staging, ) + def _get_copy_index_cached(self, request_ids, num_contexts, beam_width): + """`IndexMapper.get_copy_index` memoized for one prepare() step. + + Keyed on the request-id list *contents*, not its identity: the caller may + reuse a single list object and mutate it between steps. The memo is reset + at the top of every step by ``compute_sliding_block_tables``, which always + runs before the other users. + """ + key = (tuple(request_ids), num_contexts, beam_width) + if getattr(self, "_copy_idx_memo_key", None) == key: + return self._copy_idx_memo_val + val = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + self._copy_idx_memo_key = key + self._copy_idx_memo_val = val + return val + @nvtx_range_debug("dsv4_compute_sliding_block_tables") def compute_sliding_block_tables( self, @@ -1152,7 +1199,16 @@ def compute_sliding_block_tables( num_contexts: int, ) -> None: """Compute all per-layer sliding-window block tables for this batch.""" - copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, 1) + # One get_copy_index walk per step instead of three. This function, + # copy_batch_compress_block_tables (once per compression ratio) and + # copy_batch_indexer_compress_block_tables all call it with the same + # (request_ids, num_contexts, beam_width). The mapping is deterministic in + # those arguments and returns a view of one shared pinned buffer, so the + # later calls were recomputing a result equal to the one already there. + # Each walk costs O(num_seqs) ATen dispatches (select + as_strided + fill_ + # per entry). This is the first call of the step, so it seeds the memo. + self._copy_idx_memo_key = None + copy_idx = self._get_copy_index_cached(request_ids, num_contexts, 1) num_tables = copy_idx.size(0) self._num_tables = num_tables @@ -1167,10 +1223,12 @@ def compute_sliding_block_tables( ] device_copy_idx = self._copy_idx_to_device(copy_idx) - self._device_kv_cache_block_offsets_input.copy_( - self.host_kv_cache_block_offsets, - non_blocking=True, - ) + # compute_sliding_block_tables runs first in + # DeepseekV4TrtllmAttentionMetadata.prepare, so this is where the + # per-iteration upload normally happens; the per-ratio and indexer tables + # then reuse it instead of re-uploading. + self._block_offsets_uploaded = False + self._ensure_device_block_offsets() if scratch_descs_by_pool is not None: scratch_begs, scratch_ends, scratch_slots = self._copy_scratch_metadata_to_device( @@ -1223,7 +1281,14 @@ def copy_batch_block_offsets( """ assert beam_width == 1, "DSV4 only supports beam width 1 now" assert dst_tensor.is_cuda, "copy_batch_block_offsets expects a CUDA destination" - dst_tensor.fill_(BAD_PAGE_INDEX) + # Fill only what the copy below does not overwrite. The copy fully + # populates [:, :_num_tables, 0, :], so pre-filling those entries was + # wasted work; the rows past _num_tables and the beams past 0 still have + # to carry BAD_PAGE_INDEX because padded CUDA-graph token slots can index + # them through req_idx_per_token. + if self._num_tables < dst_tensor.size(1): + dst_tensor[:, self._num_tables :, :, :].fill_(BAD_PAGE_INDEX) + dst_tensor[:, : self._num_tables, 1:, :].fill_(BAD_PAGE_INDEX) dst_tensor[:, : self._num_tables, 0, :].copy_( self._precomputed_sliding_block_tables[ :, DeepseekV4AttentionType.SWA.value, : self._num_tables, : @@ -1243,7 +1308,14 @@ def copy_batch_sliding_block_tables( Copy the per-layer block tables for attentions managed in sliding-window mode to the GPU tensor. """ assert dst_tensor.is_cuda, "copy_batch_sliding_block_tables expects a CUDA destination" - dst_tensor.fill_(BAD_PAGE_INDEX) + # Fill only the rows the copy below does not overwrite. The tail past + # _num_tables must stay BAD_PAGE_INDEX because padded CUDA-graph token + # slots can still index it through req_idx_per_token; the head is fully + # overwritten by the copy, so filling the whole + # [num_local_layers, num_sliding_types, capacity, max_blocks] tensor first + # was redundant. + if self._num_tables < dst_tensor.size(2): + dst_tensor[:, :, self._num_tables :, :].fill_(BAD_PAGE_INDEX) dst_tensor[:, :, : self._num_tables, :].copy_( self._precomputed_sliding_block_tables[:, :, : self._num_tables, :], non_blocking=True, @@ -1261,7 +1333,7 @@ def copy_batch_compress_block_tables( ) -> None: """Build the COMPRESS block table for one compression ratio and copy it to the destination.""" assert beam_width == 1, "DSV4 only supports beam width 1 now" - copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + copy_idx = self._get_copy_index_cached(request_ids, num_contexts, beam_width) staging = self._host_compress_block_tables_staging[compress_ratio] if compress_ratio == 4: pool_id = self._csa_compress_pool_id @@ -1278,8 +1350,16 @@ def copy_batch_compress_block_tables( raise RuntimeError( f"Missing COMPRESS pool metadata for compress ratio {compress_ratio}" ) - staging[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) - dst_tensor[:num_seqs].copy_(staging[:num_seqs], non_blocking=True) + if dst_tensor.is_cuda: + # Build the table on the GPU. The CPU fallback below gathers a few + # hundred KB out of host_kv_cache_block_offsets and runs torch.where on + # it once per compression ratio, on the decode critical path; the device + # op does the same arithmetic (BAD_PAGE_INDEX passthrough, otherwise + # base * scale) and also removes the host->device staging copy. + self._compute_shared_block_table_device(pool_id, scale, copy_idx, dst_tensor[:num_seqs]) + else: + staging[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) + dst_tensor[:num_seqs].copy_(staging[:num_seqs], non_blocking=True) @nvtx_range_debug("dsv4_copy_batch_indexer_compress_block_tables") def copy_batch_indexer_compress_block_tables( @@ -1289,15 +1369,27 @@ def copy_batch_indexer_compress_block_tables( beam_width: int, num_contexts: int, num_seqs: int, + device_block_table: Optional[torch.Tensor] = None, ) -> None: - """Build the shared INDEXER_COMPRESS compatibility block table.""" + """Build the shared INDEXER_COMPRESS compatibility block table. + + When `device_block_table` is given the table is produced directly on the + device and `host_block_table` is left untouched. + """ assert beam_width == 1, "DSV4 only supports beam width 1 now" - copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + copy_idx = self._get_copy_index_cached(request_ids, num_contexts, beam_width) pool_id = self._csa_indexer_compress_pool_id scale = self._csa_indexer_compress_scale if pool_id is None or scale is None: raise RuntimeError("Missing INDEXER_COMPRESS pool metadata") - host_block_table[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) + if device_block_table is not None: + # Same rationale as copy_batch_compress_block_tables: build on device + # and skip both the host gather and the staging copy. + self._compute_shared_block_table_device( + pool_id, scale, copy_idx, device_block_table[:num_seqs] + ) + else: + host_block_table[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) @staticmethod def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs): diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 1f76e5a73e51..6a98b86d6687 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -423,6 +423,17 @@ def __post_init__(self): ) self.empty_topk_indices_buffer.fill_(-1) + # token_positions: absolute position of each token in its sequence. + # Persistent so the fused token-position kernel writes in place instead + # of allocating a fresh tensor per iteration. + self.token_positions_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="token_positions_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # SWA local indices self.swa_local_indices_cuda = self.get_empty( self.cuda_graph_buffers, @@ -526,10 +537,7 @@ def prepare_for_indexer_k_cache(self): beam_width=self.beam_width, num_contexts=self.num_contexts, num_seqs=self.num_seqs, - ) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], - non_blocking=True, + device_block_table=self.indexer_k_cache_block_offsets, ) # Columns beyond each sequence's allocated indexer blocks contain BAD_PAGE_INDEX (-1). # CUDA-graph padded token slots may still compute scatter addresses from those columns @@ -554,35 +562,80 @@ def prepare_for_block_tables(self): num_seqs=self.num_seqs, ) + def prepare_for_indices_conversion(self): + """Device-side req_idx_per_token (overrides the CPU staging path). + + The base class builds this with a CPU ``repeat_interleave`` plus a + pinned H2D memcpy. DeepSeek-V4 already keeps ``seq_lens`` on device, so + one kernel computes ``cu_seq_lens`` and the per-token request index + together -- and ``prepare()`` reuses that ``cu_seq_lens`` instead of + recomputing it on the host. + """ + num_requests = self.num_seqs + if num_requests <= 0: + return + if self._seq_lens_cuda is None: + # seq_lens has not been set yet; fall back to the base implementation. + super().prepare_for_indices_conversion() + return + torch.ops.trtllm.compute_token_positions( + self._seq_lens_cuda[:num_requests], + None, + self.cu_seq_lens_cuda, + self.req_idx_per_token, + None, + self.num_tokens, + True, + ) + def prepare_for_deepseek_v4_indices(self, token_positions=None): """Prepare SWA/compressed local indices and sparse_mla_topk_lens.""" window_size = self.window_size - device = self.swa_local_indices_cuda.device if token_positions is None: - # Initial prepare() path — build token_positions from CPU data. + # Initial prepare() path. cu_seq_lens_cuda is already populated by + # prepare_for_indices_conversion, so reuse the scan and only run the + # per-token phase (one kernel instead of arange + searchsorted + + # two gathers + two adds). num_tokens = self.num_tokens num_requests = self.num_seqs + token_positions = self.token_positions_cuda[:num_tokens] + torch.ops.trtllm.compute_token_positions( + self._seq_lens_cuda[:num_requests], + self.cached_token_lens_cuda[:num_requests], + self.cu_seq_lens_cuda, + self.req_idx_per_token, + token_positions, + num_tokens, + False, + ) - # cu_seq_lens_cuda must already be populated before this call - cu_seq_lens = self.cu_seq_lens_cuda[: num_requests + 1] - cached_tokens = self.cached_token_lens_cuda[:num_requests] - - token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) - req_idx = torch.searchsorted(cu_seq_lens[1:].to(torch.int32), token_idx, right=True) - offsets = token_idx - cu_seq_lens[req_idx].to(torch.int32) - token_positions = cached_tokens[req_idx].to(torch.int32) + offsets - - self._prepare_deepseek_v4_indices_compiled( - token_positions, - window_size, - self.max_compressed_indices[128], - self.sparse_mla_topk, - self.swa_local_indices_cuda, - self.compressed_local_indices_cuda, - self.sparse_mla_topk_lens, - self._compress_ratios_sorted, - ) + if token_positions.is_cuda: + # Single fused CUDA kernel: replaces ~15 ATen launches (arange x2, + # unsqueeze/expand, clamp, add, where x2, floordiv, full, plus one + # clamp/add/copy per compress_ratio) and their temporaries. + torch.ops.trtllm.deepseek_v4_compute_indices( + token_positions, + window_size, + self.max_compressed_indices[128], + self.sparse_mla_topk, + self.swa_local_indices_cuda, + self.compressed_local_indices_cuda, + self.sparse_mla_topk_lens.get(1), + self.sparse_mla_topk_lens.get(4), + self.sparse_mla_topk_lens.get(128), + ) + else: + self._prepare_deepseek_v4_indices_compiled( + token_positions, + window_size, + self.max_compressed_indices[128], + self.sparse_mla_topk, + self.swa_local_indices_cuda, + self.compressed_local_indices_cuda, + self.sparse_mla_topk_lens, + self._compress_ratios_sorted, + ) @staticmethod @maybe_compile(dynamic=True, options={"max-autotune": True}) @@ -778,13 +831,13 @@ def prepare(self): cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests] self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda, ctx_output_sizes) - self._compute_compressed_mask( - self.new_comp_kv_lens_cuda, - self.cu_new_comp_kv_cuda, - self.compressed_mask_cuda, + ratios = self._compress_ratios_sorted + torch.ops.trtllm.deepseek_v4_compute_compressed_mask( + [self.new_comp_kv_lens_cuda[r] for r in ratios], + [self.cu_new_comp_kv_cuda[r] for r in ratios], + [self.compressed_mask_cuda[r] for r in ratios], + [int(self.num_total_compressed_tokens[r]) for r in ratios], num_requests, - self.num_total_compressed_tokens, - self._compress_ratios_sorted, ) def prepare_compressed_kv_metadata( @@ -811,27 +864,54 @@ def prepare_compressed_kv_metadata( num_contexts = self.num_contexts num_generations = self.num_generations - self._compute_per_ratio_kv_lens( - kv_lens, - cached_tokens, - batch_size, - self.compressed_kv_lens_cuda, - self.past_kv_lens_cuda, - self.new_comp_kv_lens_cuda, - self.cu_new_comp_kv_cuda, - self._compress_ratios_sorted, - ) - - if num_contexts > 0: - self._compute_ctx_compressed_position_ids( + ratios = self._compress_ratios_sorted + if kv_lens.is_cuda: + # One launch for all ratios: replaces 4 ATen ops + a cumsum + a pad + # per ratio (~18 dispatches at 3 ratios) with a single fused kernel. + torch.ops.trtllm.deepseek_v4_compute_per_ratio_kv_lens( + kv_lens, + cached_tokens, + ratios, + [self.compressed_kv_lens_cuda[r] for r in ratios], + [self.past_kv_lens_cuda[r] for r in ratios], + [self.new_comp_kv_lens_cuda[r] for r in ratios], + [self.cu_new_comp_kv_cuda[r] for r in ratios], + ) + else: + self._compute_per_ratio_kv_lens( + kv_lens, + cached_tokens, + batch_size, + self.compressed_kv_lens_cuda, self.past_kv_lens_cuda, + self.new_comp_kv_lens_cuda, self.cu_new_comp_kv_cuda, - self.compressed_position_ids_cuda, - num_contexts, - self._compress_ratios_sorted, - ctx_output_sizes, + ratios, ) + if num_contexts > 0: + if ctx_output_sizes is not None: + # Host-side counts are already available, so the whole per-ratio + # loop (arange + searchsorted + gather + mul per ratio) collapses + # into one launch. + torch.ops.trtllm.deepseek_v4_compute_ctx_compressed_position_ids( + [self.past_kv_lens_cuda[r] for r in ratios], + [self.cu_new_comp_kv_cuda[r] for r in ratios], + [self.compressed_position_ids_cuda[r] for r in ratios], + ratios, + [int(ctx_output_sizes[r]) for r in ratios], + num_contexts, + ) + else: + self._compute_ctx_compressed_position_ids( + self.past_kv_lens_cuda, + self.cu_new_comp_kv_cuda, + self.compressed_position_ids_cuda, + num_contexts, + ratios, + ctx_output_sizes, + ) + if self.num_gen_tokens_per_seq > 0 and num_generations > 0: # Extract output_offset as Python int per ratio to avoid # tensor-scalar slice inside compiled function. @@ -840,15 +920,18 @@ def prepare_compressed_kv_metadata( r: self.cu_new_comp_kv_cuda[r][num_contexts].item() if num_contexts > 0 else 0 for r in self._compress_ratios_sorted } - self._compute_gen_compressed_position_ids( - self.past_kv_lens_cuda, - self.cu_new_comp_kv_cuda, - self.compressed_position_ids_cuda, + gen_comp_counts = [ + num_generations * ((self.num_gen_tokens_per_seq + r - 1) // r) for r in ratios + ] + torch.ops.trtllm.deepseek_v4_compute_gen_compressed_position_ids( + [self.past_kv_lens_cuda[r] for r in ratios], + [self.cu_new_comp_kv_cuda[r] for r in ratios], + [self.compressed_position_ids_cuda[r] for r in ratios], + ratios, + gen_comp_counts, + [int(gen_output_offsets[r]) for r in ratios], num_contexts, - num_generations, - self.num_gen_tokens_per_seq, - self._compress_ratios_sorted, - gen_output_offsets, + num_contexts + num_generations, ) def on_update_kv_lens(self): @@ -874,13 +957,13 @@ def on_update_kv_lens(self): ) self.prepare_compressed_kv_metadata(kv_lens, cached_tokens, ctx_output_sizes) - self._compute_compressed_mask( - self.new_comp_kv_lens_cuda, - self.cu_new_comp_kv_cuda, - self.compressed_mask_cuda, + ratios = self._compress_ratios_sorted + torch.ops.trtllm.deepseek_v4_compute_compressed_mask( + [self.new_comp_kv_lens_cuda[r] for r in ratios], + [self.cu_new_comp_kv_cuda[r] for r in ratios], + [self.compressed_mask_cuda[r] for r in ratios], + [int(self.num_total_compressed_tokens[r]) for r in ratios], batch_size, - self.num_total_compressed_tokens, - self._compress_ratios_sorted, ) token_positions = self._compute_token_positions( @@ -890,6 +973,7 @@ def on_update_kv_lens(self): num_tokens, self.cu_seq_lens_cuda, self.req_idx_per_token, + self.token_positions_cuda, ) self.prepare_for_deepseek_v4_indices(token_positions) @@ -930,8 +1014,25 @@ def _compute_token_positions( num_tokens: int, cu_seq_lens_buf: torch.Tensor, req_idx_per_token_buf: torch.Tensor, + token_positions_buf: torch.Tensor, ) -> torch.Tensor: - """Compute cu_seq_lens, req_idx_per_token, and token_positions (eager).""" + """Compute cu_seq_lens, req_idx_per_token, and token_positions.""" + if seq_lens.is_cuda: + # One kernel: scan + per-token binary search + position math, + # replacing ~8 ATen dispatches (pad/cumsum/arange/searchsorted/two + # gathers/two adds) and their temporaries. + token_positions = token_positions_buf[:num_tokens] + torch.ops.trtllm.compute_token_positions( + seq_lens, + cached_tokens, + cu_seq_lens_buf, + req_idx_per_token_buf, + token_positions, + num_tokens, + True, + ) + return token_positions + device = seq_lens.device # cu_seq_lens diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1bb0243f12ca..88e90df92144 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -91,6 +91,7 @@ l0_b200: - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15) + - unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py TIMEOUT (15) # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_mla_helix.py - unittest/_torch/modules/test_fused_add_rms_norm_quant.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4e7347edd644..a3117e32f62a 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -20,6 +20,7 @@ l0_b300: - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py TIMEOUT (15) - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/thop/serial - unittest/_torch/executor # 250s diff --git a/tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py b/tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py new file mode 100644 index 000000000000..c704a27a79e3 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host-path tests for the DeepSeek-V4 block-table preparation. + +These cover two behaviours that are easy to break and expensive to notice: + +1. ``_get_copy_index_cached`` must return the same mapping the uncached + ``IndexMapper.get_copy_index`` call would have produced, must recompute when + the request set changes (including in-place mutation of a reused list), and + must be invalidated at the start of every step. +2. The block-table destination buffers are only partially overwritten by the + copy that follows, so the untouched padding must still carry + ``BAD_PAGE_INDEX``. Padded CUDA-graph token slots can index those rows + through ``req_idx_per_token``, so leaving stale values there is a silent + correctness bug rather than a crash. +""" + +import torch + +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, +) +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX + + +class _RecordingIndexMapper: + """Stands in for the C++ IndexMapper, counting how often it is walked.""" + + def __init__(self): + self.calls = 0 + + def get_copy_index(self, request_ids, num_contexts, beam_width): + self.calls += 1 + # Mirror the real mapper's contract closely enough for these tests: one + # row per request, deterministic in the arguments. + return torch.tensor( + [(r * 10 + num_contexts + beam_width) for r in request_ids], + dtype=torch.int32, + ) + + +def _manager_with_mapper(): + mgr = DeepseekV4CacheManager.__new__(DeepseekV4CacheManager) + mgr.index_mapper = _RecordingIndexMapper() + return mgr + + +def test_copy_index_memo_reuses_one_walk_per_step(): + mgr = _manager_with_mapper() + request_ids = [3, 1, 4, 1, 5] + + mgr._copy_idx_memo_key = None + first = mgr._get_copy_index_cached(request_ids, 2, 1) + again = mgr._get_copy_index_cached(request_ids, 2, 1) + third = mgr._get_copy_index_cached(request_ids, 2, 1) + + assert mgr.index_mapper.calls == 1, "memo should collapse three walks into one" + torch.testing.assert_close(first, again, rtol=0, atol=0) + torch.testing.assert_close(first, third, rtol=0, atol=0) + + +def test_copy_index_memo_matches_uncached_result(): + """The memo must not change the value, only how often it is computed.""" + cached_mgr = _manager_with_mapper() + plain_mgr = _manager_with_mapper() + request_ids = [7, 2, 9] + + cached_mgr._copy_idx_memo_key = None + got = cached_mgr._get_copy_index_cached(request_ids, 1, 1) + expected = plain_mgr.index_mapper.get_copy_index(request_ids, 1, 1) + + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_copy_index_memo_recomputes_on_changed_arguments(): + mgr = _manager_with_mapper() + mgr._copy_idx_memo_key = None + + mgr._get_copy_index_cached([1, 2], 0, 1) + assert mgr.index_mapper.calls == 1 + # different request set + mgr._get_copy_index_cached([1, 3], 0, 1) + assert mgr.index_mapper.calls == 2 + # different num_contexts + mgr._get_copy_index_cached([1, 3], 1, 1) + assert mgr.index_mapper.calls == 3 + # different beam width + mgr._get_copy_index_cached([1, 3], 1, 2) + assert mgr.index_mapper.calls == 4 + + +def test_copy_index_memo_keys_on_contents_not_identity(): + """The caller may reuse one list object and mutate it between steps.""" + mgr = _manager_with_mapper() + reused = [1, 2, 3] + + mgr._copy_idx_memo_key = None + before = mgr._get_copy_index_cached(reused, 0, 1).clone() + assert mgr.index_mapper.calls == 1 + + reused[1] = 99 # same object, different contents + after = mgr._get_copy_index_cached(reused, 0, 1) + + assert mgr.index_mapper.calls == 2, "must not serve a stale mapping" + assert not torch.equal(before, after) + + +def test_copy_index_memo_is_reset_each_step(): + """Resetting the key is what makes the memo safe across steps.""" + mgr = _manager_with_mapper() + request_ids = [4, 5] + + mgr._copy_idx_memo_key = None + mgr._get_copy_index_cached(request_ids, 0, 1) + assert mgr.index_mapper.calls == 1 + + # start of the next step + mgr._copy_idx_memo_key = None + mgr._get_copy_index_cached(request_ids, 0, 1) + assert mgr.index_mapper.calls == 2 + + +def test_sliding_block_table_padding_keeps_bad_page_index(): + """Only [:_num_tables] is overwritten; the tail must stay BAD_PAGE_INDEX.""" + layers, types, capacity, max_blocks, num_tables = 2, 5, 9, 4, 3 + src = torch.arange(layers * types * num_tables * max_blocks, dtype=torch.int32).reshape( + layers, types, num_tables, max_blocks + ) + + # Reference: the original code filled the whole tensor first. + reference = torch.empty((layers, types, capacity, max_blocks), dtype=torch.int32) + reference.fill_(BAD_PAGE_INDEX) + reference[:, :, :num_tables, :].copy_(src) + + # Optimized: pre-poison everything, fill only the tail, then copy the head. + optimized = torch.full((layers, types, capacity, max_blocks), -999999, dtype=torch.int32) + if num_tables < optimized.size(2): + optimized[:, :, num_tables:, :].fill_(BAD_PAGE_INDEX) + optimized[:, :, :num_tables, :].copy_(src) + + torch.testing.assert_close(optimized, reference, rtol=0, atol=0) + assert not (optimized == -999999).any(), "head was not fully overwritten" + + +def test_sliding_block_table_padding_when_full(): + """With num_tables == capacity there is no tail to fill.""" + layers, types, capacity, max_blocks = 1, 2, 6, 3 + src = torch.arange(layers * types * capacity * max_blocks, dtype=torch.int32).reshape( + layers, types, capacity, max_blocks + ) + + reference = torch.empty((layers, types, capacity, max_blocks), dtype=torch.int32) + reference.fill_(BAD_PAGE_INDEX) + reference[:, :, :capacity, :].copy_(src) + + optimized = torch.full((layers, types, capacity, max_blocks), -999999, dtype=torch.int32) + if capacity < optimized.size(2): # false: guard must skip the fill + optimized[:, :, capacity:, :].fill_(BAD_PAGE_INDEX) + optimized[:, :, :capacity, :].copy_(src) + + torch.testing.assert_close(optimized, reference, rtol=0, atol=0) + + +def test_block_offsets_padding_keeps_bad_page_index(): + """copy_batch_block_offsets writes only beam 0 of [:_num_tables].""" + layers, capacity, beams, max_blocks, num_tables = 2, 7, 2, 3, 4 + src = torch.arange(layers * num_tables * max_blocks, dtype=torch.int32).reshape( + layers, num_tables, max_blocks + ) + + reference = torch.empty((layers, capacity, beams, max_blocks), dtype=torch.int32) + reference.fill_(BAD_PAGE_INDEX) + reference[:, :num_tables, 0, :].copy_(src) + + optimized = torch.full((layers, capacity, beams, max_blocks), -999999, dtype=torch.int32) + if num_tables < optimized.size(1): + optimized[:, num_tables:, :, :].fill_(BAD_PAGE_INDEX) + optimized[:, :num_tables, 1:, :].fill_(BAD_PAGE_INDEX) + optimized[:, :num_tables, 0, :].copy_(src) + + torch.testing.assert_close(optimized, reference, rtol=0, atol=0) + assert not (optimized == -999999).any(), "some region was left unwritten" diff --git a/tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py b/tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py new file mode 100644 index 000000000000..8c8539680845 --- /dev/null +++ b/tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Correctness tests for the DeepSeek-V4 metadata-preparation CUDA ops. + +Each op replaces a python/ATen reference that used to run in +``attn_metadata.prepare``. The references are reproduced verbatim here so the +tests pin the exact semantics rather than the current implementation. +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + +RATIO_SETS = [[1, 4, 128], [1, 128], [128], [1, 4], [1]] + + +# ---------------------------------------------------------------- references + + +def _ref_indices(token_positions, window_size, max_comp_128, sparse_mla_topk, ratios): + device = token_positions.device + num_tokens = token_positions.shape[0] + + positions = token_positions.unsqueeze(1) + swa_offsets = torch.arange(window_size, dtype=torch.int32, device=device) + swa_start = (positions - window_size + 1).clamp(min=0) + swa_indices = swa_start + swa_offsets + swa = torch.where(swa_indices > positions, -1, swa_indices).to(torch.int32) + + num_valid = (token_positions + 1) // 128 + comp_col = torch.arange(max_comp_128, dtype=torch.int32, device=device) + valid_mask = comp_col.unsqueeze(0) < num_valid.unsqueeze(1) + comp = torch.where( + valid_mask, + comp_col.unsqueeze(0).expand(num_tokens, -1), + torch.full((num_tokens, max_comp_128), -1, dtype=torch.int32, device=device), + ) + + kv_lens = token_positions + 1 + lens = {} + for r in ratios: + if r == 1: + total = kv_lens.clamp(max=window_size) + elif r == 4: + total = window_size + (kv_lens // r).clamp(max=sparse_mla_topk) + elif r == 128: + total = window_size + kv_lens // r + else: + raise ValueError(r) + lens[r] = total.to(torch.int32) + return swa, comp, lens + + +def _ref_per_ratio(kv_lens, cached_tokens, ratios): + out = {} + for r in ratios: + compressed = (kv_lens // r).to(torch.int32) + past = (cached_tokens // r).to(torch.int32) + new_comp = compressed - past + cu = torch.nn.functional.pad(torch.cumsum(new_comp, dim=0), (1, 0)).to(torch.int32) + out[r] = (compressed, past, new_comp, cu) + return out + + +def _ref_mask(new_comp, cu, total_tokens, batch_size, device): + token_idx = torch.arange(total_tokens, dtype=torch.int32, device=device) + seq_idx = torch.searchsorted(cu[1:], token_idx, right=True).clamp_(max=batch_size - 1) + return (token_idx - cu[seq_idx]) < new_comp[seq_idx] + + +def _ref_ctx_position_ids(past_kv, cu, total, ratio, num_contexts, device): + ctx_idx = torch.arange(total, dtype=torch.int32, device=device) + ctx_cu = cu[: num_contexts + 1].to(torch.int32) + ctx_req = torch.searchsorted(ctx_cu[1:], ctx_idx, right=True) + return ((past_kv[:num_contexts][ctx_req] + (ctx_idx - ctx_cu[ctx_req])) * ratio).to(torch.int32) + + +def _ref_gen_position_ids(past_kv, cu, gen_comp, offset, ratio, num_contexts, batch_size, device): + output_idx = torch.arange(gen_comp, dtype=torch.int32, device=device) + offset + req_idx = torch.searchsorted(cu[: batch_size + 1], output_idx, right=True) - 1 + req_idx = req_idx.clamp(min=num_contexts, max=batch_size - 1) + return ((past_kv[req_idx] + (output_idx - cu[req_idx])) * ratio).to(torch.int32) + + +def _ref_token_positions(seq_lens, cached_tokens, batch_size, num_tokens, device): + cu = torch.nn.functional.pad(torch.cumsum(seq_lens.to(torch.int), dim=0), (1, 0)).to( + torch.int32 + ) + token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) + req_idx = torch.searchsorted(cu[1 : batch_size + 1].to(torch.int32), token_idx, right=True) + positions = cached_tokens[req_idx].to(torch.int32) + (token_idx - cu[req_idx].to(torch.int32)) + return cu, req_idx.to(torch.int32), positions + + +# -------------------------------------------------------------------- tests + + +@pytest.mark.parametrize("ratios", RATIO_SETS) +@pytest.mark.parametrize( + "positions_spec", + [ + "zeros", + "boundaries", + "window_edge", + "topk_saturating", + "random_long", + "decode_like", + ], +) +@pytest.mark.parametrize("window_size,max_comp,topk", [(128, 256, 2048), (128, 8, 64), (64, 4, 16)]) +def test_compute_indices(ratios, positions_spec, window_size, max_comp, topk): + device = "cuda" + torch.manual_seed(0) + if positions_spec == "zeros": + tp = torch.zeros(5, dtype=torch.int32, device=device) + elif positions_spec == "boundaries": + tp = torch.tensor([0, 1, 127, 128, 255, 256, 511, 512], dtype=torch.int32, device=device) + elif positions_spec == "window_edge": + tp = torch.tensor( + [window_size - 2, window_size - 1, window_size, window_size + 1], + dtype=torch.int32, + device=device, + ) + elif positions_spec == "topk_saturating": + tp = torch.tensor( + [topk * 4 - 1, topk * 4, topk * 4 + 1, 100000], dtype=torch.int32, device=device + ) + elif positions_spec == "random_long": + tp = torch.randint(0, 40000, (129,), dtype=torch.int32, device=device) + else: + tp = torch.randint(1000, 4000, (256,), dtype=torch.int32, device=device) + + num_tokens = tp.shape[0] + ref_swa, ref_comp, ref_lens = _ref_indices(tp, window_size, max_comp, topk, ratios) + + # Deliberately over-allocate rows so the op must honour num_tokens, not the + # buffer shape (mirrors the CUDA-graph padded buffers in production). + swa = torch.full((num_tokens + 7, window_size), 123, dtype=torch.int32, device=device) + comp = torch.full((num_tokens + 7, max_comp), 123, dtype=torch.int32, device=device) + lens = {r: torch.full((num_tokens + 7,), 123, dtype=torch.int32, device=device) for r in ratios} + + torch.ops.trtllm.deepseek_v4_compute_indices( + tp, + window_size, + max_comp, + topk, + swa, + comp, + lens.get(1), + lens.get(4), + lens.get(128), + ) + + torch.testing.assert_close(swa[:num_tokens], ref_swa, rtol=0, atol=0) + torch.testing.assert_close(comp[:num_tokens], ref_comp, rtol=0, atol=0) + for r in ratios: + torch.testing.assert_close(lens[r][:num_tokens], ref_lens[r], rtol=0, atol=0) + # Padding rows must be untouched. + assert (swa[num_tokens:] == 123).all() + assert (comp[num_tokens:] == 123).all() + + +@pytest.mark.parametrize("ratios", RATIO_SETS) +@pytest.mark.parametrize("batch_size", [1, 2, 7, 37, 64, 146, 257, 300]) +def test_compute_per_ratio_kv_lens(ratios, batch_size): + device = "cuda" + torch.manual_seed(batch_size) + kv_lens = torch.randint(1, 40000, (batch_size,), dtype=torch.int32, device=device) + cached = (kv_lens * torch.rand(batch_size, device=device)).to(torch.int32) + ref = _ref_per_ratio(kv_lens, cached, ratios) + + pad = 5 + compressed = { + r: torch.zeros(batch_size + pad, dtype=torch.int32, device=device) for r in ratios + } + past = {r: torch.zeros(batch_size + pad, dtype=torch.int32, device=device) for r in ratios} + new_comp = {r: torch.zeros(batch_size + pad, dtype=torch.int32, device=device) for r in ratios} + cu = {r: torch.zeros(batch_size + 1 + pad, dtype=torch.int32, device=device) for r in ratios} + + torch.ops.trtllm.deepseek_v4_compute_per_ratio_kv_lens( + kv_lens, + cached, + ratios, + [compressed[r] for r in ratios], + [past[r] for r in ratios], + [new_comp[r] for r in ratios], + [cu[r] for r in ratios], + ) + + for r in ratios: + exp_c, exp_p, exp_n, exp_cu = ref[r] + torch.testing.assert_close(compressed[r][:batch_size], exp_c, rtol=0, atol=0) + torch.testing.assert_close(past[r][:batch_size], exp_p, rtol=0, atol=0) + torch.testing.assert_close(new_comp[r][:batch_size], exp_n, rtol=0, atol=0) + torch.testing.assert_close(cu[r][: batch_size + 1], exp_cu, rtol=0, atol=0) + + +@pytest.mark.parametrize("ratios", RATIO_SETS) +@pytest.mark.parametrize("batch_size", [1, 3, 37, 146, 257]) +def test_compute_compressed_mask(ratios, batch_size): + device = "cuda" + torch.manual_seed(batch_size + 11) + kv_lens = torch.randint(1, 8000, (batch_size,), dtype=torch.int32, device=device) + cached = (kv_lens * torch.rand(batch_size, device=device)).to(torch.int32) + ref = _ref_per_ratio(kv_lens, cached, ratios) + + totals = {r: int(ref[r][3][batch_size].item()) for r in ratios} + masks = {r: torch.zeros(max(totals[r], 1) + 8, dtype=torch.bool, device=device) for r in ratios} + torch.ops.trtllm.deepseek_v4_compute_compressed_mask( + [ref[r][2] for r in ratios], + [ref[r][3] for r in ratios], + [masks[r] for r in ratios], + [totals[r] for r in ratios], + batch_size, + ) + for r in ratios: + if totals[r] == 0: + continue + expected = _ref_mask(ref[r][2], ref[r][3], totals[r], batch_size, device) + torch.testing.assert_close(masks[r][: totals[r]], expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("ratios", RATIO_SETS) +@pytest.mark.parametrize("num_contexts", [1, 4, 33]) +def test_compute_ctx_compressed_position_ids(ratios, num_contexts): + device = "cuda" + torch.manual_seed(num_contexts + 7) + kv_lens = torch.randint(1, 6000, (num_contexts,), dtype=torch.int32, device=device) + cached = (kv_lens * torch.rand(num_contexts, device=device)).to(torch.int32) + ref = _ref_per_ratio(kv_lens, cached, ratios) + + counts = {r: int(ref[r][3][num_contexts].item()) for r in ratios} + out = {r: torch.zeros(max(counts[r], 1) + 8, dtype=torch.int32, device=device) for r in ratios} + torch.ops.trtllm.deepseek_v4_compute_ctx_compressed_position_ids( + [ref[r][1] for r in ratios], + [ref[r][3] for r in ratios], + [out[r] for r in ratios], + ratios, + [counts[r] for r in ratios], + num_contexts, + ) + for r in ratios: + if counts[r] == 0: + continue + expected = _ref_ctx_position_ids(ref[r][1], ref[r][3], counts[r], r, num_contexts, device) + torch.testing.assert_close(out[r][: counts[r]], expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("ratios", RATIO_SETS) +@pytest.mark.parametrize( + "num_contexts,num_generations,gen_tokens_per_seq", + [ + (0, 8, 1), + (0, 16, 4), + (2, 6, 1), + (5, 11, 4), + (0, 1, 1), + ], +) +def test_compute_gen_compressed_position_ids( + ratios, num_contexts, num_generations, gen_tokens_per_seq +): + device = "cuda" + batch_size = num_contexts + num_generations + torch.manual_seed(batch_size + gen_tokens_per_seq) + kv_lens = torch.randint(1, 6000, (batch_size,), dtype=torch.int32, device=device) + cached = (kv_lens * torch.rand(batch_size, device=device)).to(torch.int32) + ref = _ref_per_ratio(kv_lens, cached, ratios) + + counts, offsets = {}, {} + for r in ratios: + counts[r] = num_generations * ((gen_tokens_per_seq + r - 1) // r) + offsets[r] = int(ref[r][3][num_contexts].item()) if num_contexts > 0 else 0 + + out = { + r: torch.zeros(offsets[r] + counts[r] + 8, dtype=torch.int32, device=device) for r in ratios + } + torch.ops.trtllm.deepseek_v4_compute_gen_compressed_position_ids( + [ref[r][1] for r in ratios], + [ref[r][3] for r in ratios], + [out[r] for r in ratios], + ratios, + [counts[r] for r in ratios], + [offsets[r] for r in ratios], + num_contexts, + batch_size, + ) + for r in ratios: + if counts[r] == 0: + continue + expected = _ref_gen_position_ids( + ref[r][1], ref[r][3], counts[r], offsets[r], r, num_contexts, batch_size, device + ) + torch.testing.assert_close( + out[r][offsets[r] : offsets[r] + counts[r]], expected, rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("seq_lens_spec", ["uniform_1", "mixed", "single", "long_ctx"]) +@pytest.mark.parametrize("batch_size", [1, 2, 8, 64, 146, 257]) +def test_compute_token_positions(seq_lens_spec, batch_size): + device = "cuda" + torch.manual_seed(batch_size + 3) + if seq_lens_spec == "uniform_1": + seq_lens = torch.ones(batch_size, dtype=torch.int32, device=device) + elif seq_lens_spec == "mixed": + seq_lens = torch.randint(1, 40, (batch_size,), dtype=torch.int32, device=device) + elif seq_lens_spec == "single": + seq_lens = torch.full((batch_size,), 7, dtype=torch.int32, device=device) + else: + seq_lens = torch.randint(1, 600, (batch_size,), dtype=torch.int32, device=device) + + num_tokens = int(seq_lens.sum().item()) + cached = torch.randint(0, 30000, (batch_size,), dtype=torch.int32, device=device) + ref_cu, ref_req, ref_pos = _ref_token_positions( + seq_lens, cached, batch_size, num_tokens, device + ) + + cu = torch.zeros(batch_size + 1 + 4, dtype=torch.int32, device=device) + req = torch.zeros(num_tokens + 4, dtype=torch.int32, device=device) + pos = torch.zeros(num_tokens + 4, dtype=torch.int32, device=device) + + torch.ops.trtllm.compute_token_positions(seq_lens, cached, cu, req, pos, num_tokens, True) + torch.testing.assert_close(cu[: batch_size + 1], ref_cu, rtol=0, atol=0) + torch.testing.assert_close(req[:num_tokens], ref_req, rtol=0, atol=0) + torch.testing.assert_close(pos[:num_tokens], ref_pos, rtol=0, atol=0) + + # req_idx must equal the CPU repeat_interleave form the base class used. + expected_req = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens + ) + torch.testing.assert_close(req[:num_tokens], expected_req, rtol=0, atol=0) + + # reuse mode: cu already populated, only the per-token phase runs + req2 = torch.zeros(num_tokens + 4, dtype=torch.int32, device=device) + pos2 = torch.zeros(num_tokens + 4, dtype=torch.int32, device=device) + torch.ops.trtllm.compute_token_positions(seq_lens, cached, cu, req2, pos2, num_tokens, False) + torch.testing.assert_close(req2[:num_tokens], ref_req, rtol=0, atol=0) + torch.testing.assert_close(pos2[:num_tokens], ref_pos, rtol=0, atol=0) + + # req-only mode: token_positions omitted + req3 = torch.zeros(num_tokens + 4, dtype=torch.int32, device=device) + torch.ops.trtllm.compute_token_positions(seq_lens, None, cu, req3, None, num_tokens, True) + torch.testing.assert_close(req3[:num_tokens], ref_req, rtol=0, atol=0) + + +def test_compute_indices_rejects_bad_shapes(): + device = "cuda" + tp = torch.zeros(4, dtype=torch.int32, device=device) + swa = torch.zeros((4, 8), dtype=torch.int32, device=device) + comp = torch.zeros((4, 4), dtype=torch.int32, device=device) + # window_size larger than the buffer's columns + with pytest.raises(RuntimeError): + torch.ops.trtllm.deepseek_v4_compute_indices(tp, 16, 4, 8, swa, comp, None, None, None) + # wrong dtype + tp_f = torch.zeros(4, dtype=torch.float32, device=device) + with pytest.raises(RuntimeError): + torch.ops.trtllm.deepseek_v4_compute_indices(tp_f, 8, 4, 8, swa, comp, None, None, None) + + +def test_per_ratio_rejects_mismatched_lists(): + device = "cuda" + kv = torch.ones(4, dtype=torch.int32, device=device) + cached = torch.zeros(4, dtype=torch.int32, device=device) + buf = torch.zeros(4, dtype=torch.int32, device=device) + cu = torch.zeros(5, dtype=torch.int32, device=device) + # two ratios but only one buffer per list + with pytest.raises(RuntimeError): + torch.ops.trtllm.deepseek_v4_compute_per_ratio_kv_lens( + kv, cached, [1, 4], [buf], [buf], [buf], [cu] + ) + + +def _ref_shared_block_table(block_offsets, pool_id, copy_idx, scale, bad_page_index=-1): + """Reference for ``DeepseekV4CacheManager._compute_shared_block_table``. + + The python path gathered ``block_offsets[pool_id, copy_idx, 0, :]`` on the host + and mapped it with ``where(base == BAD_PAGE_INDEX, BAD_PAGE_INDEX, base * scale)``. + """ + base = block_offsets[pool_id, copy_idx, 0, :] + return torch.where(base == bad_page_index, bad_page_index, base * scale) + + +@pytest.mark.parametrize("scale", [1, 4, 128]) +@pytest.mark.parametrize("num_pools,capacity,max_blocks", [(1, 8, 4), (3, 64, 16), (7, 130, 129)]) +@pytest.mark.parametrize("num_seqs", [1, 5, 64]) +def test_compute_shared_block_table(scale, num_pools, capacity, max_blocks, num_seqs): + device = "cuda" + torch.manual_seed(num_seqs * 31 + max_blocks) + if num_seqs > capacity: + pytest.skip("num_seqs must fit the mapper capacity") + + block_offsets = torch.randint( + 0, 1 << 20, (num_pools, capacity, 2, max_blocks), dtype=torch.int32, device=device + ) + # Sprinkle BAD_PAGE_INDEX so the passthrough branch is exercised, including a + # fully-invalid row and a fully-valid row. + block_offsets[block_offsets % 5 == 0] = -1 + block_offsets[:, 0, 0, :] = -1 + if capacity > 1: + block_offsets[:, 1, 0, :] = 7 + + copy_idx = torch.randperm(capacity, device=device)[:num_seqs].to(torch.int32) + pool_id = num_pools - 1 + + out = torch.full((num_seqs, max_blocks), 12345, dtype=torch.int32, device=device) + torch.ops.trtllm.compute_shared_block_table(block_offsets, copy_idx, pool_id, scale, out) + + expected = _ref_shared_block_table(block_offsets, pool_id, copy_idx.long(), scale) + torch.testing.assert_close(out, expected.to(torch.int32), rtol=0, atol=0) + + +def test_compute_shared_block_table_leaves_padding_untouched(): + """Rows past num_seqs must not be written: padded CUDA-graph slots read them.""" + device = "cuda" + block_offsets = torch.ones((2, 16, 2, 8), dtype=torch.int32, device=device) + copy_idx = torch.arange(4, dtype=torch.int32, device=device) + out = torch.full((16, 8), -7, dtype=torch.int32, device=device) + + torch.ops.trtllm.compute_shared_block_table(block_offsets, copy_idx, 1, 4, out[:4]) + + torch.testing.assert_close( + out[4:], torch.full((12, 8), -7, dtype=torch.int32, device=device), rtol=0, atol=0 + ) + torch.testing.assert_close( + out[:4], torch.full((4, 8), 4, dtype=torch.int32, device=device), rtol=0, atol=0 + ) + + +# The shared-memory prefix scans pick a template instantiation from a size ladder +# (<=512, <=2048, else the compile-time bound). Batch sizes that straddle those +# boundaries take different code paths, and the largest tier is only safe because +# the op layer rejects anything above the bound. Both need coverage. + +_SCAN_TIER_BATCHES = [511, 512, 513, 2047, 2048, 2049, 4096] + + +@pytest.mark.parametrize("batch_size", _SCAN_TIER_BATCHES) +def test_token_positions_across_scan_tiers(batch_size): + device = "cuda" + torch.manual_seed(batch_size) + seq_lens = torch.randint(1, 5, (batch_size,), dtype=torch.int32, device=device) + num_tokens = int(seq_lens.sum().item()) + cached = torch.randint(0, 1000, (batch_size,), dtype=torch.int32, device=device) + + ref_cu, ref_req, ref_pos = _ref_token_positions( + seq_lens, cached, batch_size, num_tokens, device + ) + cu = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + req = torch.zeros(num_tokens, dtype=torch.int32, device=device) + pos = torch.zeros(num_tokens, dtype=torch.int32, device=device) + + torch.ops.trtllm.compute_token_positions(seq_lens, cached, cu, req, pos, num_tokens, True) + torch.testing.assert_close(cu, ref_cu, rtol=0, atol=0) + torch.testing.assert_close(req, ref_req, rtol=0, atol=0) + torch.testing.assert_close(pos, ref_pos, rtol=0, atol=0) + + +@pytest.mark.parametrize("batch_size", _SCAN_TIER_BATCHES) +def test_per_ratio_kv_lens_across_scan_tiers(batch_size): + device = "cuda" + torch.manual_seed(batch_size + 7) + ratios = [1, 4, 128] + kv_lens = torch.randint(1, 9000, (batch_size,), dtype=torch.int32, device=device) + cached = (kv_lens * torch.rand(batch_size, device=device)).to(torch.int32) + ref = _ref_per_ratio(kv_lens, cached, ratios) + + comp = {r: torch.zeros(batch_size, dtype=torch.int32, device=device) for r in ratios} + past = {r: torch.zeros(batch_size, dtype=torch.int32, device=device) for r in ratios} + new_comp = {r: torch.zeros(batch_size, dtype=torch.int32, device=device) for r in ratios} + cu = {r: torch.zeros(batch_size + 1, dtype=torch.int32, device=device) for r in ratios} + + torch.ops.trtllm.deepseek_v4_compute_per_ratio_kv_lens( + kv_lens, + cached, + ratios, + [comp[r] for r in ratios], + [past[r] for r in ratios], + [new_comp[r] for r in ratios], + [cu[r] for r in ratios], + ) + for r in ratios: + exp_c, exp_p, exp_n, exp_cu = ref[r] + torch.testing.assert_close(comp[r], exp_c, rtol=0, atol=0) + torch.testing.assert_close(past[r], exp_p, rtol=0, atol=0) + torch.testing.assert_close(new_comp[r], exp_n, rtol=0, atol=0) + torch.testing.assert_close(cu[r], exp_cu, rtol=0, atol=0) + + +def test_token_positions_rejects_batch_above_scan_bound(): + """Above the compile-time bound the op must fail loudly, not scribble.""" + device = "cuda" + too_big = 4097 + seq_lens = torch.ones(too_big, dtype=torch.int32, device=device) + cached = torch.zeros(too_big, dtype=torch.int32, device=device) + cu = torch.zeros(too_big + 1, dtype=torch.int32, device=device) + req = torch.zeros(too_big, dtype=torch.int32, device=device) + pos = torch.zeros(too_big, dtype=torch.int32, device=device) + + with pytest.raises(RuntimeError): + torch.ops.trtllm.compute_token_positions(seq_lens, cached, cu, req, pos, too_big, True) + + # With the scan skipped (cu_seq_lens supplied by the caller) the bound does + # not apply, so the same batch must go through. + cu_ref = torch.zeros(too_big + 1, dtype=torch.int32, device=device) + cu_ref[1:] = torch.cumsum(seq_lens, 0, dtype=torch.int32) + torch.ops.trtllm.compute_token_positions(seq_lens, cached, cu_ref, req, pos, too_big, False) + expected_req = torch.arange(too_big, dtype=torch.int32, device=device) + torch.testing.assert_close(req, expected_req, rtol=0, atol=0) + + +def test_per_ratio_kv_lens_rejects_batch_above_scan_bound(): + device = "cuda" + too_big = 4097 + kv = torch.ones(too_big, dtype=torch.int32, device=device) + cached = torch.zeros(too_big, dtype=torch.int32, device=device) + buf = torch.zeros(too_big, dtype=torch.int32, device=device) + cu = torch.zeros(too_big + 1, dtype=torch.int32, device=device) + + with pytest.raises(RuntimeError): + torch.ops.trtllm.deepseek_v4_compute_per_ratio_kv_lens( + kv, cached, [1], [buf], [buf], [buf], [cu] + )