Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 5 additions & 27 deletions cpp/tensorrt_llm/kernels/IndexerTopK.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
285 changes: 0 additions & 285 deletions cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu

This file was deleted.

Loading