Skip to content

[None][perf] DSv4 indexer Top-K: align kFTarget=kK, uniform kSeqSmall=4096, diagnostic env knobs #14413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lfr-0531 merged 1 commit into from
May 25, 2026
Merged

[None][perf] DSv4 indexer Top-K: align kFTarget=kK, uniform kSeqSmall=4096, diagnostic env knobs #14413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
28 changes: 22 additions & 6 deletions cpp/tensorrt_llm/kernels/heuristic_topk.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,15 +224,27 @@ struct GvrParams; // primary undefined → compile-time error for bad combos
template <>
struct GvrParams<float, 512>
{
static constexpr int kFTarget = 384;
// 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<float, 1024>
{
static constexpr int kFTarget = 2560;
// kFTarget = kK (see GvrParams<float, 512> 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;
};
Expand All @@ -251,15 +263,17 @@ struct GvrParams<float, 2048>
template <>
struct GvrParams<__nv_bfloat16, 512>
{
static constexpr int kFTarget = 384;
// kFTarget aligned to kK — see GvrParams<float, 512> rationale.
static constexpr int kFTarget = 512;
static constexpr int kC = 5120;
static constexpr int kNumBins = 512;
};

template <>
struct GvrParams<__nv_bfloat16, 1024>
{
static constexpr int kFTarget = 2560;
// kFTarget = kK — see GvrParams<float, 1024> rationale.
static constexpr int kFTarget = 1024;
static constexpr int kC = 5120;
static constexpr int kNumBins = 512;
};
Expand All @@ -275,15 +289,17 @@ struct GvrParams<__nv_bfloat16, 2048>
template <>
struct GvrParams<__half, 512>
{
static constexpr int kFTarget = 384;
// kFTarget aligned to kK — see GvrParams<float, 512> rationale.
static constexpr int kFTarget = 512;
static constexpr int kC = 5120;
static constexpr int kNumBins = 512;
};

template <>
struct GvrParams<__half, 1024>
{
static constexpr int kFTarget = 2560;
// kFTarget = kK — see GvrParams<float, 1024> rationale.
static constexpr int kFTarget = 1024;
static constexpr int kC = 5120;
static constexpr int kNumBins = 1024;
};
Expand Down
90 changes: 72 additions & 18 deletions cpp/tensorrt_llm/kernels/indexerTopK.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,29 +732,71 @@ struct SchemeXBounds
int kSeqSmall;
};

inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem)
// 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)
{
static std::once_flag sOnce;
static int sSm = 0;
static int sL2 = 0;
static int sNMin = 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);
constexpr int kSeqSmallDefault = 12288;
char const* env = std::getenv("TRTLLM_HEURISTIC_NMIN");
if (env != nullptr)
{
int const v = std::atoi(env);
sNMin = (v >= 1024 && v <= 200000) ? v : kSeqSmallDefault;
sNMinEnv = (v >= 1024 && v <= 200000) ? v : 0;
}
else
char const* env_bsmax = std::getenv("TRTLLM_HEURISTIC_BSMAX");
if (env_bsmax != nullptr)
{
sNMin = kSeqSmallDefault;
int const v = std::atoi(env_bsmax);
sBsMax = (v >= 1 && v <= 65536) ? v : 0;
}
});

Expand All @@ -766,7 +808,14 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem)
? static_cast<int>(static_cast<int64_t>(sL2) * 9 / 10 / (static_cast<int64_t>(numColumns) * bytesPerElem))
: b.kBsWave;
b.kBsLarge = std::min(b.kBsWave, b.kBsL2 > 0 ? b.kBsL2 : b.kBsWave);
b.kSeqSmall = sNMin;
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;
}

Expand All @@ -789,7 +838,10 @@ 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, …).
auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4);
// 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 maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock);
Expand Down Expand Up @@ -887,15 +939,16 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic
// queries hardware attrs). At N≈70K both bounds produce ~426, so the
// L2 axis is a no-op there; for larger N it auto-tightens the threshold.
//
// Small-N lower bound `kSeqSmall` (default 12288) lets the Heuristic
// axis take over wherever the original Radix-radix branch would have
// triggered. Random-data benchmarks suggest the crossover is 16384,
// but workloads with strongly preIdx-correlated logits make P1 stats
// accurate and P2 converge in 1-2 iterations, shifting the real
// crossover into the [12288, 16384] band. Configurable via
// TRTLLM_HEURISTIC_NMIN env (>=1024).
// 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);
auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, topK);
int const kBsWave = bounds.kBsWave;
int const kBsL2 = bounds.kBsL2;
int const kBsLarge = bounds.kBsLarge;
Expand Down Expand Up @@ -1035,7 +1088,8 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int*
int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold;

// bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32.
auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast<int>(sizeof(InputT)));
// K-aware kSeqSmall — see fp32 dispatcher for rationale.
auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast<int>(sizeof(InputT)), topK);
int const kBsLarge = bounds.kBsLarge;
int const kSeqSmall = bounds.kSeqSmall;

Expand Down Expand Up @@ -1152,7 +1206,7 @@ bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytes
{
return false;
}
auto const bounds = getSchemeXBounds(numColumns, bytesPerElem);
auto const bounds = getSchemeXBounds(numColumns, bytesPerElem, topK);
return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge;
}

Expand Down
29 changes: 22 additions & 7 deletions tests/unittest/_torch/thop/parallel/test_indexer_topk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1456,7 +1456,11 @@ def run_fn(logits, seq_lens):
@pytest.mark.parametrize("batch_size", [1, 64, 128])
@pytest.mark.parametrize("next_n", [1, 2, 3])
@pytest.mark.parametrize("index_topk", [512, 1024, 2048])
@pytest.mark.parametrize("num_tokens", [8192, 16384])
# 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],
Expand Down Expand Up @@ -1593,11 +1597,19 @@ def _run_indexer_topk_decode_v4_gvr_check(
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 (= 12288) for every
# row; the kernel will divide actual_kv_len by compress_ratio internally,
# so a floor of (kSeqSmall + 1) * compress_ratio + next_n on the
# uncompressed seq_len guarantees compressed N stays in the GVR window.
min_uncompressed = (12288 + 1) * compress_ratio + next_n
# 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)

Expand Down Expand Up @@ -1682,7 +1694,10 @@ def _run_indexer_topk_decode_v4_gvr_check(
@pytest.mark.parametrize("batch_size", [1, 64])
@pytest.mark.parametrize("next_n", [1, 2, 3])
@pytest.mark.parametrize("index_topk", [512, 1024, 2048])
@pytest.mark.parametrize("num_tokens", [65536, 131072])
# 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],
Expand Down
Loading