diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 6e6e7b29b059..656675e641ad 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -34,20 +34,16 @@ namespace kernels // (a value <= 0 selects the internal default). int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0); -/// 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}) +/// fp32 indexer TopK decode — three dispatch tiers: /// - 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, + int const stride1, int const next_n, int const topK = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); -/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except -/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and +/// bf16 indexer TopK decode — same dispatch tiers as the fp32 entry, except /// 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 @@ -57,35 +53,17 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic /// 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, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const next_n, int const topK = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); /// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, - int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); + int const topK = 2048, 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); -/// 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 TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu deleted file mode 100644 index 839ae5483c05..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu +++ /dev/null @@ -1,285 +0,0 @@ -/* - * 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/kernels/heuristicTopKDecode.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/envUtils.h" - -// Import gvrTopKJob (__device__ __noinline__, the GVR micro-kernel) and -// all helpers. gvrTopKJob is independently optimized by ptxas, matching standalone -// SASS quality regardless of the caller's prologue code. -#include "tensorrt_llm/kernels/heuristic_topk.cuh" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -using heuristic_topk::BLOCK_SIZE; -using heuristic_topk::GvrDtypeTraits; -using heuristic_topk::GvrParams; -using heuristic_topk::gvrTopKJob; -using heuristic_topk::gvrTopKJobDtype; -using heuristic_topk::KernelSmemTplK; - -// Templated on TopK so the launcher can dispatch K=512/1024/2048 to the -// 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, int compressRatio) -{ - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - - int const rowIdx = blockIdx.x; - int const seq_len = seqLens[rowIdx / next_n]; - // 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; - float* __restrict__ outputValues = scratchValues + static_cast(rowIdx) * topK; - int* __restrict__ outputIndices = outIndices + static_cast(rowIdx) * topK; - - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - if (N <= topK) - { - int const tid = threadIdx.x; - for (int i = tid; i < N; i += BLOCK_SIZE) - { - outputValues[i] = input[i]; - outputIndices[i] = i; - } - for (int i = N + tid; i < topK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - return; - } - - // 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(); -#endif -} - -// ============================================================================ -// Multi-dtype path (bf16 / fp16) -// ============================================================================ -// Mirrors heuristicTopKMultiRowKernel for bf16/fp16 inputs. The kernel body -// is structurally identical; only the input/output dtype, the smem-key -// dtype, and the GVR job (gvrTopKJobDtype) differ. - -// 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, 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 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; - InputT* __restrict__ outputValues = scratchValues + static_cast(rowIdx) * topK; - int* __restrict__ outputIndices = outIndices + static_cast(rowIdx) * topK; - - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - if (N <= topK) - { - int const tid = threadIdx.x; - for (int i = tid; i < N; i += BLOCK_SIZE) - { - outputValues[i] = input[i]; - outputIndices[i] = i; - } - InputT const neg_max = GvrDtypeTraits::from_fp32(-FLT_MAX); - for (int i = N + tid; i < topK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - return; - } - - // 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)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// 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, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>( - __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, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>( - __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, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>( - __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, int); -template __global__ void heuristicTopKMultiRowKernel<1024>( - 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, int); - -// Dispatch on topK at runtime — each TopK-instantiation gets its own smem -// size (driven by GvrParams::kC/kNumBins) and own kfn pointer -// (cudaFuncSetAttribute / cudaLaunchKernelEx target the right kernel). -// -// fp32 routes to heuristicTopKMultiRowKernel; bf16/fp16 route to -// heuristicTopKMultiRowKernelDtype. Vector-load alignment -// requirement is 4 elements for fp32 (float4) and 8 elements for bf16/fp16 -// (int4 of 16-bit). In TRT-LLM the logits stride is always a multiple of -// tokens_per_block (≥64), so the alignment check is never hit at runtime -// — it's an assert against caller misuse. -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, - int compressRatio, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO( - topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}"); - - constexpr int kAlign = std::is_same_v ? 4 : 8; - TLLM_CHECK_WITH_INFO(stride0 % kAlign == 0 || numRows <= 1, - "heuristicTopKDecode requires logits stride0 divisible by %d for multi-row launch", kAlign); - - auto launchOne = [&]() - { - // bf16/fp16 path also uses fp32 keys[] in smem (down-conversion deferred). - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - size_t const smemSize = sizeof(SmemT); - - auto kfn = []() - { - if constexpr (std::is_same_v) - return heuristicTopKMultiRowKernel; - else - return heuristicTopKMultiRowKernelDtype; - }(); - - if (smemSize > 48u * 1024u) - { - cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemSize)); - } - - cudaLaunchConfig_t config; - config.gridDim = numRows; - config.blockDim = BLOCK_SIZE; - config.dynamicSmemBytes = smemSize; - 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, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK, - preIdxStride, preIdxCount, compressRatio); - }; - - switch (topK) - { - case 512: launchOne.template operator()<512>(); break; - case 1024: launchOne.template operator()<1024>(); break; - case 2048: launchOne.template operator()<2048>(); break; - default: TLLM_THROW("heuristicTopKDecode: topK validated above; unreachable"); - } -} - -} // anonymous namespace - -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, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - 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, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, - 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, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h deleted file mode 100644 index 0d2330f76545..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h +++ /dev/null @@ -1,61 +0,0 @@ -/* - * 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -inline constexpr int kHeuristicTopK = 2048; -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, - 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, - 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, - int compressRatio, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh deleted file mode 100644 index d75933bd5d38..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ /dev/null @@ -1,2006 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -// ============================================================================ -// heuristic_topk.cuh — Heuristic-Guided Top-K (Sort-Free, Histogram-Based) -// -// Outer name: "heuristic" (algorithm family + public dispatcher / launchers). -// Inner name: "gvr" (Guess-Verify-Refine) — the single-CTA single-row -// micro-kernel implementing the algorithm of: -// "Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding -// on Blackwell via Temporal Correlation" -// -// Optimised for NVIDIA B200 (Blackwell, sm_100), single thread-block kernel. -// -// GVR phase mapping: -// P1 (preIdx stats) ┐ Guess: estimate the K-th-value -// P2 (secant threshold search) ┘ threshold from previous-step top-K -// indices, then refine the guess via -// count-only secant iterations. -// P3 (collect) — Verify: scatter the elements that -// pass the guessed threshold into -// shared memory and confirm the -// candidate count is in the safe band. -// P4 (histogram snap + partition) — Refine: 2048-bin histogram snap to -// the exact K-th value, then partition -// the candidates into the output set. -// ============================================================================ - -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace heuristic_topk -{ - -// ============================================================================ -// Multi-dtype Trait Layer -// ============================================================================ -// Encapsulates dtype-specific cvt intrinsics + vector load width so the -// kernel body can be templated cleanly. For fp32 the trait is identity. -// -// Arithmetic (threshold, accumulators, bin index) is always fp32; only -// the HBM input container, smem keys, and output values follow InputT. -// GVR is HBM-bandwidth-bound, so fp32 ALU has no measurable cost. - -template -struct GvrDtypeTraits; - -template <> -struct GvrDtypeTraits -{ - using SmemKey = float; - static constexpr int VEC_W = 4; // int4 = 4 × fp32 - static constexpr int SMEM_KEY_BYTES = 4; - - __device__ static __forceinline__ float to_fp32(float v) - { - return v; - } - - __device__ static __forceinline__ float from_fp32(float v) - { - return v; - } - - __device__ static __forceinline__ void unpack4(int4 raw, float* out) - { - out[0] = __int_as_float(raw.x); - out[1] = __int_as_float(raw.y); - out[2] = __int_as_float(raw.z); - out[3] = __int_as_float(raw.w); - } -}; - -template <> -struct GvrDtypeTraits<__nv_bfloat16> -{ - using SmemKey = __nv_bfloat16; - static constexpr int VEC_W = 8; // int4 = 8 × bf16 = 4 × bf162 - static constexpr int SMEM_KEY_BYTES = 2; - - __device__ static __forceinline__ float to_fp32(__nv_bfloat16 v) - { - return __bfloat162float(v); - } - - __device__ static __forceinline__ __nv_bfloat16 from_fp32(float v) - { - return __float2bfloat16_rn(v); - } - - __device__ static __forceinline__ void unpack8(int4 raw, float* out) - { - auto* p = reinterpret_cast<__nv_bfloat162*>(&raw); -#pragma unroll - for (int j = 0; j < 4; j++) - { - out[2 * j] = __low2float(p[j]); - out[2 * j + 1] = __high2float(p[j]); - } - } -}; - -template <> -struct GvrDtypeTraits<__half> -{ - using SmemKey = __half; - static constexpr int VEC_W = 8; // int4 = 8 × fp16 = 4 × half2 - static constexpr int SMEM_KEY_BYTES = 2; - - __device__ static __forceinline__ float to_fp32(__half v) - { - return __half2float(v); - } - - __device__ static __forceinline__ __half from_fp32(float v) - { - return __float2half_rn(v); - } - - __device__ static __forceinline__ void unpack8(int4 raw, float* out) - { - auto* p = reinterpret_cast<__half2*>(&raw); -#pragma unroll - for (int j = 0; j < 4; j++) - { - out[2 * j] = __low2float(p[j]); - out[2 * j + 1] = __high2float(p[j]); - } - } -}; - -// ============================================================================ -// Configuration Constants -// ============================================================================ - -constexpr int BLOCK_SIZE = 512; -constexpr int WARP_SIZE = 32; -constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; - -constexpr int TOP_K = 2048; -constexpr int HEURISTIC_SIZE = 2048; -constexpr int SAFETY_MARGIN = 2048; -constexpr int MAX_CANDIDATES = TOP_K + SAFETY_MARGIN * 2; // 6144 - -constexpr int MAX_REFINE_ITERS = 15; -// Phase-3 repair budget: bisecting on the uint32 key image collapses any -// bracket to adjacent floats in <= 32 steps; 40 adds slack. -constexpr int MAX_REPAIR_ITERS = 40; -constexpr int NUM_BINS = 2048; - -static_assert(TOP_K % BLOCK_SIZE == 0); -static_assert(MAX_CANDIDATES % BLOCK_SIZE == 0); - -// ============================================================================ -// Multi-K Trait Layer -// ============================================================================ -// Per-(InputT, TopK) compile-time trait encoding the secant-search target -// `kFTarget`, candidate-buffer cap `kC`, and Phase-4 histogram bin count -// `kNumBins`. -// -// kFTarget under V3.2-decode preIdx semantics (preIdx = top-K of prev row): -// Phase-1 pmean lands near the right tail of the prev-row top-K, so the -// Phase-2 secant initial bracket is biased high. A tighter K-proportional -// target converges faster than the M=2048 era's flat target. -// -// `kFTarget` is the secant solver's **soft steering target**, not the -// convergence condition. The Phase-2 loop converges whenever the -// candidate count falls within `[kK, kCC]` (see the `done = 1` check at -// the end of `gvrTopKJob`'s P1+P2 scope). A `kFTarget` below `kK` is -// intentional and useful for small K — with preIdx-seeded P1 landing -// the initial threshold near the right tail of the prev-row top-K, the -// secant's first interpolation often overshoots; biasing the target -// below kK pulls the next iteration's threshold down more aggressively -// and reaches the legal `[kK, kCC]` band in fewer secant steps. The -// concrete multipliers below were tuned empirically over V4 M=K cells -// (commit 8 in this PR): -// K=512: 0.75K = 384 -// K=1024: 2.5K = 2560 -// K=2048: kept at 1.5K (3072) for fp32 to preserve SASS byte-identity -// with the V2e production hot path. bf16/fp16 K=2048 use 2K -// (4096), where there is no prior production baseline to honor. -// -// kC = 5120 for K=512/1024 across all dtypes — drops smem footprint enough -// to leave headroom for 4-5 CTA/SM theoretical occupancy when register -// allocation permits. fp32 K=2048 keeps kC=6144 to preserve V2e SASS -// byte-identity (production fp32 K=2048 hot path is the correctness -// floor; smem layout change would alter SASS). -// -// kNumBins varies per (T, K) by atomic-contention vs Phase-4 setup-cost -// trade-off: -// - Below ~1024 bins, atomicAdd contention on the bin-counter array -// dominates Phase-4 cost. -// - Above 1024 bins, the Phase-4 histogram clear+scan setup cost -// dominates over the contention savings. -// The optimum varies because (a) candidate-buffer size kC scales the -// atomic-contention denominator, and (b) bf16/fp16 paths read smem keys -// through Trait::to_fp32, shifting the Phase-3 ↔ Phase-4 ratio. Hence -// per-(T, K) tabulation rather than a closed-form rule. -// -// Primary template intentionally left undefined: any unsupported (T, K) -// combination triggers a compile-time error rather than a runtime fall- -// through. -template -struct GvrParams; // primary undefined → compile-time error for bad combos - -template <> -struct GvrParams -{ - // kFTarget=kK aligns the secant's soft steering target with the band's - // lower edge; eliminates the upper-clamp saturation on tight-σ + high-A2 - // layers (L36/L42/L28). Cross-prompt simulator validation on swe-bench - // 32k/64k/100k showed 2.19× / 1.77× / 1.51× total P2-iter reduction with - // zero cap-hits, zero per-layer regression vs the prior kFTarget=384. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -template <> -struct GvrParams -{ - // kFTarget = kK (see GvrParams rationale). Q9k Pro 32k - // K=1024 native sweep (M=K=1024) finds kFT=1024 reduces sum_mean - // P2 iters from 35.33 (kFT=2560) → 30.21 (1.17× speedup) with zero - // per-layer regression and zero cap-hits. The prior kFT=2560 setting - // was tuned with M=512 K=1024 (sparse_attention_config default - // index_topk=512 inherited Flash's K), which does not represent - // production Pro behavior (production: M = K). - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -// fp32 K=2048 preserves V2e SASS byte-identity with the production hot path. -// Changing kFTarget or kC would alter ptxas output for the existing -// `gvrTopKJob<2048>` / `heuristicTopKMultiRowKernel<2048>` instantiations. -template <> -struct GvrParams -{ - static constexpr int kFTarget = 3072; - static constexpr int kC = 6144; - static constexpr int kNumBins = NUM_BINS; -}; - -template <> -struct GvrParams<__nv_bfloat16, 512> -{ - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__nv_bfloat16, 1024> -{ - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__nv_bfloat16, 2048> -{ - static constexpr int kFTarget = 4096; - static constexpr int kC = 5120; - static constexpr int kNumBins = NUM_BINS; -}; - -template <> -struct GvrParams<__half, 512> -{ - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__half, 1024> -{ - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -template <> -struct GvrParams<__half, 2048> -{ - static constexpr int kFTarget = 4096; - static constexpr int kC = 5120; - static constexpr int kNumBins = NUM_BINS; -}; - -// kC must remain divisible by BLOCK_SIZE (vector loads). -static_assert(GvrParams::kC % BLOCK_SIZE == 0); -static_assert(GvrParams::kC % BLOCK_SIZE == 0); -static_assert(GvrParams::kC % BLOCK_SIZE == 0); - -// ============================================================================ -// Shared Memory Layout -// ============================================================================ -// Templated on (SmemKey, candidate-cap, num-bins). Default -// (MAX_CANDIDATES=6144, NUM_BINS=2048) sizes: -// fp32 : ~59 KB -// bf16/fp16 : ~47 KB -// K=512/1024 instantiations cap candidates at kC=5120 (~51 KB fp32 / -// ~41 KB bf16/fp16) per GvrParams::kC. - -template -struct KernelSmemTplK -{ - alignas(16) SmemKey keys[CCap]; // CCap × sizeof(SmemKey) (4B fp32 / 2B bf16/fp16) - alignas(16) int vals[CCap]; // CCap × 4B - - int warp_counts[NUM_WARPS]; // 64 B - int histogram[NumBinsT]; // NumBinsT × 4B (default 2048 → 8 KB) - int per_thread_counts[BLOCK_SIZE]; // cached from the most recent blockCountGE call (Phase-3 reuse) - - float threshold; - int cand_count; - int done; - - float val_lo, val_hi; - int cnt_lo, cnt_hi; - - float pmax_saved; - int out_count; -}; - -// Convenience alias for the default-cap layout (kC=6144, kNumBins=2048), -// used by the K=2048 instantiations. -template -using KernelSmemTpl = KernelSmemTplK; - -using KernelSmem = KernelSmemTpl; - -// ============================================================================ -// Warp-Level Reduction Primitives -// ============================================================================ - -#if __CUDA_ARCH__ >= 800 - -__device__ __forceinline__ int warpReduceSum(int val) -{ - return __reduce_add_sync(0xffffffffu, val); -} - -__device__ __forceinline__ unsigned floatToOrderedUint(float f) -{ - unsigned u = __float_as_uint(f); - return (u & 0x80000000u) ? ~u : (u | 0x80000000u); -} - -__device__ __forceinline__ float orderedUintToFloat(unsigned u) -{ - return __uint_as_float((u & 0x80000000u) ? (u & ~0x80000000u) : ~u); -} - -__device__ __forceinline__ float warpReduceMin(float val) -{ - unsigned u = floatToOrderedUint(val); - u = __reduce_min_sync(0xffffffffu, u); - return orderedUintToFloat(u); -} - -__device__ __forceinline__ float warpReduceMax(float val) -{ - unsigned u = floatToOrderedUint(val); - u = __reduce_max_sync(0xffffffffu, u); - return orderedUintToFloat(u); -} - -#else - -__device__ __forceinline__ int warpReduceSum(int val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val += __shfl_down_sync(0xffffffffu, val, off); - return val; -} - -__device__ __forceinline__ float warpReduceMin(float val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val = fminf(val, __shfl_xor_sync(0xffffffffu, val, off)); - return val; -} - -__device__ __forceinline__ float warpReduceMax(float val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val = fmaxf(val, __shfl_xor_sync(0xffffffffu, val, off)); - return val; -} - -#endif - -// ============================================================================ -// Order-preserving float <-> uint32 map (arch-independent) -// ============================================================================ -// Same bijection as floatToOrderedUint above but defined for every -// __CUDA_ARCH__; the Phase-3 repair bisects on this key image so the -// bracket provably collapses (a float-average midpoint has no such bound). -__device__ __forceinline__ unsigned gvrOrderKey(float f) -{ - unsigned u = __float_as_uint(f); - return (u & 0x80000000u) ? ~u : (u | 0x80000000u); -} - -__device__ __forceinline__ float gvrOrderKeyToFloat(unsigned u) -{ - return __uint_as_float((u & 0x80000000u) ? (u & ~0x80000000u) : ~u); -} - -// ============================================================================ -// Device: Block count ≥ threshold in GLOBAL memory (1-sync pattern) -// ============================================================================ - -// Templated on SmemT so K=512/1024 paths can pass a kC=5120 layout. -// Default (KernelSmem) targets the K=2048 kC=6144 layout. -template -__device__ __forceinline__ void blockCountGE( - float const* __restrict__ input, int N, float threshold, SmemT* smem, int tid, int warp_id, int lane) -{ - int c = 0; - for (int i = tid * 4; i + 3 < N; i += BLOCK_SIZE * 4) - { - float4 v4 = __ldg(reinterpret_cast(input + i)); - c += (v4.x >= threshold) + (v4.y >= threshold) + (v4.z >= threshold) + (v4.w >= threshold); - } - for (int i = (N & ~3) + tid; i < N; i += BLOCK_SIZE) - c += (__ldg(&input[i]) >= threshold); - - // cache per-thread count for Phase 3 sub-pass 1 reuse - smem->per_thread_counts[tid] = c; - - c = warpReduceSum(c); - - if (lane == 0) - smem->warp_counts[warp_id] = c; - __syncthreads(); - - if (tid == 0) - { - int t = 0; - for (int w = 0; w < NUM_WARPS; w++) - t += smem->warp_counts[w]; - smem->cand_count = t; - } -} - -// ============================================================================ -// Fused snap iteration (2 syncs per call) -// ============================================================================ - -// Templated on (TopK, SmemT) so K=512/1024 paths reuse the same helper. -template -__device__ __forceinline__ void blockFusedSnapIter(SmemT* smem, int count, int tid, int warp_id, int lane) -{ - float const thr = smem->threshold; - - int lge = 0, lgt = 0; - float s_up = FLT_MAX, s_down = -FLT_MAX; - - for (int i = tid; i < count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; - lge += (v >= thr); - lgt += (v > thr); - if (v > thr) - s_up = fminf(s_up, v); - if (v < thr) - s_down = fmaxf(s_down, v); - } - - int packed = (lge << 16) | lgt; - packed = warpReduceSum(packed); - s_up = warpReduceMin(s_up); - s_down = warpReduceMax(s_down); - - if (lane == 0) - { - smem->warp_counts[warp_id] = packed; - smem->histogram[warp_id] = __float_as_int(s_up); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(s_down); - } - __syncthreads(); - - if (tid == 0) - { - int tp = 0; - float total_up = FLT_MAX, total_down = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - tp += smem->warp_counts[w]; - total_up = fminf(total_up, __int_as_float(smem->histogram[w])); - total_down = fmaxf(total_down, __int_as_float(smem->histogram[NUM_WARPS + w])); - } - smem->cnt_lo = tp >> 16; - smem->cnt_hi = tp & 0xFFFF; - - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - - if (cgt >= TopK) - { - if (total_up < FLT_MAX) - smem->threshold = total_up; - } - else if (cge < TopK) - { - if (total_down > -FLT_MAX) - smem->threshold = total_down; - } - } - __syncthreads(); -} - -// ============================================================================ -// Dtype-templated helpers (bf16 / fp16) -// ============================================================================ -// Mirror of `blockCountGE` for bf16/fp16 inputs (8-wide vector load via -// `Trait::unpack8` + fp32 up-cast before threshold compare). The Phase-4 -// snap iter is NOT mirrored: `gvrTopKJobDtype` stores smem `keys[]` as -// fp32 even on bf16/fp16 paths (deferred-conversion optimization, see -// the `gvrTopKJobDtype` comment block below), so it reuses the fp32 -// `blockFusedSnapIter` helper directly. - -// Templated on SmemT so K=512/1024 dtype paths can pass a kC=5120 layout. -// Default targets the K=2048 kC=6144 layout. -template ::SmemKey>> -__device__ __forceinline__ void blockCountGEDtype( - InputT const* __restrict__ input, int N, float threshold, SmemT* smem, int tid, int warp_id, int lane) -{ - using Trait = GvrDtypeTraits; - static_assert(Trait::VEC_W == 8, "blockCountGEDtype is for bf16/fp16 (8-wide vector); use blockCountGE for fp32"); - - int c = 0; - for (int i = tid * 8; i + 7 < N; i += BLOCK_SIZE * 8) - { - int4 raw = __ldg(reinterpret_cast(input + i)); - float v[8]; - Trait::unpack8(raw, v); -#pragma unroll - for (int j = 0; j < 8; j++) - c += (v[j] >= threshold); - } - for (int i = (N & ~7) + tid; i < N; i += BLOCK_SIZE) - c += (Trait::to_fp32(__ldg(&input[i])) >= threshold); - - smem->per_thread_counts[tid] = c; - - c = warpReduceSum(c); - - if (lane == 0) - smem->warp_counts[warp_id] = c; - __syncthreads(); - - if (tid == 0) - { - int t = 0; - for (int w = 0; w < NUM_WARPS; w++) - t += smem->warp_counts[w]; - smem->cand_count = t; - } -} - -// ============================================================================ -// Device function: algorithm body (independently optimized by ptxas) -// __noinline__ ensures ptxas allocates registers and schedules instructions -// for this function independently from the caller, matching standalone SASS. -// ============================================================================ - -// Templated on TopK so K=512/1024/2048 fp32 paths share this body. The -// runtime `topK` parameter is kept for header-API compatibility with -// callers that pass it; the kernel asserts at entry that it matches -// the template instantiation. -template -__device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int const N, int const* __restrict__ preIdx, - int const M, int const topK, float* __restrict__ outputValues, int* __restrict__ outputIndices, - KernelSmemTplK::kC, GvrParams::kNumBins>* smem, - int const preIdxOffset = 0) -{ - using Params = GvrParams; - constexpr int kK = TopK; - constexpr int kCC = Params::kC; - constexpr int kBins = Params::kNumBins; - constexpr int kFTarget = Params::kFTarget; - - int const tid = threadIdx.x; - int const warp_id = tid / WARP_SIZE; - int const lane = tid & (WARP_SIZE - 1); - unsigned const full_mask = 0xffffffffu; - - { - // ================================================================ - // Phase 1 (GVR Guess, part 1) — Min/Max/Mean of pre-indexed values - // ================================================================ - - float local_min = FLT_MAX; - float local_max = -FLT_MAX; - float local_sum = 0.0f; - int local_cnt = 0; - for (int i = tid; i < M; i += BLOCK_SIZE) - { - int idx = __ldg(&preIdx[i]) + preIdxOffset; - if (idx >= 0 && idx < N) - { - float v = __ldg(&input[idx]); - local_min = fminf(local_min, v); - local_max = fmaxf(local_max, v); - local_sum += v; - local_cnt++; - } - } - - float wmin = warpReduceMin(local_min); - float wmax = warpReduceMax(local_max); - float wsum = local_sum; -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - wsum += __shfl_down_sync(0xffffffffu, wsum, off); - int wcnt = warpReduceSum(local_cnt); - - if (lane == 0) - { - smem->histogram[warp_id] = __float_as_int(wmin); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(wmax); - smem->histogram[NUM_WARPS * 2 + warp_id] = __float_as_int(wsum); - smem->histogram[NUM_WARPS * 3 + warp_id] = wcnt; - } - __syncthreads(); - - if (tid == 0) - { - float pmin = FLT_MAX, pmax = -FLT_MAX, psum = 0.0f; - int pcnt = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - pmin = fminf(pmin, __int_as_float(smem->histogram[w])); - pmax = fmaxf(pmax, __int_as_float(smem->histogram[NUM_WARPS + w])); - psum += __int_as_float(smem->histogram[NUM_WARPS * 2 + w]); - pcnt += smem->histogram[NUM_WARPS * 3 + w]; - } - float pmean = (pcnt > 0) ? psum / (float) pcnt : (pmin + pmax) * 0.5f; - - smem->pmax_saved = pmax; - smem->threshold = pmean; - smem->val_lo = pmin; - smem->val_hi = pmax; - smem->cnt_lo = M + M / 4; - smem->cnt_hi = 1; - smem->done = 0; - } - __syncthreads(); - - // Degenerate hint (all gathered values identical or out of range): - // reset to a trusted bracket instead of emitting row[0:K]; done = 2 - // skips the secant (it cannot converge on a full-range bracket) and - // hands the row to the Phase-3 repair. The hint only affects speed. - if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) - { - if (tid == 0) - { - float const seed = (smem->val_hi <= -FLT_MAX) ? 0.0f : smem->pmax_saved; - smem->val_lo = -FLT_MAX; - smem->val_hi = FLT_MAX; - smem->cnt_lo = N; - smem->cnt_hi = 0; - smem->threshold = seed; - smem->done = 2; - } - __syncthreads(); - } - - // ================================================================ - // Phase 2 (GVR Guess, part 2) — Secant-interpolation threshold search - // ================================================================ - - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - - for (int iter = 0; iter < MAX_REFINE_ITERS; iter++) - { - if (smem->done) - break; - if (tid == 0) - { - float vlo = smem->val_lo, vhi = smem->val_hi; - int clo = smem->cnt_lo, chi = smem->cnt_hi; - constexpr int target = kFTarget; - float range = vhi - vlo; - float nv; - if (clo > chi && range > 1e-10f) - { - float f = (float) (clo - target) / (float) (clo - chi); - f = fmaxf(0.05f, fminf(0.95f, f)); - if (iter == 0) - f = fminf(f, 0.50f); - nv = vlo + range * f; - } - else - nv = (vlo + vhi) * 0.5f; - if (nv <= vlo) - nv = vlo + range * 0.05f; - if (nv >= vhi) - nv = vhi - range * 0.05f; - if (nv == vlo || nv == vhi) - { - nv = (vlo + vhi) * 0.5f; - if (nv == vlo || nv == vhi) - { - smem->threshold = vlo; - smem->done = 2; - } - else - smem->threshold = nv; - } - else - smem->threshold = nv; - } - __syncthreads(); - if (smem->done) - break; - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - } - - if (tid == 0 && !smem->done) - { - if (smem->cnt_lo <= kCC * 2) - smem->threshold = smem->val_lo; - else - smem->threshold = smem->val_hi; - smem->done = 2; - } - __syncthreads(); - } // end of P1+P2 scope - - // ================================================================ - // Phase 3 (GVR Verify) — Ballot-free candidate collect - // ================================================================ - - // done==1: Phase 2 verified cand_count in [kK, kCC]; skip the re-check. - // Otherwise the secant did not converge and `threshold` carries no - // guarantee: repair BOTH sides (the old loop only handled overflow, so - // an undershooting threshold shipped a -1-padded, silently wrong top-K). - if (smem->done != 1) - { - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - // Anchor the untested bracket end at a float extreme: Phase 1 seeds - // both ends from HINTED values with invented counts, so they can sit - // on the same side of the K-th value. count(-FLT_MAX) >= kK, - // count(FLT_MAX) = 0. - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->val_hi = FLT_MAX; - } - else if (c < kK) - { - smem->val_hi = smem->threshold; - smem->val_lo = -FLT_MAX; - } - } - __syncthreads(); - - // Invariant maintained below: count(val_lo) >= kK. - for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) - { - unsigned const klo = gvrOrderKey(smem->val_lo); - unsigned const khi = gvrOrderKey(smem->val_hi); - if (khi <= klo + 1u) - break; // bracket collapsed to adjacent representable values - if (tid == 0) - smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); - __syncthreads(); - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - smem->val_lo = smem->threshold; - else if (c < kK) - smem->val_hi = smem->threshold; - } - __syncthreads(); - } - - // Still short of kK: the bracket collapsed; val_lo admits >= kK by - // the anchor invariant (or the row has < kK finite entries and the - // -1 tail pad is the correct answer). - if (smem->cand_count < kK) - { - if (tid == 0) - smem->threshold = smem->val_lo; - __syncthreads(); - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - } - // blockCountGE publishes cand_count from tid 0 only; the branch below - // must be uniform across the block. - __syncthreads(); - - // Collapsed bracket still over kCC = a tie plateau wider than the - // candidate buffer: emit everything strictly above val_lo (< kK by - // construction) plus arbitrary ties — a valid tie-aware top-K. The - // adjacency guard keeps the emit sound if the loop ever ran dry. - if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) - { - float const thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - for (int i = tid; i < N; i += BLOCK_SIZE) - { - float const v = __ldg(&input[i]); - if (v > thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = v; - outputIndices[p] = i; - } - } - } - __syncthreads(); - int const n_gt = min(smem->out_count, kK); - if (tid == 0) - smem->out_count = n_gt; - __syncthreads(); - for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) - { - float const v = __ldg(&input[i]); - if (v == thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = v; - outputIndices[p] = i; - } - } - } - __syncthreads(); - for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } - return; - } - } - - // Reuse per-thread counts cached by the last blockCountGE call (saves - // one full N-scan; blockCountGE's __syncthreads guarantees visibility). - int my_total_qual = smem->per_thread_counts[tid]; - - int thread_prefix = my_total_qual; -#pragma unroll - for (int off = 1; off < WARP_SIZE; off *= 2) - { - int other = __shfl_up_sync(full_mask, thread_prefix, off); - if (lane >= off) - thread_prefix += other; - } - int my_excl_offset = thread_prefix - my_total_qual; - int warp_total_qual = __shfl_sync(full_mask, thread_prefix, WARP_SIZE - 1); - - if (lane == 0) - smem->warp_counts[warp_id] = warp_total_qual; - __syncthreads(); - - if (tid == 0) - { - int total = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - int cnt = smem->warp_counts[w]; - smem->warp_counts[w] = total; - total += cnt; - } - smem->cand_count = total; - } - __syncthreads(); - - int my_write_pos = smem->warp_counts[warp_id] + my_excl_offset; - - { - float const thr = smem->threshold; - for (int i = tid * 4; i + 3 < N; i += BLOCK_SIZE * 4) - { - float4 v4 = __ldg(reinterpret_cast(input + i)); -#pragma unroll - for (int j = 0; j < 4; j++) - { - float val = (&v4.x)[j]; - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; - smem->vals[my_write_pos] = i + j; - my_write_pos++; - } - } - } - for (int i = (N & ~3) + tid; i < N; i += BLOCK_SIZE) - { - float val = __ldg(&input[i]); - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; - smem->vals[my_write_pos] = i; - my_write_pos++; - } - } - } - __syncthreads(); - - // ================================================================ - // Phase 4 (GVR Refine) — Histogram-based selection + partition - // ================================================================ - - int const cand_count = min(smem->cand_count, kCC); - - if (cand_count == kK) - { - for (int i = tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = smem->keys[i]; - outputIndices[i] = smem->vals[i]; - } - return; - } - - if (cand_count > kK) - { - float cmin = FLT_MAX, cmax = -FLT_MAX; - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; - cmin = fminf(cmin, v); - cmax = fmaxf(cmax, v); - } - cmin = warpReduceMin(cmin); - cmax = warpReduceMax(cmax); - if (lane == 0) - { - smem->warp_counts[warp_id] = __float_as_int(cmin); - smem->histogram[warp_id] = __float_as_int(cmax); - } - __syncthreads(); - - float block_min = FLT_MAX, block_max = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - block_min = fminf(block_min, __int_as_float(smem->warp_counts[w])); - block_max = fmaxf(block_max, __int_as_float(smem->histogram[w])); - } - if (block_max <= block_min) - block_max = block_min + 1e-6f; - - for (int i = tid; i < kBins; i += BLOCK_SIZE) - smem->histogram[i] = 0; - __syncthreads(); - - float range1 = block_max - block_min; - float inv1 = (range1 > 0.0f) ? ((float) (kBins - 1) + 0.99f) / range1 : 0.0f; - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - int bin = (int) ((smem->keys[i] - block_min) * inv1); - bin = min(max(bin, 0), kBins - 1); - atomicAdd(&smem->histogram[bin], 1); - } - __syncthreads(); - - // Parallel K-th bin search (3-step). - // Step 1: each warp sums BINS_PER_WARP consecutive bins (high→low). - // Step 2: tid=0 locates the target warp in NUM_WARPS steps. - // Step 3: one thread in that warp scans its BINS_PER_WARP bins. - // Total serial depth: NUM_WARPS + BINS_PER_WARP steps vs full kBins. - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - static_assert(kBins % NUM_WARPS == 0, "kBins must be divisible by NUM_WARPS"); - // Step 1: each warp accumulates its slice of bins (high→low) - int warp_bin_sum = 0; - for (int j = 0; j < BINS_PER_WARP; j++) - warp_bin_sum += smem->histogram[kBins - 1 - warp_id * BINS_PER_WARP - j]; - if (lane == 0) - smem->warp_counts[warp_id] = warp_bin_sum; - } - __syncthreads(); // S-4b3a - - // Step 2: tid=0 finds which warp contains the K-th element - if (tid == 0) - { - int cum = 0, tw = NUM_WARPS - 1; - for (int w = 0; w < NUM_WARPS; w++) - { - cum += smem->warp_counts[w]; - if (cum >= kK) - { - tw = w; - break; - } - } - // Recompute prefix before target warp for step 3 - cum = 0; - for (int w = 0; w < tw; w++) - cum += smem->warp_counts[w]; - smem->cnt_lo = cum; // prefix count before target warp - smem->cnt_hi = tw; // target warp index - } - __syncthreads(); // S-4b3b - - // Step 3: one thread in target warp scans its BINS_PER_WARP bins - if (warp_id == smem->cnt_hi && lane == 0) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - int base_cum = smem->cnt_lo; - float thr = block_min; - for (int j = 0; j < BINS_PER_WARP; j++) - { - int b = kBins - 1 - smem->cnt_hi * BINS_PER_WARP - j; - base_cum += smem->histogram[b]; - if (base_cum >= kK) - { - thr = block_min + (float) b * range1 / (float) kBins; - break; - } - } - smem->threshold = thr; - } - __syncthreads(); // S-4b3c - - // snap_limit must equal cand_count to guarantee convergence: each - // iteration either strictly decreases cgt (raise thr to next - // distinct value above) or strictly increases cge (lower thr to - // next distinct value below) by >= 1, so worst-case convergence - // takes cand_count - kK + 1 iters. The older bound `cand_count/4` - // silently accepted a non-converged threshold; Pass 1 then picked - // K elements in scan order from `cgt > kK` candidates, missing - // some true top-K members (~0.09 % intermittent at small-kNumBins - // mean-zero distributions). Common path still converges in 1-3 - // iters; the higher upper bound only affects the long-tail cells. - bool snap_converged = false; - int snap_limit = cand_count; - for (int si = 0; si < snap_limit; si++) - { - blockFusedSnapIter(smem, cand_count, tid, warp_id, lane); - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - if (cgt < kK && cge >= kK) - { - snap_converged = true; - break; - } - } - (void) snap_converged; - - float sel_thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - - // Two-pass selection: pass 1 emits strictly-greater-than-threshold - // candidates, pass 2 fills remaining slots with tie-values. An - // interleaved single-pass implementation would be unstable across - // rows with many ties at the K-th rank — equal-valued candidates - // could displace strictly-greater ones depending on their relative - // order in the candidate buffer. Splitting the passes makes the - // selection deterministic regardless of buffer ordering. - - // Pass 1: strictly greater than sel_thr - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; - - bool emit_gt = (i < cand_count) && (v > sel_thr); - unsigned mask_gt = __ballot_sync(full_mask, emit_gt); - if (mask_gt) - { - int cnt = __popc(mask_gt); - int moff = __popc(mask_gt & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_gt && bp + moff < kK) - { - outputValues[bp + moff] = v; - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - // Pass 2: equal to sel_thr (fills remaining slots) - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; - - bool emit_eq = (i < cand_count) && (v == sel_thr); - unsigned mask_eq = __ballot_sync(full_mask, emit_eq); - if (mask_eq) - { - int cnt = __popc(mask_eq); - int moff = __popc(mask_eq & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_eq && bp + moff < kK) - { - outputValues[bp + moff] = v; - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - int filled = min(smem->out_count, kK); - for (int i = filled + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } - return; - } - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - outputValues[i] = smem->keys[i]; - outputIndices[i] = smem->vals[i]; - } - for (int i = cand_count + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } -} - -// ============================================================================ -// gvrTopKKernel — single-row global wrapper (1 CTA, 1 row). -// Calls gvrTopKJob (independently-optimized device function). -// For multi-row decode launches, see heuristicTopKMultiRowKernel in -// heuristicTopKDecode.cu — both share the same micro-kernel job. -// ============================================================================ - -// Templated on TopK so the launcher can dispatch K=512/1024/2048 to the -// same kernel template. -// -// __launch_bounds__ uses the single-arg form (no minBlocksPerSM hint) so -// nvcc applies the same register heuristic as `heuristicTopKMultiRowKernel` -// in heuristicTopKDecode.cu. Adding `, 1` would lower theoretical occupancy -// from 75% (REG=40) to 50% (REG=64) for the K=2048 fp32 path. -template -__global__ void __launch_bounds__(BLOCK_SIZE) - gvrTopKKernel(float const* __restrict__ input, int const N, int const* __restrict__ preIdx, int const M, - int const topK, float* __restrict__ outputValues, int* __restrict__ outputIndices) -{ - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - gvrTopKJob(input, N, preIdx, M, topK, outputValues, outputIndices, smem, /*preIdxOffset=*/0); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// ============================================================================ -// gvrTopKJobDtype — bf16/fp16 device function -// ============================================================================ -// Mirror of gvrTopKJob with trait-driven dtype substitutions: -// - HBM input read : Trait::to_fp32(__ldg(&input[i])) -// - outputValues : InputT (Trait::from_fp32 at writeback) -// - vector load : 8-wide via Trait::unpack8 -// - blockCountGE : blockCountGEDtype -// `blockFusedSnapIter` is reused directly from the fp32 path — -// smem `keys[]` are stored as fp32 here (see deferred-conversion note -// below) so no dtype-specialized snap helper is needed. -// All arithmetic (threshold, accumulators, bin index) stays fp32. The fp32 -// dtype path uses gvrTopKJob (above), not this template; instantiated only -// for bf16 and fp16. -// -// smem `keys[]` are stored as fp32 even on the bf16/fp16 paths: deferring -// the down-conversion out of the Phase-3 collect loop saves more than the -// extra ~10 KB of smem costs. The fp32 keys move conversion to the output -// writeback (one cvt per surviving candidate) instead of every smem store. -template -__device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, int const N, - int const* __restrict__ preIdx, int const M, int const topK, InputT* __restrict__ outputValues, - int* __restrict__ outputIndices, - KernelSmemTplK::kC, GvrParams::kNumBins>* smem, - int const preIdxOffset = 0) -{ - using Trait = GvrDtypeTraits; - using SmemKey = float; // keys stay fp32; conversion deferred to output writeback - using Params = GvrParams; - constexpr int kK = TopK; - constexpr int kCC = Params::kC; - constexpr int kBins = Params::kNumBins; - constexpr int kFTarget = Params::kFTarget; - static_assert(Trait::VEC_W == 8, "gvrTopKJobDtype is for bf16/fp16 (8-wide); fp32 uses gvrTopKJob"); - - int const tid = threadIdx.x; - int const warp_id = tid / WARP_SIZE; - int const lane = tid & (WARP_SIZE - 1); - unsigned const full_mask = 0xffffffffu; - - { - // ================================================================ - // Phase 1 — Min/Max/Mean of pre-indexed values - // ================================================================ - - float local_min = FLT_MAX; - float local_max = -FLT_MAX; - float local_sum = 0.0f; - int local_cnt = 0; - for (int i = tid; i < M; i += BLOCK_SIZE) - { - int idx = __ldg(&preIdx[i]) + preIdxOffset; - if (idx >= 0 && idx < N) - { - float v = Trait::to_fp32(__ldg(&input[idx])); - local_min = fminf(local_min, v); - local_max = fmaxf(local_max, v); - local_sum += v; - local_cnt++; - } - } - - float wmin = warpReduceMin(local_min); - float wmax = warpReduceMax(local_max); - float wsum = local_sum; -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - wsum += __shfl_down_sync(0xffffffffu, wsum, off); - int wcnt = warpReduceSum(local_cnt); - - if (lane == 0) - { - smem->histogram[warp_id] = __float_as_int(wmin); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(wmax); - smem->histogram[NUM_WARPS * 2 + warp_id] = __float_as_int(wsum); - smem->histogram[NUM_WARPS * 3 + warp_id] = wcnt; - } - __syncthreads(); - - if (tid == 0) - { - float pmin = FLT_MAX, pmax = -FLT_MAX, psum = 0.0f; - int pcnt = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - pmin = fminf(pmin, __int_as_float(smem->histogram[w])); - pmax = fmaxf(pmax, __int_as_float(smem->histogram[NUM_WARPS + w])); - psum += __int_as_float(smem->histogram[NUM_WARPS * 2 + w]); - pcnt += smem->histogram[NUM_WARPS * 3 + w]; - } - float pmean = (pcnt > 0) ? psum / (float) pcnt : (pmin + pmax) * 0.5f; - - smem->pmax_saved = pmax; - smem->threshold = pmean; - smem->val_lo = pmin; - smem->val_hi = pmax; - smem->cnt_lo = M + M / 4; - smem->cnt_hi = 1; - smem->done = 0; - } - __syncthreads(); - - // Degenerate hint (all gathered values identical or out of range): - // reset to a trusted bracket instead of emitting row[0:K]; done = 2 - // skips the secant (it cannot converge on a full-range bracket) and - // hands the row to the Phase-3 repair. The hint only affects speed. - if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) - { - if (tid == 0) - { - float const seed = (smem->val_hi <= -FLT_MAX) ? 0.0f : smem->pmax_saved; - smem->val_lo = -FLT_MAX; - smem->val_hi = FLT_MAX; - smem->cnt_lo = N; - smem->cnt_hi = 0; - smem->threshold = seed; - smem->done = 2; - } - __syncthreads(); - } - - // ================================================================ - // Phase 2 — Secant-interpolation threshold search - // ================================================================ - - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - - for (int iter = 0; iter < MAX_REFINE_ITERS; iter++) - { - if (smem->done) - break; - if (tid == 0) - { - float vlo = smem->val_lo, vhi = smem->val_hi; - int clo = smem->cnt_lo, chi = smem->cnt_hi; - constexpr int target = kFTarget; - float range = vhi - vlo; - float nv; - if (clo > chi && range > 1e-10f) - { - float f = (float) (clo - target) / (float) (clo - chi); - f = fmaxf(0.05f, fminf(0.95f, f)); - if (iter == 0) - f = fminf(f, 0.50f); - nv = vlo + range * f; - } - else - nv = (vlo + vhi) * 0.5f; - if (nv <= vlo) - nv = vlo + range * 0.05f; - if (nv >= vhi) - nv = vhi - range * 0.05f; - if (nv == vlo || nv == vhi) - { - nv = (vlo + vhi) * 0.5f; - if (nv == vlo || nv == vhi) - { - smem->threshold = vlo; - smem->done = 2; - } - else - smem->threshold = nv; - } - else - smem->threshold = nv; - } - __syncthreads(); - if (smem->done) - break; - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - } - - if (tid == 0 && !smem->done) - { - if (smem->cnt_lo <= kCC * 2) - smem->threshold = smem->val_lo; - else - smem->threshold = smem->val_hi; - smem->done = 2; - } - __syncthreads(); - } // end of P1+P2 scope - - // ================================================================ - // Phase 3 — Ballot-free candidate collect - // ================================================================ - - // Mirror of the fp32 Phase-3 repair in gvrTopKJob (see comments there). - if (smem->done != 1) - { - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - // See the fp32 path: anchor the untested bracket end at a float extreme - // so count(val_lo) >= kK > count(val_hi) holds by construction. - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->val_hi = FLT_MAX; - } - else if (c < kK) - { - smem->val_hi = smem->threshold; - smem->val_lo = -FLT_MAX; - } - } - __syncthreads(); - - // Invariant maintained below: count(val_lo) >= kK. - for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) - { - unsigned const klo = gvrOrderKey(smem->val_lo); - unsigned const khi = gvrOrderKey(smem->val_hi); - if (khi <= klo + 1u) - break; // bracket collapsed to adjacent representable values - if (tid == 0) - smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); - __syncthreads(); - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - smem->val_lo = smem->threshold; - else if (c < kK) - smem->val_hi = smem->threshold; - } - __syncthreads(); - } - - if (smem->cand_count < kK) - { - if (tid == 0) - smem->threshold = smem->val_lo; - __syncthreads(); - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - } - // blockCountGEDtype publishes cand_count from tid 0 only; the branch - // below must be uniform across the block. - __syncthreads(); - - // Collapsed bracket with > kCC elements at the threshold: emit the - // strictly-greater set plus arbitrary ties directly (see fp32 path). - // The direct emit below is only valid once the bracket has collapsed: - // it assumes count(> thr) < kK, which is exactly "val_hi is the next - // representable value above val_lo and count(val_hi) < kK". If the - // loop ran out of iterations without collapsing (it cannot, given - // MAX_REPAIR_ITERS >= 32, but the guard keeps that an invariant rather - // than an assumption) fall through to the ordinary collect. - if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) - { - float const thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - for (int i = tid; i < N; i += BLOCK_SIZE) - { - float const v = Trait::to_fp32(__ldg(&input[i])); - if (v > thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = Trait::from_fp32(v); - outputIndices[p] = i; - } - } - } - __syncthreads(); - int const n_gt = min(smem->out_count, kK); - if (tid == 0) - smem->out_count = n_gt; - __syncthreads(); - for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) - { - float const v = Trait::to_fp32(__ldg(&input[i])); - if (v == thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = Trait::from_fp32(v); - outputIndices[p] = i; - } - } - } - __syncthreads(); - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } - return; - } - } - - int my_total_qual = smem->per_thread_counts[tid]; - - int thread_prefix = my_total_qual; -#pragma unroll - for (int off = 1; off < WARP_SIZE; off *= 2) - { - int other = __shfl_up_sync(full_mask, thread_prefix, off); - if (lane >= off) - thread_prefix += other; - } - int my_excl_offset = thread_prefix - my_total_qual; - int warp_total_qual = __shfl_sync(full_mask, thread_prefix, WARP_SIZE - 1); - - if (lane == 0) - smem->warp_counts[warp_id] = warp_total_qual; - __syncthreads(); - - if (tid == 0) - { - int total = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - int cnt = smem->warp_counts[w]; - smem->warp_counts[w] = total; - total += cnt; - } - smem->cand_count = total; - } - __syncthreads(); - - int my_write_pos = smem->warp_counts[warp_id] + my_excl_offset; - - { - float const thr = smem->threshold; - // 8-wide vector load (int4 = 8 × bf16/fp16) - for (int i = tid * 8; i + 7 < N; i += BLOCK_SIZE * 8) - { - int4 raw = __ldg(reinterpret_cast(input + i)); - float v[8]; - Trait::unpack8(raw, v); -#pragma unroll - for (int j = 0; j < 8; j++) - { - float val = v[j]; - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; // P0: defer convert to output - smem->vals[my_write_pos] = i + j; - my_write_pos++; - } - } - } - // Tail loop (N % 8) - for (int i = (N & ~7) + tid; i < N; i += BLOCK_SIZE) - { - float val = Trait::to_fp32(__ldg(&input[i])); - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; // P0: defer convert to output - smem->vals[my_write_pos] = i; - my_write_pos++; - } - } - } - __syncthreads(); - - // ================================================================ - // Phase 4 — Histogram-based selection + partition - // ================================================================ - - int const cand_count = min(smem->cand_count, kCC); - - if (cand_count == kK) - { - for (int i = tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = Trait::from_fp32(smem->keys[i]); // P0: convert at output - outputIndices[i] = smem->vals[i]; - } - return; - } - - if (cand_count > kK) - { - float cmin = FLT_MAX, cmax = -FLT_MAX; - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; // P0: keys already fp32 - cmin = fminf(cmin, v); - cmax = fmaxf(cmax, v); - } - cmin = warpReduceMin(cmin); - cmax = warpReduceMax(cmax); - if (lane == 0) - { - smem->warp_counts[warp_id] = __float_as_int(cmin); - smem->histogram[warp_id] = __float_as_int(cmax); - } - __syncthreads(); - - float block_min = FLT_MAX, block_max = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - block_min = fminf(block_min, __int_as_float(smem->warp_counts[w])); - block_max = fmaxf(block_max, __int_as_float(smem->histogram[w])); - } - if (block_max <= block_min) - block_max = block_min + 1e-6f; - - for (int i = tid; i < kBins; i += BLOCK_SIZE) - smem->histogram[i] = 0; - __syncthreads(); - - float range1 = block_max - block_min; - float inv1 = (range1 > 0.0f) ? ((float) (kBins - 1) + 0.99f) / range1 : 0.0f; - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - int bin = (int) ((smem->keys[i] - block_min) * inv1); // P0: keys fp32 - bin = min(max(bin, 0), kBins - 1); - atomicAdd(&smem->histogram[bin], 1); - } - __syncthreads(); - - // Parallel K-th bin search (2-step) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - static_assert(kBins % NUM_WARPS == 0, "kBins must be divisible by NUM_WARPS"); - int warp_bin_sum = 0; - for (int j = 0; j < BINS_PER_WARP; j++) - warp_bin_sum += smem->histogram[kBins - 1 - warp_id * BINS_PER_WARP - j]; - if (lane == 0) - smem->warp_counts[warp_id] = warp_bin_sum; - } - __syncthreads(); - - if (tid == 0) - { - int cum = 0, tw = NUM_WARPS - 1; - for (int w = 0; w < NUM_WARPS; w++) - { - cum += smem->warp_counts[w]; - if (cum >= kK) - { - tw = w; - break; - } - } - cum = 0; - for (int w = 0; w < tw; w++) - cum += smem->warp_counts[w]; - smem->cnt_lo = cum; - smem->cnt_hi = tw; - } - __syncthreads(); - - if (warp_id == smem->cnt_hi && lane == 0) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - int base_cum = smem->cnt_lo; - float thr = block_min; - for (int j = 0; j < BINS_PER_WARP; j++) - { - int b = kBins - 1 - smem->cnt_hi * BINS_PER_WARP - j; - base_cum += smem->histogram[b]; - if (base_cum >= kK) - { - thr = block_min + (float) b * range1 / (float) kBins; - break; - } - } - smem->threshold = thr; - } - __syncthreads(); - - // snap_limit must equal cand_count to guarantee convergence; see - // detailed rationale in the fp32 `gvrTopKJob` path. - bool snap_converged = false; - int snap_limit = cand_count; - for (int si = 0; si < snap_limit; si++) - { - blockFusedSnapIter(smem, cand_count, tid, warp_id, lane); // P0: keys fp32 → use fp32 snap helper - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - if (cgt < kK && cge >= kK) - { - snap_converged = true; - break; - } - } - (void) snap_converged; - - float sel_thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - - // Pass 1: strictly greater than sel_thr - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; // P0: keys fp32 - - bool emit_gt = (i < cand_count) && (v > sel_thr); - unsigned mask_gt = __ballot_sync(full_mask, emit_gt); - if (mask_gt) - { - int cnt = __popc(mask_gt); - int moff = __popc(mask_gt & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_gt && bp + moff < kK) - { - outputValues[bp + moff] = Trait::from_fp32(v); - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - // Pass 2: equal to sel_thr (fills remaining slots) - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; // P0: keys fp32 - - bool emit_eq = (i < cand_count) && (v == sel_thr); - unsigned mask_eq = __ballot_sync(full_mask, emit_eq); - if (mask_eq) - { - int cnt = __popc(mask_eq); - int moff = __popc(mask_eq & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_eq && bp + moff < kK) - { - outputValues[bp + moff] = Trait::from_fp32(v); - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - int filled = min(smem->out_count, kK); - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = filled + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } - return; - } - - // cand_count < kK fallback - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - outputValues[i] = Trait::from_fp32(smem->keys[i]); // P0: convert at output - outputIndices[i] = smem->vals[i]; - } - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = cand_count + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } -} - -// ============================================================================ -// gvrTopKKernelDtype — bf16/fp16 single-row global wrapper -// ============================================================================ -// Templated on (InputT, TopK). __launch_bounds__ uses the single-arg form -// so nvcc applies the same register heuristic as the multi-row dtype kernel -// in heuristicTopKDecode.cu. See `gvrTopKKernel` note above. - -template -__global__ void __launch_bounds__(BLOCK_SIZE) - gvrTopKKernelDtype(InputT const* __restrict__ input, int const N, int const* __restrict__ preIdx, int const M, - int const topK, InputT* __restrict__ outputValues, int* __restrict__ outputIndices) -{ - using SmemKey = typename GvrDtypeTraits::SmemKey; - using SmemT - = KernelSmemTplK::kC, GvrParams::kNumBins>; // dtype keys fp32 - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - gvrTopKJobDtype(input, N, preIdx, M, topK, outputValues, outputIndices, smem, - /*preIdxOffset=*/0); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// ============================================================================ -// Explicit kernel instantiations — 9 (T × K) combos. -// ============================================================================ -// Mirrors the same pattern used in heuristicTopKDecode.cu for the multi-row -// kernels. Forces nvcc to emit each `gvrTopKKernel` / `gvrTopKKernelDtype -// ` host-side wrapper stub *before* `launchHeuristicTopK` takes their -// address via `cudaLaunchKernelEx`. Without these declarations, certain nvcc -// stubgen versions (CI containers on sm_89 / sm_120f) emit the implicit stub -// inside the kernel body and then conflict with their own subsequent -// "explicit specialization of __wrapper__device_stub_*" pass — surfacing -// as a "specialization after instantiation" build error against -// cudafe1.stub.c. Header is included only by heuristicTopKDecode.cu (one TU) -// so no ODR concern. -template __global__ void gvrTopKKernel<512>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernel<1024>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernel<2048>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__half, 512>(__half const*, int, int const*, int, int, __half*, int*); -template __global__ void gvrTopKKernelDtype<__half, 1024>(__half const*, int, int const*, int, int, __half*, int*); -template __global__ void gvrTopKKernelDtype<__half, 2048>(__half const*, int, int const*, int, int, __half*, int*); - -// ============================================================================ -// Launch Wrapper -// ============================================================================ - -namespace detail -{ -// Per-(T, TopK) launcher implementation. Hoisted out of `launchHeuristicTopK` -// (was a C++20 templated lambda `[&]()`) because that pattern -// confuses nvcc's cudafe1 stub generator: taking the address of -// `gvrTopKKernel` / `gvrTopKKernelDtype` from inside a -// templated capturing lambda triggers an "explicit specialization of -// `__wrapper__device_stub_gvrTopKKernel` after instantiation" error -// against the auto-generated host wrapper stub. A regular function template -// avoids the quirk and stays in C++17 (no templated-lambda extension warning). -// -// Kernel body, GvrParams traits, kfn selection, opt-in smem, PDL attr, and -// cudaLaunchKernelEx call are byte-identical to the previous lambda body — -// SASS is unchanged for all 9 (T, K) instantiations. -template -cudaError_t launchHeuristicTopKImpl(T const* input, int N, int const* preIdx, int M, int topK, T* outputValues, - int* outputIndices, cudaStream_t stream, bool enablePDL) -{ - // dtype path uses fp32 smem keys (deferred convert). Launcher - // allocates smem with float keys regardless of input dtype. - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - size_t const smemSize = sizeof(SmemT); - - // Resolve target kernel function pointer at compile time. - auto kfn = []() - { - if constexpr (std::is_same_v) - return gvrTopKKernel; - else - return gvrTopKKernelDtype; - }(); - - if (smemSize > 48u * 1024u) - { - int device; - cudaGetDevice(&device); - int maxSmem; - cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); - if (smemSize > static_cast(maxSmem)) - return cudaErrorInvalidConfiguration; - cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemSize)); - } - - cudaLaunchConfig_t config{}; - config.gridDim = dim3(1); - config.blockDim = dim3(BLOCK_SIZE); - config.dynamicSmemBytes = smemSize; - config.stream = stream; - - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enablePDL ? 1 : 0; - config.attrs = attrs; - config.numAttrs = 1; - - cudaLaunchKernelEx(&config, kfn, input, N, preIdx, M, topK, outputValues, outputIndices); - return cudaGetLastError(); -} -} // namespace detail - -template -cudaError_t launchHeuristicTopK(T const* input, int N, IdxT const* preIdx, int M, int topK, T* outputValues, - IdxT* outputIndices, cudaStream_t stream = 0) -{ - static_assert(sizeof(IdxT) == sizeof(int), "launchHeuristicTopK only supports 32-bit indices"); - static_assert(std::is_same_v || std::is_same_v || std::is_same_v, - "launchHeuristicTopK supports only fp32 / bf16 / fp16"); - - // GvrParams specializations cover K ∈ {512, 1024, 2048}; reject others. - if (topK != 512 && topK != 1024 && topK != 2048) - return cudaErrorInvalidValue; - - // Dispatch on (T, topK) → 9 distinct kernel-pointer paths. Each - // instantiation captures its own (kFTarget, kC, kNumBins) tuple via - // GvrParams so all values are compile-time constants inside - // the kernel body. Opt-in smem + cudaLaunchKernelEx + PDL handling is - // shared across all 9 paths via `detail::launchHeuristicTopKImpl`. - - // Honor the standard TRTLLM_ENABLE_PDL env var (default on; set "0" to - // disable). - bool enablePDL = true; - if (char const* env = std::getenv("TRTLLM_ENABLE_PDL")) - { - if (env[0] == '0' && env[1] == '\0') - enablePDL = false; - } - - switch (topK) - { - case 512: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - case 1024: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - case 2048: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - default: return cudaErrorInvalidValue; - } -} - -// Explicit instantiations — fp32 + bf16/fp16 -template cudaError_t launchHeuristicTopK( - float const*, int, int const*, int, int, float*, int*, cudaStream_t); -template cudaError_t launchHeuristicTopK<__nv_bfloat16, int>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*, cudaStream_t); -template cudaError_t launchHeuristicTopK<__half, int>( - __half const*, int, int const*, int, int, __half*, int*, cudaStream_t); - -} // namespace heuristic_topk diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index d63f203a048c..b1bce6f2bf67 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -19,7 +19,6 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/heuristicTopKDecode.h" #include "tensorrt_llm/kernels/noAuxTcKernels.h" #include #include @@ -718,105 +717,19 @@ constexpr int kMaxBlocksPerRowDecode = 10; // 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; - int l2Bytes; - int kBsWave; - int kBsL2; - int kBsLarge; - int kSeqSmall; -}; - -// Uniform small-N lower bound for the Heuristic GVR path across all K. -// Aligns the GVR routing boundary with the Radix multi-CTA split-work -// threshold (maxByCols = N / kDecodeMinColsPerSubBlock(=2048) ≥ 2 at -// N ≥ 4096), so the dispatcher's algorithmic-handoff point is consistent: -// below 4096 the Radix path resolves to single-CTA insertion-sort and GVR -// is not attempted; at or above 4096 GVR may be considered. -// DSv4 swe-bench synth sweeps on B200/B300 (V3.2-Q19c protocol, May 2026): -// N=4K cells across K ∈ {512, 1024, 2048} all win — GVR R/H bf16 = 3.07× -// (K=512) / 2.57× (K=1024) / 1.34× (K=2048). -// N=2K cells across the same 9 (K × dtype) combos all show GVR R/H < 1 -// (0.55× – 0.84×), justifying 4K as the floor. -inline int kSeqSmallDefaultForK(int /*topK*/) -{ - return 4096; -} - -inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem, int topK) +// Cached device SM count for the wave-aware blocks-per-row dispatch below. +inline int getDeviceSmCount() { static std::once_flag sOnce; static int sSm = 0; - static int sL2 = 0; - // ----------------------------------------------------------------------- - // Diagnostic / tuning escape-hatch env overrides. Both are OFF by default - // and the K-aware / hardware-derived defaults below are expected to be - // optimal for production. Use only for microbenchmarks, regression - // bisection, or workload-specific tuning where the defaults are clearly - // suboptimal. - // - // TRTLLM_HEURISTIC_NMIN (valid range [1024, 200000]) - // Overrides `kSeqSmall` (Heuristic small-N threshold) for ALL K. - // Lower risk: only shifts a perf threshold; the kernel still - // produces an exact top-K either way. Setting it too low routes - // more N → Heuristic and may be slower than the fallback for - // small N; correctness is preserved. - // - // TRTLLM_HEURISTIC_BSMAX (valid range [1, 65536]) - // Overrides `kBsLarge` (BS upper bound for Heuristic) past the - // hardware-derived min(kBsWave, kBsL2). Higher risk: bypasses - // L2/occupancy safety bounds, so heuristic may run in working-set - // ranges where it has not been tuned (L2 thrash, suboptimal grid - // configs). Primary use is indexer microbenchmarks that need a - // BS-scaling comparison against the Radix path on identical inputs. - // ----------------------------------------------------------------------- - // sNMinEnv > 0 iff TRTLLM_HEURISTIC_NMIN is set to a valid value. When set, - // it overrides the per-K default for ALL K. - static int sNMinEnv = 0; - static int sBsMax = 0; std::call_once(sOnce, []() { int dev = 0; cudaGetDevice(&dev); cudaDeviceGetAttribute(&sSm, cudaDevAttrMultiProcessorCount, dev); - cudaDeviceGetAttribute(&sL2, cudaDevAttrL2CacheSize, dev); - char const* env = std::getenv("TRTLLM_HEURISTIC_NMIN"); - if (env != nullptr) - { - int const v = std::atoi(env); - sNMinEnv = (v >= 1024 && v <= 200000) ? v : 0; - } - char const* env_bsmax = std::getenv("TRTLLM_HEURISTIC_BSMAX"); - if (env_bsmax != nullptr) - { - int const v = std::atoi(env_bsmax); - sBsMax = (v >= 1 && v <= 65536) ? v : 0; - } }); - - SchemeXBounds b; - b.smCount = sSm; - b.l2Bytes = sL2; - b.kBsWave = (sSm > 0) ? (sSm * 3 - sSm / 8) : 426; - b.kBsL2 = (sL2 > 0 && numColumns > 0) - ? static_cast(static_cast(sL2) * 9 / 10 / (static_cast(numColumns) * bytesPerElem)) - : b.kBsWave; - b.kBsLarge = std::min(b.kBsWave, b.kBsL2 > 0 ? b.kBsL2 : b.kBsWave); - if (sBsMax > 0) - { - // BSMAX env override bypasses the hardware-derived L2/occupancy bound - // (see the BSMAX section in the call_once block above for risk notes). - b.kBsLarge = sBsMax; - } - // NMIN env override (if set) wins over the per-K default for ALL K. - b.kSeqSmall = (sNMinEnv > 0) ? sNMinEnv : kSeqSmallDefaultForK(topK); - return b; + return sSm; } } // namespace @@ -838,12 +751,8 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW // Query the actual SM count from the driver so the dispatch tracks the // hardware rather than a baked-in target (H100=132, B200=148, …). - // topK=0: blocks-per-row computation is K-agnostic; kSeqSmall is uniform - // 4096 across K, so the topK arg is unused for the kSeqSmall lookup as - // well, and only smCount/kBsWave/kBsL2 are consumed here. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, /*topK=*/0); - TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count"); - int const smCount = bounds.smCount; + int const smCount = getDeviceSmCount(); + TLLM_CHECK_WITH_INFO(smCount > 0, "indexerTopK: failed to query device SM count"); int const maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock); int const maxBp = std::min(maxByCols, kMaxBlocksPerRowDecode); @@ -889,110 +798,9 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, - int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, - int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) + int const stride1, int const next_n, int const topK, int const compressRatio, cudaStream_t const stream) { 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` is uniform 4096 across all K (see - // kSeqSmallDefaultForK). 4K is the dispatcher's algorithmic-handoff - // point: below 4096 the Radix path resolves to single-CTA insertion-sort - // (maxByCols = N/2048 = 1 → bp=1; useRadixSort = N≥12288 = false), and - // GVR is empirically slower than insertion-sort below 4K across all - // K ∈ {512, 1024, 2048} × dtype ∈ {fp32, bf16, fp16} (R/H ∈ [0.55, 0.84] - // at N=2K; DSv4 V3.2-Q19c synth sweeps May 2026). Configurable via - // TRTLLM_HEURISTIC_NMIN env (>=1024), which overrides the default. - // ======================================================================== - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, topK); - int const kBsWave = bounds.kBsWave; - int const kBsL2 = bounds.kBsL2; - int const kBsLarge = bounds.kBsLarge; - int const kSeqSmall = bounds.kSeqSmall; - - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - // 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; - std::call_once(sDebugOnceFlag, - []() - { - char const* env = std::getenv("TRTLLM_SCHEMEX_DEBUG"); - sDebug = (env != nullptr && env[0] == '1'); - }); - if (sDebug) - { - 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, compressRatio, stream); - sync_check_cuda_error(stream); - return; - } - int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); cudaLaunchAttribute attrs[1]; @@ -1052,13 +860,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic // ============================================================================ // 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): +// Dispatch chain: // numColumns < kSortingAlgorithmThreshold (12288) → insertion sort // kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort // numColumns ≥ splitWorkThreshold (200K default) → unsupported @@ -1078,8 +880,7 @@ namespace 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) + int const compressRatio, cudaStream_t const stream) { static_assert(std::is_same_v || std::is_same_v, "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); @@ -1087,25 +888,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* constexpr int kNumThreadsPerBlock = 512; int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. - // K-aware kSeqSmall — see fp32 dispatcher for rationale. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT)), topK); - 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) + if (numColumns < kSortingAlgorithmThreshold) { // Insertion sort path — InputT propagated; histogram/sort run on float keys. auto* kernel_instance = &topKPerRowDecode(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride0, stride1, next_n, topK, 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, int const compressRatio, - cudaStream_t const stream) + int const compressRatio, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride1, next_n, topK, compressRatio, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, @@ -1199,17 +980,6 @@ void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int con sync_check_cuda_error(stream); } -bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem) -{ - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - if (!isSupportedTopK) - { - return false; - } - auto const bounds = getSchemeXBounds(numColumns, bytesPerElem, topK); - return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; -} - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 7e9f3bd7070d..2352f82def77 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -36,9 +36,8 @@ 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, int64_t compress_ratio, - std::optional const& radix_aux_indices, std::optional const& radix_aux_logits) + int64_t next_n, int64_t index_topk, 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"); @@ -70,53 +69,18 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t TORCH_CHECK(logits_stride_0 >= 0, "logits_stride_0 must be greater than or equal to 0"); TORCH_CHECK(logits_stride_1 >= 0, "logits_stride_1 must be greater than or equal to 0"); - int32_t const* preIdxPtr = nullptr; - int32_t preIdxStride = 0; - int32_t preIdxCount = 0; - if (pre_idx.has_value()) - { - auto const& preIdxTensor = pre_idx.value(); - TORCH_CHECK(preIdxTensor.is_cuda(), "pre_idx must be a CUDA tensor"); - TORCH_CHECK(preIdxTensor.device() == logits.device(), "pre_idx must be on the same device as logits"); - TORCH_CHECK(preIdxTensor.is_contiguous(), "pre_idx must be contiguous"); - TORCH_CHECK(preIdxTensor.dim() == 2, "pre_idx must be a 2D Tensor"); - TORCH_CHECK(preIdxTensor.size(0) * next_n == numRows64, - "pre_idx first dimension must equal logits.size(0)/next_n (one hint row per batch element)"); - preIdxPtr = preIdxTensor.data_ptr(); - preIdxStride = static_cast(preIdxTensor.stride(0)); - preIdxCount = static_cast(preIdxTensor.size(1)); - } - - // Caller-owned scratch buffer for heuristic TopK output values. - // Must be pre-allocated with stable address for CUDA Graph compatibility. - // scratch dtype must match input dtype. auto const logits_dtype = logits.scalar_type(); TORCH_CHECK(logits_dtype == at::ScalarType::Float || logits_dtype == at::ScalarType::BFloat16 || logits_dtype == at::ScalarType::Half, "indexer_topk_decode: logits dtype must be float32, bfloat16, or float16; got ", logits_dtype); - void* heuristicScratchPtr = nullptr; - if (heuristic_scratch.has_value()) - { - auto const& scratchTensor = heuristic_scratch.value(); - TORCH_CHECK(scratchTensor.is_cuda(), "heuristic_scratch must be a CUDA tensor"); - TORCH_CHECK( - scratchTensor.device() == logits.device(), "heuristic_scratch must be on the same device as logits"); - TORCH_CHECK(scratchTensor.is_contiguous(), "heuristic_scratch must be contiguous"); - TORCH_CHECK(scratchTensor.numel() >= static_cast(num_rows) * index_topk, - "heuristic_scratch must have at least numRows * index_topk elements"); - TORCH_CHECK(scratchTensor.scalar_type() == logits_dtype, - "heuristic_scratch dtype must match logits dtype (got scratch=", scratchTensor.scalar_type(), - ", logits=", logits_dtype, ")"); - heuristicScratchPtr = scratchTensor.data_ptr(); - } int32_t splitWorkThreshold = 200 * 1000; auto stream = at::cuda::getCurrentCUDAStream(logits.get_device()); 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 + // fp32 path — Insertion / Radix / Radix-split-work dispatcher. + // 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 @@ -126,8 +90,7 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t 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). The + // Caller-owned scratch with stable address (CUDA Graph safe). The // Python TopK module supplies these from its reusable buffer arena. auto const& ai = radix_aux_indices.value(); auto const& al = radix_aux_logits.value(); @@ -160,23 +123,22 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t } tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), 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); + logits_stride_1, static_cast(next_n), static_cast(index_topk), + 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), + logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), 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), static_cast(compress_ratio), stream); + static_cast(next_n), static_cast(index_topk), static_cast(compress_ratio), + stream); } } @@ -226,8 +188,7 @@ 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, int compress_ratio=1, " - "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); + "int compress_ratio=1, Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 177ba04fc1f2..1f7a7ee8a739 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -38,7 +38,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from .params import DSAParams +from .params import DSAParams, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -635,44 +635,76 @@ def __init__( self._enable_heuristic_topk = ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ) - # Opt-in self-sampling GVR top-K decode (CuTeDSL, env-gated: - # TRTLLM_GVR_SELF_SAMPLING=1). Same operator contract as the tiered - # heuristic path (per-request device kv_lens, raw prev-top-K hints, - # per-row MTP window, in-kernel n <= topK short path); tuning is - # frozen from indexer_max_seq_len at capture time, so the launch is - # CUDA-graph-replay safe. The TopK module's hardware-format gate - # falls through to the CUDA GVR path with a one-time warning; - # contract violations inside the engine raise. - self._use_self_sampling_topk = ( - os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" - and IS_CUTLASS_DSL_AVAILABLE - # datacenter Blackwell only; consumer Blackwell (sm_120/121) - # lacks thread-block clusters - and get_sm_version() in (100, 103) - and sparse_params.index_topk in (512, 1024, 2048) - and compress_ratio in (1, 4) + # Two-level GVR dispatch: enable_heuristic_topk selects the GVR + # family over the exact radix path; use_self_sampling_topk (default + # True) selects the hint-free self-sampling engine over the + # temporal-hint engines. Tuning is frozen from indexer_max_seq_len + # at capture time, so the launch is CUDA-graph-replay safe. The + # TopK module's hardware-format gate falls back to the exact + # insertion/radix path with a one-time warning; contract violations + # inside the engine raise. + self._use_self_sampling_topk = use_self_sampling_gvr( + enable_heuristic_topk=self._enable_heuristic_topk, + use_self_sampling_topk=sparse_params.use_self_sampling_topk, + index_topk=sparse_params.index_topk, + compress_ratio=compress_ratio, + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + sm_version=get_sm_version(), ) + if os.environ.get("TRTLLM_GVR_SELF_SAMPLING") is not None: + logger.warning_once( + "TRTLLM_GVR_SELF_SAMPLING is retired and ignored: the " + "self-sampling GVR engine is selected by the " + "use_self_sampling_topk sparse-attention config field " + "(default True) when enable_heuristic_topk is set.", + key="gvr_self_sampling_env_retired", + ) + if ( + self._enable_heuristic_topk + and sparse_params.use_self_sampling_topk + and not self._use_self_sampling_topk + ): + logger.warning_once( + "use_self_sampling_topk=True but the self-sampling GVR " + "prerequisites are not met " + f"(cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " + f"sm={get_sm_version()}, " + f"index_topk={sparse_params.index_topk}, " + f"compress_ratio={compress_ratio}); falling back to the " + "temporal GVR path (exact radix when the DSL engine is " + "unavailable).", + key="gvr_self_sampling_prereq_fallback", + ) self.mtp_index_share = sparse_params.mtp_index_share - if self.use_cute_dsl_topk: + if ( + self._enable_heuristic_topk + and IS_CUTLASS_DSL_AVAILABLE + # datacenter Blackwell only; consumer Blackwell (sm_120/121) + # lacks the thread-block clusters both GVR engines use + and get_sm_version() in (100, 103) + ): + decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR + else: + if self._enable_heuristic_topk: + logger.warning_once( + "enable_heuristic_topk=True but the DSL GVR engine is " + f"unavailable (cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " + f"sm={get_sm_version()}); using the exact radix decode " + "top-K instead.", + key="gvr_prereq_radix_fallback", + ) decode_top_k_implementation = ( - TopKImplementation.CUTE_DSL_GVR - if self._enable_heuristic_topk - else TopKImplementation.CUTE_DSL_RADIX + TopKImplementation.CUTE_DSL_RADIX + if self.use_cute_dsl_topk + else TopKImplementation.CUDA_RADIX ) - elif self._enable_heuristic_topk: - decode_top_k_implementation = TopKImplementation.CUDA_GVR - else: - decode_top_k_implementation = TopKImplementation.CUDA_RADIX - if self._use_self_sampling_topk and self._enable_heuristic_topk: - # env opt-in overrides the decode implementation; the GVR prior - # contract is identical - decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR_V2 self.top_k = TopK( self.index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, + gvr_self_sampling=self._use_self_sampling_topk, ) # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM @@ -1376,7 +1408,8 @@ def sparse_attn_indexer( num_gen_tokens = num_tokens - num_ctx_tokens gvr_prior_indices = None - if self._enable_heuristic_topk: + if self.top_k.needs_gvr_prior: + assert metadata.gvr_prior_indices is not None local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] gvr_prior_indices = metadata.gvr_prior_indices[local_layer] if is_generation is None: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index b69d856d743b..601e184d7289 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -4,7 +4,6 @@ from __future__ import annotations -import os from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional @@ -29,7 +28,7 @@ _pick_dsl_expand, _select_indexer_compress_ratio, ) -from .params import DSAMetadataParams +from .params import DSAMetadataParams, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -103,6 +102,11 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # Number of compressed KV tokens for context requests num_ctx_kv_tokens: int = 0 gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None + # Temporal-GVR prior state: allocated only when the two-level dispatch + # selects the temporal engine (the self-sampling engine keeps no + # cross-step state). + needs_gvr_prior: bool = field(default=False, init=False) + gvr_prior_indices: Optional[torch.Tensor] = field(default=None, init=False) def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -176,6 +180,23 @@ def __post_init__(self): if hasattr(self.kv_cache_manager, "compressed_block_sizes"): tpb = tpb // _effective_compress_ratio_divisor(self._indexer_compress_ratio) self._tokens_per_block = tpb + # Mirror the indexer's two-level GVR decision. The representative + # compress ratio matches every live indexer: the DeepSeek-V4 backend + # only builds indexers on cr=4 layers and plain DSA uses cr=1. + self.use_self_sampling_topk = use_self_sampling_gvr( + enable_heuristic_topk=self.enable_gvr_topk, + use_self_sampling_topk=sparse_metadata_params.use_self_sampling_topk, + index_topk=self.num_sparse_topk, + compress_ratio=self._indexer_compress_ratio, + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + sm_version=get_sm_version(), + ) + self.needs_gvr_prior = ( + self.enable_gvr_topk + and IS_CUTLASS_DSL_AVAILABLE + and get_sm_version() in (100, 103) + and not self.use_self_sampling_topk + ) self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) self.create_buffers_for_indexer(capture_graph=capture_graph) @@ -354,18 +375,19 @@ def warmup_selfsampling_topk( warmed keys are the ones dispatch actually looks up. Batches outside this set still compile lazily on first touch. The helper enumerates one representative row per distinct engine compile key, so large - batch lists warm in bounded time and memory. No-op unless the opt-in - gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + batch lists warm in bounded time and memory. No-op unless the + two-level dispatch (enable_heuristic_topk + use_self_sampling_topk) + selects the self-sampling engine. """ - if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": - return - # same hardware gates as the dispatch flag (indexer __init__): never - # compile these kernels on unsupported stacks during warmup + # same two-level dispatch and hardware gates as the indexer __init__: + # never compile these kernels on unsupported stacks during warmup if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() not in (100, 103): return if not self.enable_gvr_topk or self.kv_cache_manager is None: return - top_k = getattr(self.sparse_metadata_params, "index_topk", None) + if not self.use_self_sampling_topk: + return + top_k = self.sparse_mla_topk if not top_k or int(top_k) not in (512, 1024, 2048): return cr = int(self._indexer_compress_ratio) if self._indexer_compress_ratio else 1 @@ -397,8 +419,9 @@ def warmup_selfsampling_topk( row_stride = msl_c if row_stride % 4: return - # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len - # is compressed — same multiply-back as the dispatch seam) + # The helper takes max_seq_len in KV-token space; + # get_indexer_max_seq_len is compressed, so multiply it back as at + # the dispatch seam. try: _ss_host.warmup_varlen( int(top_k), @@ -542,11 +565,7 @@ def on_update_kv_lens(self) -> None: def _compute_kv_lens_row_reorder(self) -> None: """Prepare the longest-job-first GVR row order once per forward step.""" next_n = 1 + self.max_draft_tokens - if ( - self.enable_gvr_topk - and self.use_cute_dsl_topk - and self.num_generations * next_n >= 2 * self.num_sms - ): + if self.needs_gvr_prior and self.num_generations * next_n >= 2 * self.num_sms: gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32) self.kv_lens_row_reorder_buffer[: self.num_generations].copy_(order) @@ -824,7 +843,7 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - if self.enable_gvr_topk: + if self.needs_gvr_prior: self.gvr_prior_indices = self.get_empty( self.cuda_graph_buffers, ( @@ -837,14 +856,13 @@ def create_buffers_for_indexer(self, capture_graph=False): capture_graph=capture_graph, ) self.gvr_prior_indices.zero_() - if self.use_cute_dsl_topk: - self.kv_lens_row_reorder_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), - cache_name="kv_lens_row_reorder_buffer", - dtype=torch.int32, - capture_graph=capture_graph, - ) + self.kv_lens_row_reorder_buffer = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + cache_name="kv_lens_row_reorder_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index d72f15c45b2d..fbea670a5978 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -20,6 +20,31 @@ pass +def use_self_sampling_gvr( + *, + enable_heuristic_topk: bool, + use_self_sampling_topk: bool, + index_topk: int | None, + compress_ratio: int, + is_cute_dsl_available: bool, + sm_version: int, +) -> bool: + """Return whether the two-level dispatch picks the self-sampling engine. + + Shared by the indexer (per-layer TopK construction) and the attention + metadata (prior-state allocation and warmup) so both sides of the + dispatch agree. + """ + return ( + enable_heuristic_topk + and use_self_sampling_topk + and is_cute_dsl_available + and sm_version in (100, 103) + and index_topk in (512, 1024, 2048) + and compress_ratio in (1, 4) + ) + + @dataclass(kw_only=True, slots=True) class DSABackendForwardArgs(SparseBackendForwardArgs): """DSA inputs passed from the MLA module to its backend.""" @@ -41,6 +66,7 @@ class DSAMetadataParams(SparseMetadataParams): q_split_threshold: int has_shared_indexer_layers: bool = False mtp_index_share: bool = False + use_self_sampling_topk: bool = True @dataclass(frozen=True) @@ -58,6 +84,10 @@ class DSAParams(SparseParams): q_split_threshold: int = 8192 indexer_rope_interleave: bool = False enable_heuristic_topk: bool = False + # Second-level GVR dispatch: hint-free self-sampling engine (True) vs + # temporal previous-step-hint engines (False). Only meaningful when + # enable_heuristic_topk is set. + use_self_sampling_topk: bool = True indexer_k_dtype: Literal["fp8", "fp4"] = "fp8" # Shared layers reuse the preceding full layer's top-k. is_full_indexer_layer: bool = True diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 31cec1dfc264..4b0eee7bd734 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -336,8 +336,6 @@ def _(logits, indices, next_n, index_topk, - pre_idx=None, - heuristic_scratch=None, compress_ratio=1, radix_aux_indices=None, radix_aux_logits=None): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 36c2d67ff965..a43d39d2e59a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -29,6 +29,7 @@ import contextlib import sys +from typing import Any import cutlass import cutlass.cute as cute @@ -1326,7 +1327,8 @@ def __init__( next_n: int = 1, cr_shift: int = 0, r_const: int = 1, - ): + hint_free: bool = False, + ) -> None: assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) assert kpt in (1, 2, 4, 8) and minb in (1, 2, 4) @@ -1346,6 +1348,8 @@ def __init__( self.next_n = int(next_n) self.cr_shift = int(cr_shift) self.r_const = int(r_const) + # hint-free: gather_hint sites compiled out (sentinel pass-through) + self.hint_free = bool(hint_free) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) and self.r_const >= 1 # TSH-floor staging arm. SPLIT-only compile-time key; the CUDA form @@ -1973,9 +1977,10 @@ def kern( if T > cutlass.Float32(_NEG_INF): needg = cutlass.Int32(0) if needg != cutlass.Int32(0): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT - ) # 2 barriers inside + if cutlass.const_expr(not self.hint_free): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) # 2 barriers inside T = GMIN if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: @@ -2373,9 +2378,10 @@ def kern( else: # LAZY GATHER (sentinel equality flag) if GMIN == cutlass.Float32(C.SENT_LO): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT - ) + if cutlass.const_expr(not self.hint_free): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) @@ -2960,19 +2966,21 @@ def __call__( _COMPILE_CACHE = {} -def get_compiled(tpl, options_extra: str = ""): +def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) -> Any: """Compile (or fetch) the gvr_main variant for constexpr tuple tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG) — legacy, or tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG, NEXT_N, CR_SHIFT, R_CONST) — per-row varlen mode (TSHG slot is ignored: varlen compiles the TSH machinery in whenever SPLIT and gates it per row at runtime).""" - key = (tuple(tpl), options_extra) + key = (tuple(tpl), options_extra, bool(hint_free)) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit if len(tpl) == 7: blk, u, minb, nbs, kpt, split, tshg = tpl - kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg)) + kern = GvrMainKernel( + blk, u, minb, nbs, kpt, bool(split), bool(tshg), hint_free=bool(hint_free) + ) else: blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl kern = GvrMainKernel( @@ -2987,6 +2995,7 @@ def get_compiled(tpl, options_extra: str = ""): next_n=next_n, cr_shift=cr_shift, r_const=r_const, + hint_free=bool(hint_free), ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -3309,7 +3318,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - ): + hint_free: bool = False, + ) -> None: assert blk in (256, 512, 1024) and vpt in (1, 2, 4) assert nbh in (256, 512, 1024, 2048) assert nbh % blk == 0 or blk % nbh == 0 @@ -3334,8 +3344,11 @@ def __init__( # derived compile-time constants self.S = vpt * 4 self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) - self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 - self.use_img = img and vpt == 1 + # hint-free: bracket = min/max fold of the first k row values + # (already in registers); the hint-gather bracket arms are forced off + self.hint_free = bool(hint_free) + self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hint_free) + self.use_img = img and vpt == 1 and (not hint_free) self.brl = (minb * blk <= 1024) or (vpt == 1) # ------------------------------------------------------------------ @@ -3551,7 +3564,7 @@ def kern( # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any # dependent gather; compiled out under DEG. pvs = [] - if cutlass.const_expr(not self.deg): + if cutlass.const_expr(not (self.deg or self.hint_free)): for t in cutlass.range_constexpr(KPT): pv = cutlass.Int32(-1) j = tid + cutlass.Int32(t * self.blk) @@ -3646,6 +3659,19 @@ def kern( lmin = fkey(lmn) lmax = fkey(lmx) # monotone cute.arch.barrier() # bm dies + elif cutlass.const_expr(self.hint_free and not self.deg): + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for s in cutlass.range_constexpr(S): + pos = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if pos < k: + v = _val(frags, s) + lmn = fmin_f32(lmn, v) + lmx = fmax_f32(lmx, v) + lmin = fkey(lmn) + lmax = fkey(lmx) elif cutlass.const_expr(self.deg): lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) @@ -4207,10 +4233,18 @@ def __call__( _COMPILE_CACHE__reg: dict = {} -def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): +def get_compiled__reg( + tpl: tuple, + dump_dir: str | None = None, + pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + hint_free: bool = False, +) -> Any: """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, MINB, KPT, CUR, DEG, IMG, NBH).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free)) compiled = _COMPILE_CACHE__reg.get(key) if compiled is None: from cutlass.cute import runtime as _crt @@ -4229,6 +4263,7 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_ varlen=varlen, next_n=next_n, cr_shift=cr_shift, + hint_free=hint_free, ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() @@ -4377,7 +4412,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - ): + hint_free: bool = False, + ) -> None: assert blk == 1024, "gvr_clus is always BLK=1024" assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1)" assert nbs == 256, "SNB must stay 256" @@ -4392,6 +4428,7 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) + self.hint_free = bool(hint_free) # hint-free: gather_hint sites compiled out self.lcs = cs.bit_length() - 1 # log2(CS) for the per-row Q shift self.blk = blk self.u = u @@ -4947,9 +4984,10 @@ def kern( needg = cutlass.Int32(0) if needg != cutlass.Int32(0): # degenerate sample: identical on every rank of the cluster - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers + if cutlass.const_expr(not self.hint_free): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers T = GMIN if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: @@ -5178,9 +5216,10 @@ def kern( if tshtaken == cutlass.Int32(0): # LAZY GATHER — every rank computes identical GMIN if GMIN == cutlass.Float32(C.SENT_LO): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers inside + if cutlass.const_expr(not self.hint_free): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers inside floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) @@ -5552,24 +5591,44 @@ def __call__( def get_compiled__clus( - tpl, + tpl: tuple, scap: int = 8192, cmp_: int = 2048, options_extra: str = "", varlen: bool = False, next_n: int = 1, cr_shift: int = 0, -): + hint_free: bool = False, +) -> Any: """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every reachable route has 8192/2048 — asserted by run__clus()).""" - key = (tuple(tpl), scap, cmp_, options_extra, bool(varlen), int(next_n), int(cr_shift)) + key = ( + tuple(tpl), + scap, + cmp_, + options_extra, + bool(varlen), + int(next_n), + int(cr_shift), + bool(hint_free), + ) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: return hit blk, u, minb, nbs, cs = tpl kern = GvrClusKernel( - blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_, varlen=varlen, next_n=next_n, cr_shift=cr_shift + blk, + u, + minb, + nbs, + cs, + scap=scap, + cmp_=cmp_, + varlen=varlen, + next_n=next_n, + cr_shift=cr_shift, + hint_free=hint_free, ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -5787,7 +5846,8 @@ def __init__( varlen: bool = False, next_n: int = 1, cr_shift: int = 0, - ): + hint_free: bool = False, + ) -> None: assert blk == BLKC, "all instantiations BLK=BLKC=1024" assert vpt in (1, 2, 4) and cs in (2, 4, 8) self.blk = blk @@ -5803,6 +5863,8 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) + # hint-free: P0 samples the first k row elements (coalesced) instead of the hint + self.hint_free = bool(hint_free) self.S = vpt * 4 self.span = blk * vpt # float4 per CTA @@ -5958,8 +6020,12 @@ def kern( # ---- P0: redundant hint gather, EVERY CTA (k<=BLK by dispatch # gate). One coalesced word per thread, NO cluster barrier — # GMIN/GMAX identical everywhere by construction. - if tid < k: - pv0 = ld_g_i32(p_addr, tid) + if cutlass.const_expr(self.hint_free): + if tid < k: + pv0 = tid + else: + if tid < k: + pv0 = ld_g_i32(p_addr, tid) # ---- P1: row load — predicated flat float4[VPT] batch (the CUDA # has NO exact-fit peel here, guard is per-load). Issue all loads @@ -6461,16 +6527,31 @@ def __call__( _COMPILE_CACHE__regclus: dict = {} -def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): +def get_compiled__regclus( + tpl: tuple, + dump_dir: str | None = None, + pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + hint_free: bool = False, +) -> Any: """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" - key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free)) compiled = _COMPILE_CACHE__regclus.get(key) if compiled is None: from cutlass.cute import runtime as _crt blk, vpt, cs = tpl kernel = GvrRegClusKernel( - blk, vpt, cs, pdl=pdl, varlen=varlen, next_n=next_n, cr_shift=cr_shift + blk, + vpt, + cs, + pdl=pdl, + varlen=varlen, + next_n=next_n, + cr_shift=cr_shift, + hint_free=hint_free, ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 165048c8dcfe..cea15314e1f4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -679,7 +679,14 @@ def _main(blk_, minb_, u_, split_): _VARLEN_CACHE = {} -def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): +def _varlen_launcher( + num_rows: int, + npad: int, + k: int, + n_env: int, + next_n: int, + cr: int, +) -> tuple: """Capture-time varlen plan + compiled launcher. The gvr_main port is the universally correct fallback; specialist family tiers below. Every choice here is a function of capture-stable quantities only — mirroring @@ -699,7 +706,11 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): plan_free = route(num_rows, n_eff, npad, k) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + tuple(plan_free["tpl"]), + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + hint_free=True, ) lc = ("reg_clus", fn, n_eff) _VARLEN_CACHE[key] = lc @@ -713,7 +724,11 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # / short-row handling lives in-kernel. if plan_free["kernel"] in ("reg", "regimg"): fn = dev.get_compiled__reg( - tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + tuple(plan_free["tpl"]), + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + hint_free=True, ) rt_f = plan_free["rt"] lc = ( @@ -739,6 +754,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): varlen=True, next_n=next_n, cr_shift=cr_shift, + hint_free=True, ) lc = ( "clus", @@ -754,7 +770,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # TSHG (tpl[6]) is dead under varlen (the ctor compiles the TSH # machinery in whenever SPLIT); normalize it out of the compile key so # row counts differing only in that slot share one engine - fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const)) + fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const), hint_free=True) big = num_rows * r_const <= 148 aim_base = ( ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) @@ -1234,17 +1250,15 @@ def run_ws( def run_varlen( logits: torch.Tensor, - pre_idx: torch.Tensor, kv_lens: torch.Tensor, indices: torch.Tensor, next_n: int = 1, compress_ratio: int = 1, values: torch.Tensor | None = None, max_seq_len: int | None = None, - engine: str = "auto", workspace: torch.Tensor | None = None, ) -> None: - """Production-contract varlen entry (per-row device kv_lens). + """Run hint-free self-sampling Top-K with per-request device KV lengths. Row semantics (mirror of ``heuristicTopKDecode.cu`` and the in-tree ``cute_dsl_gvr_topk_decode`` runner): @@ -1255,19 +1269,18 @@ def run_varlen( not new-token seq_lens); row ``r`` uses ``n_r = (kv_lens[r // next_n] - next_n + (r % next_n) + 1) // compress_ratio`` valid entries (cr 1 = DSv3.2, 4 = DSv4 Flash/Pro); - ``pre_idx`` ``[batch, k]`` is REQUEST-level raw prev-step top-K, - shared by a request's ``next_n`` rows (offset-free hint contract); + the bracket is derived from the current row itself (register families: + min/max fold of the first k row values; streaming families do not + consume a temporal hint on the accept path); ``k`` comes from + ``indices.shape[1]``; per-row ``n_r <= k`` takes the short path (identity + ``-1`` tail). - ENGINES: ``engine="auto"`` (default) launches the per-row IN-KERNEL - gvr_main varlen port — ONE launch for the whole batch; each CTA reads its - row's kv_len on device and re-derives the sampling ladder (route_dynamic + The per-row in-kernel engine launches once for the whole batch. Each CTA + reads its row's kv_len on device and re-derives the sampling ladder (route_dynamic formula mirror), so with ``max_seq_len`` given (a capture-stable engine constant, e.g. dsa.py's ``indexer_max_seq_len``) the call performs NO host reads. Without ``max_seq_len`` the envelope comes from ONE ``kv_lens.max()`` host read (documented sync, refused under capture). - ``engine="reference"`` keeps the b=1 host-loop reference implementation — - the differential oracle the in-kernel engine is validated against. KNOWN LIMITATION: on rows containing NaN logits the selected index SET can differ from ``heuristicTopKDecode.cu`` (both kernels order NaNs @@ -1306,10 +1319,6 @@ def run_varlen( batch = num_rows // nn if kv_lens.shape[0] != batch: raise RuntimeError(f"kv_lens length {kv_lens.shape[0]} != num_rows/next_n = {batch}") - if len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch: - raise RuntimeError( - f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" - ) d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: raise RuntimeError(f"device index out of range: {d}") @@ -1323,142 +1332,103 @@ def run_varlen( if ws is None: ws = default_workspace(logits) - if engine == "auto": - # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- - # Full validation battery (the engine bypasses _run_impl — every - # check the batch-uniform path enforces is replayed here; the - # batch-dim check is CRITICAL: the kernel grid comes from - # logits.shape[0], so a short indices/values tensor would be written - # out of bounds). - if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): - raise RuntimeError("all tensors must be CUDA") - if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: - raise RuntimeError("logits must be float32; pre_idx/indices int32") - if len(indices.shape) != 2 or indices.shape[0] != num_rows: - raise RuntimeError( - f"indices must be [num_rows={num_rows}, >=k], got {tuple(indices.shape)}" - ) - k = pre_idx.shape[1] - if indices.shape[1] < k: - raise RuntimeError(f"indices width {indices.shape[1]} < k={k}") - if not (pre_idx.is_contiguous() and indices.is_contiguous() and kv_lens.is_contiguous()): - raise RuntimeError("pre_idx/indices/kv_lens must be contiguous") - # logits: accept row-major views with a wider row stride (the DSL - # paged-MQA logits arena is 256-aligned and column-sliced — a legal - # NON-contiguous view). The kernel only needs (base, row stride): - # widen back to a compact [rows, stride] view over the same storage; - # the tail columns are never classified (per-row n gates all reads). - if logits.stride(1) != 1: - raise RuntimeError("logits inner stride must be 1") - npad = logits.stride(0) if num_rows > 1 else logits.shape[1] - lg = logits - if not logits.is_contiguous(): - need = logits.storage_offset() + num_rows * npad - if logits.untyped_storage().size() // 4 < need: - raise RuntimeError("logits view storage too small to widen to its row stride") - lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) - if npad & 3: - raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") - if lg.data_ptr() & 15: - raise RuntimeError("logits base must be 16-byte aligned") - if values is not None: - if not values.is_cuda or values.dtype is not _F32: - raise RuntimeError("values must be CUDA float32") - if ( - len(values.shape) != 2 - or values.shape[0] != num_rows - or values.shape[1] < k - or not values.is_contiguous() - ): - raise RuntimeError( - f"values must be contiguous [num_rows={num_rows}, >=k], " - f"got {tuple(values.shape)}" - ) - cshift = 0 if cr == 1 else 2 - if max_seq_len is not None: - n_env = int(max_seq_len) >> cshift - else: - if _is_capturing(): - raise RuntimeError( - "run_varlen without max_seq_len reads kv_lens.max() on " - "host — pass max_seq_len (a capture-stable engine " - "constant) under CUDA graph capture" - ) - n_env = int(kv_lens.max().item()) >> cshift - # eager mode: quantize the data-dependent envelope up to the next - # power of two so a growing decode does not recompile at every - # R increment (bounded plans, bounded _VARLEN_CACHE) - n_env = 1 << max(n_env - 1, 1).bit_length() - n_env = min(max(n_env, 1), npad) - key = (num_rows, npad, k, n_env, nn, cr) - lc = _VARLEN_CACHE.get(key) - if lc is None: - if _is_capturing(): - raise RuntimeError( - "varlen launcher not compiled for this shape — warm up " - "before CUDA graph capture" - ) - lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) - idx = indices - if idx.shape[1] != k: - idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) - vals = values - if vals is not None and vals.shape[1] != k: - vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - if lc[0] == "reg_clus": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) - lc[1](lg, pre_idx, kv_lens, idx, lc[2]) - elif lc[0] == "reg": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) - lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) - elif lc[0] == "clus": - # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, - # SCAP, CMP, dead DYN x5) - lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) - else: - _, fn, pre, tail = lc - fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) - if vals is not None: - idx64 = idx.to(torch.int64) - vals.copy_(lg.gather(1, idx64.clamp_min(0))) - vals.masked_fill_(idx < 0, torch.finfo(_F32).min) - return - if engine != "reference": - raise RuntimeError(f"engine must be 'auto' or 'reference', got {engine!r}") - - # ---- reference engine (differential oracle): b=1 host loop -------------- - if _is_capturing(): - raise RuntimeError( - "run_varlen reference engine reads kv_lens on host, illegal under CUDA graph capture" - ) - # match the engine's flat-packed output convention for wider-than-k - # buffers (pack ONCE from the tensor base, then slice per row) - k = pre_idx.shape[1] - idx = indices - if len(idx.shape) != 2 or idx.shape[0] != num_rows: + # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- + # Full validation battery (the engine bypasses _run_impl — every + # check the batch-uniform path enforces is replayed here; the + # batch-dim check is CRITICAL: the kernel grid comes from + # logits.shape[0], so a short indices/values tensor would be written + # out of bounds). + if not (logits.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if logits.dtype is not _F32 or indices.dtype is not _I32: + raise RuntimeError("logits must be float32; indices must be int32") + if len(indices.shape) != 2 or indices.shape[0] != num_rows: raise RuntimeError( f"indices must be [num_rows={num_rows}, >=k], got {tuple(indices.shape)}" ) + k = indices.shape[1] + if not (indices.is_contiguous() and kv_lens.is_contiguous()): + raise RuntimeError("indices/kv_lens must be contiguous") + # logits: accept row-major views with a wider row stride (the DSL + # paged-MQA logits arena is 256-aligned and column-sliced — a legal + # NON-contiguous view). The kernel only needs (base, row stride): + # widen back to a compact [rows, stride] view over the same storage; + # the tail columns are never classified (per-row n gates all reads). + if logits.stride(1) != 1: + raise RuntimeError("logits inner stride must be 1") + npad = logits.stride(0) if num_rows > 1 else logits.shape[1] + lg = logits + if not logits.is_contiguous(): + need = logits.storage_offset() + num_rows * npad + if logits.untyped_storage().size() // 4 < need: + raise RuntimeError("logits view storage too small to widen to its row stride") + lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) + if npad & 3: + raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") + if lg.data_ptr() & 15: + raise RuntimeError("logits base must be 16-byte aligned") + if values is not None: + if not values.is_cuda or values.dtype is not _F32: + raise RuntimeError("values must be CUDA float32") + if ( + len(values.shape) != 2 + or values.shape[0] != num_rows + or values.shape[1] < k + or not values.is_contiguous() + ): + raise RuntimeError( + f"values must be contiguous [num_rows={num_rows}, >=k], got {tuple(values.shape)}" + ) + cshift = 0 if cr == 1 else 2 + if max_seq_len is not None: + n_env = int(max_seq_len) >> cshift + else: + if _is_capturing(): + raise RuntimeError( + "run_varlen without max_seq_len reads kv_lens.max() on " + "host — pass max_seq_len (a capture-stable engine " + "constant) under CUDA graph capture" + ) + n_env = int(kv_lens.max().item()) >> cshift + # eager mode: quantize the data-dependent envelope up to the next + # power of two so a growing decode does not recompile at every + # R increment (bounded plans, bounded _VARLEN_CACHE) + n_env = 1 << max(n_env - 1, 1).bit_length() + n_env = min(max(n_env, 1), npad) + key = (num_rows, npad, k, n_env, nn, cr) + lc = _VARLEN_CACHE.get(key) + if lc is None: + if _is_capturing(): + raise RuntimeError( + "varlen launcher not compiled for this shape — warm up before CUDA graph capture" + ) + lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) + idx = indices if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) vals = values if vals is not None and vals.shape[1] != k: vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - kl = kv_lens.tolist() # the ONE documented D2H sync of this engine - for r in range(num_rows): - # production graph slots can carry kv_len < next_n (padded / evicted - # requests): clamp to the empty row, emitting all -1 — the same - # contract the in-kernel engine implements - actual = max(kl[r // nn] - nn + (r % nn) + 1, 0) - req = r // nn - _run_impl( - logits[r : r + 1], - pre_idx[req : req + 1], - actual // cr, - idx[r : r + 1], - ws, - None if vals is None else vals[r : r + 1], - ) + # Hint-free engines do not read the compiled kernel's pre_idx ABI slot. + pre_arg = idx + if lc[0] == "reg_clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) + lc[1](lg, pre_arg, kv_lens, idx, lc[2]) + elif lc[0] == "reg": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) + elif lc[0] == "clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, + # SCAP, CMP, dead DYN x5) + lc[1](lg, pre_arg, kv_lens, idx, *lc[2]) + else: + _, fn, pre, tail = lc + fn(lg, pre_arg, idx, ws, *pre, kv_lens, *tail) + if vals is not None: + idx64 = idx.to(torch.int64) + vals.copy_(lg.gather(1, idx64.clamp_min(0))) + vals.masked_fill_(idx < 0, torch.finfo(_F32).min) + return __all__ = [ @@ -1511,6 +1481,7 @@ def warmup_varlen( producer layout (e.g. the DSL paged-MQA arena's 256-element rounding) must pass it; the 64-element default only matches producers that round the same way. + """ dev = torch.cuda.current_device() nn = max(1, int(next_n)) @@ -1560,7 +1531,15 @@ def warmup_varlen( raise RuntimeError( f"row_stride must be a float4-multiple >= n_env={n_env}, got {row_stride}" ) - key = (dev, int(top_k), int(max_seq_len), int(compress_ratio), nn, tuple(rows_list), npad) + key = ( + dev, + int(top_k), + int(max_seq_len), + int(compress_ratio), + nn, + tuple(rows_list), + npad, + ) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: return @@ -1569,20 +1548,18 @@ def warmup_varlen( # contiguous prefix views (compile keys depend on shapes only) logits = torch.zeros((rows_max, npad), dtype=torch.float32, device=dev) kv_lens = torch.full((rows_max // nn,), int(max_seq_len), dtype=torch.int32, device=dev) - pre_idx = torch.zeros((rows_max // nn, int(top_k)), dtype=torch.int32, device=dev) out = torch.empty((rows_max, int(top_k)), dtype=torch.int32, device=dev) for rows in rows_list: batch = rows // nn run_varlen( logits[:rows], - pre_idx[:batch], kv_lens[:batch], out[:rows], next_n=nn, compress_ratio=int(compress_ratio), max_seq_len=int(max_seq_len), ) - del logits, kv_lens, pre_idx, out + del logits, kv_lens, out torch.cuda.synchronize() # band launches compiled every ENGINE; now populate the per-row-count # LAUNCHER cache entries for the exact requested row counts (pure host diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 680e530d6c2f..71915e3acf3a 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -994,6 +994,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = sparse_attention_config.q_split_threshold indexer_rope_interleave = sparse_attention_config.indexer_rope_interleave enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk + use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk indexer_k_dtype = sparse_attention_config.indexer_k_dtype else: index_n_heads = pretrained_config.index_n_heads @@ -1007,6 +1008,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = 8192 indexer_rope_interleave = False enable_heuristic_topk = False + use_self_sampling_topk = True default_sparse_attention_config = DeepSeekV4SparseAttentionConfig( ) indexer_k_dtype = default_sparse_attention_config.indexer_k_dtype @@ -1023,6 +1025,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_config['q_split_threshold'] = q_split_threshold indexer_config['indexer_rope_interleave'] = indexer_rope_interleave indexer_config['enable_heuristic_topk'] = enable_heuristic_topk + indexer_config['use_self_sampling_topk'] = use_self_sampling_topk indexer_config['indexer_k_dtype'] = indexer_k_dtype return indexer_config @@ -1060,6 +1063,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): use_cute_dsl_paged_mqa_logits = sparse_attention_config.use_cute_dsl_paged_mqa_logits q_split_threshold = sparse_attention_config.q_split_threshold enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk + use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk indexer_k_dtype = sparse_attention_config.indexer_k_dtype index_share_for_mtp_iteration = sparse_attention_config.index_share_for_mtp_iteration else: @@ -1072,6 +1076,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): use_cute_dsl_paged_mqa_logits = False q_split_threshold = 8192 enable_heuristic_topk = False + use_self_sampling_topk = True indexer_k_dtype = "fp8" index_share_for_mtp_iteration = None kwargs[ @@ -1088,6 +1093,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold=q_split_threshold, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk, + use_self_sampling_topk=use_self_sampling_topk, indexer_k_dtype=indexer_k_dtype, index_share_for_mtp_iteration= index_share_for_mtp_iteration) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 950b02c30f93..609447938fa2 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -20,15 +20,11 @@ class TopKImplementation(str, Enum): TORCH = "torch" CUDA_RADIX = "cuda_radix" CUTE_DSL_RADIX = "cute_dsl_radix" - CUDA_GVR = "cuda_gvr" CUTE_DSL_GVR = "cute_dsl_gvr" - CUTE_DSL_GVR_V2 = "cute_dsl_gvr_v2" _GVR_IMPLEMENTATIONS = { - TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, - TopKImplementation.CUTE_DSL_GVR_V2, } _MAX_RADIX_BLOCKS_PER_ROW = 10 @@ -49,6 +45,7 @@ def __init__( prefill_implementation: TopKImplementation | None = None, decode_implementation: TopKImplementation | None = None, compress_ratio: int = 1, + gvr_self_sampling: bool = True, ) -> None: super().__init__() self.top_k = top_k @@ -59,6 +56,17 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio + # Second-level GVR dispatch for CUTE_DSL_GVR: True selects the + # hint-free self-sampling engine, False the temporal-hint engine. + self.gvr_self_sampling = gvr_self_sampling + + @property + def needs_gvr_prior(self) -> bool: + """Return whether decode consumes previous-step Top-K indices.""" + return ( + self.decode_implementation == TopKImplementation.CUTE_DSL_GVR + and not self.gvr_self_sampling + ) def forward( self, @@ -87,8 +95,11 @@ def forward( next_n: Number of decode rows per request. max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` - is the required caller-owned int32 previous selection with - shape ``[num_requests, top_k]`` on ``scores.device``. + is required by the temporal GVR path (``CUTE_DSL_GVR`` with + ``gvr_self_sampling=False``). It is + caller-owned int32 previous selection with shape + ``[num_requests, top_k]`` on ``scores.device``. The + self-sampling engine does not consume this state. ``gvr_row_order`` is an optional int32 request ordering with shape ``[num_requests]`` on the same device. @@ -196,8 +207,6 @@ def _forward_decode_radix( output_indices, next_n, self.top_k, - pre_idx=None, - heuristic_scratch=None, compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, @@ -261,8 +270,7 @@ def _forward_decode_gvr( gvr_prior_indices: torch.Tensor | None = None, gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: - assert gvr_prior_indices is not None - if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR_V2: + if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR and self.gvr_self_sampling: assert max_seq_len is not None if ( # engine hardware-format gate (falls through otherwise): @@ -279,24 +287,25 @@ def _forward_decode_gvr( and scores.data_ptr() % 16 == 0 and (scores.shape[0] > 1 or scores.shape[1] % 4 == 0) ): + # hint-free k derives from the output width; pin it to the module's k + assert output_indices.shape[1] == self.top_k from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen logger.info_once( "self-sampling GVR top-K engaged " f"(K={self.top_k}, cr={self.compress_ratio}, " - f"next_n={next_n}).", + f"next_n={next_n}, hint-free).", key="selfsampling_topk_engaged", ) - # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): + # Self-sampling GVR varlen engine: # one launch for the batch; per-row n from device kv_lens, # capture-stable tuning from the max-seq-len engine constant - # (no host reads — CUDA-graph safe). Hints are consumed raw - # (offset-free contract). The module receives max_seq_len in - # COMPRESSED index space; run_varlen's max_seq_len is in - # kv-token space like sequence_lengths — multiply back. + # (no host reads — CUDA-graph safe). The module receives + # max_seq_len in compressed index space; run_varlen's value + # is in KV-token space like sequence_lengths, so multiply it + # back by the compression ratio. selfsampling_topk_run_varlen( scores, - gvr_prior_indices, sequence_lengths, output_indices, next_n=next_n, @@ -305,20 +314,12 @@ def _forward_decode_gvr( ) return output_indices logger.warning_once( - "TRTLLM_GVR_SELF_SAMPLING=1 but the decode scores do not " + "self-sampling GVR is selected but the decode scores do not " "satisfy the engine's hardware-format gate " f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " - "falling through to the CUDA GVR top-K path.", + "falling back to the CUDA insertion/radix Top-K path.", key="selfsampling_topk_fallthrough", ) - if self.decode_implementation != TopKImplementation.CUTE_DSL_GVR: - # CUDA_GVR, or the V2 hardware-format fall-through above - workspace = self._get_workspace( - scores, - (scores.shape[0], self.top_k), - scores.dtype, - "top_k_cuda_gvr_workspace", - ) radix_indices, radix_values = self._get_radix_workspace(scores) torch.ops.trtllm.indexer_topk_decode( scores, @@ -326,25 +327,25 @@ def _forward_decode_gvr( output_indices, next_n, self.top_k, - pre_idx=gvr_prior_indices, - heuristic_scratch=workspace, compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) - else: - assert max_seq_len is not None - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - scores, - gvr_prior_indices, - sequence_lengths, - output_indices, - self.top_k, - next_n=next_n, - compress_ratio=self.compress_ratio, - max_seq_len=max_seq_len, - order_row=gvr_row_order, - ) + return output_indices + + assert gvr_prior_indices is not None + assert max_seq_len is not None + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + scores, + gvr_prior_indices, + sequence_lengths, + output_indices, + self.top_k, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_seq_len, + order_row=gvr_row_order, + ) return output_indices def update_gvr_prior_from_prefill( @@ -366,7 +367,7 @@ def update_gvr_prior_from_prefill( The slice starting at ``request_offset`` is updated in place. request_offset: First request row to update in the prior state. """ - if self.decode_implementation not in _GVR_IMPLEMENTATIONS: + if not self.needs_gvr_prior: return assert gvr_prior_indices is not None last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3772b100aaab..eb2a0b77b93a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -941,11 +941,20 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): default=False, description= "Whether to enable Guess-Verify-Refine (GVR) Top-K for the DSA decode " - "indexer. GVR reuses previous-step Top-K indices as hints to reduce " - "threshold search iterations. Currently supported for index_topk ∈ " - "{512, 1024, 2048} on Blackwell (SM100+), with compress_ratio ∈ {1, 4} " - "(DSv3.2 + DSv4 indexers). Falls back to the production insertion/" - "radix Top-K path when prerequisites are not met.") + "indexer instead of the exact insertion/radix Top-K path. Currently " + "supported for index_topk ∈ {512, 1024, 2048} on Blackwell (SM100+), " + "with compress_ratio ∈ {1, 4} (DSv3.2 + DSv4 indexers). Falls back to " + "the production insertion/radix Top-K path when prerequisites are not " + "met. `use_self_sampling_topk` selects the GVR engine generation.") + use_self_sampling_topk: bool = Field( + default=True, + description= + "Select the GVR engine generation when enable_heuristic_topk is set: " + "True (default) runs the hint-free self-sampling engine, which derives " + "its search bracket from the current row and keeps no cross-step " + "state; False runs the temporal-hint engines, which reuse the " + "previous decode step's Top-K indices as hints. Ignored when " + "enable_heuristic_topk is False.") indexer_k_dtype: Literal["fp8", "fp4"] = Field( default="fp8", description= @@ -1087,6 +1096,7 @@ def _value(name: str, default=None): q_split_threshold=self.q_split_threshold, indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, indexer_k_dtype=self.indexer_k_dtype, is_full_indexer_layer=self._is_full_indexer_layer( pretrained_config, kwargs.get("layer_idx")), @@ -1119,6 +1129,7 @@ def _value(name: str, default=None): index_head_dim=_value("index_head_dim", 128), enable_indexer_skip=self.skip_indexer_for_short_seqs, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, use_cute_dsl_topk=self.use_cute_dsl_topk, use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits), q_split_threshold=self.q_split_threshold, @@ -1206,6 +1217,7 @@ def _value(name: str, default=None): q_split_threshold=self.q_split_threshold, indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, indexer_k_dtype=self.indexer_k_dtype, compress_ratios=self.compress_ratios, window_size=self.window_size, @@ -1231,6 +1243,7 @@ def _value(name: str, default=None): index_head_dim=_value("index_head_dim", 128), enable_indexer_skip=self.skip_indexer_for_short_seqs, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, use_cute_dsl_topk=self.use_cute_dsl_topk, use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits), q_split_threshold=self.q_split_threshold, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 0295fdbef3c1..8cbd43261284 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -56,6 +56,7 @@ from tensorrt_llm._torch.attention_backend.sparse.dsa.cache_manager import ( _resolve_fp8_ds_mla_head_dim, ) +from tensorrt_llm._torch.attention_backend.sparse.dsa.params import use_self_sampling_gvr from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation @@ -121,11 +122,15 @@ def _set_torch_top_k(indexer: Indexer) -> None: ) -def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): +@pytest.mark.parametrize("use_self_sampling", [True, False]) +def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_sampling): sparse_config = DeepSeekV4SparseAttentionConfig( compress_ratios=[1, 4, 128], index_head_dim=96, + index_topk=512, indexer_k_dtype="fp8", + enable_heuristic_topk=True, + use_self_sampling_topk=use_self_sampling, ) sparse_metadata_params = sparse_config.to_sparse_metadata_params() metadata = object.__new__(DSAtrtllmAttentionMetadata) @@ -139,8 +144,18 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): metadata.create_buffers_for_mla_rope_append = Mock() metadata.create_buffers_for_indexer = Mock() - with patch( - "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", + return_value=100, + ), ): DSAtrtllmAttentionMetadata.__post_init__(metadata) @@ -148,6 +163,35 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): assert metadata.compress_ratios == [1, 4, 128] assert metadata._indexer_compress_ratio == 4 assert metadata._tokens_per_block == 64 + # The metadata mirror of the two-level dispatch drives prior allocation. + assert metadata.use_self_sampling_topk == use_self_sampling + assert metadata.needs_gvr_prior == (not use_self_sampling) + + +@pytest.mark.parametrize( + "kwargs,expected", + [ + (dict(), True), + (dict(sm_version=103, index_topk=2048, compress_ratio=4), True), + (dict(enable_heuristic_topk=False), False), + (dict(use_self_sampling_topk=False), False), + (dict(is_cute_dsl_available=False), False), + (dict(sm_version=120), False), + (dict(index_topk=256), False), + (dict(compress_ratio=2), False), + ], +) +def test_use_self_sampling_gvr(kwargs, expected): + base = dict( + enable_heuristic_topk=True, + use_self_sampling_topk=True, + index_topk=512, + compress_ratio=1, + is_cute_dsl_available=True, + sm_version=100, + ) + base.update(kwargs) + assert use_self_sampling_gvr(**base) is expected @pytest.mark.parametrize( @@ -213,8 +257,7 @@ def make_mock(num_generations, kv_lens_list): kv_lens_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda") row_order_buffer = torch.zeros(64, dtype=torch.int32, device="cuda") return SimpleNamespace( - enable_gvr_topk=True, - use_cute_dsl_topk=True, + needs_gvr_prior=True, num_generations=num_generations, num_sms=num_sms, max_draft_tokens=next_n - 1, @@ -273,7 +316,9 @@ def test_gvr_prior_writeback_uses_aux_stream(): enable_indexer_skip=True, ) indexer = create_indexer(sparse_config) - indexer._enable_heuristic_topk = True + # temporal GVR consumes the prior; force it independent of hardware + indexer.top_k.decode_implementation = TopKImplementation.CUTE_DSL_GVR + indexer.top_k.gvr_self_sampling = False indexer.aux_stream = torch.cuda.Stream() metadata.gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), @@ -332,6 +377,7 @@ def test_shared_topk_lifecycle(): metadata.enable_context_mla_with_cached_kv = False metadata.enable_indexer_skip = False metadata.enable_gvr_topk = False + metadata.needs_gvr_prior = False metadata.get_empty = Mock( side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) ) @@ -428,7 +474,7 @@ def test_indexer_post_load_weights_caches_fused_weight(): [ (False, False, TopKImplementation.CUDA_RADIX), (True, False, TopKImplementation.CUTE_DSL_RADIX), - (False, True, TopKImplementation.CUDA_GVR), + (False, True, TopKImplementation.CUTE_DSL_GVR), (True, True, TopKImplementation.CUTE_DSL_GVR), ], ) @@ -460,6 +506,56 @@ def test_indexer_configures_one_top_k_module( assert isinstance(indexer.top_k, TopK) assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert indexer.top_k.decode_implementation == expected_decode + if enable_heuristic: + # index_topk=128 misses the self-sampling prerequisites, so the + # default use_self_sampling_topk=True falls back to the temporal path. + assert not indexer.top_k.gvr_self_sampling + assert indexer.top_k.needs_gvr_prior + + +@skip_pre_hopper +@pytest.mark.parametrize( + "use_self_sampling,use_cute_dsl,expected_decode", + [ + (True, False, TopKImplementation.CUTE_DSL_GVR), + (True, True, TopKImplementation.CUTE_DSL_GVR), + (False, True, TopKImplementation.CUTE_DSL_GVR), + (False, False, TopKImplementation.CUTE_DSL_GVR), + ], +) +def test_indexer_two_level_gvr_dispatch( + monkeypatch, + use_self_sampling, + use_cute_dsl, + expected_decode, +): + # The retired TRTLLM_GVR_SELF_SAMPLING env must be ignored: with + # use_self_sampling_topk=False the temporal path must win regardless. + monkeypatch.setenv("TRTLLM_GVR_SELF_SAMPLING", "1") + sparse_config = DeepSeekSparseAttentionConfig( + index_head_dim=128, + index_n_heads=32, + index_topk=512, + use_cute_dsl_topk=use_cute_dsl, + enable_heuristic_topk=True, + use_self_sampling_topk=use_self_sampling, + ) + + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.get_sm_version", + return_value=100, + ), + ): + indexer = create_indexer(sparse_config) + + assert indexer.top_k.decode_implementation == expected_decode + assert indexer.top_k.gvr_self_sampling == use_self_sampling + assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) def _ceil_to_ue8m0(x: torch.Tensor): @@ -3523,8 +3619,8 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, ) - indexer._enable_heuristic_topk = True metadata_skip.gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), device="cuda", diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 9682d7981ab5..34fa45a5e0e5 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the reusable sparse index-selection Top-K module.""" -from contextlib import nullcontext +import sys +from types import SimpleNamespace from unittest.mock import Mock, call import pytest @@ -127,8 +128,6 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: output, 2, 2, - pre_idx=None, - heuristic_scratch=None, compress_ratio=4, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, @@ -142,6 +141,7 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: 2, decode_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4, + gvr_self_sampling=False, ) scores = torch.randn(1, 8) logical_lengths = torch.tensor([32], dtype=torch.int32) @@ -179,7 +179,11 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) next_n = 2 lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) row_order = torch.tensor([2, 0, 3, 1], dtype=torch.int32) @@ -201,8 +205,105 @@ def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: assert gvr.call_args.kwargs["order_row"] is row_order +def _install_fake_selfsampling_runner(monkeypatch) -> Mock: + """Replace the lazily imported self-sampling varlen entry with a Mock.""" + runner = Mock() + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k", + SimpleNamespace(selfsampling_topk_run_varlen=runner), + ) + return runner + + +def _run_gvr_v2_decode(top_k: TopK, out_width: int = 2) -> None: + scores = torch.randn(1, 8) # satisfies the V2 hardware-format gate + top_k( + scores, + torch.empty(1, out_width, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.tensor([32], dtype=torch.int32), + scan_lengths=torch.tensor([8], dtype=torch.int32), + next_n=1, + max_seq_len=16, + ) + + +def test_gvr_v2_decode_is_hint_free(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + + _run_gvr_v2_decode(top_k) + + args, kwargs = runner.call_args + assert len(args) == 3 + assert args[0].shape == (1, 8) + assert args[1].tolist() == [32] + assert args[2].shape == (1, 2) + assert kwargs == {"next_n": 1, "compress_ratio": 4, "max_seq_len": 64} + assert not top_k.needs_gvr_prior + + +def test_gvr_v2_hardware_gate_falls_back_without_prior(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + decode = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + scores = torch.randn(1, 8, dtype=torch.bfloat16) + lengths = torch.tensor([32], dtype=torch.int32) + output = torch.empty(1, 2, dtype=torch.int32) + + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=torch.tensor([8], dtype=torch.int32), + max_seq_len=16, + ) + + runner.assert_not_called() + decode.assert_called_once_with( + scores, + lengths, + output, + 1, + 2, + compress_ratio=4, + radix_aux_indices=None, + radix_aux_logits=None, + ) + + +def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None: + """Hint-free k derives from the output width; a scratch wider than + top_k must be rejected before launch, not silently become the k.""" + runner = _install_fake_selfsampling_runner(monkeypatch) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + with pytest.raises(AssertionError): + _run_gvr_v2_decode(top_k, out_width=3) + + runner.assert_not_called() + + def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) prior_indices = torch.zeros(3, 2, dtype=torch.int32) @@ -214,6 +315,31 @@ def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: ) assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + assert top_k.needs_gvr_prior + + +def test_gvr_v2_does_not_update_prior_from_prefill() -> None: + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + + top_k.update_gvr_prior_from_prefill( + torch.tensor([[4, 5]], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + prior_indices, + ) + + assert prior_indices.tolist() == [[0, 0]] + assert not top_k.needs_gvr_prior + + +def test_needs_gvr_prior_follows_two_level_dispatch() -> None: + assert not TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR).needs_gvr_prior + assert TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ).needs_gvr_prior + assert not TopK(2).needs_gvr_prior def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: @@ -247,76 +373,12 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: output, 1, 1, - pre_idx=None, - heuristic_scratch=None, compress_ratio=1, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) -def test_cuda_gvr_reserves_workspace_during_capture(monkeypatch) -> None: - decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) - device_context = Mock(side_effect=lambda _: nullcontext()) - monkeypatch.setattr(torch.cuda, "device", device_context) - - top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = Mock( - shape=(1, 8), - dtype=torch.float32, - is_cuda=True, - device=torch.device("cuda", 3), - ) - lengths = torch.tensor([8], dtype=torch.int32) - output = torch.empty(1, 2, dtype=torch.int32) - radix_indices = torch.empty(1, 10, 2, dtype=torch.int32) - radix_values = torch.empty(1, 10, 2) - workspace = torch.empty(1, 2) - prior_indices = torch.zeros(1, 2, dtype=torch.int32) - buffers = Mock() - buffers.get_buffer.side_effect = [workspace, radix_indices, radix_values] - monkeypatch.setattr(TopK, "_memory_buffers", buffers) - - top_k( - scores, - output, - is_prefill=False, - sequence_lengths=lengths, - scan_lengths=lengths, - gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, - ) - - assert buffers.get_buffer.call_args_list == [ - call( - (scores.shape[0], 2), - dtype=scores.dtype, - buffer_name="top_k_cuda_gvr_workspace_cuda:3", - reserve_buffer=True, - ), - call( - (scores.shape[0], 10, 2), - dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace_cuda:3", - reserve_buffer=True, - ), - call( - (scores.shape[0], 10, 2), - dtype=torch.float32, - buffer_name="top_k_radix_values_workspace_cuda:3", - reserve_buffer=True, - ), - ] - assert device_context.call_args_list == [call(scores.device)] * 3 - runtime_call = decode.call_args_list[-1] - assert runtime_call.kwargs["pre_idx"] is prior_indices - assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr() - assert runtime_call.kwargs["radix_aux_indices"] is radix_indices - assert runtime_call.kwargs["radix_aux_logits"] is radix_values - assert prior_indices.tolist() == [[0, 0]] - - def test_unsupported_prefill_implementation_raises() -> None: top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_RADIX) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 44fbb6396b0c..ad3a2dadf5b5 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -228,20 +228,16 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) -def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="auto"): +def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" - batch, rows = len(kv), len(kv) * next_n + rows = len(kv) * next_n n_r = [(kv[r // next_n] - next_n + (r % next_n) + 1) // cr for r in range(rows)] npad = (max(n_r) + 63) // 64 * 64 gen = torch.Generator(device=_DEV).manual_seed(seed) logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 for r in range(rows): logits[r, n_r[r] :] = 3e38 # poison beyond each row's OWN n_r - pre_idx = torch.empty((batch, top_k), dtype=torch.int32, device=_DEV) - for q in range(batch): - nmin = max(min(n_r[q * next_n : (q + 1) * next_n]), 1) - pre_idx[q] = torch.randint(0, nmin, (top_k,), generator=gen, dtype=torch.int32, device=_DEV) indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) values = ( torch.full((rows, top_k), 7.0, dtype=torch.float32, device=_DEV) if with_values else None @@ -249,13 +245,11 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) ss_host.run_varlen( logits, - pre_idx, kv_lens, indices, next_n=next_n, compress_ratio=cr, values=values, - engine=engine, ) torch.cuda.synchronize() fmin = torch.finfo(torch.float32).min @@ -281,7 +275,28 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut assert torch.equal(values[r], torch.gather(logits[r], 0, idx)) -@pytest.mark.parametrize("engine", ["auto", "reference"]) +def _reference_varlen_indices(logits, kv_lens, next_n, compress_ratio, top_k): + """Build a simple torch reference for the hint-free varlen contract.""" + reference = torch.full((logits.shape[0], top_k), -1, dtype=torch.int32, device=logits.device) + lengths = kv_lens.tolist() + for row in range(logits.shape[0]): + valid = ( + max( + lengths[row // next_n] - next_n + row % next_n + 1, + 0, + ) + // compress_ratio + ) + valid = min(valid, logits.shape[1]) + if valid <= 0: + continue + if valid <= top_k: + reference[row, :valid] = torch.arange(valid, dtype=torch.int32, device=logits.device) + else: + reference[row] = torch.topk(logits[row, :valid], top_k).indices.to(torch.int32) + return reference + + @pytest.mark.parametrize( "kv,next_n,cr,top_k", [ @@ -294,50 +309,9 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut ], ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], ) -def test_selfsampling_topk_varlen(kv, next_n, cr, top_k, engine): - """run_varlen production contract: per-row n from device kv_lens with the - MTP window formula, request-level hints, per-row short path — on BOTH the - per-row in-kernel engine ("auto") and the b=1 reference loop.""" - _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, engine=engine) - - -def test_selfsampling_topk_varlen_engine_matches_reference(): - """Differential: the in-kernel engine's per-row value multisets must - equal the reference loop's on a mixed batch (deep SPLIT rows, tsh band, - short rows, compressed space).""" - kv = [524288, 131075, 32800, 2000, 65540, 8192, 262144, 900] - top_k, cr = 1024, 4 - rows = len(kv) - n_r = [(v - 1 + 1) // cr for v in kv] - npad = (max(n_r) + 63) // 64 * 64 - gen = torch.Generator(device=_DEV).manual_seed(77) - logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - for r in range(rows): - logits[r, n_r[r] :] = 3e38 - pre_idx = torch.empty((rows, top_k), dtype=torch.int32, device=_DEV) - for q in range(rows): - pre_idx[q] = torch.randint( - 0, max(n_r[q], 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV - ) - kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) - out_a = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) - out_r = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre_idx, kv_lens, out_a, compress_ratio=cr) - ss_host.run_varlen(logits, pre_idx, kv_lens, out_r, compress_ratio=cr, engine="reference") - torch.cuda.synchronize() - for r in range(rows): - if n_r[r] <= top_k: - assert torch.equal(out_a[r], out_r[r]) or torch.equal( - torch.sort(out_a[r]).values, torch.sort(out_r[r]).values - ) - else: - ga = torch.sort( - torch.gather(logits[r], 0, out_a[r].to(torch.int64)) + 0.0, descending=True - ).values - gr = torch.sort( - torch.gather(logits[r], 0, out_r[r].to(torch.int64)) + 0.0, descending=True - ).values - assert torch.equal(ga, gr), f"row {r}: engine != reference" +def test_selfsampling_topk_varlen(kv, next_n, cr, top_k): + """Validate the hint-free per-row KV-length and MTP window contract.""" + _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr) def test_selfsampling_topk_varlen_values(): @@ -361,90 +335,46 @@ def test_selfsampling_topk_varlen_launch_modes(rows, base, step, top_k): def test_selfsampling_topk_varlen_zero_kv_slot(): """Padded / evicted CUDA-graph request slots can carry kv_len < next_n - (even 0): both engines must emit the empty short row (all -1), not raise — + (even 0): the engine must emit the empty short row (all -1), not raise — mixed with live MTP rows of another request in the same launch.""" - for engine in ("auto", "reference"): - gen = torch.Generator(device=_DEV).manual_seed(9) - logits = torch.randn((8, 8192), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) - kv_lens = torch.tensor([0, 8192], dtype=torch.int32, device=_DEV) - indices = torch.full((8, 512), -7, dtype=torch.int32, device=_DEV) - for r in range(4, 8): - n = 8192 - 4 + (r - 4) + 1 - logits[r, n:] = 3e38 - ss_host.run_varlen( - logits, pre_idx, kv_lens, indices, next_n=4, compress_ratio=1, engine=engine - ) - torch.cuda.synchronize() - assert bool((indices[:4] == -1).all()), f"{engine}: kv=0 rows must be all -1" - for r in range(4, 8): - n = 8192 - 4 + (r - 4) + 1 - idx = indices[r].to(torch.int64) - assert int(idx.min()) >= 0 and int(idx.max()) < n - ref = torch.topk(logits[r, :n], 512).values - got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values - assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values) - - -def test_selfsampling_topk_varlen_wide_buffers_flat_packed(): - """indices wider than k follow the CUDA flat-packed contract (rows at - stride k from the tensor base) IDENTICALLY on both engines.""" - kv = [9000, 300] - top_k, width = 512, 512 + 64 - gen = torch.Generator(device=_DEV).manual_seed(21) - logits = torch.randn((2, 9024), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 - logits[0, 9000:] = 3e38 - logits[1, 300:] = 3e38 - pre_idx = torch.randint(0, 300, (2, top_k), generator=gen, dtype=torch.int32, device=_DEV) - kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) - outs = {} - for engine in ("auto", "reference"): - wide = torch.full((2, width), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre_idx, kv_lens, wide, compress_ratio=1, engine=engine) - torch.cuda.synchronize() - outs[engine] = wide.reshape(-1)[: 2 * top_k].view(2, top_k).clone() - packed_a, packed_r = outs["auto"], outs["reference"] - # row 0 (kernel row): same value multiset via the SAME packed convention - ga = torch.sort( - torch.gather(logits[0], 0, packed_a[0].to(torch.int64)) + 0.0, descending=True - ).values - gr = torch.sort( - torch.gather(logits[0], 0, packed_r[0].to(torch.int64)) + 0.0, descending=True - ).values - assert torch.equal(ga, gr), "wide-buffer packing convention diverged between engines" - # row 1 (short row): identity + -1 tail at the packed location, bit-equal - expect = torch.cat( - [ - torch.arange(300, dtype=torch.int32, device=_DEV), - torch.full((top_k - 300,), -1, dtype=torch.int32, device=_DEV), - ] - ) - assert torch.equal(torch.sort(packed_a[1, :300]).values, expect[:300]) - assert bool((packed_a[1, 300:] == -1).all()) - assert torch.equal(torch.sort(packed_r[1, :300]).values, expect[:300]) - assert bool((packed_r[1, 300:] == -1).all()) + gen = torch.Generator(device=_DEV).manual_seed(9) + logits = torch.randn((8, 8192), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + kv_lens = torch.tensor([0, 8192], dtype=torch.int32, device=_DEV) + indices = torch.full((8, 512), -7, dtype=torch.int32, device=_DEV) + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + logits[r, n:] = 3e38 + ss_host.run_varlen(logits, kv_lens, indices, next_n=4, compress_ratio=1) + torch.cuda.synchronize() + assert bool((indices[:4] == -1).all()), "kv=0 rows must be all -1" + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + idx = indices[r].to(torch.int64) + assert int(idx.min()) >= 0 and int(idx.max()) < n + ref = torch.topk(logits[r, :n], 512).values + got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values) def test_selfsampling_topk_varlen_guards(): logits = torch.randn((2, 8192), dtype=torch.float32, device=_DEV) - pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) indices = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) kv = torch.tensor([8192, 8192], dtype=torch.int32, device=_DEV) with pytest.raises(RuntimeError, match="kv_lens length"): - ss_host.run_varlen(logits, pre_idx, kv[:1], indices) + ss_host.run_varlen(logits, kv[:1], indices) with pytest.raises(RuntimeError, match="not divisible"): - ss_host.run_varlen(logits, pre_idx, kv, indices, next_n=3) + ss_host.run_varlen(logits, kv, indices, next_n=3) with pytest.raises(RuntimeError, match="compress_ratio"): - ss_host.run_varlen(logits, pre_idx, kv, indices, compress_ratio=2) + ss_host.run_varlen(logits, kv, indices, compress_ratio=2) with pytest.raises(RuntimeError, match="CUDA tensor"): - ss_host.run_varlen(logits, pre_idx, kv.cpu(), indices) + ss_host.run_varlen(logits, kv.cpu(), indices) with pytest.raises(RuntimeError, match="num_rows"): # request-level-shaped indices under MTP: MUST be rejected (the # kernel grid comes from logits rows — silent OOB writes otherwise) - ss_host.run_varlen(logits, pre_idx[:1], kv[:1], indices[:1], next_n=2) + ss_host.run_varlen(logits, kv[:1], indices[:1], next_n=2) with pytest.raises(RuntimeError, match="contiguous"): strided = torch.zeros((2, 2), dtype=torch.int32, device=_DEV)[:, 0] - ss_host.run_varlen(logits, pre_idx, strided, indices) + ss_host.run_varlen(logits, strided, indices) def test_selfsampling_topk_varlen_cuda_graph(): @@ -456,7 +386,6 @@ def test_selfsampling_topk_varlen_cuda_graph(): rows, top_k, msl = 4, 512, 262144 npad = msl logits = torch.randn((rows, npad), dtype=torch.float32, device=_DEV) - 2.0 - pre_idx = torch.zeros((rows, top_k), dtype=torch.int32, device=_DEV) kv_lens = torch.tensor([100, 4099, 131070, 200000], dtype=torch.int32, device=_DEV) indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) @@ -470,17 +399,14 @@ def refresh(step): for r in range(rows): n = min(kv[r], npad) logits[r, n:] = 3e38 - pre_idx[r] = torch.randint( - 0, max(n, 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV - ) return kv refresh(0) - ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + ss_host.run_varlen(logits, kv_lens, indices, compress_ratio=1, max_seq_len=msl) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + ss_host.run_varlen(logits, kv_lens, indices, compress_ratio=1, max_seq_len=msl) for step in range(1, 6): kv = refresh(step) indices.fill_(-7) @@ -603,10 +529,9 @@ def test_selfsampling_topk_varlen_rejects_non_fp32(): the loud contract message instead of a CuTe typing failure).""" logits = torch.randn(1, 8192, device=_DEV, dtype=torch.bfloat16) kv = torch.tensor([8000], dtype=torch.int32, device=_DEV) - pre = torch.zeros(1, 512, dtype=torch.int32, device=_DEV) out = torch.empty(1, 512, dtype=torch.int32, device=_DEV) with pytest.raises(RuntimeError, match="float32"): - ss_host.run_varlen(logits, pre, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) + ss_host.run_varlen(logits, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) def test_selfsampling_topk_varlen_zero_window_rows(): @@ -620,9 +545,8 @@ def test_selfsampling_topk_varlen_zero_window_rows(): rows = kv.numel() * nn npad = (msl // cr + 63) // 64 * 64 logits = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.zeros(kv.numel(), k, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + ss_host.run_varlen(logits, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) torch.cuda.synchronize() assert (out[nn:] == -1).all().item(), "n<=0 rows must be fully -1-padded" for r in range(nn): @@ -647,11 +571,10 @@ def test_selfsampling_warmup_row_stride_matches_arena(): arena = torch.randn(rows, stride, dtype=torch.float32, device=_DEV) logits = arena[:, :msl] # non-contiguous column slice, like serving kv = torch.full((rows,), msl, dtype=torch.int32, device=_DEV) - pre = torch.zeros(rows, k, dtype=torch.int32, device=_DEV) out = torch.empty(rows, k, dtype=torch.int32, device=_DEV) g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): - ss_host.run_varlen(logits, pre, kv, out, max_seq_len=msl) + ss_host.run_varlen(logits, kv, out, max_seq_len=msl) g.replay() torch.cuda.synchronize() ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values @@ -672,24 +595,13 @@ def test_selfsampling_varlen_regclus_parity_and_oracle(): rows = batch * nn torch.manual_seed(7) lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.tensor([msl_c * cr, 900, nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg_clus", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl_c * cr, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -708,15 +620,14 @@ def test_selfsampling_varlen_regclus_cuda_graph(): rows = 8 torch.manual_seed(11) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -749,24 +660,13 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): assert fam == want, (fam, want, k, msl_c) msl = msl_c * cr lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.tensor([msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -797,27 +697,16 @@ def test_selfsampling_varlen_clus_parity_and_oracle(): assert plan["kernel"] == "clus" and plan["cluster"] == want_cs, plan batch = rows // nn lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) lens = [msl_c * cr, 900, nn - 1, 20000, 40000, 300000, msl_c * cr // 2, 5000] kv = torch.tensor( [lens[i % len(lens)] for i in range(batch)], dtype=torch.int32, device=_DEV ) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) key = (rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] - ss_host.run_varlen( - lg, - pre, - kv, - ref, - next_n=nn, - compress_ratio=cr, - max_seq_len=msl_c * cr, - engine="reference", - ) torch.cuda.synchronize() + ref = _reference_varlen_indices(lg, kv, nn, cr, k) for r in range(rows): if (ref[r] >= 0).any(): row = lg[r].float() @@ -837,17 +726,16 @@ def test_selfsampling_varlen_clus_cuda_graph(): rows = 32 torch.manual_seed(29) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() key = (rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -864,17 +752,16 @@ def test_selfsampling_varlen_reg_cuda_graph(): rows = 16 torch.manual_seed(17) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() key = (rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) with torch.cuda.graph(g): - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values for _ in range(2): out.fill_(-7) @@ -893,10 +780,9 @@ def test_selfsampling_varlen_full_row_range(): torch.manual_seed(3) for rows in (304, 1024): lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() key = (rows, msl_c, k, msl_c, 1, cr) assert key in ss_host._VARLEN_CACHE, "row count must dispatch in-engine" @@ -935,14 +821,13 @@ def test_selfsampling_varlen_heterogeneous_lengths_main(): rows = batch * nn # 304 -> route_streaming main family torch.manual_seed(5) lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) - pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) kv = torch.randint(1, msl_c * cr, (batch,), dtype=torch.int32, device=_DEV) kv[0] = msl_c * cr # full length kv[1] = 900 # short (n <= k) kv[2] = nn - 1 # zero-window (every row of the request empty) kv[3] = k * cr + nn # just above the short path out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() kl = kv.tolist() for r in range(rows): diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 33ebdb2c82ca..e065db98ac59 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -7,38 +7,12 @@ # http://www.apache.org/licenses/LICENSE-2.0 """ -Distribution-parameterized correctness tests for the heuristic indexer_topk_decode. - -Logits are sampled from four distribution families that characterise -negative-shifted decode-phase logit spaces (means −0.5 to −4.5): - - beta — bounded, bell-shaped; near-zero / moderate / deep negative mean - logistic — heavy-tailed symmetric (leptokurtic) - lognorm — positively skewed, wide support - weibull_min — right-skewed extreme-value; narrow and wide spread variants - -Two extensions beyond the baseline test_indexer_topk.py: - - pre_idx (heuristic candidates) - Shape [batch_size, index_topk]. For each batch element b, pre_idx[b] is - built from the base row's actual top-K: - - pre_idx[b, 0] = argmax (kernel invariant) - - floor(index_topk * success_ratio) slots drawn from actual top-K indices - - remaining slots filled with random valid indices - success_ratio is a pytest parameter (>= 0.4). - - MTP structure (next_n > 1) - When next_n > 1, consecutive rows within each batch element share most of - their logit values. For batch element b with valid_base = row_ends[b*next_n] - and MTP offset nni = 1…next_n-1: - logits[b*next_n + nni, nni : nni+valid_base] = logits[b*next_n, 0 : valid_base] - Positions 0..nni-1 are independently sampled (new token positions); - positions >= nni+valid_base remain -inf. - -Logit shapes (batch_size, next_n, num_tokens) match test_indexer_topk.py. +Correctness tests for the indexer Top-K custom ops: the CUDA +`indexer_topk_decode` / `indexer_topk_prefill` dispatchers +(insertion / radix / radix-split-work tiers) and the CuTe DSL +radix / filtered Top-K kernels. """ -import numpy as np import pytest import torch from utils.util import getSMVersion, skip_pre_blackwell, skip_pre_hopper @@ -54,15 +28,6 @@ if not torch.cuda.is_available(): pytest.skip("CUDA is required for indexer_topk tests", allow_module_level=True) -try: - import scipy.stats as _scipy_stats - from scipy.special import gamma as _gamma - - _HAS_SCIPY = True -except ImportError: - _HAS_SCIPY = False - - # --------------------------------------------------------------------------- # Prefill parameter helpers (unchanged from test_indexer_topk.py) # --------------------------------------------------------------------------- @@ -380,7 +345,7 @@ def test_indexer_topk_decode_launch_policy_transitions( # # 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). +# buffers with stable addresses (CUDA-graph safe). # # These tests verify: # (a) caller-owned-aux output matches the default (th::empty) path, @@ -870,460 +835,6 @@ def test_filtered_topk_varlen_odd_k(top_k, dtype_name): ) -# --------------------------------------------------------------------------- -# Distribution configs for heuristic decode correctness tests -# -# Each entry is a dict with keys: -# dist — distribution family -# mean — target mean (negative; typical decode logit range −0.5 to −4.5) -# std — target standard deviation -# full_range — support width (high − low), used as the bounding interval -# c — Weibull shape parameter (weibull_min only; c≈14 for moderate skew) -# -# Parameter derivation (all analytical, no external data dependencies): -# -# beta: -# low = mean − full_range/2, high = mean + full_range/2 -# mu01 = (mean − low) / full_range -# conc = mu01*(1−mu01) / (std/full_range)² − 1 -# α = conc*mu01, β = conc*(1−mu01) -# -# logistic: -# scale = std * √3 / π [std(logistic) = scale*π/√3] -# CDF inversion: x = mean + scale * ln(u/(1−u)), u ~ U(0,1) -# -# lognorm (left-shifted to loc = mean − full_range/2): -# pos_mean = full_range / 2 -# σ = √(log(1 + (std/pos_mean)²)) → matches target std exactly -# scale = exp(log(pos_mean) − σ²/2) → matches target mean exactly -# -# weibull_min (shape c fixed; loc/scale solved from mean/std): -# scale = std / √(Γ(1+2/c) − Γ²(1+1/c)) -# loc = mean − scale * Γ(1+1/c) -# --------------------------------------------------------------------------- - -_DECODE_DIST_CONFIGS = [ - # --- Beta: bounded, bell-shaped --- - # shallow negative mean, wide spread - dict(dist="beta", mean=-0.75, std=1.90, full_range=13.60), - # moderate negative mean - dict(dist="beta", mean=-2.96, std=1.68, full_range=12.85), - # deep negative mean, narrow spread - dict(dist="beta", mean=-4.51, std=1.75, full_range=11.24), - # --- Logistic: heavy-tailed symmetric (leptokurtic) --- - dict(dist="logistic", mean=-0.47, std=1.46, full_range=12.32), - # --- Lognorm: positively skewed, wide support --- - dict(dist="lognorm", mean=-4.12, std=2.55, full_range=17.28), - # --- Weibull minimum: right-skewed extreme-value --- - # wider spread - dict(dist="weibull_min", mean=-3.04, std=1.57, full_range=12.30, c=14.0), - # narrower spread - dict(dist="weibull_min", mean=-2.26, std=1.28, full_range=9.71, c=14.0), -] - -# Human-readable pytest IDs: dist_mean_std -_DECODE_DIST_IDS = [ - f"{c['dist']}_m{abs(c['mean']):.2f}_s{c['std']:.2f}" for c in _DECODE_DIST_CONFIGS -] - - -# --------------------------------------------------------------------------- -# Distribution-aware logit generator -# --------------------------------------------------------------------------- - - -def _fit_beta_params(mean: float, std: float, low: float, high: float): - """Fit Beta(α, β) on [low, high] to match target mean and std.""" - r = high - low - mu = (mean - low) / r - var = min((std / r) ** 2, mu * (1 - mu) * 0.99) - conc = mu * (1 - mu) / var - 1 - return conc * mu, conc * (1 - mu) - - -def _fit_weibull_params(mean: float, std: float, c: float): - """Fit Weibull_min(c, loc, scale) to match target mean and std.""" - g1 = _gamma(1 + 1 / c) - g2 = _gamma(1 + 2 / c) - scale = std / np.sqrt(g2 - g1**2) - return c, mean - scale * g1, scale - - -def create_distributed_logits( - cfg: dict, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - dtype: torch.dtype, - seed: int, -) -> torch.Tensor: - """ - Generate a logits tensor sampled from the distribution specified by *cfg*. - - Values outside [row_start, row_end) are set to -inf. All distribution - parameters are derived analytically from (mean, std, full_range). - - Args: - cfg: One entry from _DECODE_DIST_CONFIGS - row_starts: (num_rows,) inclusive start column per row - row_ends: (num_rows,) exclusive end column per row - dtype: Target torch dtype - seed: NumPy RNG seed - - Returns: - Tensor (num_rows, max_len) with sampled values and -inf padding. - """ - rng = np.random.default_rng(seed) - num_rows = int(row_starts.shape[0]) - max_len = int(row_ends.cpu().max().item()) - # Pad to multiple of 8 so stride0 satisfies the alignment requirement of - # launchHeuristicTopKDecode for both fp32 (float4 = 4 elements) and - # bf16/fp16 (int4 = 16 B = 8 elements) in multi-row mode (matches TRT-LLM - # runtime where strides are always multiples of tokens_per_block >= 64). - max_len = (max_len + 7) & ~7 - size = (num_rows, max_len) - - dist = cfg["dist"] - mean, std, full_range = cfg["mean"], cfg["std"], cfg["full_range"] - low = mean - full_range / 2 - high = mean + full_range / 2 - - if dist == "beta": - alpha, beta_p = _fit_beta_params(mean, std, low, high) - samples = (rng.beta(alpha, beta_p, size=size) * (high - low) + low).astype(np.float32) - - elif dist == "logistic": - s = std * np.sqrt(3) / np.pi - u = rng.uniform(1e-7, 1 - 1e-7, size=size) - samples = (mean + s * np.log(u / (1 - u))).astype(np.float32) - - elif dist == "lognorm": - loc = low - pos_mean = max(mean - loc, 1e-6) # = full_range / 2 - sigma = float(np.sqrt(np.log(1 + (std / pos_mean) ** 2))) - scale = np.exp(np.log(pos_mean) - sigma**2 / 2) - samples = _scipy_stats.lognorm.rvs( - s=sigma, - loc=loc, - scale=scale, - size=size, - random_state=int(rng.integers(2**31 - 1)), - ).astype(np.float32) - - elif dist == "weibull_min": - c, loc, scale = _fit_weibull_params(mean, std, cfg.get("c", 14.0)) - samples = _scipy_stats.weibull_min.rvs( - c, - loc=loc, - scale=scale, - size=size, - random_state=int(rng.integers(2**31 - 1)), - ).astype(np.float32) - - else: - raise ValueError(f"Unknown distribution: {dist!r}") - - # Clip to [low, high] to bound the effective value range to exactly full_range. - # Unbounded distributions (lognorm, logistic, weibull_min) can produce outliers - # above `high` that inflate the histogram bin width in the kernel's 256-bin - # threshold search, causing boundary-element misidentification. - samples = np.clip(samples, low, high).astype(np.float32) - - logits = torch.from_numpy(samples).to(dtype=dtype, device="cuda") - col_idx = torch.arange(max_len, device="cuda").unsqueeze(0) - mask = (col_idx < row_starts.unsqueeze(1)) | (col_idx >= row_ends.unsqueeze(1)) - logits[mask] = float("-inf") - return logits - - -# --------------------------------------------------------------------------- -# MTP structure: make consecutive rows within a batch correlated -# --------------------------------------------------------------------------- - - -def apply_mtp_structure( - logits: torch.Tensor, - batch_size: int, - next_n: int, - row_ends: torch.Tensor, -) -> torch.Tensor: - """ - Enforce MTP (Multi-Token Prediction) logit correlation within each batch. - - For batch element b with base row valid length valid_base = row_ends[b*next_n], - each MTP offset nni = 1…next_n-1 satisfies: - - logits[b*next_n + nni, nni : nni+valid_base] = logits[b*next_n, 0 : valid_base] - - Positions 0..nni-1 of each MTP row remain independently sampled (new token - positions). Positions nni+valid_base.. are already -inf from - create_distributed_logits (since row_ends[b*next_n+nni] = valid_base+nni), - so no additional masking is required after this function. - - Args: - logits: (batch_size*next_n, max_len) float tensor; modified in-place - batch_size: number of batch elements - next_n: MTP factor; returns logits unchanged when next_n == 1 - row_ends: (batch_size*next_n,) exclusive end column per row - - Returns: - Same logits tensor with MTP segments overwritten. - """ - if next_n == 1: - return logits - - for b in range(batch_size): - base = b * next_n - valid_base = int(row_ends[base].item()) # valid length of base row - for nni in range(1, next_n): - # Copy base row positions [0:valid_base] → MTP row positions [nni:nni+valid_base] - # Positions 0..nni-1 stay independently sampled; positions >=nni+valid_base stay -inf. - logits[base + nni, nni : nni + valid_base] = logits[base, :valid_base] - - return logits - - -# --------------------------------------------------------------------------- -# pre_idx generator: heuristic candidate indices for the indexer kernel -# --------------------------------------------------------------------------- - - -def generate_pre_idx( - 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: - """ - Build the heuristic pre-prediction index tensor for each batch element. - - The V3.2 multi-row kernel (`heuristicTopKMultiRowKernel{,Dtype}` in - cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu) internally adds - `preIdxOffset = (rowIdx % next_n) + 1` to every pre_idx slot during its - Phase-1 stats reduction (heuristic_topk.cuh:654/1209). Production V3.2 - callers therefore pass pre_idx in PREVIOUS-step coordinates so the - kernel's +1 / +2 / +3 shift maps prev positions to current-step - positions correctly. - - This test builds pre_idx from `current_logits.topk()`, then applies a - `-1` shift before returning, so the kernel's internal `+1` brings every - hint back to its intended current-step position (preserves the kernel's - argmax invariant: kernel reads `input[(argmax_pos - 1) + 1] = input[argmax_pos]`). - - For batch element b (base row = b*next_n): - - pre_idx[b, 0] = argmax of the base row (kernel invariant) - - n_hit slots = floor(index_topk * success_ratio) indices drawn - WITHOUT replacement from the actual top-K - - n_fill = index_topk - n_hit slots - = indices drawn WITHOUT replacement from the - non-top-K pool (all valid indices except top-K) - - No element appears more than once in pre_idx[b]. The hit and fill pools - are disjoint by construction, so cross-pool duplicates are impossible. - - Edge case: when valid_len < index_topk (short sequences, ~10% of batches), - the non-top-K pool may be smaller than n_fill. In that case all available - non-top-K indices are used first; any remaining slots are filled from the - unused top-K tail (topk_idx[n_hit:]) to preserve the no-duplicate guarantee - as far as possible. - - Args: - logits: (batch_size*next_n, max_len) logits tensor - row_ends: (batch_size*next_n,) valid lengths per row - batch_size: number of batch elements - next_n: MTP factor; base row index is b*next_n - index_topk: number of pre-predicted candidates (K) - success_ratio: fraction of pre_idx drawn from actual top-K (>= 0.4) - seed: torch manual seed for reproducibility - - Returns: - pre_idx: int32 tensor of shape (batch_size, index_topk), no duplicates - """ - 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 (no duplicates); index 0 = argmax (kernel invariant) - _, topk_idx = logits[base, :valid_len].topk(k) - - # --- Hit slots: sample n_hit from top-K without replacement --- - # Always include argmax at position 0. - 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() - - # --- Fill slots: sample n_fill from non-top-K pool without replacement --- - # The non-top-K pool is disjoint from topk_idx, so no cross-pool duplicates. - n_fill = index_topk - n_hit - if n_fill > 0: - # Build non-top-K pool: all valid indices that are NOT in topk_idx - topk_mask = torch.zeros(valid_len, dtype=torch.bool, device=logits.device) - topk_mask[topk_idx] = True - non_topk = torch.where(~topk_mask)[0] # shape: (valid_len - k,) - - if len(non_topk) >= n_fill: - # Normal case: enough non-top-K candidates - perm = torch.randperm(len(non_topk), device=logits.device)[:n_fill] - pre_idx[b, n_hit:] = non_topk[perm].int() - else: - # Edge case (valid_len ≈ index_topk): use all non-top-K first, - # then fill remaining from the unused top-K tail (topk_idx[n_hit:]) - 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:] # not yet in pre_idx[b] - 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() - - # V3.2 compensation: kernel adds `(rowIdx % next_n) + 1` to every pre_idx - # entry during P1 stats reduction. Shifting by -1 here means that for the - # base row (rowIdx % next_n == 0, offset = +1) the kernel reads the exact - # current-step positions our `topk()` selected. Negative entries are - # silently dropped by the kernel's `idx >= 0 && idx < N` range check. - pre_idx -= 1 - 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 - - # radix filter single-cta test. @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @@ -1643,283 +1154,6 @@ def run_fn(logits, seq_lens): ) -# ============================================================================ -# Heuristic Decode Distribution-Parameterised Tests -# ============================================================================ - - -@skip_pre_blackwell -@pytest.mark.skipif(not _HAS_SCIPY, reason="scipy required for distribution tests") -@pytest.mark.parametrize("success_ratio", [0.5, 0.9]) -@pytest.mark.parametrize("dist_cfg", _DECODE_DIST_CONFIGS, ids=_DECODE_DIST_IDS) -@pytest.mark.parametrize("batch_size", [1, 64, 128]) -@pytest.mark.parametrize("next_n", [1, 2, 3]) -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -# num_tokens=4096 added to cover the new uniform kSeqSmall=4096 boundary -# across all K (indexerTopK.cu kSeqSmallDefaultForK). num_tokens=4096 sits -# right at the GVR routing threshold for every K ∈ {512, 1024, 2048} so the -# assertion validates the just-inside-GVR path correctness. -@pytest.mark.parametrize("num_tokens", [4096, 8192, 16384]) -@pytest.mark.parametrize( - "dtype", - [torch.float32, torch.bfloat16, torch.float16], - ids=["fp32", "bf16", "fp16"], -) -def test_indexer_topk_decode_dist( - dist_cfg, batch_size, next_n, index_topk, num_tokens, success_ratio, dtype -): - """ - Correctness test for the heuristic indexer_topk_decode across realistic - logit distributions, MTP correlation structures, pre_idx accuracy levels, - GVR-supported K values, and supported logit dtypes. - """ - torch.manual_seed(24) - torch.cuda.manual_seed(24) - - 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) - # Clamp so that every base row has valid_len >= 1 (i.e., seq_len >= next_n). - # Without this, seq_len < next_n produces non-positive row_ends for offset 0. - seq_lens = seq_lens.clamp(min=next_n) - row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 - - # 1. Sample logits from the target distribution - logits = create_distributed_logits(dist_cfg, row_starts, row_ends, dtype, seed=42) - - # 2. Apply MTP correlation: consecutive rows share their tail logits - if next_n > 1: - logits = apply_mtp_structure(logits, batch_size, next_n, row_ends) - - # 3. Build heuristic pre-prediction indices - pre_idx = generate_pre_idx( - logits, - row_ends, - batch_size, - next_n, - index_topk, - success_ratio=success_ratio, - seed=7, - ) - - # 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, - radix_aux_indices=radix_aux_indices, - radix_aux_logits=radix_aux_logits, - ) - torch.cuda.synchronize() - - # 5. Reference: exact torch.topk masked to valid range - max_row_len = int(row_ends.max().item()) - 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 - torch_indices = torch_indices.masked_fill(~(mask_lo & mask_hi), -1) - - # GVR Top-K is an exact algorithm: with same-dtype `logits.topk` as the - # reference, the sorted output values must be bit-identical (bf16 -> fp32 - # promotion inside the kernel is lossless and order-preserving, so the - # K-th cutoff is identical in both comparison spaces). Any value gap is - # a real kernel bug, not algorithmic noise — keep the default 1e-5 - # tolerance and let CI surface regressions. - assert compare_top_k_results( - logits, - indices, - torch_indices, - row_starts, - row_ends, - index_topk, - ), ( - f"heuristic indexer_topk_decode mismatch: dist={dist_cfg['dist']}, " - f"mean={dist_cfg['mean']}, std={dist_cfg['std']}, " - f"next_n={next_n}, success_ratio={success_ratio}, dtype={dtype}" - ) - - -# ============================================================================ -# 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") - row_indices = torch.arange(num_gen_tokens, device="cuda") // next_n - next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n - - # Uncompressed seq_lens are what the kernel receives in `seq_lens`. - # Clamp so that compressed_actual_kv_len > kSeqSmall 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. - # kSeqSmall is uniform 4096 across K (matches indexerTopK.cu - # kSeqSmallDefaultForK). - ksmall = 4096 - min_uncompressed = (ksmall + 1) * compress_ratio + next_n - if min_uncompressed >= num_tokens: - pytest.skip( - f"num_tokens={num_tokens} too small to clamp into the GVR window for " - f"K={index_topk} (needs uncompressed > {min_uncompressed})" - ) - 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 = 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, - 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) - torch_indices = torch_indices.masked_fill(~mask, -1) - - assert compare_top_k_results( - logits, indices, torch_indices, row_starts, row_ends, index_topk - ), ( - 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]) -# num_tokens=32768 added so that all K ∈ {512, 1024, 2048} hit the uniform -# kSeqSmall=4096 boundary at compress_ratio=4 (helper's clamp floor = -# (4096+1)*4+next_n ≈ 16389 < 32768, so no skip is triggered for any K). -@pytest.mark.parametrize("num_tokens", [32768, 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 - ) - - # ============================================================================ # CuTE DSL Prefill Top-K Tests # ============================================================================ @@ -2350,100 +1584,3 @@ def test_prefill_overflow_policy_overflow( dtype, row_start_offset=row_start_offset, ) - - -# ============================================================================ -# GVR Phase-3 threshold-repair regressions: hints that defeat the threshold -# search (undershoot / degenerate hint / tie plateau wider than kC) used to -# produce a silently wrong top-K (-1 pads or row[0:K]). -# ============================================================================ - - -def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): - """Run indexer_topk_decode (cr=4, BS=1) and assert a tie-aware exact top-K.""" - n = logits_row.shape[-1] - dtype = logits_row.dtype - logits = logits_row.view(1, n).contiguous() - pre_idx = pre_idx_row.view(1, index_topk).to(torch.int32).contiguous() - seq_lens = torch.full((1,), n * 4, dtype=torch.int32, device="cuda") - # -1 sentinel so unwritten slots trip the assertions below. - indices = torch.full((1, index_topk), -1, dtype=torch.int32, device="cuda") - scratch = torch.empty(index_topk, dtype=dtype, device="cuda") - aux_indices, aux_logits = _build_radix_aux_buffers(1, index_topk) - torch.ops.trtllm.indexer_topk_decode( - logits, - seq_lens, - indices, - 1, - index_topk, - pre_idx, - scratch, - compress_ratio=4, - radix_aux_indices=aux_indices, - radix_aux_logits=aux_logits, - ) - torch.cuda.synchronize() - - assert int((indices < 0).sum()) == 0, ( - f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" - ) - # Distinctness: a duplicate+omission pair on a tie plateau would leave - # the sorted value multiset below unchanged. - n_unique = int(torch.unique(indices[0]).numel()) - assert n_unique == index_topk, ( - f"{tag}: only {n_unique} of {index_topk} output indices are distinct" - ) - flat = logits[0].float() - got = flat[indices[0].long()].sort().values - ref = flat.topk(index_topk).values.sort().values - assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk" - - -@skip_pre_blackwell -@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"] -) -@pytest.mark.parametrize("hint", ["bottom_k", "uniform_max", "random"]) -def test_indexer_topk_decode_gvr_hostile_hint(index_topk, num_tokens, dtype, hint): - """A hint that points away from the top-K must not change the result. - - ``uniform_max`` (every slot = argmax) additionally collapses Phase 1's - min/max bracket to a point, which used to short-circuit the kernel into - emitting row[0:K]. - """ - torch.manual_seed(1234) - logits = torch.randn(num_tokens, dtype=torch.float32, device="cuda").to(dtype) - flat = logits.float() - if hint == "bottom_k": - pre = flat.topk(index_topk, largest=False).indices - elif hint == "uniform_max": - pre = flat.argmax().repeat(index_topk) - else: - pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") - _gvr_decode_exact_check(logits, pre, index_topk, f"hint={hint}") - - -@skip_pre_blackwell -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -@pytest.mark.parametrize("n_tie", [6000, 20000, 100000]) -@pytest.mark.parametrize( - "dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"] -) -def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie, dtype): - """More ties at the K-th value than the candidate buffer can hold: no - threshold lands in [K, kC], so the repair must emit the strictly-greater - set plus arbitrary ties. bf16/fp16 cover the reduced-precision driver's - separate direct-emit block.""" - torch.manual_seed(1234) - num_tokens = 131072 - n_above = index_topk // 2 - logits = torch.full((num_tokens,), -1.0, dtype=torch.float32, device="cuda") - logits[:n_above] = torch.linspace(2.0, 3.0, n_above, device="cuda") - logits[n_above : n_above + n_tie] = 1.0 - # Plateau (1.0) and floor (-1.0) are exact in every dtype; casting can - # only merge strictly-greater values with each other, which is tolerated. - logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous().to(dtype) - pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") - _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}")