[None][feat] DSv4: enable GVR Heuristic Top-K for compress_ratio=4#14219
Conversation
|
/bot run --disable-fail-fast |
Widens the canUseHeuristic gate in indexerTopK.cu to allow
compressRatio == 4 in addition to 1, wires
sparse_attention_config.enable_heuristic_topk through model_config.py,
and adds a parameterized cr=4 unit test
(test_indexer_topk_decode_dist_v4_cr4).
Refactor:
- heuristicTopKDecode.{cu,h}: take compressRatio as argument; force
preIdxOffset = 0 when compressRatio != 1 (V4 indexer operates in
compressed-token-index space where new entries are appended at the
end, so prev-step indices in [0, c_prev-1] remain valid as-is in
[0, c_curr-1] — no +1/+2/+3 shift needed).
- indexerTopK.cu: bool compressRatioOk = (compressRatio == 1 ||
compressRatio == 4).
- model_config.py: pass-through sparse_attention_config fields
(use_cute_dsl_topk, q_split_threshold, indexer_rope_interleave,
enable_heuristic_topk, indexer_k_dtype, …) so user-set values
survive the V4 sparse_attention_config rebuild. Previously the V4
rebuild silently dropped these back to subclass defaults, so a
user setting enable_heuristic_topk: true in
--extra_llm_api_options had no effect on the V4 model.
- test_indexer_topk.py: dist-parameterized correctness test for V4
(cr=4) covering MTP windows {1,2,3}, all GVR-supported K
{512,1024,2048}, and fp32/bf16/fp16 dtypes; plus
apply_mtp_structure_compressed and generate_pre_idx_v4 helpers
that honor the cr=4 compressed-index space semantics.
Perf evidence (8x B300 SXM6 AC, indexer kernel-isolated microbench,
numColumns=16642 = §4 production [Scheme X] gate hit, K and pre_idx
hit-rate aligned with V4 HF config & test convention):
Flash (index_topk=512, success_ratio=0.6):
bf16: R/H 1.55-1.67x over nR ∈ {4,8,12,16}
fp32: R/H 2.01-2.08x
Pro (index_topk=1024, success_ratio=0.6):
bf16: R/H 1.40-1.57x
fp32: R/H 2.02-2.17x
Attribution (sr=0 pessimistic baseline; sr = prev-step pre_idx hit
rate, 0.6 matches generate_pre_idx_v4 test convention):
K axis (K=2048 → 512/1024 @ sr=0): +0.27x bf16 / +0.42x fp32 (dominant)
sr axis (sr=0 → 0.6 @ V4 K): +0.04 to +0.17x (bf16 K=512 strongest)
NCU SOL/occupancy at nR=16 bf16 (legacy K=2048 capture; SOL/occupancy
conclusions apply equally at V4 K): both kernels latency-bound
(Compute SOL 2.7-3.8%, Achieved Occupancy ~25%, 0.05 full waves on
148 SMs). Heuristic wins by doing less work per row (pre_idx hints
skip the Radix histogram), not by better HW utilization.
End-to-end trtllm-bench (Flash, ISL=65536 BS=4 MTP=3, autotuner-off
greedy-seeded "clean" conditions): Heuristic +6.4% req/s vs Radix.
The earlier autotuner-on result (-4.4%) was fully attributed to
autotuner tactic interaction + sampling-induced DAR drift.
Accuracy: Flash GSM8K 96.85 (matches prior baseline; Heuristic is an
exact algorithm, no precision loss).
Full report and reproduction scripts:
perf_logs/long_isl/REPORT.md (NCU report, V4-faithful microbench JSONs
under microbench_v4_K/, YAML configs, bench scripts all archived).
Made-with: Claude Code (Opus 4.7, 1M context)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
9a30058 to
4a84ea5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48757 [ run ] triggered by Bot. Commit: |
|
PR_Github #48759 [ run ] triggered by Bot. Commit: |
|
PR_Github #48757 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #48767 [ run ] triggered by Bot. Commit: |
|
PR_Github #48759 [ run ] completed with state |
|
PR_Github #48767 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48776 [ run ] triggered by Bot. Commit: |
|
PR_Github #48776 [ run ] completed with state |
Adds a one-shot write-back at the end of the prefill block that seeds
`metadata.heuristic_prev_topk` with each finishing-prefill sequence's
last-context-token top-K, before the first decode step of that sequence
reads it as `preIdx`. Mirrors the design intent of the original GVR /
DeepSeek-V3.2 reference: the previous step's top-K is the warm-start
for the next step's secant, and there is no good reason for "previous"
to skip the prefill→decode boundary specifically.
Without this, decode step 0 of every sequence runs from the all-zero
init left behind by `heuristic_prev_topk.zero_()`, which after the
kernel's +1 temporal offset resolves to all-index-1 — a benign but
uninformative seed that forces full P2 secant bisection from scratch.
Captured real-loop data on DSv4-Flash (1 prompt × BS=1 × K=512 ×
TP=8 EP=8 MTP=0) shows the prefill last-row top-K shares
~60-75 % of indices with the eventual decode step-0 top-K per layer
(layers 20 / 22 / 40 / 42 measured); this is comparable to the
steady-state preIdx ∩ topk hit-rate on subsequent decode steps
(0.61-0.84 mean per layer), so the bridge gives step 0 a warm start
indistinguishable in quality from a "normal" mid-decode step.
Slot convention mirrors the existing decode write-back at the bottom
of the decode block: newly-promoted gens append after currently-active
gens, i.e., `[num_generations : num_generations + num_contexts]`.
Gated on:
* `self._enable_heuristic_topk` (no-op if GVR disabled)
* `has_prefill` (no-op for pure-decode iters)
* `not metadata.skip_indexer_for_ctx_reqs` (no-op when the indexer
is being short-circuited for ctx requests)
Numerical safety: same dtype, same buffer storage, same compressed-
index space as the existing decode write-back. No new collectives, no
new graph nodes for CUDA Graph compatibility.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Made-with: Claude Code
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #48887 [ run ] triggered by Bot. Commit: |
Apply the exact reformatting that pre-commit/yapf produced on b6291ef in CI, so the GitHub Actions Pre-commit Check passes. Made-with: Claude Code (Opus 4.7) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #48907 [ run ] triggered by Bot. Commit: |
|
PR_Github #48887 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #48933 [ run ] triggered by Bot. Commit: |
|
PR_Github #48907 [ run ] completed with state |
|
PR_Github #48933 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49030 [ run ] triggered by Bot. Commit: |
|
PR_Github #49030 [ run ] completed with state |
Mirrors heuristicTopKDecode.cu PR NVIDIA#14219 cr-aware branch in the DSL GVR Top-K kernel. compress_ratio=1 (default) preserves DSv3.2 behavior exactly; compress_ratio=4 enables the DSv4 (overlap-compressor) indexer path: * pre_idx_offset = 0 (vs (row % next_n) + 1 for cr=1) — in compressed- index space, new entries append at the end so prev-step indices remain valid as-is. * N = actual_kv_len / cr — logits/preIdx live in compressed-token- index space when cr > 1. GvrParams TABLE is also keyed by (dtype, K, cr) so V3.2 and V4 use their respectively tuned kFTarget values: cr=1 (V3.2): kFTarget = 384 (K=512) / 2560 (K=1024), pre-NVIDIA#14413. cr=4 (V4): kFTarget = kK = 512 (K=512) / 1024 (K=1024), PR NVIDIA#14413. K=2048: identical across cr (V4 doesn't natively use K=2048). Cache key includes compress_ratio so different cr settings compile separate kernels. assert restricts compress_ratio in {1, 4}. Verified: 1152/1152 pytest configs pass on cr=1 default path. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
…VIDIA#14219) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit bef645d) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Summary
canUseHeuristicgate incpp/tensorrt_llm/kernels/indexerTopK.cufromcompressRatio == 1tocompressRatio ∈ {1, 4}, enabling GVR Heuristic Top-K on the DSv4 indexer (which operates in compressed-token-index space).compressRatiothroughlaunchHeuristicTopKDecode(3 dtype overloads). Kernel forcespreIdxOffset = 0whencr != 1(compressed entries are appended at the end, so prev-step indices in[0, c_prev-1]remain valid in[0, c_curr-1]), and dividesseqLensbycompressRatiointernally to derive per-row compressedN.DeepSeekSparseAttentionConfigsibling fields intensorrt_llm/_torch/model_config.py:681-731so user-setenable_heuristic_topk: true(and friends) survive the V4sparse_attention_configrebuild. Previously the V4 rebuild silently dropped these back to subclass defaults — a user settingenable_heuristic_topk: truein--extra_llm_api_optionshad no effect on V4.test_indexer_topk_decode_dist_v4_cr4intests/unittest/_torch/thop/parallel/test_indexer_topk.py(cr=4 × MTP {1,2,3} × K {512,1024,2048} × {fp32,bf16,fp16} × batch_size {1,64} × num_tokens {65536,131072}) plusapply_mtp_structure_compressed/generate_pre_idx_v4helpers.See the commit message for full kernel-isolated perf evidence and NCU SOL/occupancy analysis.
Perf evidence (8× B300 SXM6 AC)
Indexer kernel-isolated microbench at production shape (
numColumns=16642,sr=0.6matchinggenerate_pre_idx_v4):End-to-end trtllm-bench (Flash ISL=65536 BS=4 MTP=3, autotuner-off greedy, clean conditions): +6.4% req/s / −4.8% TPOT / −5.4% TTFT vs Radix.
E2E accuracy validation (8× B300)
Canonical accuracy judge: 4 independent GVR-ON GSM8K runs (strict-match, multi-token output)
All Δ ≤ 0.5 pp (within 1× stderr). Pro K=1024 + CUDA Graph capture + sparse-MLA path validated with no OOB regression.
Brittle-benchmark coverage (informational)
Pro MMLU 5-shot drift is NOT caused by GVR. The published baseline (67.03) was measured with
autotuner-ON; the GVR-ON matrix flipped TWO variables simultaneously (GVR-ON+autotuner-OFF). A 2×2 control withGVR-OFF + autotuner-OFFyielded 62.94 (Δ −4.09 pp), proving autotuner-OFF alone causes the dominant drift; the GVR-ON delta vs the autotuner-OFF baseline is actually +2.56 pp (recovery). MMLU 5-shot chat-mode forces 1-token outputs (14 017/14 042 Pro outputs are a single A/B/C/D letter), making it pp-level sensitive to floating-point reduction order changes from any config flip — unsuitable for sparse-attention feature regression detection. GSM8K strict-match (rows above) is the canonical judge.How to enable GVR — V3.2 vs V4
The user-facing flag (
enable_heuristic_topk: true) is identical in both, but the dispatch chain differs in two ways under the hood:algorithmvalue (Pydantic discriminator atllm_args.py:318, 410)dsa→DeepSeekSparseAttentionConfigdeepseek_v4→DeepSeekV4SparseAttentionConfig(subclass)enable_heuristic_topkfield originllm_args.py:350-358)index_topkdefaultcompressRatioin indexernumColumns = ISL / 4)canUseHeuristic(indexerTopK.cu:910, 1045)compressRatio == 1compressRatio ∈ {1, 4}(PR change)model_config.pypass-through of user-setenable_heuristic_topkmodel_config.py:743-779)update_sparse_attention_indexer_config(...)(model_config.py:681-731) — added by this PR; previously V4 rebuild silently dropped the user override back to subclass defaultFalseinvokeIndexerTopKPrefill,indexerTopK.cu:1126-1146)YAML — V3.2 (reference, unchanged)
YAML — V4 (new; requires this PR)
Without this PR, the V4 YAML above silently has no effect (V4 rebuild path drops the user override back to subclass
False); with this PR, V4 honors it the same way V3.2 always has.Verification (both models)
Run with
TRTLLM_SCHEMEX_DEBUG=1set in the environment and grep stderr / log for[Scheme X]:If the dispatcher prints
-> Radix pathfor a workload you expect to use GVR, one of the gate conditions is unmet (typical causes:index_topk ∉ {512, 1024, 2048},numColumns < kSeqSmall = 16384, ornumRows ≥ kBsLarge). The gate falls back to Radix safely — never wrong output, just no GVR speedup.V3.2 backward compatibility
preIdxOffset = (rowIdx % next_n) + 1andN = actual_kv_len(sincecr==1).int compress_ratio=1as default in both C++m.def(cpp/tensorrt_llm/thop/IndexerTopKOp.cpp:199) and Python fake op (tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py:212).test_indexer_topk_decode_dist(no cr parameter) auto-usescr=1→ no test modification needed; cr=1 production callers indsa.py:102similarly unchanged.git merge-base origin/feat/deepseek_v4 HEAD == 37c2e05ecf— HEAD is a strict descendant of upstream (no rebase / sibling conflict).Test plan
pytest tests/unittest/_torch/thop/parallel/test_indexer_topk.py::test_indexer_topk_decode_dist_v4_cr4— 108 cases,@skip_pre_blackwelltest_indexer_topk_decode_diststill passes (cr=1 default)/bot run --disable-fail-fast— triggered 2026-05-17T13:51:24Z (initial run failed on stale clang-format + ruff-format; re-triggered after auto-fix amend)Follow-up (separate PR): add
test_heuristic_topktoTestDeepSeekV4Flashintests/integration/defs/accuracy/test_llm_api_pytorch.py, mirroring V3.2'sheuristic_topk_mtp1parametrization at line 3344,3348; register intests/integration/test_lists/test-db/l0_dgx_b300_ds.yml.🤖 Generated with Claude Code