diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 324e597dc779..6e6e7b29b059 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,47 +27,63 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -/// Indexer TopK decode. Three tiers: -/// - GVR Heuristic (preIdx provided, K in {512,1024,2048}, numColumns in -/// [kSeqSmall, splitWorkThreshold), numRows below the -/// architecture-derived wave/L2 bound). -/// - Single-block (numColumns < split-work threshold) -/// - Multi-pass radix (numColumns >= split-work threshold; requires -/// `scratch` sized via indexerTopKDecodeScratchBytes, -/// zero-init on first call and may be reused). -/// -/// `is_prefill = true` forces single-block (split-work suppressed). -void invokeIndexerTopKDecode(float 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, - float* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); +// Number of blocks-per-row used by the multi-block split + merge dispatch path of +// invokeIndexerTopKDecode. Returns 1 when the single-block path is preferred. +// Callers that allocate aux buffers must use this same helper to size them, and +// must pass the same splitWorkThreshold they will pass to invokeIndexerTopKDecode +// (a value <= 0 selects the internal default). +int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0); -/// Size of the multi-pass radix `scratch` buffer for these shapes. -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK); +/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four +/// fallback tiers: +/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048}) +/// - Insertion sort (N < kSortingAlgorithmThreshold) +/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) +/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) +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); -/// bf16 overload; same contract. +/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except +/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and +/// the split-work tier is unsupported (the bf16/fp16 entry does not expose +/// the float aux buffers required for split-work). Insertion + radix tiers +/// share topKPerRowDecode with fp32 — histogram and sort run on float keys +/// after a static_cast(InputT) at HBM-read sites. +/// +/// Aborts with TLLM_CHECK if numColumns ≥ splitWorkThreshold; callers in +/// that regime must use the fp32 entry. 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 = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, cudaStream_t const stream = 0, - void* scratch = nullptr, size_t scratchBytes = 0, bool is_prefill = false); + int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, + cudaStream_t const stream = 0); -/// fp16 overload; same contract. +/// 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, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); + __half* heuristicScratch = nullptr, int const compressRatio = 1, 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, cudaStream_t const stream = 0); -/// True iff invokeIndexerTopKDecode would pick the GVR tier for this shape: -/// K in {512,1024,2048}, numColumns in [kSeqSmall, splitWorkThreshold), and -/// numRows below the architecture-derived wave/L2 bound. Lets callers -/// provision preIdx / heuristicScratch only when needed. +/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic +/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx +/// is provided and stride1 == 1. Useful for callers that need to provision a +/// preIdx tensor or heuristicScratch buffer only when GVR will be selected. +/// +/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048}, +/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where +/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem. +/// +/// @param numRows logits rows (batch · next_n) +/// @param numColumns logits columns (max sequence length) +/// @param topK requested output size +/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16) bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu index 25b0de3ef0a5..839ae5483c05 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -47,15 +47,19 @@ using heuristic_topk::KernelSmemTplK; // same kernel template. Smem layout is derived from GvrParams // at compile time. template -__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float const* __restrict__ logits, - int const* __restrict__ seqLens, int const* __restrict__ preIdx, float* __restrict__ scratchValues, - int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) +__global__ void __launch_bounds__(BLOCK_SIZE) + heuristicTopKMultiRowKernel(float const* __restrict__ logits, int const* __restrict__ seqLens, + int const* __restrict__ preIdx, float* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, + int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) { using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - int const N = seq_len - next_n + (rowIdx % next_n) + 1; + // seqLens is in uncompressed token space; the logits/preIdx live in + // compressed-index space when compressRatio > 1 (DSv4 indexer). + int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int const N = actual_kv_len / compressRatio; float const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -84,9 +88,19 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float return; } - // +1 accounts for the temporal shift: prev_topk indices were computed at - // seq_len-1, but the current step has one additional KV token appended. - int const preIdxOffset = (rowIdx % next_n) + 1; + // Temporal-shift offset to map prev-step's top-K indices into this step's + // KV index space. + // compressRatio == 1 (DSv3.2): +1 — KV grew by exactly 1 token per + // decode step; prev indices were at seq_len-1 so a uniform +1 maps + // them to the equivalent positions under the indexer's "newest-first" + // layout. The (rowIdx % next_n) addend extends this to MTP windows. + // compressRatio == 4 (DSv4): 0 — in compressed-index space new + // compressed entries are appended at the end; prev indices in + // [0, c_prev-1] remain valid as-is. Per-row Δc varies (0 or 1) with + // prev kv_len mod 4 alignment, but a uniform offset of 0 stays + // within-bounds for all rows and preserves the temporal-correlation + // hint (vertical top-K consistency validated offline). + int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; gvrTopKJob(input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -103,16 +117,18 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float // Templated on (InputT, TopK). Smem layout is derived from // GvrParams. template -__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, - int const* __restrict__ seqLens, int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, - int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) +__global__ void __launch_bounds__(BLOCK_SIZE) + heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, int const* __restrict__ seqLens, + int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, + int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) { // dtype path uses fp32 keys[] in smem (down-conversion deferred to writeback). using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - int const N = seq_len - next_n + (rowIdx % next_n) + 1; + int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int const N = actual_kv_len / compressRatio; InputT const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -142,7 +158,8 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I return; } - int const preIdxOffset = (rowIdx % next_n) + 1; + // See fp32 path: cr==1 → (rowIdx % next_n)+1; cr!=1 (DSv4) → 0. + int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; gvrTopKJobDtype( input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -152,24 +169,25 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I // Explicit instantiations — 6 (dtype × K) combos. Launchers dispatch on // runtime topK via switch, so all 6 must be available at link time. +// Trailing `int` is the compressRatio parameter (1 = V3.2, 4 = V4 indexer). template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 1024>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<512>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<1024>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<2048>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); // Dispatch on topK at runtime — each TopK-instantiation gets its own smem // size (driven by GvrParams::kC/kNumBins) and own kfn pointer @@ -184,7 +202,7 @@ template __global__ void heuristicTopKMultiRowKernel<2048>( template void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int const* preIdx, int* outIndices, InputT* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { TLLM_CHECK_WITH_INFO( topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}"); @@ -224,7 +242,7 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int config.attrs = attrs; cudaLaunchKernelEx(&config, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK, - preIdxStride, preIdxCount); + preIdxStride, preIdxCount, compressRatio); }; switch (topK) @@ -240,26 +258,26 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); } void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, - topK, preIdxStride, preIdxCount, numRows, stream); + topK, preIdxStride, preIdxCount, numRows, compressRatio, stream); } void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h index 5e38c64cba06..0d2330f76545 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -33,21 +33,28 @@ inline constexpr int kHeuristicSize = 2048; /// Launch heuristic TopK decode kernel — fp32 input. /// @param scratchValues Caller-owned buffer of size [numRows * topK] floats. /// Required for CUDA Graph compatibility — must have a stable device address. +/// @param compressRatio KV compression ratio (1 = V3.2 indexer; 4 = V4 indexer +/// whose logits/preIdx live in compressed-token-index space). For +/// compressRatio != 1, preIdxOffset is forced to 0 (append-at-end in +/// compressed space → prev-step indices remain valid as-is); the +/// existing (rowIdx % next_n)+1 shift is used only when compressRatio==1. void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); /// Launch heuristic TopK decode kernel — bf16 input. /// scratchValues is [numRows * topK] of bf16 (matches input dtype). +/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); /// Launch heuristic TopK decode kernel — fp16 input. /// scratchValues is [numRows * topK] of fp16 (matches input dtype). +/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index ad1e728b3298..b35531a480a2 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -39,6 +39,13 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +// Radix histogram bin count, used by topKPerRowJob's per-step distribution +// pass. 2048 = 2^11 maps directly to the fp16 fast path's 11-bit bin index +// (sign + exponent + top mantissa bits, mask 0x7FF). All file-scope dispatch +// thresholds below derive from this so that a future change to the histogram +// width does not require re-tuning the heuristics. +constexpr int kNumBins = 2048; + namespace { @@ -169,6 +176,20 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx, int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues, SmemFinalType& smemFinal, int stride1, int rowStart, int topK) { + // Step 0 is the fp16 fast path; if it could not resolve top-K (threshold bin + // exceeded kNumFinalItems) we restart the fp32 radix from scratch in steps 1-3. + // Discard any candidates step 0 wrote into smemOutput so step 1 doesn't + // double-count valid entries that fall under both fp16 and fp32 thresholds + // (this is what produced 2x duplicated indices when many -FLT_MAX padding + // entries dominated the threshold bin in the multi-block merge path). + if constexpr (step == 1) + { + if (threadIdx.x == 0) + { + smemFoundTopKValues[0] = 0; + smemFinalDstIdx[0] = 0; + } + } // Clear the histogram. #pragma unroll for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) @@ -271,18 +292,13 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; - // Skip auto-promote at step 0 when we'll continue: half-precision bins - // don't align with step 2's full-precision bit-pattern filter, so a - // step-0 promote would be double-counted at step 2. - bool const step0WillContinue = (step == 0) && (smemFinalBinSize[0] > kNumFinalItems); - auto processBins = [&](InputT logitIn, int idx) { float const logit = static_cast(logitIn); if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - if (binIdx < thresholdBinIdx && !step0WillContinue) + if (binIdx < thresholdBinIdx) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -368,21 +384,28 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i return smemFinalBinSize[0] > kNumFinalItems; } -// Smem holder for topKPerRowJob's final-sort. Always reserves -// BlockRadixSort::TempStorage so the sort algorithm can be picked at runtime -// (the union's size is dominated by FinalItems = 16 KB at our shapes). -template -struct TopKSmem +// Follows half - 11 - 11 - 10 bit iterations +template +static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, + int* outIndices, float* outLogits, int stride1, int topK) { + // The number of slots for the final pass. static constexpr int kNumFinalItems = 2048; + // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; + // The class to sort the elements during the final pass. using FinalSort = cub::BlockRadixSort; - using FinalSortTempStorage = typename FinalSort::TempStorage; + using FinalSortTempStorage = std::conditional_t; + // The class to compute the inclusive prefix-sum over the histogram. using Scan = cub::BlockScan; + // The structure to store the final items (for the final pass). struct FinalItems { + // Shared memory to store the indices for the final pass. int indices[kNumFinalItems]; + // Shared memory to store the logits for the final pass. float logits[kNumFinalItems]; }; @@ -392,47 +415,34 @@ struct TopKSmem int data[kNumBins]; }; - union Final + // Shared memory to compute the block sort. + __shared__ union { FinalItems items; FinalSortTempStorage finalSort; Histogram histo; - }; - - Final smemFinal; - int smemThresholdBinIdx[1]; - int smemFinalDstIdx[1]; - int smemFinalBinSize[1]; - int smemFoundTopKValues[1]; -}; - -// Follows half - 11 - 11 - 10 bit iterations -template -static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, - int* outIndices, float* outLogits, int stride1, int topK, TopKSmem& smem) -{ - static constexpr int kNumFinalItems = TopKSmem::kNumFinalItems; - static constexpr int kNumFinalItemsPerThread = TopKSmem::kNumFinalItemsPerThread; - using FinalSort = typename TopKSmem::FinalSort; - - auto& smemFinal = smem.smemFinal; - int* smemThresholdBinIdx = smem.smemThresholdBinIdx; - int* smemFinalDstIdx = smem.smemFinalDstIdx; - int* smemFinalBinSize = smem.smemFinalBinSize; - int* smemFoundTopKValues = smem.smemFoundTopKValues; + } smemFinal; // Shared memory to store the selected indices. // If we are processing using multiple blocks, we need to store the logits and // indices. extern __shared__ int32_t smemOutput[]; + // Shared memory to store the threshold bin. + __shared__ int smemThresholdBinIdx[1]; + // Shared memory counter to register the candidates for the final phase. + __shared__ int smemFinalDstIdx[1]; + // Shared memory to determine if the threshold bin fits in the final items. + __shared__ int smemFinalBinSize[1]; + // Shared memory to keep track of the top-k values found so far by the + // previous iterations + __shared__ int smemFoundTopKValues[1]; + // The length of the row. int rowLen = rowEnd - rowStart; // Shortcut if the length of the row is smaller than Top-K. Indices are not - // sorted by their corresponding logit. Unreachable when mergeBlocks=true: - // both merge callers pass rowLen = numBlocksPerRow * topK > topK. + // sorted by their corresponding logit. if (rowLen <= topK) { for (int rowIt = threadIdx.x; rowIt < rowLen; rowIt += kNumThreadsPerBlock) @@ -502,12 +512,10 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i if (!continueToNextStep) { - // Sort the threshold-bin candidates. Insertion sort wins below ~512 - // items (O(n^2/T) with no fixed cost); BlockRadixSort wins above - // (constant cost padded to kNumFinalItems = 2048). - constexpr int kInsertionSortBranchThreshold = 512; - int const finalCount = smemFinalDstIdx[0]; - if (finalCount > kInsertionSortBranchThreshold) + // The histogram did not proceed to the final 10 bits, therefore we need to + // sort the final items The logits of the elements to be sorted in the final + // pass. + if constexpr (useRadixSort) { // Sorting with radix sort float finalLogits[kNumFinalItemsPerThread]; @@ -518,6 +526,12 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { finalLogits[ii] = -FLT_MAX; + // Indices must be paired with the -FLT_MAX sentinel logit so unused + // slots survive SortDescendingBlockedToStriped and the post-sort copy + // as -1 rather than garbage. Without this, when the threshold bin is + // dominated by -FLT_MAX padding (multi-block merge path) the trailing + // top-K slots receive arbitrary register contents. + finalIndices[ii] = -1; } // Read the elements from SMEM. @@ -525,7 +539,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < finalCount) + if (srcIdx < smemFinalDstIdx[0]) { finalLogits[ii] = smemFinal.items.logits[srcIdx]; finalIndices[ii] = smemFinal.items.indices[srcIdx]; @@ -558,13 +572,13 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } else { - // Sorting with insertion sort. + // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; - for (int i = threadIdx.x; i < finalCount; i += kNumThreadsPerBlock) + for (int i = threadIdx.x; i < smemFinalDstIdx[0]; i += kNumThreadsPerBlock) { int outIndex = 0; auto logit = smemFinal.items.logits[i]; - for (int j = 0; j < finalCount; j++) + for (int j = 0; j < smemFinalDstIdx[0]; j++) { auto otherLogit = smemFinal.items.logits[j]; if (logit < otherLogit || (logit == otherLogit && i < j)) @@ -572,6 +586,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndex++; } } + // Store if outIndex is in bounds if (outIndex + baseIdx < topK) { smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; @@ -593,10 +608,6 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndices[i] = smemOutput[i]; outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; } - else if constexpr (mergeBlocks) - { - outIndices[i] = smemOutput[i]; - } else { if (stride1 == 1) @@ -613,7 +624,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } } // namespace -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(InputT const* logits, int const* rowStarts, int const* rowEnds, int* outIndices, int stride0, int stride1, int const topK, int const offsetIndex) @@ -621,9 +632,6 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - // The number of bins in the histogram. - static constexpr int kNumBins = 2048; - // The row computed by this block. int rowIdx = blockIdx.x + offsetIndex; @@ -635,32 +643,30 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( outIndices += static_cast(rowIdx) * topK; logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK, smem); + topKPerRowJob( + nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif } -template +template 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, float* outLogits = nullptr, - int const numBlocksToMerge = 0, int const* indices = nullptr) + 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) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - // The number of bins in the histogram. - static constexpr int kNumBins = 2048; - // The row computed by this block. int rowIdx = blockIdx.x; // The range of logits within the row. int rowStart = 0; int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1; + int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int rowEnd = actual_kv_len / compressRatio; // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) @@ -683,9 +689,8 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I } logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK, smem); + topKPerRowJob( + indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif @@ -694,415 +699,29 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I namespace { -// Multi-pass radix: 4 launches (half-precision top-11 bits, then float bits -// [21..32), [10..21), [0..10)). Per-row state lives in DRAM scratch; candidate -// data uses two ping-pong buffers. Pass 3's last block emits the final top-K. -static constexpr int kRadixBins = 2048; - -// Per-row scratch state. The last block of each pass (picked via -// `finishedBlocks`) scans the global histogram and writes the next pass's -// threshold. -struct alignas(64) RadixState -{ - int candCount; // candidates entering current pass - int outIdx; // running outIndices write position - int kRemaining; // topK - outIdx - int filterCnt; // running candidate-buf write position - int thresholdBin; // prior pass's threshold; overwritten this pass's last block - int finishedBlocks; // last-block atomic, reset between passes - int thresholdLess; // count of bins below thresholdBin; pass-3 emit routes - // ties (bin == threshold) into a disjoint slot range. - int padding[1]; -}; - -// Common last-block trailer: prefix-scan the merged global histogram, locate -// the bin where the running count crosses kRemaining, stash to state for the -// next pass. step < 3 also resets the per-row global histogram. In step 1 -// the trailer writes st.candCount = full row length for pass 2 to consume -// (pass 1 reads the length inline from seqLens, so no init kernel is needed). -template -__device__ __forceinline__ void radixLastBlockTrailer(int* gHist, RadixState& st, int topK, int rowFullLen) -{ - using Scan = cub::BlockScan; - __shared__ typename Scan::TempStorage scanStorage; - __shared__ int s_thresholdBin; - __shared__ int s_runningBefore; - __shared__ int s_thresholdCount; - - if (threadIdx.x == 0) - { - s_thresholdBin = -1; - s_runningBefore = 0; - s_thresholdCount = 0; - } - __syncthreads(); - - // kRemaining for THIS pass: - // step 1: topK (no auto-promotes have happened yet). - // step 2/3: topK - outIdx (where outIdx was atomic-incremented during - // the filter loop). The histogram for the next pass's - // threshold pick must target this fresh value, not the - // stale state.kRemaining left over from the previous pass. - int const kRem = (step == 1) ? topK : (topK - st.outIdx); - if (kRem <= 0) - { - // All top-k slots already filled by auto-promotes in this pass — - // no need for a next pass to emit anything. - if (threadIdx.x == 0) - { - st.thresholdBin = -1; - st.candCount = 0; - st.kRemaining = 0; - st.finishedBlocks = 0; - st.filterCnt = 0; - (void) rowFullLen; - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } - return; - } - constexpr int kRoundsPerScan = kRadixBins / kThreads; - int running = 0; - for (int r = 0; r < kRoundsPerScan; ++r) - { - int bin = r * kThreads + threadIdx.x; - int c = gHist[bin]; - int prefix, total; - Scan(scanStorage).ExclusiveSum(c, prefix, total); - prefix += running; - int next = prefix + c; - if (prefix < kRem && next >= kRem && s_thresholdBin == -1) - { - atomicCAS(&s_thresholdBin, -1, bin); - if (s_thresholdBin == bin) - { - s_runningBefore = prefix; - s_thresholdCount = c; - } - } - running += total; - __syncthreads(); - if (s_thresholdBin != -1) - break; - } - __syncthreads(); - - if (threadIdx.x == 0) - { - // If the cumsum over the whole histogram never reached kRem the row - // has fewer items than we still need (e.g. decode rows shorter than - // topK). Set thresholdBin to a sentinel above any valid bin so the - // next pass's `bin < thresholdBin` test accepts every surviving - // candidate as auto-promote. - st.thresholdBin = (s_thresholdBin == -1) ? kRadixBins : s_thresholdBin; - // s_runningBefore is the count of histogram items in bins < threshold. - // In the sentinel case it stays at its init 0; that's fine because - // pass-3's inline final-emit only uses thresholdLess to position the - // ties (bin == threshold) write base, and the sentinel branch has - // no items in the threshold bin (the sentinel is above all valid bins). - st.thresholdLess = s_runningBefore; - st.finishedBlocks = 0; - if constexpr (step == 1) - { - // Pass 2 also scans the full row from `logits`, so it reads - // st.candCount = full row length. Pass 1 did not write to it - // and the state struct started at zero (cudaMemsetAsync). - st.candCount = rowFullLen; - (void) s_thresholdCount; - (void) kRem; - } - else - { - st.candCount = st.filterCnt; - int newKRem = topK - st.outIdx; - if (newKRem < 0) - newKRem = 0; - st.kRemaining = newKRem; - st.filterCnt = 0; - } - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } -} - -// Fused histogram + filter pass kernel. -// -// step == 1: pass 1, no filter, just histogram top-11 bits of the row. -// step == 2: pass 2 reads `logits` (pass 1 didn't write a candidate buffer), -// for each item: -// bin1 < thresholdBin1 → write to outIndices (auto-promote) -// bin1 == thresholdBin1 → append to candBufOut, count its bin2 -// in the histogram for pass 2 -// else → drop -// step == 3: same as step 2 but reads `candBufIn` (pass 2's output) and -// uses extractBinIdx<2> for the prior-bits check, extractBinIdx<3> -// for the histogram. -// -// Last block of every pass runs `radixLastBlockTrailer` to compute the -// next pass's threshold and reset cross-pass state. -template -static __global__ __launch_bounds__(kThreads) void radixPassKernel(InputT const* logits, int const* seqLens, - int* outIndices, int const* candBufIn, int* candBufOut, int* histograms, RadixState* state, int stride0, int next_n, - int topK) -{ - int rowIdx = blockIdx.y; - int blockInRow = blockIdx.x; - int blocksPerRow = gridDim.x; - - RadixState& st = state[rowIdx]; - int* gHist = histograms + static_cast(rowIdx) * kRadixBins; - - __shared__ int sHist[kRadixBins]; - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - sHist[i] = 0; - __syncthreads(); - - if constexpr (step == 1) - { - // Read seqLen inline so pass 1 does not depend on st.candCount being - // pre-initialised by a separate init kernel; the cudaMemsetAsync that - // zeroes state+histograms together is enough. Pass-1 trailer below - // writes st.candCount = seqLens[rowIdx] for pass 2 to consume. - int const rowEnd = seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1; - InputT const* in = logits + static_cast(rowIdx) * stride0; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t /*idx*/) - { - float const v = static_cast(vIn); - uint32_t bin = extractBinIdx(v); - atomicAdd(&sHist[bin], 1); - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else if constexpr (step == 2) - { - int const rowEnd = st.candCount; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArr = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t i) - { - float const v = static_cast(vIn); - int bin1 = static_cast(extractBinIdx<1>(v)); - if (bin1 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = static_cast(i); - } - else if (bin1 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArr[pos] = static_cast(i); - uint32_t bin2 = extractBinIdx(v); - atomicAdd(&sHist[bin2], 1); - } - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else // step == 3 - { - int const candCnt = st.candCount; - int const* candArrIn = candBufIn + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArrOut = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - for (int i = blockInRow * kThreads + threadIdx.x; i < candCnt; i += blocksPerRow * kThreads) - { - int srcIdx = candArrIn[i]; - float v = static_cast(in[srcIdx]); - int bin2 = static_cast(extractBinIdx<2>(v)); - if (bin2 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = srcIdx; - } - else if (bin2 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArrOut[pos] = srcIdx; - uint32_t bin3 = extractBinIdx(v); - atomicAdd(&sHist[bin3], 1); - } - } - } - __syncthreads(); - - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - { - int c = sHist[i]; - if (c) - atomicAdd(&gHist[i], c); - } - - __threadfence(); - __shared__ int isLast; - if (threadIdx.x == 0) - { - int prev = atomicAdd(&st.finishedBlocks, 1); - isLast = (prev == blocksPerRow - 1) ? 1 : 0; - } - __syncthreads(); - if (!isLast) - return; - - int const rowFullLen = (step == 1) ? (seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1) : 0; - radixLastBlockTrailer(gHist, st, topK, rowFullLen); - - if constexpr (step == 3) - { - // Final emit, folded into the last block of pass 3. Scan candBufOut - // (top 22 bits == thresholdBin2) and route items into outIndices by - // bin3 vs thresholdBin3. Two-counter scheme so threshold-bin ties - // don't race definite top-k items on the same atomic: - // bin3 < thresh3 → slots [ltBase, ltBase + prefix3) - // bin3 == thresh3 → slots [ltBase + prefix3, topK) - __syncthreads(); - int const filterCnt = st.candCount; - int const thresh3 = st.thresholdBin; - int const prefix3 = st.thresholdLess; - int const* candArr = candBufOut + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int const ltBase = st.outIdx; // already at outBase here - int const eqBase = ltBase + prefix3; - int const eqCap = topK - eqBase; // ≥ 0 by trailer invariant - __shared__ int sEqEmitted; - if (threadIdx.x == 0) - sEqEmitted = 0; - __syncthreads(); - for (int i = threadIdx.x; i < filterCnt; i += kThreads) - { - int srcIdx = candArr[i]; - float v = static_cast(in[srcIdx]); - int bin3 = static_cast(extractBinIdx<3>(v)); - if (bin3 < thresh3) - { - // atomicAdd on st.outIdx is safe: by construction exactly - // prefix3 items fall in this branch, so pos stays in - // [ltBase, eqBase) which is strictly inside [0, topK). - int pos = atomicAdd(&st.outIdx, 1); - outIdxArr[pos] = srcIdx; - } - else if (bin3 == thresh3) - { - int pos = atomicAdd(&sEqEmitted, 1); - if (pos < eqCap) - outIdxArr[eqBase + pos] = srcIdx; - } - } - __syncthreads(); - if (threadIdx.x == 0) - { - int eq = sEqEmitted < eqCap ? sEqEmitted : eqCap; - int filled = eqBase + eq; - if (filled > topK) - filled = topK; - for (int i = filled; i < topK; ++i) - outIdxArr[i] = -1; - } - // Reset the per-row global histogram and st.outIdx so the next call - // sees a clean state without a per-call cudaMemsetAsync. Caller must - // zero-initialize the scratch buffer before the first call. - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - if (threadIdx.x == 0) - st.outIdx = 0; - } -} - -// Scratch layout (uint8 buffer, 64-byte aligned regions): -// RadixState[numRows] -// int histograms[numRows * kRadixBins] (zeroed on first call by -// torch::zeros allocator; pass-3 -// trailer zeroes for subsequent -// calls) -// int candBuf1[numRows * stride0] (pass 2 → pass 3 input) -// int candBuf2[numRows * stride0] (pass 3 → fused final filter) -static size_t radixScratchBytes(int numRows, int numColumns) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - size_t s = 0; - s += roundUp(sizeof(RadixState) * numRows); - s += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - return s; -} - -template -static void launchMultiPassRadix(void* scratch, InputT const* logits, int const* seqLens, int* outIndices, int numRows, - int numColumns, int topK, int stride0, int next_n, cudaLaunchAttribute const* attrs, cudaStream_t stream) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - char* base = static_cast(scratch); - RadixState* state = reinterpret_cast(base); - base += roundUp(sizeof(RadixState) * numRows); - int* histograms = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - int* candBuf1 = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - int* candBuf2 = reinterpret_cast(base); - - int sm_cnt = 132; - { - int dev = 0; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&sm_cnt, cudaDevAttrMultiProcessorCount, dev); - } - // Block fan-out heuristic: target ~4 active blocks/SM (one wave at the - // achievable occupancy of radixPassKernel<512, 1>), with a per-block work - // floor of 2048 items (4 items/thread at 512-wide). - int targetTotalBlocks = sm_cnt * 4; - int numBlocksPerRow = (targetTotalBlocks + numRows - 1) / numRows; - int maxByCols = numColumns / 2048; - if (numBlocksPerRow > maxByCols) - numBlocksPerRow = maxByCols; - if (numBlocksPerRow < 1) - numBlocksPerRow = 1; - - constexpr int kPassThreads = 512; - - auto launchPass = [&](void const* kernel, int const* candIn, int* candOut) - { - cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(numBlocksPerRow, numRows); - cfg.blockDim = kPassThreads; - cfg.dynamicSmemBytes = 0; - cfg.stream = stream; - cfg.numAttrs = 1; - cfg.attrs = const_cast(attrs); - void* args[] = {(void*) &logits, (void*) &seqLens, (void*) &outIndices, (void*) &candIn, (void*) &candOut, - (void*) &histograms, (void*) &state, (void*) &stride0, (void*) &next_n, (void*) &topK}; - cudaLaunchKernelExC(&cfg, kernel, args); - }; - - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, (int*) nullptr); - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, candBuf1); - // Pass 3 emits the final top-K inline in its last-block trailer (see - // radixPassKernel) instead of requiring a separate filter launch. - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) candBuf1, candBuf2); -} - -// Architecture-derived GVR eligibility bounds (cached per-process). +// Insertion vs radix crossover for the single-block path. Radix always pays a +// histogram-clear + per-step scan over kNumBins bins (~4 refinement passes); +// insertion only maintains a topK SMEM array. The crossover is where insertion's +// O(numColumns * lg topK) catches up to radix's O(numColumns + kNumBins) per +// pass — measured empirically around 6 histograms of work. +constexpr int kSortingAlgorithmThreshold = 6 * kNumBins; +// Force the multi-block split-and-merge path above this column count: per-block +// work would otherwise dwarf the merge-pass cost. Callers may override via the +// splitWorkThreshold parameter. Tuned empirically; below this width the merge +// launch overhead isn't worth saving the per-block radix cost. +constexpr int kDefaultSplitWorkThreshold = 200 * 1000; +// Cap blocks-per-row in the multi-block path. Bounds the aux-buffer size +// (numRows * blocksPerRow * topK * sizeof(int32)) and the merge-step input width +// so the second-pass merge kernel stays SMEM-resident. +constexpr int kMaxBlocksPerRowDecode = 10; +// Each sub-block must amortize the radix histogram overhead, i.e. cover at least +// one full histogram pass worth of columns. +constexpr int kDecodeMinColsPerSubBlock = kNumBins; + +// Scheme X bound calculator — shared between fp32 and bf16/fp16 dispatchers. +// Caches hardware attrs (SM count, L2 capacity) and the small-N threshold +// once per process via std::call_once. Per-call cost is just two reads +// from cached static variables plus a small arithmetic block, no syscalls. struct SchemeXBounds { int smCount; @@ -1151,33 +770,148 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) return b; } -// Unified dispatcher (fp32 / bf16 / fp16). Each tier's kernel is templated on -// InputT and casts to float at HBM-read sites; the scratch buffer is uint8. -template -void invokeIndexerTopKDecodeImpl(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, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +} // namespace + +int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold) +{ + if (numRows <= 0) + { + return 1; + } + + int const forceSplitThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // Preserve original behavior for very long sequences. + if (numColumns >= forceSplitThreshold) + { + return kMaxBlocksPerRowDecode; + } + + // Query the actual SM count from the driver so the dispatch tracks the + // hardware rather than a baked-in target (H100=132, B200=148, …). + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4); + TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count"); + int const smCount = bounds.smCount; + int const maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock); + int const maxBp = std::min(maxByCols, kMaxBlocksPerRowDecode); + + int blocksPerRow; + if (numRows < smCount / 2) + { + // Sub-half-wave band: bp=2 by itself leaves SMs idle (numRows × 2 < smCount), + // so sweep bp ∈ [2, maxBp] for the choice that minimizes waves(bp) / bp, + // where waves(bp) = ceil(numRows * bp / smCount). The wave-quantization-aware + // sweep avoids spilling one extra block per row across a wave boundary, which + // is what a naive ceil(smCount / numRows) target would do. + int bestBp = 1; + int bestWaves = 1; // numRows < smCount/2 → bp=1 always fits in a single wave + for (int bp = 2; bp <= maxBp; ++bp) + { + int const totalBlocks = numRows * bp; + int const waves = (totalBlocks + smCount - 1) / smCount; + // waves / bp < bestWaves / bestBp ⇔ waves * bestBp < bestWaves * bp + if (waves * bestBp < bestWaves * bp) + { + bestWaves = waves; + bestBp = bp; + } + } + blocksPerRow = bestBp; + } + else + { + // numRows >= smCount/2: bp=2 saturates SMs with one wave (numRows*2 >= smCount) + // and stays on the multi-block split+merge path. Crucially this path uses + // a different kernel instantiation than bp=1, and only the bp=1 single-block + // radix kernel pays the wave-scheduling cliff when gridDim.x approaches + // smCount (measured on B200, cols=196608, topK=2048: BS=131=125us, + // BS=132=312us, BS=148=390us — the cliff disappears entirely with bp=2). + // bp=2 is also the cheapest split (smallest merge input); larger bp piles + // on merge-pass overhead without proportional gain on the shapes measured. + // Falls back to 1 only when maxByCols caps it at 1 for very narrow rows + // (numColumns < kDecodeMinColsPerSubBlock). + blocksPerRow = std::min(2, maxByCols); + } + return std::max(1, blocksPerRow); +} + +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) { - // Split-work cutoff: matches main's 200k default. is_prefill forces - // single-block via a 1<<30 threshold no shape can reach: prefill chunks are - // bounded by max_num_tokens, well below the multi-pass radix crossover at - // any practical setting. - int const adaptiveSplitWorkThreshold = is_prefill ? (1 << 30) : 200 * 1000; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : adaptiveSplitWorkThreshold; constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // ======================================================================== + // Small-N dispatch axis. + // + // GVR Heuristic Top-K has a *fixed* per-launch overhead from Phase-1 + // (preIdx stats reduction over M=2048) and Phase-4 (2048-bin histogram + // snap), totaling ~11 µs regardless of N. For small N (≤16K), this + // fixed cost dominates and the kernel loses to the existing + // insertion-sort/radix path. Empirically (random data, B200 BS=1): + // N=8192 : Heuristic 16.5 µs vs Radix 11.2 µs (radix 1.47× faster) + // N=16384 : Heuristic 21.9 µs vs Radix 22.0 µs (parity) + // N=32768 : Heuristic 26.1 µs vs Radix 32.9 µs (heuristic 1.26× faster) + // N=131072 : Heuristic 43.4 µs vs Radix 76.1 µs (heuristic 1.75× faster) + // + // Route N < kSeqSmall to the existing Radix/Insertion path (which itself + // splits at kSortingAlgorithmThreshold=12288). kSeqSmall is set at the + // empirical crossover point. + // + // ======================================================================== + // Architecture-derived BS-threshold dispatch — jointly bounded by + // occupancy AND L2 cache capacity. + // + // Two physical constraints bound when the per-row heuristic kernel + // remains faster than a radix streaming kernel: + // + // (A) Occupancy bound — 3·SM − SM/8 (wave geometry + setup margin) + // Each CTA uses ~58 KB SMEM (fixed, independent of N), so B200's + // 228 KB dynamic SMEM allows max 3 CTA/SM. Above 3·SM rows per + // launch, tail-wave imbalance causes stragglers. The -SM/8 margin + // (~1/8 wave) covers CTA setup + L2 ingestion overhead. + // On B200(148 SM): 3×148 − 18 = 426. + // + // (B) L2 cache bound — 0.9·L2 / (4·N) per-CTA logits fit + // Each CTA streams its row (N×4B) through L2 per Phase-2 iter. + // With num_concurrent_CTAs × N × 4B > L2, eviction dominates. + // On B200(126 MB L2) with N=70K: 0.9·126MB/(4·70690) ≈ 440, + // which is ~ equal to (A)=426 — the two constraints cross over + // near the SWE-Bench data point. + // For N > 73K the L2 bound tightens below (A) and must take + // over; e.g. N=128K → kBsL2=238, N=196K → kBsL2=155. + // + // Dispatch threshold = min(kBsWave, kBsL2), still data-agnostic (only + // queries hardware attrs). At N≈70K both bounds produce ~426, so the + // L2 axis is a no-op there; for larger N it auto-tightens the threshold. + // + // Small-N lower bound `kSeqSmall` (default 12288) lets the Heuristic + // axis take over wherever the original Radix-radix branch would have + // triggered. Random-data benchmarks suggest the crossover is 16384, + // but workloads with strongly preIdx-correlated logits make P1 stats + // accurate and P2 converge in 1-2 iterations, shifting the real + // crossover into the [12288, 16384] band. Configurable via + // TRTLLM_HEURISTIC_NMIN env (>=1024). + // ======================================================================== + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4); + int const kBsWave = bounds.kBsWave; + int const kBsL2 = bounds.kBsL2; + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; - // GVR eligibility (matches main's rule): supported K, stride1 contiguous, - // preIdx + scratch provided, numColumns in [kSeqSmall, splitWorkThreshold), - // and numRows below the architecture-derived wave/L2 bound. is_prefill - // suppresses GVR through effectiveSplitWorkThreshold being huge. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - bool const canUseHeuristic = preIdx != nullptr && stride1 == 1 && isSupportedTopK && preIdxCount == topK - && preIdxStride >= preIdxCount && heuristicScratch != nullptr && numColumns >= bounds.kSeqSmall - && numColumns < effectiveSplitWorkThreshold && numRows < bounds.kBsLarge; - - // Env-gated dispatch trace (TRTLLM_SCHEMEX_DEBUG=1). + // compressRatio == 1: DSv3.2 indexer (no compressor). + // compressRatio == 4: DSv4 indexer (overlap compressor); logits/preIdx in + // compressed-token-index space. Kernel handles N = actual_kv_len/cr and + // forces preIdxOffset=0 internally for cr != 1. + bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); + bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; + + // Optional env-gated dispatch trace (set TRTLLM_SCHEMEX_DEBUG=1 to enable) { static std::once_flag sDebugOnceFlag; static bool sDebug = false; @@ -1189,87 +923,204 @@ void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* }); if (sDebug) { - fprintf(stderr, "[Scheme X] numRows=%d numColumns=%d kBsLarge=%d kSeqSmall=%d -> %s path\n", numRows, - numColumns, bounds.kBsLarge, bounds.kSeqSmall, canUseHeuristic ? "Heuristic" : "Radix"); + fprintf(stderr, + "[Scheme X] numRows=%d numColumns=%d kBsWave=%d kBsL2=%d kBsLarge=%d kSeqSmall=%d smCount=%d " + "L2=%dMB -> %s path%s\n", + numRows, numColumns, kBsWave, kBsL2, kBsLarge, kSeqSmall, bounds.smCount, + bounds.l2Bytes / (1024 * 1024), canUseHeuristic ? "Heuristic" : "Radix", + (numColumns < kSeqSmall) ? " (small-N route)" : ""); } } if (canUseHeuristic) { launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); + sync_check_cuda_error(stream); + return; } - else if (numColumns < effectiveSplitWorkThreshold) + + int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); + + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + + if (blocksPerRow == 1) { - // Single-block tier: one CTA per row. - auto* kernel_instance = &topKPerRowDecode; + // Single block per row. Below kSortingAlgorithmThreshold use insertion sort, + // above use the histogram-radix path. + bool const useRadixSort = numColumns >= kSortingAlgorithmThreshold; + auto* kernel_instance = useRadixSort ? &topKPerRowDecode + : &topKPerRowDecode; + cudaLaunchConfig_t config; config.gridDim = numRows; config.blockDim = kNumThreadsPerBlock; config.dynamicSmemBytes = topK * sizeof(int32_t); config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - /*outLogits=*/nullptr, /*numBlocksToMerge=*/0, /*indices=*/nullptr); + compressRatio, nullptr, 0, nullptr); } else { - // Multi-pass radix. radixPassKernel reads logits contiguously, so - // strided inputs would rank the wrong values — gate on stride1 == 1. - // (The single-block tier handles stride1 != 1 via topKPerRowJob's - // strided fallback.) - TLLM_CHECK_WITH_INFO(stride1 == 1, "indexer top-k split-work tier (multi-pass radix) requires stride1 == 1."); - TLLM_CHECK_WITH_INFO(scratch != nullptr && scratchBytes >= radixScratchBytes(numRows, numColumns), - "indexer top-k split-work tier: scratch buffer missing or too small."); - cudaLaunchAttribute radixAttrs[1]; - radixAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - radixAttrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - launchMultiPassRadix( - scratch, logits, seqLens, indices, numRows, numColumns, topK, stride0, next_n, radixAttrs, stream); + // Split each row across `blocksPerRow` blocks, then merge with a second pass. + auto* kernel_instance_part1 = &topKPerRowDecode; + cudaLaunchConfig_t config_part1; + config_part1.gridDim = dim3(numRows, blocksPerRow); + config_part1.blockDim = kNumThreadsPerBlock; + config_part1.dynamicSmemBytes = 2 * topK * sizeof(int32_t); + config_part1.stream = stream; + config_part1.numAttrs = 1; + config_part1.attrs = attrs; + + cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, + next_n, compressRatio, outLogitsAux, 0, nullptr); + + constexpr int kNumThreadsPerBlockMerge = 1024; + auto* kernel_instance_part2 = &topKPerRowDecode; + cudaLaunchConfig_t config_part2; + config_part2.gridDim = numRows; + config_part2.blockDim = kNumThreadsPerBlockMerge; + config_part2.dynamicSmemBytes = topK * sizeof(int32_t); + config_part2.stream = stream; + config_part2.numAttrs = 1; + config_part2.attrs = attrs; + + cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, + topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); } sync_check_cuda_error(stream); } -} // anonymous namespace +// ============================================================================ +// bf16 / fp16 dispatcher overloads +// ============================================================================ +// Reuses the BS-threshold + small-N dispatch axes (kBsLarge, kSeqSmall) from +// the fp32 dispatcher, except kBsL2 uses sizeof(InputT) bytes/element instead +// of 4 — L2 footprint is half, so bf16/fp16 path remains valid for larger BS +// than fp32 at the same N. +// +// Fallback chain when GVR-Heuristic preconditions are not met (preIdx +// missing, BS too large, or numColumns < kSeqSmall): +// numColumns < kSortingAlgorithmThreshold (12288) → insertion sort +// kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort +// numColumns ≥ splitWorkThreshold (200K default) → unsupported +// +// Insertion + radix tiers use the same topKPerRowDecode kernel as fp32 with +// InputT propagated through; the histogram and sort steps operate on float +// keys after a static_cast(InputT) at HBM-read sites, so accuracy is +// identical to casting input to fp32 before the kernel. +// +// The split-work tier requires float aux buffers (outLogitsAux / +// outIndicesAux) that the bf16/fp16 entry does not expose; callers in that +// regime must use the fp32 entry. -void invokeIndexerTopKDecode(float 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, float* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +namespace { - invokeIndexerTopKDecodeImpl(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); -} -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int /*topK*/) +template +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) { - return radixScratchBytes(numRows, numColumns); + static_assert(std::is_same_v || std::is_same_v, + "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); + + constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; + + bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); + // See fp32 path: cr==1 (V3.2) and cr==4 (V4 indexer) are both supported. + bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); + bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; + + if (canUseHeuristic) + { + launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, + preIdxStride, preIdxCount, numRows, compressRatio, stream); + } + else if (numColumns < kSortingAlgorithmThreshold) + { + // Insertion sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else if (numColumns < effectiveSplitWorkThreshold) + { + // Radix sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else + { + TLLM_CHECK_WITH_INFO(false, + "indexer_topk_decode bf16/fp16 path does not support numColumns >= splitWorkThreshold " + "(split-work path requires float aux buffers not exposed in the bf16/fp16 entry). " + "Got numColumns=%d splitWorkThreshold=%d. Use the fp32 entry for this regime.", + numColumns, effectiveSplitWorkThreshold); + } + + sync_check_cuda_error(stream); } +} // anonymous namespace + 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, cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, - scratchBytes, is_prefill); + invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, 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, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, + cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); + invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, @@ -1278,10 +1129,18 @@ void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int con { constexpr int kNumThreadsPerBlock = 512; - // One launch over all rows; the per-row sort algorithm is picked at - // runtime inside topKPerRowJob. - topKPerRowPrefill<<>>( - logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + int numInsertionBlocks = std::min(numRows, kSortingAlgorithmThreshold); + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + + if (numRows > kSortingAlgorithmThreshold) + { + int numRadixBlocks = numRows - kSortingAlgorithmThreshold; + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, kSortingAlgorithmThreshold); + } sync_check_cuda_error(stream); } @@ -1293,8 +1152,6 @@ bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytes { return false; } - // Mirrors the dispatcher's effectiveSplitWorkThreshold default. - constexpr int kDefaultSplitWorkThreshold = 200 * 1000; auto const bounds = getSchemeXBounds(numColumns, bytesPerElem); return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; } diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 696b8d471043..1981c417dbae 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -38,9 +38,10 @@ namespace torch_ext void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, th::Tensor const& indices, int64_t next_n, int64_t index_topk, std::optional const& pre_idx, - std::optional const& heuristic_scratch, std::optional const& done_counter_scratch, - std::optional const& scratch, bool is_prefill) + std::optional const& heuristic_scratch, int64_t compress_ratio, + std::optional const& radix_aux_indices, std::optional const& radix_aux_logits) { + TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); TORCH_CHECK(logits.is_cuda() && seq_lens.is_cuda() && indices.is_cuda(), "logits, seq_lens, and indices must be CUDA tensors"); @@ -110,70 +111,74 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t heuristicScratchPtr = scratchTensor.data_ptr(); } - // Multi-pass radix scratch. Pre-allocate it for CUDA-graph stable - // addresses, or let the wrapper allocate internally below. - void* multiPassScratchPtr = nullptr; - size_t multiPassScratchBytes = 0; - if (scratch.has_value()) - { - auto const& t = scratch.value(); - TORCH_CHECK(t.is_cuda(), "scratch must be a CUDA tensor"); - TORCH_CHECK(t.device() == logits.device(), "scratch must be on the same device as logits"); - TORCH_CHECK(t.is_contiguous(), "scratch must be contiguous"); - TORCH_CHECK(t.scalar_type() == at::ScalarType::Byte, "scratch must be uint8"); - multiPassScratchPtr = t.data_ptr(); - multiPassScratchBytes = static_cast(t.numel()); - } - - // 0 lets invokeIndexerTopKDecode pick its numRows-aware threshold. - int32_t const splitWorkThreshold = 0; + int32_t splitWorkThreshold = 200 * 1000; auto stream = at::cuda::getCurrentCUDAStream(logits.get_device()); - // Mirror the kernel's threshold so we only allocate when split-work fires. - int32_t const adaptiveSplitWorkThreshold = 200 * 1000; - th::Tensor scratch_internal; - if (num_columns >= adaptiveSplitWorkThreshold && multiPassScratchPtr == nullptr) - { - size_t const bytes = tk::indexerTopKDecodeScratchBytes(num_rows, num_columns, static_cast(index_topk)); - // Zero-init: the radix kernel reads the per-row state on first call. - scratch_internal - = th::zeros({static_cast(bytes)}, th::TensorOptions().dtype(th::kByte).device(logits.device())); - multiPassScratchPtr = scratch_internal.data_ptr(); - multiPassScratchBytes = bytes; - } - - // `done_counter_scratch` is kept on the op signature for source compat - // but unused since the fused split-work tier was removed. Warn once so - // callers still passing it know to drop the argument. - if (done_counter_scratch.has_value()) - { - TORCH_WARN_ONCE( - "indexer_topk_decode: `done_counter_scratch` is deprecated and ignored since the fused split-work tier " - "was removed; drop this argument. The multi-pass radix path allocates its scratch internally, or you may " - "pass a pre-allocated buffer via `scratch` for CUDA-graph capture."); - } - if (logits_dtype == at::ScalarType::Float) { + // fp32 path — full Scheme X v1.2 dispatcher (GVR / Insertion / Radix / + // Radix-split-work). Caller-owned radix_aux_{indices,logits} are the + // split-work scratch buffers and are dereferenced only when + // blocksPerRow > 1; for blocksPerRow == 1 the dispatcher passes + // nullptr through to the kernel and never touches them (see + // indexerTopK.cu's invokeIndexerTopKDecode bp==1 branch). + int const blocks_per_row = tk::computeIndexerTopKDecodeBlocksPerRow(num_rows, num_columns, splitWorkThreshold); + int32_t* aux_indices_ptr = nullptr; + float* aux_logits_ptr = nullptr; + if (radix_aux_indices.has_value() && radix_aux_logits.has_value()) + { + // Caller-owned scratch with stable address (CUDA Graph safe; + // matches the heuristic_scratch convention noted above). All + // in-tree callers under CUDA Graph capture go through dsa.py's + // DSAtrtllmAttentionMetadata which always pre-allocates these. + auto const& ai = radix_aux_indices.value(); + auto const& al = radix_aux_logits.value(); + TORCH_CHECK(ai.is_cuda() && al.is_cuda(), "radix_aux_{indices,logits} must be CUDA tensors"); + TORCH_CHECK(ai.device() == logits.device() && al.device() == logits.device(), + "radix_aux_{indices,logits} must be on the same device as logits"); + TORCH_CHECK(ai.is_contiguous() && al.is_contiguous(), "radix_aux_{indices,logits} must be contiguous"); + TORCH_CHECK(ai.scalar_type() == th::kInt32, "radix_aux_indices must be int32"); + TORCH_CHECK(al.scalar_type() == th::kFloat32, "radix_aux_logits must be float32"); + if (blocks_per_row > 1) + { + int64_t const needed_elts = static_cast(num_rows) * blocks_per_row * index_topk; + TORCH_CHECK(ai.numel() >= needed_elts && al.numel() >= needed_elts, + "radix_aux_{indices,logits} must hold at least num_rows*blocks_per_row*index_topk elements (got ", + ai.numel(), " / ", al.numel(), ", need ", needed_elts, ")"); + } + aux_indices_ptr = ai.data_ptr(); + aux_logits_ptr = al.data_ptr(); + } + else + { + // No caller-supplied scratch. Allowed only when blocksPerRow == 1 + // (kernel never dereferences nullptr aux pointers in that case). + // blocksPerRow > 1 without caller-owned scratch is the original + // G4 CUDA-Graph stale-pointer hazard — reject early. + TORCH_CHECK(blocks_per_row == 1, + "radix_aux_{indices,logits} must be pre-allocated by the caller when blocks_per_row > 1 " + "(got blocks_per_row=", + blocks_per_row, "). Required for CUDA Graph safety."); + } tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), - splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), - static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast(heuristicScratchPtr), stream, multiPassScratchPtr, multiPassScratchBytes, is_prefill); + aux_logits_ptr, aux_indices_ptr, splitWorkThreshold, num_rows, num_columns, logits_stride_0, + logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, + preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), stream); } else if (logits_dtype == at::ScalarType::BFloat16) { tk::invokeIndexerTopKDecode(reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, - preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), stream, multiPassScratchPtr, - multiPassScratchBytes, is_prefill); + preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), + static_cast(compress_ratio), stream); } else // Half { tk::invokeIndexerTopKDecode(reinterpret_cast<__half const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast<__half*>(heuristicScratchPtr), stream, multiPassScratchPtr, multiPassScratchBytes, is_prefill); + static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), stream); } } @@ -188,6 +193,8 @@ void indexer_topk_prefill(th::Tensor const& logits, th::Tensor const& row_starts TORCH_CHECK(indices.dim() == 2, "indices must be a 2D Tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be a 2D Tensor"); + TORCH_CHECK(indices.size(1) >= index_topk, "indices.size(1) must be >= index_topk, got ", indices.size(1), " < ", + index_topk); auto const inputSize = logits.sizes(); auto const numRows64 = inputSize[0]; @@ -213,16 +220,6 @@ void indexer_topk_prefill(th::Tensor const& logits, th::Tensor const& row_starts static_cast(logits_stride_1), static_cast(index_topk), stream); } -// Returns the size in bytes of the `scratch` buffer required by -// indexer_topk_decode's multi-pass radix path for the given shape. Callers -// can allocate `torch.empty(size, dtype=torch.uint8, device='cuda')` and -// pass the result as the `scratch` argument. -int64_t indexer_topk_decode_scratch_bytes(int64_t num_rows, int64_t num_columns, int64_t index_topk) -{ - return static_cast(tk::indexerTopKDecodeScratchBytes( - static_cast(num_rows), static_cast(num_columns), static_cast(index_topk))); -} - } // end namespace torch_ext TRTLLM_NAMESPACE_END @@ -231,9 +228,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " - "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, Tensor? done_counter_scratch=None, " - "Tensor? scratch=None, bool is_prefill=False) -> ()"); - m.def("indexer_topk_decode_scratch_bytes(int num_rows, int num_columns, int index_topk) -> int"); + "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1, " + "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) @@ -241,11 +237,6 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("indexer_topk_decode", &tensorrt_llm::torch_ext::indexer_topk_decode); } -TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) -{ - m.impl("indexer_topk_decode_scratch_bytes", &tensorrt_llm::torch_ext::indexer_topk_decode_scratch_bytes); -} - TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index ab16aac4d716..592edfb4950e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -94,13 +94,21 @@ def warmup_heuristic_topk_decode(top_k: int = 2048, indices = torch.empty((1, top_k), dtype=torch.int32, device=device) pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device) scratch = torch.empty((top_k, ), dtype=torch.float32, device=device) + radix_aux_indices = torch.empty((1, 10, top_k), + dtype=torch.int32, + device=device) + radix_aux_logits = torch.empty((1, 10, top_k), + dtype=torch.float32, + device=device) torch.ops.trtllm.indexer_topk_decode(logits, seq_lens, indices, 1, top_k, pre_idx=pre_idx, - heuristic_scratch=scratch) + heuristic_scratch=scratch, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits) torch.cuda.synchronize() @@ -710,9 +718,30 @@ def __post_init__(self): capture_graph=capture_graph, ) + self._create_radix_aux_buffers(capture_graph=capture_graph) + # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) + def _create_radix_aux_buffers(self, capture_graph=False): + """Create persistent scratch for the fp32 radix split-work TopK path.""" + max_blocks_per_row = 10 + max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + self.radix_aux_indices = self.get_empty( + self.cuda_graph_buffers, + (max_gen_tokens, max_blocks_per_row, self.num_sparse_topk), + cache_name="radix_aux_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.radix_aux_logits = self.get_empty( + self.cuda_graph_buffers, + (max_gen_tokens, max_blocks_per_row, self.num_sparse_topk), + cache_name="radix_aux_logits", + dtype=torch.float32, + capture_graph=capture_graph, + ) + def _create_kv_lens_2d_buffer(self, capture_graph=False): """Pre-allocated buffer for the DeepGEMM 2D context_lens API. @@ -810,6 +839,7 @@ def update_spec_dec_param( dtype=torch.float32, capture_graph=capture_graph, ) + self._create_radix_aux_buffers(capture_graph=capture_graph) def _invalidate_pool_view_cache(self): """Invalidate the cached pool view and related step-invariant values. @@ -1387,6 +1417,7 @@ def __init__(self, self.head_dim = sparse_attention_config.index_head_dim # 128 self.index_topk = sparse_attention_config.index_topk # 2048 self.layer_idx = layer_idx + self.compress_ratio = 1 self.wq_b = Linear( self.q_lora_rank, @@ -1997,7 +2028,8 @@ def sparse_attn_indexer( logits, chunk.cu_seqlen_ks[chunk_q_start:chunk_q_end], chunk.cu_seqlen_ke[chunk_q_start:chunk_q_end], - topk_indices_buffer[global_q_start:global_q_end, :]) + topk_indices_buffer[global_q_start:global_q_end, :], + self.index_topk) else: topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), @@ -2049,7 +2081,8 @@ def sparse_attn_indexer( if use_custom_topk: torch.ops.trtllm.indexer_topk_prefill( logits, cu_seqlen_ks, cu_seqlen_ke, - topk_indices_buffer[:num_ctx_tokens, :]) + topk_indices_buffer[:num_ctx_tokens, :], + self.index_topk) else: topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), @@ -2258,7 +2291,10 @@ def sparse_attn_indexer( next_n, self.index_topk, pre_idx=pre_idx, - heuristic_scratch=heuristic_scratch) + heuristic_scratch=heuristic_scratch, + compress_ratio=self.compress_ratio, + radix_aux_indices=metadata.radix_aux_indices, + radix_aux_logits=metadata.radix_aux_logits) else: # padded positions = torch.arange( diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index fa5cbd696d8d..1c5ceed592a3 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -249,7 +249,10 @@ def _(logits, next_n, index_topk, pre_idx=None, - heuristic_scratch=None): + heuristic_scratch=None, + compress_ratio=1, + radix_aux_indices=None, + radix_aux_logits=None): # In-place operation, no return value (void function) pass diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 5d8d5dfef98a..2e9d9c75c620 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -1,3 +1,11 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + """ Distribution-parameterized correctness tests for the heuristic indexer_topk_decode. @@ -221,16 +229,13 @@ def generate_seq_lens(batch_size, min_long_seq, num_tokens): # --------------------------------------------------------------------------- -# Original random-data decode test (verbatim from test_indexer_topk.py) +# Random-data decode tests (verbatim helper from test_indexer_topk.py, plus +# regression coverage for the SM-saturation heuristic). # --------------------------------------------------------------------------- -@pytest.mark.parametrize("batch_size", [1, 64, 512, 2048]) -@pytest.mark.parametrize("next_n", [1, 2]) -@pytest.mark.parametrize("index_topk", [2048, 128]) -@pytest.mark.parametrize("num_tokens", [4096, 8192]) -def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens): - """Verify indexer_topk_decode output matches torch.topk for random logits.""" +def _run_indexer_topk_decode_check(batch_size, next_n, index_topk, num_tokens, compress_ratio): + """Run the random-data decode equivalence check against torch.topk.""" torch.manual_seed(24) torch.cuda.manual_seed(24) num_gen_tokens = batch_size * next_n @@ -239,16 +244,44 @@ def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens): next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n seq_lens = generate_seq_lens(batch_size, index_topk, num_tokens) - row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 + # Calculate actual KV lengths + actual_kv_lens = seq_lens[row_indices] - next_n + next_n_offset + 1 + + # Apply compression with floor division + row_ends = actual_kv_lens // compress_ratio + + # Generate logits with the compressed size logits = create_random_logits(row_starts, row_ends, torch.float32, 42) indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") - torch.ops.trtllm.indexer_topk_decode(logits, seq_lens, indices, next_n, index_topk) + # Pre-allocate worst-case Radix split-work aux scratch so the fp32 path is + # CUDA-Graph-safe under capture and unconditionally accepted by the cpp op + # (which rejects blocks_per_row > 1 without caller-owned scratch). For + # blocks_per_row == 1 the kernel never dereferences these — supplying them + # is harmless. See IndexerTopKOp.cpp's fp32 branch. + radix_aux_indices, radix_aux_logits = _build_radix_aux_buffers(num_gen_tokens, index_topk) + + # Run CUDA implementation with compress_ratio + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits, + ) + torch.cuda.synchronize() + # Run reference implementation on compressed row_ends max_row_len = row_ends.max().item() + if max_row_len == 0: + # All rows are empty after compression, skip comparison + return torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] mask_lo = torch_indices >= 0 mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0 @@ -260,90 +293,343 @@ def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens): ), "CUDA top_k_per_row results don't match torch.topk" +@pytest.mark.parametrize("batch_size", [1, 64, 512, 2048]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("index_topk", [2048, 512, 128]) +@pytest.mark.parametrize("num_tokens", [4096, 8192]) +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens, compress_ratio): + _run_indexer_topk_decode_check(batch_size, next_n, index_topk, num_tokens, compress_ratio) + + +# Regression coverage for the SM-saturation heuristic in indexerTopK decode. +# Exercises the multi-block split path that was newly enabled for small batches +# and moderate numColumns (commits 1c88ecd58, 38f9b59b2): +# batch_size in {16, 32, 128} bracket the smTarget transitions (>= 9, 5, 2 blocks/row) +# num_tokens 4096-32768 with cr in {1, 4} spans both the small-numColumns range +# that previously short-circuited to blocksPerRow=1 and the legacy single-block +# regime (numCols < 2048 stays single-block via the maxByCols guard). +@pytest.mark.parametrize("batch_size", [16, 32, 128]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("index_topk", [2048]) +@pytest.mark.parametrize("num_tokens", [4096, 16384, 32768]) +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_topk_decode_sm_saturation( + batch_size, next_n, index_topk, num_tokens, compress_ratio +): + _run_indexer_topk_decode_check(batch_size, next_n, index_topk, num_tokens, compress_ratio) + + +# Covers the single-block / multi-block path transition at +# batch_size == kDecodeTargetTotalBlocks (= 132). batch_size < 132 keeps the +# single-block-per-row path; at and above 132 the launch policy routes to +# the multi-block split-and-merge path with bp = 2 to avoid the wave- +# scheduling cliff that the single-block radix kernel hits once gridDim.x +# approaches smCount on B200. +@pytest.mark.parametrize("batch_size", [1, 64, 100, 132, 148, 256]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("index_topk", [2048]) +@pytest.mark.parametrize("num_tokens", [4096, 16384, 32768]) +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_topk_decode_launch_policy_transitions( + batch_size, next_n, index_topk, num_tokens, compress_ratio +): + _run_indexer_topk_decode_check(batch_size, next_n, index_topk, num_tokens, compress_ratio) + + # --------------------------------------------------------------------------- -# Split-work (multi-pass radix) decode path, opt-in via the thop wrapper +# Caller-owned Radix aux scratch (CUDA-Graph-safe path) +# +# The fp32 Radix path of indexer_topk_decode allocates its split-work aux +# buffers via th::empty per call when blocks_per_row > 1. Under CUDA Graph +# capture+replay these per-call pointers become stale when the caching +# allocator is perturbed (chunked-prefill churn at high CONC), producing +# silent corruption that surfaces as CUDA_ERROR_ILLEGAL_ADDRESS at a +# downstream sync (see DSv4-Flash pareto sweep G4). +# +# The fix added two optional kwargs `radix_aux_indices` and +# `radix_aux_logits` so the caller can supply persistent stable-address +# buffers (matching the existing `heuristic_scratch` convention). +# +# These tests verify: +# (a) caller-owned-aux output matches the default (th::empty) path, +# (b) the same caller-owned-aux op survives CUDA Graph capture+replay +# with consistent output and matches the non-graph reference, +# (c) validation rejects malformed buffers (dtype / size / device / +# contiguity). +# +# Parameters target the fp32 split-work regime: blocks_per_row > 1 +# requires numColumns / kNumBins(=2048) >= 2 (see indexerTopK.cu:775). +# For compress_ratio=4 that means num_tokens >= 16384 (numColumns = 4096). # --------------------------------------------------------------------------- -def _run_opt_in_decode(batch_size, num_tokens, *, dtype, done_counter, scratch): - """Run indexer_topk_decode on the split-work (multi-pass radix) tier and - assert output matches torch.topk. The path works for fp32 / bf16 / fp16 — - the radix scratch stays uint8 / int32 regardless of logit dtype.""" - index_topk = 2048 - next_n = 1 - torch.manual_seed(24) - torch.cuda.manual_seed(24) +_K_MAX_BLOCKS_PER_ROW_DECODE = 10 # mirror kMaxBlocksPerRowDecode in indexerTopK.cu + + +def _build_radix_aux_buffers(num_rows: int, index_topk: int): + """Allocate worst-case stable-address aux buffers for the Radix path.""" + aux_indices = torch.empty( + (num_rows, _K_MAX_BLOCKS_PER_ROW_DECODE, index_topk), + dtype=torch.int32, + device="cuda", + ) + aux_logits = torch.empty( + (num_rows, _K_MAX_BLOCKS_PER_ROW_DECODE, index_topk), + dtype=torch.float32, + device="cuda", + ) + return aux_indices, aux_logits + + +def _build_decode_inputs(batch_size, next_n, index_topk, num_tokens, compress_ratio, seed=24): + """Construct the (logits, seq_lens, row_starts, row_ends) tuple used by + the fp32 indexer_topk_decode path. + + Unlike `_run_indexer_topk_decode_check`'s helper, this one **forces + every row to span the full num_tokens window** so that the row_ends + distribution is deterministic and `row_ends.max() == num_tokens // + compress_ratio` for every parametrisation. That guarantees the test + exercises the Radix-split-work path (`blocks_per_row > 1` requires + `num_columns >= 2 * kNumBins = 4096`) instead of accidentally + short-circuiting to the single-block fallback when a random seq_len + happens to be small.""" + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) num_gen_tokens = batch_size * next_n row_starts = torch.zeros(num_gen_tokens, dtype=torch.int32, device="cuda") row_indices = torch.arange(num_gen_tokens, device="cuda") // next_n next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n - seq_lens = generate_seq_lens(batch_size, index_topk, num_tokens) - row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 + seq_lens = torch.full((batch_size,), num_tokens, dtype=torch.int32, device="cuda") + actual_kv_lens = seq_lens[row_indices] - next_n + next_n_offset + 1 + row_ends = actual_kv_lens // compress_ratio - logits = create_random_logits(row_starts, row_ends, dtype, 42) - indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") + logits = create_random_logits(row_starts, row_ends, torch.float32, 42) + return logits, seq_lens, row_starts, row_ends + +def _topk_logit_values_sorted(logits, indices, k): + """For each row, return the top-K logit VALUES (sorted descending, + invalid -1 slots replaced with -inf). Used to compare two op-call + outputs by the VALUES they selected, tolerating the kernel's + nondeterministic tie-breaking on equal-logit candidates.""" + safe_idx = indices.clone().long() + valid = safe_idx >= 0 + safe_idx[~valid] = 0 # gather requires non-negative + gathered = logits.gather(1, safe_idx) + gathered = gathered.masked_fill(~valid, float("-inf")) + sorted_vals, _ = torch.sort(gathered[:, :k], dim=-1, descending=True) + return sorted_vals + + +@pytest.mark.parametrize( + "batch_size,next_n,index_topk,num_tokens,compress_ratio", + [ + # blocks_per_row > 1 regime: numColumns >= 4096 + (8, 1, 512, 16384, 4), # mirrors DSv4 Flash CONC=8 GVR-OFF failure point + (16, 1, 512, 16384, 4), + (32, 1, 512, 16384, 4), + (8, 1, 512, 8192, 1), # compress_ratio=1 path + ], +) +def test_indexer_topk_decode_radix_aux_equivalence( + batch_size, next_n, index_topk, num_tokens, compress_ratio +): + """Caller-owned Radix aux scratch must yield allocation-pointer-independent + output: two invocations with freshly-allocated aux buffers (different + backing addresses) on identical inputs must produce the same selected + top-K logit values. Guards against the original G4 hazard at the API + boundary — the kernel must consume the supplied buffers without leaking + pointer state from any prior call.""" + logits, seq_lens, _, _ = _build_decode_inputs( + batch_size, next_n, index_topk, num_tokens, compress_ratio + ) + num_rows = logits.shape[0] + + # First invocation: caller-owned aux buffer A. + aux_a_idx, aux_a_log = _build_radix_aux_buffers(num_rows, index_topk) + indices_a = torch.empty((num_rows, index_topk), dtype=torch.int32, device="cuda") torch.ops.trtllm.indexer_topk_decode( logits, seq_lens, - indices, + indices_a, next_n, index_topk, - None, # pre_idx - None, # heuristic_scratch - done_counter, - scratch, - False, # is_prefill + compress_ratio=compress_ratio, + radix_aux_indices=aux_a_idx, + radix_aux_logits=aux_a_log, ) torch.cuda.synchronize() - max_row_len = row_ends.max().item() - torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] - mask = (torch_indices >= 0) & ((torch_indices - (row_ends - row_starts)[:, None]) < 0) - torch_indices = torch_indices.masked_fill(~mask, -1) + # Second invocation with a freshly-allocated aux buffer B at a distinct + # backing address — the only intentional difference vs invocation A. + aux_b_idx, aux_b_log = _build_radix_aux_buffers(num_rows, index_topk) + indices_b = torch.empty((num_rows, index_topk), dtype=torch.int32, device="cuda") + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices_b, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=aux_b_idx, + radix_aux_logits=aux_b_log, + ) + torch.cuda.synchronize() - assert compare_top_k_results( - logits, indices, torch_indices, row_starts, row_ends, index_topk - ), "opt-in indexer_topk_decode results don't match torch.topk" - - -@pytest.mark.parametrize("batch_size,num_tokens", [(1, 524288), (4, 262144), (16, 131072)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) -def test_indexer_topk_decode_multi_pass_radix(batch_size, num_tokens, dtype): - """Multi-pass radix path: shapes inside the low-bs / long-seq eligibility - zone with a caller-allocated uint8 scratch buffer sized by - indexer_topk_decode_scratch_bytes. Verified for fp32, bf16, and fp16.""" - index_topk = 2048 - scratch_bytes = torch.ops.trtllm.indexer_topk_decode_scratch_bytes( - batch_size, num_tokens, index_topk + # Compare by selected-logit VALUES rather than bit-identical indices — + # the Radix kernel uses histogram-radix sort with atomic ops and may + # return different (but equally valid) index orderings across calls on + # ties. Top-K logit values are the correctness invariant; existing + # `compare_top_k_results` follows the same principle. + k_effective = min(index_topk, int(logits.size(1))) + values_a = _topk_logit_values_sorted(logits, indices_a, k_effective) + values_b = _topk_logit_values_sorted(logits, indices_b, k_effective) + assert torch.allclose(values_a, values_b, rtol=0.0, atol=0.0), ( + "Top-K logit values diverged across two caller-owned-aux invocations" ) - assert scratch_bytes > 0, "scratch size must be positive" - # torch.zeros (not empty) so the first call sees a clean per-row state; - # subsequent calls are kept clean by the kernel's pass-3 trailer. - scratch = torch.zeros(scratch_bytes, dtype=torch.uint8, device="cuda") - _run_opt_in_decode(batch_size, num_tokens, dtype=dtype, done_counter=None, scratch=scratch) -@pytest.mark.parametrize("batch_size,num_tokens", [(4, 524288), (8, 262144), (16, 524288)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) -def test_indexer_topk_decode_multi_pass_radix_via_legacy_api(batch_size, num_tokens, dtype): - """Multi-pass radix tier reached through the legacy calling convention: - the caller passes the deprecated `done_counter_scratch` and no `scratch`, - so the thop wrapper allocates the radix scratch internally. `done_counter` - is accepted but ignored (its only effect now is a deprecation warning). - Verified for fp32, bf16, and fp16.""" - done_counter = torch.zeros(batch_size, dtype=torch.int32, device="cuda") - _run_opt_in_decode(batch_size, num_tokens, dtype=dtype, done_counter=done_counter, scratch=None) +@pytest.mark.parametrize( + "batch_size,next_n,index_topk,num_tokens,compress_ratio", + [ + # Same shapes as the equivalence test, exercising blocks_per_row > 1. + (8, 1, 512, 16384, 4), + (32, 1, 512, 16384, 4), + ], +) +def test_indexer_topk_decode_radix_aux_cuda_graph_replay( + batch_size, next_n, index_topk, num_tokens, compress_ratio +): + """Regression test for the original CUDA-Graph stale-pointer bug. Capture + the Radix path with caller-owned aux into a CUDA Graph, replay multiple + times, and verify the output matches a non-graph reference and stays + stable across replays. With the pre-fix per-call `th::empty` aux the + captured kernels could write to recycled memory between replays — the + persistent aux buffer eliminates that hazard.""" + logits, seq_lens, _, _ = _build_decode_inputs( + batch_size, next_n, index_topk, num_tokens, compress_ratio + ) + num_rows = logits.shape[0] + # Non-graph reference (uses the same caller-owned aux path). + aux_indices, aux_logits = _build_radix_aux_buffers(num_rows, index_topk) + indices_ref = torch.empty((num_rows, index_topk), dtype=torch.int32, device="cuda") + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices_ref, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + ) + torch.cuda.synchronize() -@skip_pre_hopper -@pytest.mark.parametrize("batch_size", _PREFILL_BATCH_SIZES) -@pytest.mark.parametrize("index_topk", [2048, 128]) -@pytest.mark.parametrize("num_tokens", _PREFILL_NUM_TOKENS) -def test_indexer_topk_prefill(batch_size, index_topk, num_tokens): - """Verify indexer_topk_prefill output matches torch.topk for variable-length rows.""" + # CUDA-Graph capture + replay using the SAME aux tensors. The captured + # kernel descriptors hold the aux pointers; persistent allocation keeps + # them valid across replays. + indices_graph = torch.empty((num_rows, index_topk), dtype=torch.int32, device="cuda") + + # Warm-up outside capture (recommended for stable allocator state). + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices_graph, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + ) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices_graph, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + ) + + # Replay several times; selected top-K logit values must match the + # non-graph reference and stay invariant across replays. (Indices + # themselves may reorder among equal-logit candidates due to atomic + # tie-breaking in the histogram-radix sort — value equality is the + # correctness invariant.) + k_effective = min(index_topk, int(logits.size(1))) + ref_values = _topk_logit_values_sorted(logits, indices_ref, k_effective) + for replay in range(8): + indices_graph.zero_() + g.replay() + torch.cuda.synchronize() + replay_values = _topk_logit_values_sorted(logits, indices_graph, k_effective) + assert torch.allclose(replay_values, ref_values, rtol=0.0, atol=0.0), ( + f"CUDA Graph replay #{replay} top-K logit values diverged from non-graph reference" + ) + + +def test_indexer_topk_decode_radix_aux_validation(): + """Caller-owned aux validation: wrong dtype, undersized buffer, wrong + device, and non-contiguous tensors must each raise.""" + batch_size, next_n, index_topk, num_tokens, compress_ratio = 8, 1, 512, 16384, 4 + logits, seq_lens, _, _ = _build_decode_inputs( + batch_size, next_n, index_topk, num_tokens, compress_ratio + ) + num_rows = logits.shape[0] + indices = torch.empty((num_rows, index_topk), dtype=torch.int32, device="cuda") + + good_idx, good_log = _build_radix_aux_buffers(num_rows, index_topk) + + def call(ai, al): + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices, + next_n, + index_topk, + compress_ratio=compress_ratio, + radix_aux_indices=ai, + radix_aux_logits=al, + ) + + # Wrong dtype for indices buffer (float32 instead of int32). + bad_idx_dtype = torch.empty_like(good_idx, dtype=torch.float32) + with pytest.raises(RuntimeError, match="radix_aux_indices must be int32"): + call(bad_idx_dtype, good_log) + + # Wrong dtype for logits buffer (int32 instead of float32). + bad_log_dtype = torch.empty_like(good_log, dtype=torch.int32) + with pytest.raises(RuntimeError, match="radix_aux_logits must be float32"): + call(good_idx, bad_log_dtype) + + # Undersized buffer (fewer elements than num_rows*blocks_per_row*index_topk). + too_small_idx = torch.empty((num_rows, 1, index_topk), dtype=torch.int32, device="cuda") + too_small_log = torch.empty((num_rows, 1, index_topk), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="at least num_rows"): + call(too_small_idx, too_small_log) + + # Non-contiguous buffer. + non_contig_idx = good_idx.transpose(0, 1) # makes it non-contiguous + assert not non_contig_idx.is_contiguous() + with pytest.raises(RuntimeError, match="must be contiguous"): + call(non_contig_idx, good_log) + + +def _run_indexer_topk_prefill_check(batch_size, index_topk, num_tokens): + """Run the prefill equivalence check against torch.topk.""" torch.manual_seed(24) torch.cuda.manual_seed(24) @@ -375,6 +661,15 @@ def test_indexer_topk_prefill(batch_size, index_topk, num_tokens): ), "CUDA top_k_per_row results don't match torch.topk" +@skip_pre_hopper +@pytest.mark.parametrize("batch_size", _PREFILL_BATCH_SIZES) +@pytest.mark.parametrize("index_topk", [2048, 128]) +@pytest.mark.parametrize("num_tokens", _PREFILL_NUM_TOKENS) +def test_indexer_topk_prefill(batch_size, index_topk, num_tokens): + """Verify indexer_topk_prefill output matches torch.topk for variable-length rows.""" + _run_indexer_topk_prefill_check(batch_size, index_topk, num_tokens) + + # ============================================================================ # CuTE DSL Top-K Tests # ============================================================================ @@ -874,6 +1169,131 @@ def generate_pre_idx( return pre_idx +def apply_mtp_structure_compressed( + logits: torch.Tensor, + batch_size: int, + next_n: int, + row_ends: torch.Tensor, +) -> torch.Tensor: + """ + cr=4-safe variant of apply_mtp_structure. + + apply_mtp_structure assumes ``row_ends[b*next_n + nni] == row_ends[b*next_n] + + nni`` (the V3.2 cr=1 invariant where each MTP draft adds exactly one KV + token). Under cr=4, ``row_ends = floor(actual_kv_len / 4)`` and that + invariant breaks: when ``actual_kv_len[base] mod 4`` lies in {1, 2, 3} + (75% of seq_lens) we get ``row_ends[base+nni] == row_ends[base]``, so the + copy ``[nni : nni+valid_base]`` overruns ``row_ends[base+nni]`` and writes + finite values into what create_distributed_logits left as -inf. The + polluted positions then leak into torch.topk's reference (which doesn't + know about the row's true compressed N), producing off-by-one counts vs. + the kernel. + + This variant clips the per-row copy length to fit within row b*next_n+nni's + valid compressed range, preserving MTP correlation where it fits and + leaving -inf positions untouched. + """ + if next_n == 1: + return logits + + for b in range(batch_size): + base = b * next_n + valid_base = int(row_ends[base].item()) + for nni in range(1, next_n): + row = base + nni + valid_row = int(row_ends[row].item()) + # Largest copy_len such that [nni, nni+copy_len) ⊆ [0, valid_row). + copy_len = max(0, min(valid_base, valid_row - nni)) + if copy_len > 0: + logits[row, nni : nni + copy_len] = logits[base, :copy_len] + + return logits + + +def generate_pre_idx_v4( + logits: torch.Tensor, + row_ends: torch.Tensor, + batch_size: int, + next_n: int, + index_topk: int, + success_ratio: float = 0.6, + seed: int = 0, +) -> torch.Tensor: + """ + DSv4 (compress_ratio=4) variant of generate_pre_idx — no `-1` shift. + + Unlike V3.2 where the kernel applies preIdxOffset = (rowIdx % next_n) + 1 + to every preIdx entry (KV grew by 1 per decode step in uncompressed space), + the V4 indexer operates in compressed-token-index space where consecutive + decode steps may add 0 or 1 compressed entries (each compressed entry + fuses 4 real tokens). Per-row Δc varies with prev kv_len mod 4 alignment, + but new compressed entries are always appended at the end so prev-step + indices in [0, c_prev-1] remain valid as-is in [0, c_curr-1]. The kernel + therefore forces preIdxOffset = 0 when compressRatio != 1, and tests must + pass preIdx in CURRENT-step coordinates (no -1 shift). + + Structure of the returned pre_idx[b]: + - pre_idx[b, 0] = argmax of the base row (kernel invariant) + - floor(K * success_ratio) slots from the actual top-K (without replace) + - remaining slots from non-top-K pool (without replace) + + Edge case (valid_len < K) handled identically to generate_pre_idx. + + Args: + logits, row_ends, batch_size, next_n, index_topk, success_ratio, seed: + See generate_pre_idx — the V4 helper mirrors its sampling logic. + + Returns: + pre_idx: int32 tensor of shape (batch_size, index_topk), entries in + the compressed current-step index space (no negative entries since + the kernel uses offset = 0). + """ + torch.manual_seed(seed) + pre_idx = torch.zeros(batch_size, index_topk, dtype=torch.int32, device=logits.device) + + for b in range(batch_size): + base = b * next_n + valid_len = int(row_ends[base].item()) + k = min(index_topk, valid_len) + + # Actual top-K of the base row; index 0 = argmax (kernel invariant). + _, topk_idx = logits[base, :valid_len].topk(k) + + n_hit = max(1, int(k * success_ratio)) + n_hit = min(n_hit, k) + + if n_hit > 1: + perm = torch.randperm(k - 1, device=logits.device)[: n_hit - 1] + hits = torch.cat([topk_idx[:1], topk_idx[1:][perm]]) + else: + hits = topk_idx[:1] + + pre_idx[b, :n_hit] = hits.int() + + n_fill = index_topk - n_hit + if n_fill > 0: + topk_mask = torch.zeros(valid_len, dtype=torch.bool, device=logits.device) + topk_mask[topk_idx] = True + non_topk = torch.where(~topk_mask)[0] + + if len(non_topk) >= n_fill: + perm = torch.randperm(len(non_topk), device=logits.device)[:n_fill] + pre_idx[b, n_hit:] = non_topk[perm].int() + else: + pre_idx[b, n_hit : n_hit + len(non_topk)] = non_topk.int() + leftover = n_fill - len(non_topk) + topk_tail = topk_idx[n_hit:] + take = min(leftover, len(topk_tail)) + if take > 0: + pre_idx[b, n_hit + len(non_topk) : n_hit + len(non_topk) + take] = topk_tail[ + :take + ].int() + + # No shift: kernel reads input[preIdx[i] + 0] = input[preIdx[i]] directly + # in compressed current-step coordinates. + return pre_idx + + @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @pytest.mark.parametrize("batch_size", [1, 4, 64, 256]) @@ -1093,8 +1513,21 @@ def test_indexer_topk_decode_dist( # 4. Run heuristic CUDA kernel — heuristic_scratch dtype must match logits. indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") heuristic_scratch = torch.empty(num_gen_tokens * index_topk, dtype=dtype, device="cuda") + # Supply Radix split-work aux scratch. For dtype=fp32 with num_columns + # below kSeqSmall the dispatcher falls through GVR to the Radix path and + # the cpp op rejects blocks_per_row > 1 without caller-owned scratch; for + # bf16/fp16 these kwargs are simply ignored. + radix_aux_indices, radix_aux_logits = _build_radix_aux_buffers(num_gen_tokens, index_topk) torch.ops.trtllm.indexer_topk_decode( - logits, seq_lens, indices, next_n, index_topk, pre_idx, heuristic_scratch + logits, + seq_lens, + indices, + next_n, + index_topk, + pre_idx, + heuristic_scratch, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits, ) torch.cuda.synchronize() @@ -1125,53 +1558,107 @@ def test_indexer_topk_decode_dist( ) -@skip_pre_blackwell -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -@pytest.mark.parametrize("batch_size", [1, 64]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) -def test_indexer_topk_decode_gvr_heuristic(index_topk, batch_size, dtype): - """Exercise the GVR Heuristic dispatcher tier directly. - - The split-work tests above always pass ``pre_idx=None`` (single-block / - multi-pass-radix tiers) at K=2048, so they never reach the GVR-eligibility - branch of ``invokeIndexerTopKDecodeImpl`` that this PR's dispatcher change - rewrote. This test pins shapes squarely inside the GVR window so the - dispatcher must select the heuristic path, and adds the missing K=512 / - K=1024 coverage: - - - pre_idx + heuristic_scratch provided, - - K in {512, 1024, 2048} (the dispatcher's isSupportedTopK set), - - full-length rows so numColumns == num_tokens == 16384 lands in - [kSeqSmall (default 12288), splitWorkThreshold (200000)) regardless of - RNG (random seq_lens could drop a short batch below kSeqSmall and - silently reroute to the single-block tier), - - numRows (<= 64) well under the architecture wave/L2 bound. - """ - next_n = 1 - num_tokens = 16384 +# ============================================================================ +# DSv4 Heuristic Decode Test (compress_ratio = 4) +# ============================================================================ +# +# Exercises the V4 indexer GVR Top-K path enabled by the +# `compressRatio == 1 || compressRatio == 4` relaxation in +# canUseHeuristic (cpp/tensorrt_llm/kernels/indexerTopK.cu). For +# compressRatio != 1 the kernel: +# 1. Computes N = (seq_len - next_n + (rowIdx % next_n) + 1) / compressRatio, +# i.e. the row's compressed-KV length (vs. uncompressed N in the V3.2 +# path). +# 2. Forces preIdxOffset = 0 (vs. (rowIdx % next_n) + 1 in V3.2), since +# compressed entries are appended at the end of the compressed KV and +# prev-step indices remain valid as-is. +# +# To reach the GVR (Heuristic) path with cr=4 we need the *compressed* +# numColumns ≥ kSeqSmall (≈12288), so the test uses num_tokens ∈ +# {65536, 131072} which gives compressed range ≈ {16K, 32K}. Smaller cr=4 +# cases (where compressed N falls below kSeqSmall) are already covered by +# test_indexer_topk_decode parametrized on compress_ratio ∈ [1, 4] — those +# exercise the Radix/Insertion fallback for the same gate. + + +def _run_indexer_topk_decode_v4_gvr_check( + batch_size: int, + next_n: int, + index_topk: int, + num_tokens: int, + dtype: torch.dtype, + dist_cfg: dict, + success_ratio: float, +): + """Run the V4 (compress_ratio=4) heuristic indexer_topk_decode check.""" torch.manual_seed(24) torch.cuda.manual_seed(24) + compress_ratio = 4 num_gen_tokens = batch_size * next_n row_starts = torch.zeros(num_gen_tokens, dtype=torch.int32, device="cuda") - # Full-length rows so numColumns deterministically lands in the GVR window. - seq_lens = torch.full((batch_size,), num_tokens, dtype=torch.int32, device="cuda") - row_ends = seq_lens # next_n == 1 => one row per batch element + row_indices = torch.arange(num_gen_tokens, device="cuda") // next_n + next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n - logits = create_random_logits(row_starts, row_ends, dtype, 42) - pre_idx = generate_pre_idx( - logits, row_ends, batch_size, next_n, index_topk, success_ratio=0.6, seed=7 + # Uncompressed seq_lens are what the kernel receives in `seq_lens`. + # Clamp so that compressed_actual_kv_len ≥ kSeqSmall (= 12288) for every + # row; the kernel will divide actual_kv_len by compress_ratio internally, + # so a floor of (kSeqSmall + 1) * compress_ratio + next_n on the + # uncompressed seq_len guarantees compressed N stays in the GVR window. + min_uncompressed = (12288 + 1) * compress_ratio + next_n + seq_lens = generate_seq_lens(batch_size, min_uncompressed, num_tokens) + seq_lens = seq_lens.clamp(min=min_uncompressed) + + # row_ends is the compressed-KV length per row (= what logits' columns + # represent in V4 — the indexer operates in compressed-token-index space). + actual_kv_lens = seq_lens[row_indices] - next_n + next_n_offset + 1 + row_ends = actual_kv_lens // compress_ratio + + # 1. Sample logits over the compressed shape. + logits = create_distributed_logits(dist_cfg, row_starts, row_ends, dtype, seed=42) + + # 2. Apply MTP correlation between rows within each batch element. + # Use the compressed-aware variant: cr=4 breaks the cr=1 invariant + # row_ends[base+nni] = row_ends[base]+nni, so the copy length must be + # clipped per-row to avoid overrunning the row's valid range. + if next_n > 1: + logits = apply_mtp_structure_compressed(logits, batch_size, next_n, row_ends) + + # 3. Build heuristic pre-prediction indices — V4 variant (no -1 shift). + pre_idx = generate_pre_idx_v4( + logits, + row_ends, + batch_size, + next_n, + index_topk, + success_ratio=success_ratio, + seed=7, ) + # 4. Run heuristic CUDA kernel with compress_ratio=4. The kernel: + # - reads logits in compressed-index space (numColumns = logits.shape[1]) + # - divides seq_lens by compress_ratio to derive per-row N + # - uses preIdxOffset = 0 (preIdx already in current-step coords) indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") - # heuristic_scratch dtype must match logits dtype; size >= numRows * index_topk. heuristic_scratch = torch.empty(num_gen_tokens * index_topk, dtype=dtype, device="cuda") - + # Supply Radix split-work aux scratch — same rationale as the V3.2 helper: + # required by the cpp op when blocks_per_row > 1, harmless otherwise. + radix_aux_indices, radix_aux_logits = _build_radix_aux_buffers(num_gen_tokens, index_topk) torch.ops.trtllm.indexer_topk_decode( - logits, seq_lens, indices, next_n, index_topk, pre_idx, heuristic_scratch + logits, + seq_lens, + indices, + next_n, + index_topk, + pre_idx, + heuristic_scratch, + compress_ratio=compress_ratio, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits, ) torch.cuda.synchronize() + # 5. Reference: torch.topk masked to the compressed row_ends. max_row_len = int(row_ends.max().item()) torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] mask = (torch_indices >= 0) & ((torch_indices - (row_ends - row_starts)[:, None]) < 0) @@ -1179,4 +1666,49 @@ def test_indexer_topk_decode_gvr_heuristic(index_topk, batch_size, dtype): assert compare_top_k_results( logits, indices, torch_indices, row_starts, row_ends, index_topk - ), f"GVR heuristic indexer_topk_decode mismatch: K={index_topk}, bs={batch_size}, dtype={dtype}" + ), ( + f"V4 heuristic indexer_topk_decode (cr=4) mismatch: dist={dist_cfg['dist']}, " + f"mean={dist_cfg['mean']}, std={dist_cfg['std']}, batch_size={batch_size}, " + f"next_n={next_n}, index_topk={index_topk}, num_tokens={num_tokens}, " + f"success_ratio={success_ratio}, dtype={dtype}" + ) + + +# Param matrix is intentionally tighter than test_indexer_topk_decode_dist: +# only one logit distribution and one success_ratio because the GVR algorithm +# is dist-/hint-quality-invariant for correctness (an exact algorithm). The +# axes that *do* differ in V4 vs V3.2 are exercised in full: +# compress_ratio = 4 (fixed — sole purpose of this test) +# next_n in {1, 2, 3} (decode + MTP windows) +# index_topk in {512, 1024, 2048} (all GVR-supported K) +# num_tokens in {65536, 131072} (compressed N ≈ 16K and 32K) +# dtype: fp32 / bf16 / fp16 (both kernel templates) +# batch_size: 1 (single-row), 64 (multi-row) +@skip_pre_blackwell +@pytest.mark.skipif(not _HAS_SCIPY, reason="scipy required for distribution tests") +@pytest.mark.parametrize("success_ratio", [0.7]) +@pytest.mark.parametrize("batch_size", [1, 64]) +@pytest.mark.parametrize("next_n", [1, 2, 3]) +@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) +@pytest.mark.parametrize("num_tokens", [65536, 131072]) +@pytest.mark.parametrize( + "dtype", + [torch.float32, torch.bfloat16, torch.float16], + ids=["fp32", "bf16", "fp16"], +) +def test_indexer_topk_decode_dist_v4_cr4( + batch_size, next_n, index_topk, num_tokens, success_ratio, dtype +): + """ + Correctness test for the DSv4 heuristic indexer_topk_decode with + compress_ratio=4 across MTP windows, all GVR-supported K, and all + supported logit dtypes. Uses one representative distribution; broader + distribution coverage is left to test_indexer_topk_decode_dist (cr=1). + """ + # Logistic chosen as the single representative distribution — its + # heavy-tailed symmetric shape produces the wide K-th-value spread that + # stresses GVR's secant threshold search most. + dist_cfg = dict(dist="logistic", mean=-0.47, std=1.46, full_range=12.32) + _run_indexer_topk_decode_v4_gvr_check( + batch_size, next_n, index_topk, num_tokens, dtype, dist_cfg, success_ratio + )