[None][perf] Block-skip fast path for self-sampling GVR top-K decode (stacked on #17821) - #18122
Draft
longcheng-nv wants to merge 42 commits into
Draft
[None][perf] Block-skip fast path for self-sampling GVR top-K decode (stacked on #17821)#18122longcheng-nv wants to merge 42 commits into
longcheng-nv wants to merge 42 commits into
Conversation
…ndalone)
Add a self-sampling variant of the GVR (Guess-Verify-Refine) heuristic
top-K decode for the DSA indexer, translated to CuTeDSL from the
optimized CUDA line (fork branch GVR-selfsampling-CuTeDSL), as two
standalone modules under cute_dsl_kernels/blackwell/top_k:
- gvr_topk_decode_self_sampling.py: merged device module — four kernel
families (sampling-ladder main / register-resident reg / cluster clus /
cluster-register reg_clus), lazily JIT-compiled per constexpr tuple.
- gvr_topk_decode_self_sampling_host.py: host companion — pure-function
dispatch route(b, n, npad, k) (bit-exact transcription of the CUDA
host dispatch, cross-checked by a 1,159,168-case boundary+fuzz sweep
plus a 300k-case parity fuzz of this merged form), per-device
workspace slab, and run/run_ws DPS entries with the CUDA binding's
hardening battery.
Contract (documented in the host module): batch-uniform host-int
n_valid in compressed index space; fp32; K in {512, 1024, 2048};
64-element-multiple row stride. Exact (tie-interchangeable) top-K.
NOT wired into the decode path: the production engine reads per-request
seq_lens on-device with per-row MTP offsets (heuristicTopKDecode.cu);
adopting that per-row contract inside these kernels is follow-up work,
so this module must not substitute for the tiered path under continuous
batching, MTP, or CUDA-graph capture.
Evidence (B200): 886-cell x 11-BS real-decode-capture grid = 9,746
cases vs the production CUDA arm: 0 INEXACT, geomean ratio 0.974.
Unit tests: tie-aware exactness (signed-zero normalized, poisoned-pad
immunity) across gate-edge and envelope shapes, run_ws with caller
workspace, guard predicates, dispatch totality; sm_100-gated, picked up
by the existing unittest/_torch/thop/parallel sweep.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… raw indices, uniform across DSv3.2/Flash/Pro The contract docstring previously instructed callers to apply the cr==1 +1 temporal shift to pre_idx (mirroring heuristicTopKDecode.cu). Drop it: hints only steer the sampling ladder — exactness never depends on them — and on real V3.2 decode captures raw prev-step hints overlap the current top-K at 0.773 vs 0.536 when +1-shifted (15 cells x 14 consecutive step-pairs, gap widening with ISL). One offset-free hint convention now serves all three models; the kernels already consume pre_idx as-is, so this is a contract-documentation fix only. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ction pad convention)
When the row has no more than topK valid entries every valid position
is in the top-K: emit identity indices and pad the tail with -1,
mirroring heuristicTopKDecode.cu:72-84. Host-level torch-op branch for
the standalone module (the CUDA-graph-safe per-row rewrite will move it
in-kernel, where a per-row fallback is impossible inside a graph).
Closes the 'n <= topK unproven' gap from the integration audit.
Tests: 8 boundary shapes (n in {64..2048}, k in {512,1024,2048},
n < k / n == k-1 / n == k) x bs {1,4} with poisoned padding, plus
kernel-path regression just above the boundary (k512_n4099,
k2048_n4111) verified exact on B200.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…lt off) Optional `values` DPS output on run()/run_ws(), production parity for the heuristicTopKDecode values writeback. Default None = OFF, matching dsa.py, which allocates the values scratch only for the non-CuTeDSL path. The indices are exact, so a gather epilogue reproduces the in-kernel writeback bit-for-bit at zero cost when disabled; the constexpr in-kernel form rides the CUDA-graph per-row rewrite. Short path pads values with -FLT_MAX (production convention). Tests: kernel path (values == gathered top-K == torch.topk multiset), short path (head copies logits, -FLT_MAX tail), wide-buffer packed re-view and dtype guard, verified on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ests Production can hand the kernel degenerate hint buffers: dsa.py initializes heuristic_prev_topk with zero_() (all-zero cold start; the prefill->decode seeding covers the common path but zero-init corners remain), and nothing forbids duplicated hints. Exactness must never depend on hint quality (hint-robustness bug class of PR NVIDIA#17550). Adds all-zero / all-same / all-max / half-duplicated hint cases on (k512, n8192) and the k2048 gate-edge n131075 — all verified exact on B200 (also probed on n131072/k1024 x bs{1,4} pre-commit, 24 configs, zero failures). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…act entry (reference engine)
Adds run_varlen(logits, pre_idx, kv_lens, indices, next_n=,
compress_ratio=, values=None) — the heuristicTopKDecode contract:
per-request device kv_lens (TOTAL cache length, uncompressed token
space), per-row n = (kv_len - next_n + row%next_n + 1)/compressRatio
(the MTP window formula, cr 1 = DSv3.2 / 4 = DSv4 Flash+Pro),
request-level raw pre_idx shared by a request's next_n rows, per-row
n <= k short path.
REFERENCE engine: one documented host read of kv_lens (raises under
CUDA-graph capture), rows driven as b=1 launches through the
batch-uniform engine. This pins the varlen/MTP contract and its test
battery; the per-row in-kernel engine (device kv_lens reads, fixed-R
thin-slicing, n-band parameter table) replaces the loop next without
changing either.
Tests: heterogeneous lengths with per-row poisoned padding, cr in
{1,4} x next_n in {1,2,4} incl. in-request n variation and
compressed-boundary-crossing rows, mixed short rows, values output,
contract guards — all verified exact on B200.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…it (route_static/route_dynamic) Factors route(b, n, npad, k) into route_static — family, compile tuple, grid/cluster/block and the rt scalars that change only at discrete n-thresholds (freezable at CUDA-graph capture time, mirroring the in-tree runner's pick_tuning(graph_capture=...) pattern) — and route_dynamic — the n-continuous scalars a per-row kernel recomputes from its own row length (CMP and the reg smem footprint; the SMP/TGT/SS2/TGT2/Q sampling ladder for the streaming families). The device-side per-row engine will mirror exactly the route_dynamic formulas. Lossless by construction and by fuzz: recombining the halves reproduces route() bit-exactly on 163,755 (b, n, k) points (threshold windows, R-boundaries, prime-stride sweep, LCG random; full result-dict equality). route_bands() enumerates the static-constant n-intervals — the (b=8, k=1024, 262144) envelope collapses to 10 contiguous bands, so the eventual in-kernel band table is tiny. Host-side groundwork only; no kernel behavior change. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…gine (gvr_main port)
The gvr_main family gains a per-row varlen mode (production
heuristicTopKDecode contract): each CTA reads its row's kv_len from a
device kv_lens tensor and re-derives n plus the whole sampling ladder
(SMP/SS2/TGT/TGT2/Q) via an exact device transcription of the
route_dynamic() host formulas — integer round(sqrt(6n)) with an isqrt
fixup (bit-parity with the host double form), Int64 TGT products
(the CUDA host math is 64-bit), and constexpr next_n / cr_shift /
r_const so every division strength-reduces. All values are pure
functions of the row index, so the R split CTAs of a row stay
grid-uniform by construction (workspace handshake unchanged).
Design points:
- No runtime return in CuTe DSL (in-tree gvr_topk_decode.py precedent):
n <= k rows run the body as a zero-work pass (n=0, TGT=INT_MAX so no
rung accepts, Q=0 empties the split slices) and an epilogue emits the
production identity/-1-pad short path.
- The TSH-floor staging gate becomes per-row runtime (tsh_en && per-row
n4 <= 32768), exactly the CUDA original's grid-uniform runtime gate;
legacy compiles keep bit-identical behavior (tsh_run == 1).
- pre_idx is REQUEST-level in varlen mode (row // next_n mapping).
- run_varlen(engine="auto") launches the batch in ONE kernel; with
max_seq_len given (capture-stable engine constant) the call performs
no host reads. engine="reference" keeps the b=1 loop as the
differential oracle. Capture-time tuning comes from route_streaming()
(the streaming half of route(), 110,003-point fuzz agreement).
- Legacy batch-uniform ABI extended with dead trailing args (dummy
kv_lens + five zeros) — one kern body, no duplication.
Validated on B200: legacy regression 5/5 exact (extended ABI), engine
vs reference differential 9/9 mixed-batch configs row-for-row equal —
deep SPLIT (b=1, n=200k), 8-row cr=4 mix to n=225k, tsh band (b=16),
MTP next_n {2,4}, all-short batches, 200-row small_dense k=2048, and
the no-host-read max_seq_len path.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…est for the varlen engine Warm up, capture one run_varlen(engine=auto, max_seq_len=...) launch — no host reads, no JIT inside capture — then replay while kv_lens grows in place: a row crossing the n <= topK short-path boundary INSIDE the graph, a row walking the 131072 band edge, and a 200k deep row. Every replay verified tie-aware exact on B200 (8-replay standalone run all green). This closes the CUDA-graph-safety design goal for the gvr_main varlen port: geometry and tuple are frozen from capture-stable quantities (route_streaming at max_seq_len), all N-dependence is per-row device arithmetic. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…TRTLLM_GVR_SELF_SAMPLING=1)
Wires the varlen engine into the DSA indexer decode top-K dispatch as
the highest-priority branch, env-gated (TRTLLM_GVR_SELF_SAMPLING=1) and
contract-gated at init (cutlass DSL present, sm100+, index_topk in
{512,1024,2048}, compress_ratio in {1,4}) — covering DSv3.2 (K=2048,
cr=1), DSv4 Flash (K=512, cr=4) and DSv4 Pro (K=1024, cr=4). The call
reuses the exact buffers of the existing cute_dsl_gvr_topk_decode
branch (request-level heuristic_prev_topk, kv_lens_cuda_runtime,
topk_indices_buffer, indexer_max_seq_len as the capture-stable tuning
constant) — no new metadata plumbing. Lazy import behind the gate;
contract violations raise loudly (explicit experiment flag, not a
silent fallback). First warmup call pays the one-time DSL JIT.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ive the TGT*2 scan target The zero-work short pass used TGT = 0x7FFFFFFF, whose TGT*2 third scan target overflows Int32 to -2, flipping every tot0 >= TGT*2 gate on the all-zero histogram (benign downstream today — empty candidate sets emit nothing — but an unnecessary cliff). Use 2^30-1 so the doubled target stays positive; 'never accepts' semantics unchanged. Full differential battery (9/9 mixed-batch configs) + graph capture/replay (8 replays) re-verified green on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… of the varlen engine and seam Fixes from a 4-dimension adversarial review of the varlen stack: - CRITICAL seam unit bug: indexer_max_seq_len is COMPRESSED index space (metadata divides by compress_ratio) but run_varlen's max_seq_len is kv-token space — the seam shifted twice, freezing a 4x-too-small tuning envelope for cr=4. Seam now multiplies back (exact identity: n_env == indexer_max_seq_len). - Seam crash on legal configs: the DSL paged-MQA logits arena is a 256-aligned buffer column-sliced to max_seq_len — a NON-contiguous view the engine used to reject. The engine now accepts row-major views (inner stride 1) and widens them back to a compact [rows, row_stride] view over the same storage (as_strided, zero copy; the tail columns are never classified — per-row n gates all reads). A dispatch-site hardware-format gate (stride %4, 16B base) falls through to the existing branches for layouts the kernel cannot address (odd-npad DeepGEMM). - OOB-write guard (three reviewers converged): the engine never validated indices/values batch dims — a request-level-shaped buffer under MTP would be silently written past its end (grid comes from logits rows). Full B1-style battery now runs on the engine path (CUDA/dtype/2-D/contiguity/batch/width/alignment), kv_lens contiguity included. - Engine/reference convention alignment: wider-than-k buffers now follow the flat-packed contract identically on both engines; the reference clamps kv_len < next_n to the empty row (all -1) exactly like the kernel — padded/evicted graph slots are a legal input, and the differential oracle can now cover them. - Eager-mode compile churn: without max_seq_len the data-dependent envelope is quantized up to the next power of two (bounded plan set and _VARLEN_CACHE; a growing decode no longer recompiles at every R increment). - Multi-stream escape hatch: run_varlen(workspace=...) (run_ws parity) so concurrent streams do not share the SPLIT publish slab. - Legacy hot path: per-device cached dummy kv_lens (no per-call allocation for the dead ABI slot). New tests: b=16 SPLIT + per-row TSH runtime gate, b=200 BLK=512 non-split, zero-kv slot mixed with live MTP rows (both engines), wide-buffer flat-packed convention (both engines), num_rows / strided-kv_lens guards. Full battery re-verified on B200 (10/10 incl. a non-contiguous arena-view differential). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Release-Check (PR_Github #67203) failed on pre-commit: ruff-format reflow on 3 files, one ruff F841 (unused next_n unpack in the varlen differential test), one codespell hit (statics -> static fields). Formatting-only plus the two mechanical fixes; kernel exactness re-verified on GPU after the reflow: full unit file 65/65 passed (sm100, standalone overlay stack). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ng-only surface Public package exports now carry ONLY the production contract (selfsampling_topk_run_varlen: per-request device kv_lens, no batch-uniformity assumption — mirrors the single-op shape of the CUDA indexer_topk_decode integration). run/run_ws keep serving as the bench DPS contract and the reference-oracle plumbing (_run_impl is what engine="reference" walks row by row), but are no longer package-exported and carry TESTING/BENCH ONLY docstring warnings. Tests already import the host module directly — zero test churn. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…y, hint/domain tests Framework-integration audit hardening (all GPU-verified, 70/70 unit file): - dispatch gate now also requires fp32 logits (falls through loudly instead of feeding a non-fp32 tensor into the fp32-typed DSL engine; production DSA logits are always fp32 today, this is belt-and-braces for future paths) - logger.info_once on first engagement + logger.warning_once on first hardware-format fall-through: operators can tell which arm served without profiling - new tests: hints containing -1 (the production short-row pad tail that flows back through heuristic_prev_topk; engine's unsigned-compare guard verified exact) and route() domain up to 8192 rows (max_batch x next_n exceeds the bench grid's b<=1024 envelope) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…elper Mirrors warmup_heuristic_topk_decode / warmup_cute_dsl_radix_topk: one tiny real launch per requested num_rows compiles the varlen engine's envelope tuples so no live request pays the first-touch DSL JIT (measured: 9.1 s cold compile for two tuples on a fresh cache; idempotent re-call 0 ms; post-warmup first real call 0.1 ms). Exposed as a module-level helper — CUDA-graph capture warmup already compiles the captured batch sizes, and wiring an automatic init hook needs max_seq_len plumbing that is not available at Indexer.__init__ time (follow-up). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…e.warmup Completes the warmup story: DSAtrtllmAttentionMetadata.warmup_selfsampling_topk mirrors warmup_cute_dsl_radix_topk (same ModelEngine hook, which is where max_seq_len is actually available — Indexer.__init__ is not). Gated on the env flag + the same init-contract conditions; compiles the eager first-touch (num_rows=next_n) tuple via warmup_varlen so no live request pays the DSL JIT. Also: seam init comment now states the two-guard reality (format gate falls through with a one-time warning; in-engine contract violations raise), and run_varlen documents the inherited NaN-ordering limitation (finite inputs incl. +/-inf are tie-aware exact). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… next_n=3 coverage run_varlen now rejects non-fp32 logits with a clear contract error (the dispatch seam falls through before this; direct callers previously hit an opaque CuTe typing failure) — regression-tested. The varlen differential gains a next_n=3 (MTP2) case: no production config uses it today, but the per-row window formula is now verified to generalize (both engines exact), closing the gap noted in the PR's MTP-coverage section. GPU-verified: 3/3 targeted tests green. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…age + robustness) Adopted from the first review round: - warmup_varlen records its done-key only after every launch succeeds, so a failed or interrupted warmup retries instead of short-circuiting to an uncompiled engine; one max-geometry allocation is reused across row counts (prefix views — compile keys depend on shapes only). - The ModelEngine warmup hook now passes the configured CUDA-graph batch sizes: the varlen launcher is keyed by the exact row count, so capture no longer depends on a prior warmup-forward having compiled its key, and warmed geometries never pay first-touch JIT. Best-effort under OOM (warn + lazy JIT) since dispatch works without warmup outside capture. - route/route_streaming reject b < 1 with a contract error instead of ZeroDivisionError. - run/_run_impl: the launch try block no longer swallows routing/compile errors as 'launch failed', and the values epilogue matches run_varlen (clamp + mask instead of a context-poisoning gather assert). - Stale module docstring refreshed (run_varlen IS the wired production entry); warmup_varlen exported in __all__; public entry points annotated. - Test SM gate aligned with the production dispatch gate (>= 100; the module previously skipped everywhere except exactly sm_100). GPU-verified: full battery 73/73; smoke 8/8 (b<1 guards, multi-rows warmup + idempotency, CUDA-graph capture at a warmed row count exact, unwarmed capture raises loudly, injected-failure warmup retries with a real recompile). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ch actually looks up The varlen launcher key includes the logits row stride (npad), and the two sides derived it differently: the DSL paged-MQA arena rounds each row up to 256 elements (CuteDSLPagedMQALogitsRunner: compute_block_kv=128, SPLIT_KV= 256), while warmup_varlen synthesized its logits at 64-element rounding. Whenever the indexer max_seq_len is not a 256 multiple, warmup compiled a variant dispatch never looks up: with CUDA graphs the pre-capture warmup forwards masked this (init just wasted a few seconds on phantom keys), but in eager serving the first live decode paid the full first-touch DSL JIT — exactly the stall the helper exists to remove. Never a correctness issue: a key miss lazily JITs the true variant. - warmup_varlen gains row_stride; the caller passes the producer's actual row stride and the done-key includes it. - metadata.warmup_selfsampling_topk mirrors the active producer: DSL paged-MQA -> round up to 256; DeepGEMM -> exact width (non-float4 widths fall through at the dispatch format gate, so nothing to warm). A future drift only degrades warmup back to unused keys, never to wrong results. - warmup now also applies the same hardware gates as the dispatch flag (CUTLASS DSL available, SM100+), so setting the env var on an unsupported stack can no longer make engine warmup compile Blackwell kernels. - New tests: capture immediately after a row_stride warmup on a column-sliced arena view (fails loudly with the old 64-rounding), and kv_lens < next_n zero-window rows (graph-padding dummies) pad-emit -1 while cohabiting normal requests stay exact. GPU-verified: full battery 75/75 (73 prior + 2 new). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…d datacenter Blackwell
The previous review round aligned the test gate with the dispatch gate at
>= 100, which silently admits consumer Blackwell (sm_120/121): those parts
lack thread-block clusters, the kernels have never been validated there,
and CI stages like RTX5080/5090 and GB10 would start collecting the suite.
Tighten all three gates (dispatch flag, warmup hook, test module) to the
validated set {sm_100, sm_103} — B200/B300-class datacenter parts — which
was the reviewer's suggested alternative. Spot-checked on sm_100: module
still collected, 4/4 targeted tests green, hooks clean.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…e (fall through above 32 rows) At high in-flight row counts the varlen streaming engine stops splitting rows across CTAs (route_streaming keeps multi-CTA R only for b <= 32, R=2 up to b <= 74, then one CTA per row) and long rows lose ~5x per call to the in-tree per-row-split kernels — measured in DSv4 Pro 1M-ISL MTP7 EP16 serving: rows=304, 2.63 ms vs 0.49 ms per top-K call, ~4x TPOT end to end. Gate the DSA dispatch on a shared MAX_VARLEN_ROWS=32 admission envelope (info_once fall-through to the in-tree path), drop over-envelope row counts from warmup_varlen so engine init never compiles keys dispatch cannot admit, and document the envelope on run_varlen. Large batches keep taking the in-tree path until a throughput tier lands (see roadmap). Two new tests: the envelope constant matches route_streaming's multi-CTA region (rows=304 collapses to R=1), and warmup drops over-envelope row counts while still warming the admitted prefix. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… single-row width gate
Two residuals from the production audit. (1) Warmup previously covered only
{next_n} and the CUDA-graph batch sizes, so eager mixed prefill+decode
batches at other in-flight generation counts paid a multi-second first-touch
JIT on the serving path (4-13 reachable unwarmed r_const bands per model
with padded graph lists). Warm every admissible row count instead — next_n
multiples up to MAX_VARLEN_ROWS — which bounds init cost at <= 32/next_n
tiny launches over the same few engine compiles and leaves nothing to
compile at serving time inside the envelope. (2) The dispatch format gate
checked only stride(0), but single-row batches derive their row window from
shape[1] (arena last-row safety), so a non-float4 indexer_max_seq_len with
next_n==1 raised inside run_varlen instead of falling through; gate that
width explicitly.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…r small batches; dedup dead TSHG key slot The varlen launcher gated TSH-floor staging on num_rows > 15, stranding small batches in SPLIT-main without the staged floor. On real deep-layer captures this is a distribution-dependent ~6x tail (v4_pro_512k L46/L52, n4=32768, rows 1-8: 142-151 us -> 25 us with staging; tie-aware exact preserved, healthy layers and n4 > 32768 rows unchanged within noise). The kernel compiles the TSH machinery whenever SPLIT and gates it per row at runtime, so this is a runtime-scalar-only change: no new engines, no compile-key change, CUDA-graph safe. Also normalize the dead TSHG slot out of the varlen compile key (the ctor overrides it under SPLIT), so row counts differing only in that slot share one engine instead of compiling twice. Varlen/envelope test subset: 22/22. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…clustered register-resident family Lift the varlen production path's force-main restriction for its first specialist family: GvrRegClusKernel gains the same per-row varlen mode already shipped in GvrMainKernel — n is re-derived per row in-kernel from device kv_lens (production heuristicTopKDecode contract), hints map to the request row, and short rows (n <= k) emit identity + (-1) tail from rank 0 and skip the body entirely (a zero-work pass would reach the degenerate crossing-overflow emitter). All derived quantities are pure functions of the row, so the whole-body guard is cluster-uniform and the cluster barriers stay aligned. The varlen launcher admits the family exactly where the free route() picks it (route parity tier 1) — a pure function of the capture-stable cache key, so CUDA-graph replay safety is unchanged; its whole admission window (n4 <= 32768) fits capture-frozen envelopes. On real captures this recovers most of the streaming-only integration tax in the family's window: the deep-layer distribution tail (v4_pro_512k L46/L52, rows 8) goes 147 us (pre-TSH-fix) / 25 us (TSH fix) -> 20.7 us, against 16-18 us for the batch-uniform standalone. Full self-sampling UT suite green (79/79) incl. two new tests: heterogeneous kv_lens + MTP row windows vs the reference oracle, and cluster-family CUDA-graph capture/replay exactness. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… row count and envelope goes through the self-sampling engines Remove the rows admission envelope: with TRTLLM_GVR_SELF_SAMPLING=1 the self-sampling engines now serve the full production range (BS 1..1024+ x next_n, envelopes to 1M kv tokens) instead of falling through to the in-tree path above 32 rows — restoring the standalone design intent. Warmup switches from row-count filtering to band-aware enumeration: one representative row per distinct engine compile key (bounded time and memory for arbitrarily large CUDA-graph batch lists), then per-row-count launcher-cache population (pure host work) so capture at any requested geometry finds its key immediately — preserving the warm-the-key-dispatch- looks-up discipline the arena-stride capture test pins. Production-lessons checklist re-validated on this change: capture-stable pure-function dispatch (family from the launcher key, per-row n in-kernel from device kv_lens), loud-fail on uncompiled keys under capture, two-pass pre-capture warmup coverage, producer row-stride key mirroring, workspace slab grow-never/keep-alive, cluster-family graph capture, heterogeneous kv_lens + MTP row-window oracle tests. Full suite 79/79. Real-capture record (vs the in-tree production op, cold-L2, all exact): rows 64..1024 wins 2.9-4.7x at N=262144, 3.1-3.7x on the former deep-layer pathology band, 1.6-3.5x on DSv3.2 163k; short-N cells (N <= ~2k) trade 0.67-0.94x at microsecond scale — accepted by design: enabled means self-sampling, everywhere. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ndependence test at throughput scale Pin the varlen row-independence contract on the streaming main family: one 304-row batch mixing random per-request kv_lens with full-length, short (n <= k), zero-window and boundary rows under MTP row windows (next_n=4) — every row must match its own per-prefix torch.topk value multiset, short rows the in-kernel identity + (-1) tail, zero-window rows all -1. No engine may assume batch-uniform indexer-logits lengths. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…register-resident family (reg/regimg) Route-parity tier 2: the varlen launcher now admits the register-resident family (and its img-window flavor) exactly where the free route picks it, instead of forcing those bands through the streaming main kernel. This family owns the small/mid-N band across all row counts in the standalone dispatch, where main pays a 1.5x tax. Kernel side (GvrRegClusKernel discipline): per-row n re-derived in-kernel from device kv_lens (request windows via next_n/cr_shift, envelope clamp); short rows (n <= k) emit identity + (-1) tail in-kernel and skip the body -- k can exceed BLK on this family, so the emit is a strided loop, not the reg_clus single predicate. kv_lens joins the ABI as a dead dummy slot in batch-uniform mode (gvr_main precedent). CMP/QC/smem stay envelope-derived launch constants: in-kernel they are pure capacity clamps, a fast-path threshold and the launch smem size -- safe upper bounds for every per-row n <= envelope. Production lessons re-checked: capture-stable dispatch (family remains a pure function of the launcher cache key), warmup band enumeration gains the reg engine key, loud-fail on uncompiled keys under capture unchanged. Tests: route-parity + heterogeneous oracle (reg, k>BLK strided short-row emit, regimg window) and CUDA-graph capture/replay; suite 82/82. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…emaining clus/reg_clus reciprocals Convert the last three IEEE-div reciprocal sites (clus sample-hist scale, clus window scale, reg_clus bin-transform scale) to cute.arch.rcp_approx, completing the fix-4 CUDA --use_fast_math parity (bare MUFU.RCP instead of the div.rn rcp+Newton+CALL chain). All three feed classify bucketing only, which is scale-invariant for any SC > 0; inputs are clamped positive finite by the existing guards. UT 82/82; paired same-GPU A/B on clus/reg_clus cells is noise-level (<=1.3%), no regression (the fix-1 spill hazard that kept main's blk!=512 arm on div.rn does not reproduce in these families). The cute.math.div(..., fastmath=True) spelling (per DSL-team guidance) is not available on cutlass 4.5.0 (this PR's validated stack); adopt it when the tree moves to DSL >= 4.6. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…cluster-split family (clus) Route-parity tier 3, completing the family coverage: the varlen production entry now admits every family exactly where the free route picks it (main / reg_clus / reg / regimg / clus) -- no force-main fallback band remains. GvrClusKernel gains the per-row varlen mode (GvrMainKernel / GvrRegClusKernel discipline): n is re-derived per row from device kv_lens; the sampling-ladder scalars (SMP/TGT/Q/SS2/TGT2, dead launch slots in this mode) are re-derived per row by the route_dynamic clus mirror (exact-integer isqrt aim ladder, Int64 target products, QUAD 16-multiplier geometry). SCAP/CMP are pure functions of (rows, CS, k) -- never of n -- so the envelope launch values are the per-row values; smem extents unchanged. The 'big' occupancy flag is launch-computed in the jit wrapper (rows * CS <= 148). One deviation from route_dynamic, documented in-kernel: the QUAD schedule is computed for every non-short row instead of only n > SCAP -- the host only ever launches this family with n > SCAP, so the SMP == 0 no-sample path is untested; per-row n <= SCAP rows get a small valid schedule (sampling only steers the rung; exactness is schedule-invariant, and sample positions stay under the row's own n). Short rows (n <= k) emit identity + (-1) tail from cluster rank 0 and skip the whole body -- 'short' is a pure function of 'row', identical across all CS ranks, so every cluster barrier stays aligned. Verification: UT 84/84 (2 new: heterogeneous parity+oracle over CS=2 and CS=4 clusters with mid rows below the standalone admission floor, and a CUDA-graph capture/replay test). Real-capture clus zone (153 cells / 255 cases from the 886-grid census): all exact, integration tax (same-session CUDA-events pairing, cmp886_r3 protocol) 1.319 -> 1.137 gm (worst 1.216), now at the same intrinsic per-row level as the other families. Graph-replay protocol on clus shapes: tax 1.096-1.124. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…mirror chain The per-row re-derivation of the sampling-ladder scalars was the dominant varlen-vs-standalone kernel tax (l1p1 ncu: inst x1.58-1.73 on main shapes, x1.26-1.28 on clus -- all of it the DSL's runtime Int32/Int64 divides and isqrt fixup loops, redundantly issued by every thread). Two fixes, chosen per family by measured mirror share: - main: warp0 alone walks the ladder and publishes SMP/SS2/TGT/TGT2 through a 4-word smem block; the other warps spend the wait issuing register-free L2 hints for this CTA's own P3 slice, then everyone reads the scalars back after one CTA barrier. The chain itself swaps runtime divides for MUFU.RCP multiplies and collapses the isqrt fixup loops to single steps (f32 sqrt of an exactly-representable int is within 1). - clus: same cheap-chain spelling, kept all-thread in place. This family's mirror share is small, so a warp0+barrier hoist EXPOSES the chain's serial latency at ~1 CTA/SM (measured 1.14->1.31 tax regression) while the redundant form hides it across warps. Schedule quantities may drift +-1 vs the host double form; exactness is schedule-invariant (same argument as the clus port's always-sample deviation). A one-step window guard keeps SMP*SS2 inside the sampled row span. Q (chunk ownership) keeps its exact form in both families. Legacy (batch-uniform) codegen is untouched: every change sits under const_expr(varlen), including the smem slots. nsys verdict (interleaved eager sa/vl, cold-L2, 200 paired reps/cell, umbriel-b200-027): v32_128k r1 1.165->1.093, pro_1024k r8 1.184->1.103, pro_1024k r32 1.135->1.083, pro_512k r64 1.131->1.071; reg_clus control cells unchanged. Treated-cell tax gm 1.154->1.087. UT 84/84; exactness smoke (main x3 + clus x2 + short/long heterogeneous kv_lens) all exact under the tie-aware value-multiset oracle. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Adopts the TopK module refactor (NVIDIA#17622): the self-sampling engine is now TopKImplementation.CUTE_DSL_GVR_V2 with the hardware-format gate living in the module and falling through to the CUDA GVR body; the env opt-in (TRTLLM_GVR_SELF_SAMPLING=1) promotes the implementation at the indexer's selection site. The module receives max_seq_len in compressed index space, so the V2 branch multiplies back to kv-token space for run_varlen. Two latent branch bugs surfaced by the merge audit are fixed in the same resolution: the warmup gate read metadata.enable_heuristic_topk, which main renamed to enable_gvr_topk, and it referenced MAX_VARLEN_ROWS, which the full-range dispatch commit removed from the host module (replaced by a local eager warm-span constant). Both would have raised AttributeError on the in-tree warmup path under the env opt-in; the e2e evidence runs never hit them because the ss arm used the pre-refactor overlay seam. Validated on the merged tree: wheel build, TopK-V2 module smoke (exact through the module path + bf16 fall-through leg), kernel UT 84/84. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Main advanced two commits during the first merge's validation; the only overlap is positional (warmup_selfsampling_topk sits next to on_update_kv_lens, whose signature gained an annotation). Python-only delta upstream — no rebuild; TopK-V2 module smoke re-passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…the TopK module edit Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…module Drop the dead per-row Q re-derivation in the clus varlen prologue (the Q launch slot has no in-kernel consumer in this family; chunk ownership derives from n4/STEPC — ruff F841), give the two torch-facing debug entries local torch imports (the DSL module is deliberately torch-free at import time — ruff F821), and take ruff-format's layout on the unit test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…in the register family Pure layout (long-line wrapping in GvrTopkRegKernel, no code change); validated against the exact pushed tree state this time — the previous lint pass ran on a local tree whose pending changes masked these lines. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…TLASS DSL 4.6.1 cute.make_fragment was removed in nvidia-cutlass-dsl 4.6.1 — the version this repo pins — so the self-sampling engines raised AttributeError at first JIT under the shipping dependency set (undetected by CI: the unit tests only collect on SM100/103 stages, which do not run this suite). cute.make_rmem_tensor is the same-signature replacement, already the convention in the in-tree DSL kernels, and exists in 4.5.0 as well, so every environment keeps working. Mechanical rename, 23 sites. Validated under BOTH pins: UT 84/84 on 4.6.1 (isolated overlay) and on 4.5.0; exactness smoke (main x3 + clus x2 + heterogeneous kv_lens) all exact on both; nsys spot-check in the normal range. A full-grid perf re-run under 4.6.1 is queued — all prior grid evidence compiled under 4.5.0, and codegen shifts measurably between compiler minors. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…o ld.global.nc.v4.f32 DSL 4.6.1's NVVM rewrites adjacent 128-bit f32 copy-atom loads into v2.b64 register-PAIR loads plus mov.b64 unpacks (PTX: v4.b32 30->12, v2.b64 0->18 on gvr_clus). The even-aligned pair constraint fragments register allocation at the 64-register wall: the SAME ptxas 13.2 gives 64 regs / 4B spill from the 4.5.0 PTX but 63 regs / 80B spill from the 4.6.1 PTX — measured +20% on clus and +8.5% on the register family at runtime (the dominant share of the 4.6.1 grid regression, clus zone tax 1.068 -> 1.223). An inline-asm boundary pins the four-scalar-f32 shape on every DSL version; the isolated copy atom still lowers correctly, so this guards against the context-dependent rewrite only. nsys verdict under 4.6.1 (200 paired cold-L2 reps, replicated): clus r32 tax 1.163 -> 1.074 (vl 16.83 -> 13.87us), reg r12 1.207 -> 1.127, main and reg_clus controls neutral. UT 84/84 on 4.6.1 and 4.5.0; exactness smoke all exact on both pins. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…time smem carveout in the register family DSL 4.6.1 moved the shared-memory carveout derivation into a compile-time MLIR attribute (smem.max_smem_per_mp, emitted when min_blocks_per_mp > 1) that only accounts for static smem. With the register family's dynamic launch smem this selects a 16 KiB shared-memory config, pinning each SM to a single resident CTA (achieved occupancy 67.6% -> 11.6%) and slowing saturating small-N grids up to 3.25x (N=1027 @ 1024 rows: 23.2us -> 7.1us varlen after the fix). Extend the existing _no_carveout() compile scope to also drop that attribute; clus/main/reg_clus launch with min_blocks_per_mp == 1 and are unaffected (verified neutral). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Condense narrative comment blocks to the technical invariant they encode and rewrite module/class docstrings as plain functional descriptions. Comment-and-docstring-only change: AST-verified equivalent (bare-string sections and assert messages normalized), tie-aware exactness smoke green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Add an opt-in block-skip row pass to the gvr_main family: a per-32-element block-maxima side tensor lets the kernel visit only logit blocks whose maximum can clear the sampled threshold (compact survivor list, warp-per- block walk with ballot-aggregated emission; capacity overflow falls back to an identity walk over the row window). Emission, verify and the retry ladder are shared with the baseline pass, so exactness semantics are unchanged. Dispatch is the shape-only bsk_gate (batch working set >= 128 MiB and N >= 72K); outside the gate run_bsk routes to the baseline. Real-capture 886-cell x 11-BS A/B (nsys pure-kernel, cold-L2, same-GPU paired): gate-open 531 pairs gm 1.735 (worst 0.942, single-cell anomaly; next-worst 1.048), gated band neutral 0.96-1.05. In production the block-maxima tensor is produced by the indexer-GEMM emission epilogue (PR NVIDIA#16953 model); this change carries the consumer side only and leaves the production route() untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Stacked on #17821 — this branch is the #17821 head plus one commit; the review scope is the top commit only. Kept as draft until #17821 merges.
Adds an opt-in block-skip row pass to the self-sampling GVR top-K decode kernel (
gvr_mainfamily):bmax), builds a compact list of the blocks whose maximum can clear the sampled threshold, and walks only those (warp-per-block, eight blocks in flight, ballot-aggregated emission, one atomic per round). If the survivor list overflows its fixed capacity, the pass falls back to an identity walk over the row window.logits[:, :n]per row.bsk_gate(b, n, npad, k) = (b * npad * 4 >= 128 MiB) and (n >= 72 * k)— no data-dependent terms. Outside the gate,run_bskroutes to the baselinerun.route()is untouched. In deployment thebmaxtensor is produced by the indexer-GEMM emission epilogue (the [None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer #16953 producer model); this PR carries the consumer side as an explicitrun_bskentry so the kernel capability can be reviewed and validated independently of the producer wiring.Performance
Real-capture grid, 886 cells x 11 batch sizes (9,746 combos; DSV3.2 + DSV4 Flash/Pro decode captures), nsys pure-kernel medians, cold-L2, same-GPU paired A/B against the #17821 head:
Gate-open speedup vs #17821, split one axis at a time:
Combined model_isl x BS (gm/worst; "—" = gated or routed to a non-main family):
Note: the commit message quotes an earlier 531-pair tally whose gate reconstruction used the unpadded row length in the working-set term; the kernel gate uses the padded row stride (64-element multiples), which additionally opens the long-ISL bs=128 column (all wins). The 582-pair numbers here are authoritative.
Test Coverage
tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk_blockskip.py:torch.topk(hostile +3e38 pads);The existing
test_gvr_selfsampling_topk.pysuite (84 cases) passes unchanged on this branch.PR Checklist
🤖 Generated with Claude Code