[None][feat] Add DSpark confidence-head verification scheduling

cpp/tensorrt_llm/kernels/IndexerTopK.h

@@ -40,11 +40,18 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW
 ///   - Insertion sort   (N < kSortingAlgorithmThreshold)
 ///   - Radix sort       (kSortingAlgorithmThreshold ≤ N < splitWork)
 ///   - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux)
+///
+/// @param rowKvLens  optional [numRows] per-row causal KV extent, in uncompressed
+///   token space. Non-null selects the ragged layout: rows no longer map to
+///   requests by `rowIdx / next_n`, so the extent every row may attend to is
+///   supplied directly instead of being reconstructed from `next_n`. Null keeps
+///   the uniform arithmetic, bit-identical to before. Mutually exclusive with
+///   `preIdx` (GVR's hint is request-indexed through `next_n` as well).
 void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux,
     int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0,
     int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0,
     int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1,
-    cudaStream_t const stream = 0);
+    int const* rowKvLens = nullptr, cudaStream_t const stream = 0);
 
 /// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except
 /// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and
@@ -59,13 +66,14 @@ void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, in
     int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1,
     int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0,
     int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1,
-    cudaStream_t const stream = 0);
+    int const* rowKvLens = nullptr, cudaStream_t const stream = 0);
 
 /// fp16 indexer TopK decode — see bf16 overload for dispatcher contract.
 void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold,
     int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n,
     int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0,
-    __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0);
+    __half* heuristicScratch = nullptr, int const compressRatio = 1, int const* rowKvLens = nullptr,
+    cudaStream_t const stream = 0);
 
 void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices,
     int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048,

cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu

@@ -297,7 +297,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
     void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv,
     int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens,
     int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks,
-    int out_elem_bytes)
+    int out_elem_bytes, int32_t const* __restrict__ new_tokens_per_seq = nullptr)
 {
     using KvScoreElemT = typename std::conditional<KV_SCORE_ELEM_BYTES == 2, __nv_bfloat16, float>::type;
     using StateElemT = typename std::conditional<STATE_ELEM_BYTES == 2, __nv_bfloat16, float>::type;
@@ -337,7 +337,16 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
     int const eff_tid = head_blk * NTHRD_INNER + tid;
 
     int const kv_len = kv_lens[batch_idx];
-    int const sp = kv_len - NEXT_N;
+    // How many tokens this request appended this step. Uniform batches append
+    // NEXT_N each, which is what the template bound encodes. Under DSpark's
+    // ragged verification each request appends its own window instead, so the
+    // count is read per request and NEXT_N degrades to a compile-time *upper*
+    // bound -- which is all the loops below need it to be, because both are
+    // already guarded (Phase 1 by token_idx < kv_len, Phase 3 by
+    // num_compressions). Keeping NEXT_N as the bound is what preserves the
+    // full unroll on the uniform path.
+    int const nn = (new_tokens_per_seq != nullptr) ? new_tokens_per_seq[batch_idx] : NEXT_N;
+    int const sp = kv_len - nn;
     int const in_off = cu_seq_lens[batch_idx];
     int const out_off = cu_kv_comp[batch_idx];
     int64_t const page_sd = static_cast<int64_t>(page_size) * STATE_DIM;
@@ -413,7 +422,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
     // ================================================================
     // Phase 2: Count how many complete compression windows finished.
     // ================================================================
-    int last_token_idx = sp + NEXT_N - 1;
+    int last_token_idx = sp + nn - 1;
     int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO;
 
     // ================================================================
@@ -663,7 +672,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
 // Generate explicit template instantiations.
 #define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW)                                                                  \
     template __global__ void pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NN, NRW>(void const*, float const*, void*, \
-        void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int);
+        void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int,   \
+        int32_t const*);
 FOREACH_DECODE_CONFIG(INST_DECODE)
 #undef INST_DECODE
 
@@ -680,6 +690,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k
     int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens,
     int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim,
     int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes,
+    int32_t const* new_tokens_per_seq,
     cudaStream_t stream)
 {
     TLLM_CHECK_WITH_INFO(
@@ -731,7 +742,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k
     {                                                                                                                  \
         pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NN, NRW><<<grid, nthreads, smem_bytes, stream>>>(kv_score, ape, \
             paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp,        \
-            page_size, max_blocks, out_elem_bytes);                                                                    \
+            page_size, max_blocks, out_elem_bytes, new_tokens_per_seq);                                                \
         return;                                                                                                        \
     }
     FOREACH_DECODE_CONFIG(TRY_LAUNCH)

cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h

@@ -47,6 +47,10 @@ void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim]  (bf16 or f
     int kv_score_elem_bytes,                     // bytes per element for kv_score (2=bf16, 4=fp32)
     int state_elem_bytes,                        // bytes per element for paged state (2=bf16, 4=fp32)
     int out_elem_bytes,                          // bytes per element for output
+    int32_t const* new_tokens_per_seq,           // [bsz] or nullptr; per-request tokens appended this
+                                                 // step. nullptr means every request appended next_n
+                                                 // (the uniform case); non-null makes next_n an upper
+                                                 // bound only, for DSpark ragged verification.
     cudaStream_t stream);
 
 // Prefill kernel: bulk compression with per-token gather/scatter + state update.

cpp/tensorrt_llm/kernels/indexerTopK.cu

@@ -654,7 +654,8 @@ template <int kNumThreadsPerBlock, bool useRadixSort, bool multipleBlocksPerRow
     typename InputT = float>
 static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens,
     int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio,
-    float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr)
+    float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr,
+    int const* rowKvLens = nullptr)
 {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     cudaGridDependencySynchronize();
@@ -664,8 +665,27 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I
 
     // The range of logits within the row.
     int rowStart = 0;
-    int seq_len = seqLens[rowIdx / next_n];
-    int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1;
+    // Two ways to learn how far back this query row may attend.
+    //
+    // Uniform (rowKvLens == nullptr): every request contributes exactly next_n
+    // rows, so the request index and the row's offset inside its verify window
+    // both fall out of integer division, and the causal extent is the request's
+    // sequence length walked back to this offset.
+    //
+    // Ragged (rowKvLens != nullptr): requests contribute different numbers of
+    // rows, so neither rowIdx / next_n nor rowIdx % next_n means anything. The
+    // caller precomputes the extent per row -- it is the only quantity this
+    // kernel ever derived from next_n, which is why one array is enough.
+    int actual_kv_len;
+    if (rowKvLens != nullptr)
+    {
+        actual_kv_len = rowKvLens[rowIdx];
+    }
+    else
+    {
+        int seq_len = seqLens[rowIdx / next_n];
+        actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1;
+    }
     int rowEnd = actual_kv_len / compressRatio;
 
     // Local pointers to this block
@@ -890,9 +910,16 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW
 void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux,
     int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0,
     int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride,
-    int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream)
+    int const preIdxCount, float* heuristicScratch, int const compressRatio, int const* rowKvLens,
+    cudaStream_t const stream)
 {
     constexpr int kNumThreadsPerBlock = 512;
+    // GVR consumes previous-step Top-K hints indexed by request, which it also
+    // recovers from next_n. Ragged batches therefore cannot take that path; the
+    // caller is expected to have disabled the hint, and this makes the
+    // requirement explicit rather than letting a stale preIdx through.
+    TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr,
+        "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint");
     int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold;
 
     // ========================================================================
@@ -1016,7 +1043,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
         config.attrs = attrs;
 
         cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n,
-            compressRatio, nullptr, 0, nullptr);
+            compressRatio, nullptr, 0, nullptr, rowKvLens);
     }
     else
     {
@@ -1031,7 +1058,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
         config_part1.attrs = attrs;
 
         cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK,
-            next_n, compressRatio, outLogitsAux, 0, nullptr);
+            next_n, compressRatio, outLogitsAux, 0, nullptr, rowKvLens);
 
         constexpr int kNumThreadsPerBlockMerge = 1024;
         auto* kernel_instance_part2 = &topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>;
@@ -1043,8 +1070,12 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
         config_part2.numAttrs = 1;
         config_part2.attrs = attrs;
 
+        // The merge pass overrides rowEnd with numBlocksToMerge * topK, so the
+        // extent is unused -- but rowKvLens still has to be forwarded, because
+        // otherwise the uniform branch would index seqLens by rowIdx / next_n,
+        // and on a ragged batch next_n is not a divisor of the row count.
         cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1,
-            topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux);
+            topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux, rowKvLens);
     }
     sync_check_cuda_error(stream);
 }
@@ -1079,10 +1110,14 @@ template <typename InputT>
 void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold,
     int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK,
     int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio,
-    cudaStream_t const stream)
+    int const* rowKvLens, cudaStream_t const stream)
 {
     static_assert(std::is_same_v<InputT, __nv_bfloat16> || std::is_same_v<InputT, __half>,
         "invokeIndexerTopKDecodeDtype is for bf16/fp16 only");
+    // See the fp32 entry: GVR's hint is request-indexed via next_n, so it and a
+    // ragged row map cannot both be live.
+    TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr,
+        "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint");
 
     constexpr int kNumThreadsPerBlock = 512;
     int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold;
@@ -1123,7 +1158,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int*
         config.attrs = attrs;
 
         cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n,
-            compressRatio, nullptr, 0, nullptr);
+            compressRatio, nullptr, 0, nullptr, rowKvLens);
     }
     else if (numColumns < effectiveSplitWorkThreshold)
     {
@@ -1143,7 +1178,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int*
         config.attrs = attrs;
 
         cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n,
-            compressRatio, nullptr, 0, nullptr);
+            compressRatio, nullptr, 0, nullptr, rowKvLens);
     }
     else
     {
@@ -1162,19 +1197,20 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int*
 void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices,
     int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1,
     int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount,
-    __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream)
+    __nv_bfloat16* heuristicScratch, int const compressRatio, int const* rowKvLens, cudaStream_t const stream)
 {
     invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns,
-        stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream);
+        stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens,
+        stream);
 }
 
 void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold,
     int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK,
     int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio,
-    cudaStream_t const stream)
+    int const* rowKvLens, cudaStream_t const stream)
 {
     invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0,
-        stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream);
+        stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens, stream);
 }
 
 void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices,