[None][perf] DeepSeek-V4: cut host overhead in generation-phase metadata preparation

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 <algorithm>
+#include <cstdint>
+
+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 <int kMaxBatch>
+__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<int32_t>(blockDim.x);
+    int32_t const rounded = ((batchSize + stride - 1) / stride) * stride;
+
+    for (int32_t i = static_cast<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(blockIdx.y);
+    if (tableId >= numTables)
+    {
+        return;
+    }
+
+    int64_t const outputOffset = static_cast<int64_t>(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<int64_t>(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq;
+
+    for (int32_t blockId
+         = static_cast<int32_t>(blockIdx.x) * static_cast<int32_t>(blockDim.x) + static_cast<int32_t>(threadIdx.x);
+         blockId < maxBlocksPerSeq; blockId += static_cast<int32_t>(gridDim.x) * static_cast<int32_t>(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<uint32_t>(kThreadsPerBlock));
+        if (batchSize <= 512)
+        {
+            computeCuSeqLensKernel<512><<<grid, block, 0, stream>>>(seqLens, cuSeqLens, batchSize);
+        }
+        else if (batchSize <= 2048)
+        {
+            computeCuSeqLensKernel<2048><<<grid, block, 0, stream>>>(seqLens, cuSeqLens, batchSize);
+        }
+        else
+        {
+            computeCuSeqLensKernel<kMaxTokenPositionScanBatch>
+                <<<grid, block, 0, stream>>>(seqLens, cuSeqLens, batchSize);
+        }
+    }
+    if (numTokens > 0)
+    {
+        int32_t const blocks = std::min((numTokens + kThreadsPerBlock - 1) / kThreadsPerBlock, 2048);
+        computeTokenPositionsKernel<<<blocks, kThreadsPerBlock, 0, stream>>>(
+            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<uint32_t>(threadsPerBlock));
+    dim3 const grid(static_cast<uint32_t>(blocksPerRow), static_cast<uint32_t>(numTables));
+    computeSharedBlockTableKernel<<<grid, block, 0, stream>>>(
+        blockOffsets, copyIdx, output, poolId, scale, copyIdxCapacity, numTables, maxBlocksPerSeq);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

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 <cstdint>
+#include <cuda_runtime_api.h>
+
+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

cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu

@@ -16,6 +16,7 @@
 
 #include "tensorrt_llm/kernels/deepseekV4BlockTable.h"
 
+#include <algorithm>
 #include <cstdint>
 
 TRTLLM_NAMESPACE_BEGIN