From a2d0a71b342e6e43776c2a7f4be10b499cbfac51 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 21 May 2026 12:44:01 +0000 Subject: [PATCH] [None][perf] DSv4 indexer Top-K: align kFTarget=kK, uniform kSeqSmall=4096, diagnostic env knobs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three correlated tuning changes for the V4 indexer Top-K Heuristic dispatch on Blackwell GPUs. All three exercise the same code path (invokeIndexerTopKDecode / Dtype) and share the same B200/B300 swe-bench synth benchmark anchor, so they ship together. 1) heuristic_topk.cuh: align kFTarget to kK for K=512/1024 (all dtypes) GvrParams<{fp32,bf16,fp16}, 512>: kFTarget = 384 -> 512 (= kK) GvrParams<{fp32,bf16,fp16}, 1024>: kFTarget = 2560 -> 1024 (= kK) GvrParams<*, 2048>: unchanged (kFTarget = 3072 fp32 / 4096 bf16,fp16). The secant solver's soft-target is now aligned with the band's lower edge. Eliminates the upper-clamp saturation observed on tight-sigma high-A2 layers (K=512) and corrects a mistune for K=1024 that was measured against M=512 instead of production M=K=1024. Measured (B200, swe-bench synth, mean over 3 cfgs, zero cap-hits, zero per-layer regression): K=512 P2-iter sum: 2.19x / 1.77x / 1.51x at N=32K/64K/100K K=1024 P2-iter sum: 1.17x at N=32K (M=K=1024) 2) indexerTopK.cu: uniform kSeqSmall = 4096 across all K kSeqSmallDefaultForK now returns 4096 for every topK (previously 4096 for K=512/1024 and 12288 for K=2048). This 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 across K: below 4096 the Radix path resolves to single-CTA insertion-sort and GVR is not attempted; at or above 4096 GVR may be considered. Measured (B300, V3.2-Q19c protocol, May 2026): N=4K cells across K ∈ {512, 1024, 2048} all win — GVR R/H bf16 = 3.07x (K=512) / 2.57x (K=1024) / 1.34x (K=2048). N=2K cells across the same 9 (K x dtype) combos all show GVR R/H < 1 (0.55x – 0.84x), justifying 4K as the floor. All four call sites (computeIndexerTopKDecodeBlocksPerRow, the fp32 dispatcher, the bf16/fp16 dispatcher, canIndexerTopKDecodeUseGvr) retain their topK arg even though the kSeqSmall lookup itself is now K-agnostic; kept for API stability and because kBsWave/kBsL2/kBsLarge are still K-aware via bytesPerElem. 3) indexerTopK.cu: diagnostic / tuning env knobs (off by default) TRTLLM_HEURISTIC_NMIN (existing): now becomes an escape-hatch that wins over the uniform default when set; previously it just shadowed a single global default. TRTLLM_HEURISTIC_BSMAX (new): overrides kBsLarge past the hardware-derived min(kBsWave, kBsL2). Higher risk than NMIN — it bypasses L2/occupancy safety bounds, so set only for indexer microbenchmarks needing a BS-scaling comparison against the Radix path on identical inputs. The hardware-derived defaults are expected to be optimal for production; both env knobs are documented in-source as diagnostic/tuning escape hatches. Tests (test_indexer_topk_decode_dist + _dist_v4_cr4): - V3.2 cr=1: num_tokens=[8192,16384] -> [4096,8192,16384] to cover the new uniform K-agnostic kSeqSmall=4096 boundary. - V4 cr=4: num_tokens=[65536,131072] -> [32768,65536,131072]; helper's hardcoded min_uncompressed=12288 floor replaced with the uniform ksmall=4096 floor + pytest.skip guard for combos where the clamp floor exceeds num_tokens. With ksmall=4096 the floor at cr=4 is ~16389 uncompressed, so num_tokens=32768 admits all K ∈ {512, 1024, 2048} without skipping. Made-with-Claude-Code: opus-4.7-1m Co-Authored-By: Claude Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/heuristic_topk.cuh | 28 ++++-- cpp/tensorrt_llm/kernels/indexerTopK.cu | 90 +++++++++++++++---- .../_torch/thop/parallel/test_indexer_topk.py | 29 ++++-- 3 files changed, 116 insertions(+), 31 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh index ed1364096e04..f21a4d7cf2c6 100644 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh @@ -224,7 +224,12 @@ struct GvrParams; // primary undefined → compile-time error for bad combos template <> struct GvrParams { - 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; }; @@ -232,7 +237,14 @@ struct GvrParams template <> struct GvrParams { - static constexpr int kFTarget = 2560; + // 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; }; @@ -251,7 +263,8 @@ struct GvrParams template <> struct GvrParams<__nv_bfloat16, 512> { - static constexpr int kFTarget = 384; + // kFTarget aligned to kK — see GvrParams rationale. + static constexpr int kFTarget = 512; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -259,7 +272,8 @@ struct GvrParams<__nv_bfloat16, 512> template <> struct GvrParams<__nv_bfloat16, 1024> { - static constexpr int kFTarget = 2560; + // kFTarget = kK — see GvrParams rationale. + static constexpr int kFTarget = 1024; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -275,7 +289,8 @@ struct GvrParams<__nv_bfloat16, 2048> template <> struct GvrParams<__half, 512> { - static constexpr int kFTarget = 384; + // kFTarget aligned to kK — see GvrParams rationale. + static constexpr int kFTarget = 512; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -283,7 +298,8 @@ struct GvrParams<__half, 512> template <> struct GvrParams<__half, 1024> { - static constexpr int kFTarget = 2560; + // kFTarget = kK — see GvrParams rationale. + static constexpr int kFTarget = 1024; static constexpr int kC = 5120; static constexpr int kNumBins = 1024; }; diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index b35531a480a2..d63f203a048c 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -732,12 +732,53 @@ 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, []() { @@ -745,16 +786,17 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) 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; } }); @@ -766,7 +808,14 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) ? 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); - 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; } @@ -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); @@ -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; @@ -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(sizeof(InputT))); + // 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; @@ -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; } diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 58a38a399e54..e731a1711f80 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -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], @@ -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) @@ -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],