[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457
[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457longcheng-nv wants to merge 50 commits into
Conversation
…nert) First of a series porting the the reference prototype histogram-ladder R0 admission fast path into the unified GvrTopKKernel. Adds the multi-threshold GE-count primitive block_count_ge_multi<M> (reuses block_count_ge's vectorized memory path with M static register counters, caching each per-thread count column) and the constructor knobs enable_r0/r0_qfracs/mt_unroll/p1b_cache/ fb_fix plus their derived rung targets. All knobs default OFF and the helper is not yet called from any kernel body, so a default-constructed kernel is byte-identical to the base; this commit has no behavioral effect. The rung placement (P1b), admission dispatch, cached-column Phase 3, log-falsi fallback and cluster-aware reduction land in follow-up commits. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ent (inert) Adds phase1b_hspace_rungs: a 256-bin SMEM histogram over the prev-topK gathered values banded by P1's pmin/pmax, with warp-0-parallel extraction of M h-space quantile rungs (crossing qneeds[m] from the top). These rungs are the candidate thresholds the single-pass R0 admission evaluates. Not yet called from any kernel body (the enable_r0 dispatch lands next), so the base kernel remains byte-identical; no behavioral effect. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…s (secant fallback) Wires the R0 histogram-ladder admission into _run_phases as a fast path: P1b places M h-space rungs, one block_count_ge_multi pass counts all M at once, and the tightest rung with count in [K, kC] is accepted, seeding Phase 3 with its cached per-thread count column (zero rescan). An R0 miss falls back to the existing phase2_secant_search (a later commit replaces that fallback with an inline log-falsi refine). Five R0 scratch buffers are allocated in run_one_row (gated on enable_r0, None otherwise) and threaded through _run_phases with the same const_expr-gate + None-propagate idiom the kernel already uses for s_cluster_partial / smem_input. Gated to enable_r0 and cluster_size==1 (the earlier R0 scans the full row in a single CTA; the slice-parallel + cluster count-merge variant for the cs>1 long-row branch lands in a follow-up commit). enable_r0 defaults False, so the compiled path is the unchanged secant search and the default build is byte-identical to the base. NOTE: unlike the two preceding commits, the enable_r0=True path executes new device code and is NOT yet silicon-validated; exactness + nsys gates on a B200 are required before enabling it by default. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
On an R0 admission miss, derive the measured [lo, hi] bracket from the rungs and, when both ends were measured, fire ONE extra count pass at a log-count regula-falsi aim (geometric center of the [K, kC] window). If it lands in-window the row is resolved with smem_ptcnt fresh from that pass; otherwise (or when an end was unmeasured) the existing secant fallback runs from the untouched P1 bracket, so the fallback state is identical to before this commit. Adds only the R1 fast-exit for near-boundary misses. Entirely inside the enable_r0 && cluster_size==1 branch, so the default build (enable_r0=False) is unchanged. Enabled path is not yet silicon-validated; gated behind the same pending B200 exactness/nsys run as the R0 admission itself. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…laces secant under R0) Replaces the R0-miss secant fallback (when fb_fix=True, the default) with the earlier bounded log-count regula-falsi refine: iterate block_count_ge with a log2_mstar-aimed threshold update over a [pmin, pmax] bracket whose end counts start unknown and get measured, converging into [K, kC] and setting done=1 so upstream Phase 3 skips its retry-shrink. fb_fix=False keeps the secant fallback. Adds fb_alpha/log2_mstar (interior aim). Inside the enable_r0 && cluster_size==1 branch, so the default build (enable_r0=False) is byte-identical. CORRECTNESS-CRITICAL: the exhaustion path lands on the measured undershoot side for value-tie plateaus; this is the case the kC >= 5K 16-bit tie contract protects and it is NOT yet silicon-validated. This commit must be gated behind the tie-exactness gate on a B200 before enable_r0 is defaulted on; treat it as the highest-risk commit in the series. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Makes the R0 admission cover the multi-CTA (cluster_size>1) path so the
load-balance kernel's long-row branch gets R0 too. block_count_ge_multi now
scans [slice_start, slice_end) and DSMEM all-reduces the M rung counts
across the cluster with a RELEASE cluster_arrive (upstream's block_count_ge
comment warns cluster_arrive_relaxed skips the release; the earlier mc used the
relaxed form). A per-cluster s_cluster_partial_m[M] scratch is allocated
(gated enable_r0 && cs>1) and the R0 / R1 / fb_fix count calls are all
slice+cluster aware. phase1b rungs stay per-CTA identical (preIdx stats are
full-row). The enable_r0 && cluster_size==1 gate is relaxed to enable_r0.
Default build (enable_r0=False) remains byte-identical: every R0 buffer and
branch is const_expr-gated on enable_r0.
Validated on B200: cluster-R0 exactness sweep 144/144 (fp32/bf16/fp16 x
cs{2,4} x K{512,1024,2048} x N up to 262144 x {real,random preIdx} + tie
plateaus), no deadlock, vs torch.topk. Cumulative exactness 330/330.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…cs + kC-diet + p1b_cache dtype-gate) Folds the earlier host dispatch policy into the ctor, all gated on enable_r0 so an OFF kernel stays byte-identical: - r0_qfracs default = (0.85, 0.35) (0.85, 0.35): the reference dispatch policy ships this M=2 ladder for every (dtype,K,N); the M=2 count pass is ~free and the R1 falsi shot covers the 3-7% bracket misses. uh4 (M=4) was silicon- falsified (mc geomean 0.956 - admission != latency). - p1b_cache default = (dtype != fp32): the gather-cache wins +0.8-2.8% on 16-bit; flat/negative on fp32. (The cached path itself lands in; the flag is inert until then.) - kC-diet: K512 single-CTA -> kC=3072 (saves 16KB SMEM). kC>=2560 is the K512 16-bit tie-safety contract, so 3072 is safe; cluster + K1024/K2048 stay stock. Validated on B200: default-dispatch exactness 48/48 (fp32/bf16/fp16, all K, incl. K512 aggressive tie plateaus at kC=3072). Corrects the prior benchmarking that used the falsified uh4 ladder. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…bit) Implements the p1b_cache fast path that 's dispatch enables for 16-bit: P1 (phase1_preidx_stats) optionally stashes each gathered preIdx value into a SMEM buffer (smem_gath, top_k f32) via a new const_expr-gated store, and a new phase1b_hspace_rungs_cached builds the rung histogram straight from that cache -- skipping P1b's second GMEM random gather (the expensive part in half precision). smem_gath is allocated in run_one_row gated on enable_r0 && p1b_cache and threaded through _run_phases; the phase1b call dispatches cached vs re-gather on the p1b_cache const_expr. The optional smem_gath param defaults to None, so with p1b_cache OFF (all fp32, and the whole base) the store branches are const_expr-eliminated and the kernel is byte-identical. Validated on B200: 16-bit default-dispatch exactness 48/48 (bf16/fp16 now exercise the cached path, incl. K512 tie plateaus at kC=3072). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…f to Phase 3/4 The cluster (cs>1) R0 path was silently WRONG on an admission/R1 hit and only produced correct output because it always missed and fell back to the secant path (exactness masked the bug). ncu root-cause: forcing a hit gave uniq 592/1024. The R0/R1 accept paths omitted the per-CTA cluster-gather state that block_count_ge sets and Phase 3/4 consume: - s_iscalars[1]=1 (done): tells Phase 3 to honor the admitted threshold and skip its retry-shrink (missing => the cluster collect diverged). - s_iscalars[5]=local slice count for the chosen rung (the per-CTA cand_count the cluster gather sums). - R1 accept now also sets done=1 (block_count_ge already sets [0]/[5]). - R0-miss persists the measured rung bracket into s_thr[1]/s_thr[2] so the fb_fix refine starts tight (the reference prototype behavior) instead of P1's wide [pmin,pmax]. Validated on B200: cluster-R0 forced-hit now exact; full grid 60/60 exact (fp32/bf16/fp16 x best/worst x K512/1024/2048 x cs1/cs4). Perf note: at large N the (0.85, 0.35) (M=2) rungs straddle [K,kC] (rung0 count > kC, rung1 < K), so cs>1 resolves via the refine (the earlier 1.14x ceiling); cs=1 small-N R0 still wins ~1.5x. enable_r0 default OFF, so base is byte-identical. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…, no re-measure) Replaces the R0-miss 'R1 single shot + fb_fix re-measure loop' with one bounded log-falsi loop SEEDED from the measured rung bracket + its known counts (clo/chi) so it does regula-falsi from iter 0 with no re-measure and no redundant R1 pass. At large N ((0.85, 0.35) rungs straddle [K,kC]) this is the resolve path; seeding trims passes. exact 5/5 spot; cs4 ncu 49.4->48.1us (the reference prototype mc-R0 41.6). enable_r0 default OFF -> base byte-identical. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…n converged snap P3 stream-write (dense+compact), P4 keys/vals accesses, and the snap-loop convergence check now go through raw integer addressing off a per-method iterator.toint() base (_smem_ld/_smem_st helpers) instead of tensor indexing, which re-derived the cluster SMEM window (S2R SR_CgaCtaId + LEA<<24) around hot accesses. P4 writeback: when the snap loop converged, the final iteration's cgt (s_iscalars[3]) is the exact strictly-greater count at sel_thr, so gt entries pack via counter[4] from 0 while ties pack via counter[5] from offset cgt — ONE candidate sweep and one barrier instead of two-pass. Same [gt | eq | pad] output partition; non-converged snap falls back to the original two-pass (its cgt would be stale). Measured (B200, real V4-Pro/Flash production shapes): P4 7.80->7.58us, wall +1-2% on top of fastsync (my-pro fs2 geomean 1.032->1.043, my-flash fs0 1.063->1.075). Full sweep 450/450 PASS. ncu follow-up: the residual snap-top S2R stall is a post-barrier pipeline-refill accounting artifact plus the un-hoisted scalar scratch accesses — see analysis notes. Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The P4 snap loop steps one distinct candidate value per iteration (~0.45us: full candidate re-scan + 2 barriers). Real logits concentrate count mass at the k-th boundary, so the selected 1/1024 bin often holds tens of values; snap stragglers of 10-20us set the wall clock at N<=32K (P4 max 15.8us at N=8K vs 7.6us mean on synth aggregate cells). When the selected level-1 bin holds >8 candidates, rebuild the histogram over just that bin (+2% slop; bin width shrinks 1024x) and re-run the 3-step k-th search before entering the snap loop. Costs one extra candidate scan (~1us), leaves snap 0-2 steps. The snap loop converges monotonically from any start, so exactness is untouched. Histogram build + 3-step search factored into _hist_build / _kth_bin_search (trace-identical to the previous inline code). Synth sweep (B300, bs64 K=1024 aggregate cells, us): 2K 12.24->10.48 4K 15.14->12.97 8K 20.32->13.60 16K 24.88->17.07 32K 32.99->21.57 64K 42.35->39.99 128K/256K flat (P2-dominated). Worst bound-gap 2.09x -> 1.48x. Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Two follow-ups to the level-2 histogram refinement (eed5eb5931): 1. Refinement chain: level-2 gate lowered 8->2 (one rebuild breaks even vs ~2 snap steps) and a level-3 pass added (gate >8, fires only when level 2 fails to split the bin: heavy ties or sub-ulp window). Unrolled 2-iteration chain over (gate, binw_cur). 2. _kth_bin_search step 1: every lane redundantly walked a bins_per_warp-deep serial LDS+IADD chain (~7% of stall samples at N=8K). Now each lane sums bins_per_warp/32 bins + one REDUX warp reduce; const_expr fallback to the serial walk when bins_per_warp is not a multiple of 32. Synth sweep (B300, bs64 K=1024 aggregate cells, us, vs eed5eb5931): 2K 10.48->9.10 4K 12.97->10.71 8K 13.60->12.42 16K 17.07->15.80 32K 21.57->21.25 64K 39.99->38.14. Worst bound-gap 1.48x -> 1.43x (128K, P2-straggler dominated); N<=32K all within 1.08-1.27x. Cumulative vs pre-refinement baseline: 8K 20.3->12.4 (-39%), 16K 24.9->15.8 (-36%), 32K 33.0->21.3 (-36%). Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…back 1. Histogram window from the P2 exit bracket: P4's candidate min/max scan is replaced (behind a runtime guard) by [vlo, vhi) from s_thr[1]/[2] — P3 collected v >= vlo and cnt(>=vhi) < kK puts the k-th inside the window. Out-of-window values clamp into edge bins (cumulative-from-top counts stay exact) and the window is P2's acceptance band, so level-1 resolution improves. Stale-bracket paths (degenerate fallback, probe variants) fail the guard and take the original scan; a plausible-but-wrong bracket costs extra snap steps only, never exactness. 2. Converged-path writeback drops its SMEM atomics: the converged snap iteration's per-warp packed(ge,gt) counts are still staged in smem_wcnt and cover exactly this scan's per-warp element sets, so warps derive deterministic [gt | eq] output bases from a register prefix instead of ~2*cand/32 serialized atomicAdd claims. Padding count becomes cge (= s_iscalars[2]) on this path. Synth sweep (B300, bs64 K=1024): 4K 10.71->10.25, 8K 12.42->12.14, 16K 15.80->15.66, 32K 21.25->21.15; large N unchanged. Both changes individually small (~1-2%) — the ncu stall attribution on the min/max and writeback-atomic regions was largely pipeline-accounting; kept for the latency-chain shortening and the documented bracket contract. Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Pre-existing hole surfaced by adversarial review: a candidate span > FLT_MAX (|v| ~ 1.7e38, fuzz-only for real logits) overflows range1 to +inf -> inv1 = +0 -> all candidates in bin 0 -> thr = NaN. NaN breaks the snap walk's ordering (all comparisons false), so it no-ops for snap_limit iterations and the row is emitted as pure padding. Clamping range1 to FLT_MAX keeps the start threshold ordered; snap's monotone walk rescues any ordered (even +/-inf) start. No perf change (8K spot 12.17us, smoke PASS). Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Adversarial review of the cs>1 path (now the recommended config for N>=64K) surfaced a pre-existing theoretical data race from the multi-CTA commit: block_count_ge uses a single s_cluster_partial slot and a single cluster rendezvous per call, so a straggler CTA's post-wait DSMEM read of a peer's slot is unfenced against that peer's NEXT-call overwrite. If it ever fired, per-CTA totals diverge -> thresholds diverge -> peer under-collection (wrong top-k) or a cluster-barrier phase mismatch (hang). The schedule window is astronomically small (peer must complete six CTA barriers + classify + a full slice re-scan while the straggler's tid0 sits between wait and its peer loads), but under the PTX memory model it is a race. Fix with zero added synchronization: call k writes/reads slot k&1 (s_cluster_partial grows [1] -> [3]; slot 2 = tid0-private call counter, zeroed in run_one_row; CTAs call in lockstep so parity stays cluster-aligned). The peer's call-(k+2) overwrite of the same slot is transitively ordered after my call-k reads by the call-(k+1) rendezvous. cs2 spot 64K/128K/256K + smoke all PASS, perf unchanged. Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…lag) Add phase4_rank_scatter as an alternative Phase-4 path behind two inert flags (enable_p4_rank_scatter / enable_p4_rank_scatter_exact, both default False). Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py (PR#15709). It replaces phase4_histogram_snap's k-th-bin search + two-pass writeback with a single fused rank-and-scatter, cutting Phase-4 barriers ~14 -> ~7; the _exact variant adds one fine-histogram recursion on the straddling coarse bin so the result stays bit-exact vs torch.topk. Production upstream currently ships ONLY the histogram-snap P4 (no rank-scatter); this brings the earlier latency win into the production kernel. All deps (llvm.bitcast, atomicAdd, _fmin_f32_inline, warp_reduce_min/max_f32, kNumBins) already present. When both flags are False the const_expr switch selects phase4_histogram_snap, so an OFF kernel is behaviour-identical to the base. Exactness verified 36/36 (K512/1024/2048, cs1/4, best/worst, N to 262144) with rank_scatter_exact ON. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…r R0 Resolve enable_p4_rank_scatter / enable_p4_rank_scatter_exact from None to ON whenever enable_r0 is set (both remain overridable). The R0 fast path now runs the earlier fused rank-and-scatter P4 instead of histogram-snap. nsys A/B over the the synthetic envelope (best/worst fp32, BS=1, seq 4k-1M, cold-L2, per-cell snap-vs-rs paired on one GPU): geomean 1.092x (K1024 1.120 / K2048 1.115 / K512 1.045), 39/52 cells win >2%, 0/52 regress >2% (worst -1.6%, noise). Largest wins at cs=1 small-N (up to 1.42x). When enable_r0 is False the default resolves to OFF, so the base kernel stays byte-identical to upstream. Exactness 36/36 (all K, cs1/4, best/worst). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…_size>1 Make the R0 p1b_cache default cs-aware. The single-CTA path keeps the (dtype != fp32) gate (fp32 occupancy regression at kC=6144), but the cluster path (cluster_size>1) now defaults p1b_cache ON for all dtypes: the SMEM gather-cache win holds and the fp32 occupancy regression does not reproduce in the latency-bound cluster kernel (different SMEM budget). Matches the reference prototype the reference cluster dispatch (unconditional ON). nsys A/B at cs=4 fp32 (the synthetic envelope best/worst BS=1, N 131k-524k, paired per-GPU): K1024 ~1.01x / K2048 ~1.02x / K512 wash, 0/18 losses, exactness 18/18. Base kernel (enable_r0=False) unaffected. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
After per-warp partials land in SMEM (one visibility barrier), EVERY warp redundantly performs the block reduce + decision a single leader used to do, in a fixed order — identical inputs and identical fp32 operation order make the result bit-identical across warps, so the publish barrier, the leader serialization, and the SMEM scalar re-reads (each costing a cluster-window S2R recompute in SASS) all disappear. Phase 4 (p4_warp_redundant, default on): - k-th bin search: 3 barriers -> 1; the serial 16-slot prefix and the 64-deep serial bin walk become lane-parallel (idx-shuffle inclusive scans + unique-crossing ballot); the selected threshold/count return in registers through level-2 refinement. - snap loop: 2 barriers/iter -> 1; threshold + convergence state in registers; warp partial staging parity double-buffered in smem_hist (a warp one iteration ahead cannot clobber a slot a slower warp still reads); the converged writeback reads the winning parity bank. Phase 2 (p2_warp_redundant, default on, cluster_size == 1): - one barrier per secant round (the staging barrier inside block_count_ge); classify + secant update replayed redundantly in registers by every warp; smem_wcnt staging parity-banked (2x num_warps); canonical exit state published once for Phase 3. Implementation notes from the falsification loop (B300, N=256K bs=64): - cute.arch.shuffle_sync_up IGNORES its offset (probed: a lane-id prefix scan returned 2^k*lane) — scans use shuffle_sync with a computed source lane instead. The interim broadcast-LDS prefix (64 predicated reads per thread per call) cost +40% instructions and erased the barrier win (ncu: barrier stall 14.1%->6.1% but 16.2M vs 11.5M inst). ncu (256K, bs64): barrier-stall 14.1% -> 6.1%; regs 64 -> 64, zero local-memory traffic both sides. A/B wall: 4K +2.6%, 32K +12%, 256K +3.3% (knobs off = leader paths, byte-identical). Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
- _hist_build: clamp the bin index in the float domain (fmax/fmin) BEFORE the int cast — fptosi is undefined for out-of-range/NaN at the IR level even though PTX cvt.rzi saturates; fmax first canonicalizes NaN to bin 0. Edge-bin semantics for in-range values are unchanged. - run_gvr_topk.gvr_topk_decode: document p4_warp_redundant / p2_warp_redundant in the Args section. Made-with-Claude-Code: fable-5 Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Pure formatting: run ruff-format over gvr_topk_decode.py so the R0 histogram-ladder / rank-scatter code added by the preceding commits is ruff-format-clean and passes the pre-commit gate. No logic change' the compiled kernel is identical (enable_r0 default stays False here; the default flip + tests land in the next commit). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…alence test Flip the GvrTopKKernel default to enable_r0=True so the R0 histogram-ladder admission + P4 rank-and-scatter become the production path, and add a direct-drive R0-vs-secant equivalence unit test. R0 wins on real DSv4/V3.2 decode-capture workloads: 24/25 cells >= 1.0 vs the secant baseline, geomean 1.330x (Pro up to 2.10x), value-set-exact 25/25, and it repairs a base-secant undershoot on one real Flash cell. The secant path is retained verbatim as the enable_r0=False fallback; on the synthetic cold-hint worst axis R0 can regress (min 0.917x), to be gated by a follow-up call-site dispatch guard. enable_r0=False remains byte-identical to the pre-R0 base. The new test drives GvrTopKKernel directly (the custom op does not expose enable_r0) and checks top-K by index set (order-independent, per the top-K contract), with a value-set fallback under half-precision boundary ties. It is also the only remaining coverage of the secant path now that the op-level tests inherit the new default. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
PR_Github #61207 [ run ] triggered by Bot. Commit: |
|
PR_Github #61102 [ run ] completed with state |
|
PR_Github #61207 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61215 [ run ] triggered by Bot. Commit: |
|
PR_Github #61215 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61302 [ run ] triggered by Bot. Commit: |
|
PR_Github #61302 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61340 [ run ] triggered by Bot. Commit: |
|
PR_Github #61340 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61367 [ run ] triggered by Bot. Commit: |
|
PR_Github #61367 [ run ] completed with state |
lfr-0531
left a comment
There was a problem hiding this comment.
Since I'm not familiar with the kernel side, I didn't dive into the details of the kernel changes. But I checked that there are no interface changes, so overall LGTM.
| @@ -353,6 +390,188 @@ def __init__( | |||
| self.FLT_MAX = 3.4028235e38 | |||
| self.NEG_FLT_MAX = -self.FLT_MAX | |||
|
|
|||
| # --- op#26 R0 histogram-ladder admission (default ON) --- | |||
There was a problem hiding this comment.
These comments are a bit verbose. Could you simplify them? Also include comments in other files. (AI always tends to over-explain things!)
There was a problem hiding this comment.
Agreed — a fair amount of the added commentary is measurement history / tuning provenance (perf numbers, audit counts, campaign refs) that belongs in the PR description rather than the code. I'll prune these down to just the invariants/contracts in the already-planned follow-up PR (the dispatch-guard one), covering all three changed files — keeping this PR's head untouched preserves the green pipeline (#49596) and your approval so it can merge as-is.
…gh-hit-rate rows) Port the in-tree R0 histogram-ladder ADMISSION concept (op#26, PR NVIDIA#16457) into the bsx throughput (tp) tier's pivot selection, closing the pr1 full-grid losses against the in-tree GVR kernel on real-capture rows at BS >= 16 while keeping the fused one-pass structure and every existing exactness invariant. Mechanism (measured on pr1 real-capture cells): the old 3-stage pivot pick targets ~3K sampled candidates inside a narrow [1.5K, 0.6kC] band. On real rows this either (a) picks a FAT rung (2-4x more P3 pushes and P4 candidates than needed - the v32/pro BS>=16 band), or (b) trusts a clustering-inflated sampled estimate and undershoots K, which fails the fused-pass reuse check and re-streams / secant-loops the full row (the flash_512k 1.6-1.8x losses; spatially clustered rows inflate a float4-sampled estimate up to 2.5x over the true count). Changes (gvr_topk_decode_bsx_tp.py only): * P2a stage-0 ADMISSION pick (R0 parity): accept the TIGHTEST ladder rung whose sampled-count confidence interval sits inside the [K, kC] acceptance window - the same "smallest exact count in [K, kC]" rule as the in-tree R0 admission, applied to the pre-pass estimates. The legacy 3-stage pick is unchanged as the fallback when no rung qualifies (cold-start / degenerate ladders take exactly the old path). * Clustering-aware sigma: sample_count now also tracks per-rung float4 OCCUPANCY, packed into the same per-thread accumulator (occ << 16 | cnt), giving the compound-Poisson sigma cnt/sqrt(occ) (equals the classic Poisson sigma on IID rows). Packing keeps registers, SMEM and the exchange at their pre-change sizes - an unpacked occ implementation measured a 14-25% whole-kernel regression (register spill in the streaming loops; A/B/bisect on v32_64k_L20). No field overflow: bsx guards npad <= 262144 => cluster-total cnt <= 8192. * Lower margin 2-sigma; K2048 uses 1.5-sigma (its [K, 4K] window is too narrow for 2-sigma to fire tight) backed by a RESCUE rung: the fused pass's second count column is now the next-fatter ladder rung instead of hmin (identical cost), so a pivot undershoot is caught with ONE collect re-stream instead of the multi-pass secant loop. Exactness machinery (accept window, P4 + tie tickets, plateau descent, ragged-N masking, degenerate emit, exit rendezvous) is untouched; admission never drops a top-K element because acceptance still requires the EXACT fused count in [K, kC]. Perf (nsys cold-L2 paired, same-rep A/B vs in-tree GVR kernel, real pr1 capture cells, B200; ratio = bsx/in-tree, lower is better): target cells before after flash_512k_L34 BS16 1.57 1.13 flash_512k_L34 BS64 1.72 1.11 flash_512k_L34 BS128 1.83 1.04 flash_512k_L34 BS256 1.54 0.73 flash_512k_L34 BS512 1.64 0.74 flash_512k_L34 BS1024 1.58 0.74 v32_32k_L50 BS16 1.40 1.21 v32_32k_L50 BS64 1.54 1.30 v32_32k_L50 BS256-1024 1.29-1.32 1.10-1.11 pro_128k_L54 BS16 1.27 1.19 pro_128k_L54 BS64 1.35 1.26 pro_128k_L54 BS128-1024 1.21-1.32 1.12-1.25 no-regression cells before after flash_512k_L36 BS1-128 0.55-0.94 0.56-0.92 flash_512k_L36 BS256-1024 0.67-0.68 0.66-0.67 v32_64k_L20 BS1-8 (reg) 0.63-0.65 0.63 v32_64k_L20 BS16-128 1.09-1.19 1.00-1.10 v32_64k_L20 BS256-1024 1.00-1.05 0.93-0.97 v32_32k_L04 BS16-1024 0.88-1.06 0.90-1.10 pro_256k_L30 BS16-1024 0.89-1.07 0.91-1.12 flash_16k_L26 BS1-1024 0.61-0.86 0.62-0.86 (reg/direct tiers untouched; the three +3-5% readings - pro_256k_L30 BS16/64, v32_32k_L04 BS64 - are inside the +-4.5% session noise floor measured on the fully untouched direct tier, e.g. flash_16k_L26 BS1 0.607 -> 0.634 with zero code change) The remaining v32_32k/pro_128k BS16-128 gap (1.19-1.30) is NOT the admission/pass-count mechanism: with admission the reuse check fires (one fused pass, admitted set slimmed 6678->2917 on v32_32k_L50) and a CS in {1,2,4} launch-shape sweep moves <= 7%. On those L2-resident shapes the in-tree kernel's whole-row-per-CTA structure is simply faster than the tp cluster split; closing it needs a tier-structure change, out of scope for this admission port. Exactness: bsx suite 56/56 (incl. new admission cases: hit-rate extremes on all three production shapes, tie plateau AT the admission threshold, count>kC overflow fallback, mixed admit/fallback ragged batch), in-tree gvr suite 671 passed / 144 skipped, 11/11 targeted screen (clustered/ties/zeros/oor/uniform/trivial), all pr1 A/B runs value-set-exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
| return GvrTopKKernel._NUM_SMS | ||
|
|
||
| @staticmethod | ||
| def pick_config( |
There was a problem hiding this comment.
pick_config duplicates the runner's existing policy, and already diverges from it
GvrTopKKernel.pick_config re-implements, heuristic for heuristic, what GvrTopKDecodeRunner._pick_tuning + the cluster_size is None auto-pick block in cute_dsl_custom_ops.py (~L6290 / ~L6626) already do — same thresholds, same tiers, same comments.
Why we duplicate the tune routine?
There was a problem hiding this comment.
Good catch — the duplication is deliberate, but I agree it should collapse to a single source of truth. Context:
Why it exists: GvrTopKKernel.launch is the direct-instantiation path — the unit tests (test_cute_dsl_gvr_topk_decode.py) and standalone/op-bench snapshots import this kernel file without the custom-op layer, so the launch-shape policy was mirrored here to keep the file self-contained.
The divergences are intentional shell behavior, not drift:
- 32B misalignment: the runner's
_pick_tuningasserts (rejects), whilelaunchsilently downgrades to 128-bit loads (documented in its docstring) — the standalone path can't assume the caller pre-validated the tensor. - cluster cap: the runner
warning_once+ clamps;launchclamps silently via a guarded import so a standalone snapshot (where_query_max_cluster_sizemay not be importable) still works. - The runner additionally owns the compile-only knobs (
enable_unroll_4/enable_phase3_unroll/use_constant_hint) and the LB-mode branch, whichpick_configintentionally doesn't cover.
Plan: fold the de-dup into the already-planned follow-up PR (the same dispatch-guard + comment-pruning one promised above): make the runner's _pick_tuning + cluster_size auto-pick delegate to GvrTopKKernel.pick_config (cute_dsl_custom_ops already imports this module), keeping the assert/warn and LB-mode differences in a thin runner-side shell. Keeping this PR's head untouched preserves the green pipeline (#49596) and the existing approval so it can merge as-is.
|
LGTM. |
| # two-level histogram, and 16-bit tie plateaus are bitwise-equal, | ||
| # where arrival order is value-exact). | ||
| if p4_exact_tail is None: | ||
| p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32 |
There was a problem hiding this comment.
Two-level linear binning does not guarantee unique fp16/bf16 values. For fp16 range [0, 65504]
with 512 bins, 1.0 and 1.25 land in the same fine bin, so arrival-order scatter can keep the
lower value at the K boundary. Please enable exact-tail handling for 16-bit or add a proven
range guard and adversarial test.
There was a problem hiding this comment.
You're right — the docstring overclaims. Binning is window-relative linear (kNumBins x 256 resolution over the Phase-2 bracket), so "16-bit keys are fully resolved" only holds when the bracket is narrow relative to the boundary-local ULP; fp16's log-dense values break it for wide brackets, as in your example. The production DSv4 path is fp32 (where p4_exact_tail handles exactly this), and the 16-bit arrival-fill predates this PR (ships with #15709's rank-scatter-exact path).
Fix will land in the already-planned follow-up PR (dispatch-guard + comment pruning + pick_config de-dup): enable exact-tail for 16-bit (2 radix levels vs fp32's 4, so cheaper), add an adversarial test, and correct this docstring.
| s_iscalars[2] = cutlass.Int32(-1) | ||
| cute.arch.barrier() | ||
| rs = rs + cutlass.Int32(1) | ||
| if s_iscalars[1] != cutlass.Int32(1): |
There was a problem hiding this comment.
s_thr[2] is the undershoot side, so its measured count can be below K. Marking done=1 skips
the Phase-3 retry and makes Phase 4 pad -1; a wide boundary plateau can trigger this. Please
route this case to an exact tie-aware fallback.
There was a problem hiding this comment.
Correct mechanics. Two things worth separating, though:
The plateau trigger: for fp32 (production path), a > kC - K (~3-4k) run of bitwise-equal values straddling the K boundary cannot realistically arise from continuously-valued logits — never observed across the 865-cell real production grid. It's reachable only for 16-bit dtypes and adversarial inputs.
The -1-pad terminal itself: inherited unchanged from the base kernel's non-convergence behavior (real-data non-convergence is always a small undershoot, via iteration exhaustion, not plateaus). -1 padding is a first-class encoding in this contract, not an error path: every decode step with seq_len < K necessarily pads with -1, so the DSA consumer already skips them on its hot path. Semantically an undershoot drops only the lowest-ranked boundary candidates; end-to-end accuracy has been validated with this behavior in production.
Agreed it should still be exact: in the follow-up PR I'll route the collapsed-bracket case to a tie-aware fallback — emit all strictly-greater candidates, top up the remaining slots from the [s_thr[1], s_thr[2]) band via the existing radix backstop (exact for both plateau and exhaustion cases) — plus plateau tests. This composes with the 16-bit exact-tail item from your other comment.
| tj = tj + cutlass.Int32(1) | ||
| cute.arch.barrier() | ||
| else: | ||
| # Persistent scalars live above the 256 digit bins |
There was a problem hiding this comment.
The exact-tail radix fallback is duplicated in the fast-path else branch and the non-fast
branch. Please extract one @cute.jit helper so future synchronization and correctness fixes
cannot drift between two copies.
There was a problem hiding this comment.
Agreed — the 4-pass radix select is textually duplicated between the p4_tail_fast else path and the non-fast branch. I'll extract a single @cute.jit helper in the follow-up PR (bundled with the items above), preserving the property that p4_tail_fast=False compiles byte-identical PTX to the original — helpers inline, and I'll re-verify the PTX diff.
mingyangHao
left a comment
There was a problem hiding this comment.
This PR can be merged now since production tests are passed. For the comments please address them in followup-PR. @longcheng-nv
Note
#16424 (the semantic dependency) has merged into
main(9095cc11d1); this branch is synced with currentmain, so the stacking is resolved — the PR's net diff vsmainis exactly the 3 files listed under Scope. CI is green at the current head (dc2bdc8a9f, pipeline #49596).Summary
Adds an R0 histogram-ladder admission fast path + a fused rank-and-scatter Phase-4 writeback to the production Blackwell cuTe-DSL GVR top-K decode kernel (
GvrTopKKernel), and makes R0 the default (enable_r0=True).Headline (real production decode captures, B200): 1.33× geomean over the shipped kernel, faster on 25/25 cells, exact on the full 2772-cell audit grid — and it repairs two correctness defects of the shipped kernel (a low-hit undershoot and a boundary-tie miss; see Correctness).
What's inside:
r0_vseed): P1's mean probe is folded into the count pass as a free extra rung (zero SMEM growth), adapting admission per row — fixes the cold-hint fat-admission regime.p4_exact_tail(fp32 default on): bit-exact boundary-tie resolution (ambiguity-gated radix select + tiny-tie fast path); 16-bit kernels byte-identical.enable_r0=False); this PR flips the kernel default only — call-site/dispatch is untouched (guard = follow-up PR).Performance
Methodology. nsys pure-kernel, cold-L2, single-GPU same-run A/B (R0 vs retained secant on identical inputs), 20 cold reps, B200. Both arms launched via the kernel's own
launch/pick_configcontract (the shapes the production runner picks). Re-measured 2026-07-20 at the current head. Inputs = per-layer indexer top-K captured from production DeepSeek-V4 Flash / Pro / V3.2 BS=1 greedy decode (9 ISL rungs 4k–1024k;N= post-compress indexer length, V4 cr=4 / V3.2 cr=1;hit= preIdx∩topK / K; V3.2preIdx= previous step's top-K). V3.2 has 7 rungs: its 160K (163,840-token) max context truncates longer prompts — the 256k row's N=163,775 is the exact valid kv length at the benched step. A calibrated synthetic envelope is used only for exactness/audit coverage (Correctness, Known limitation).BS=1 fp32, per seq-len — geomean 1.333×, 25/25 ≥ 1.0, exact 25/25
V4 Flash (K512) — geomean 1.285×
V4 Pro (K1024) — geomean 1.299×
V3.2 (K2048) — geomean 1.443×
All cells exact vs the captured reference. Flash 512k: the base kernel is additionally inexact there (undershoot) — R0 is exact and 2.10× faster.
BS scaling — R0/base geomean per BS (all captured ISL rungs × 3 dtypes; 825 cells, exact 825/825)
pick_config(the (dtype, BS, N) → launch-shape classmethod, incl. the CUDA-graphmax_seq_lencontract) andlaunch(compiled-variant cache;**kernel_overridesfor forcing knobs). cluster_size=8 (previously untested) validated: 78/78 exactness cells, beats forced cs=4 on 8/8 nsys cells.Known limitation + follow-up
enable_r0=Falseis exactly what it dispatches to.Correctness
torch.topkon the same run (plus in-range / no-duplicate / n_below=0 guards). Real 25/25, synthetic envelope 52/52, full audit grid 2772/2772 (re-verified at the current head).p4_exact_tailmakes all 12 affected audit cells bit-exact; unaffected-cell cost is noise (0.998 geomean); the tiny-tie fast path returns repair-active rows to a win (Pro 512k: 1.22×).pick_configpolicy lock;launch()autoconfig incl. forced-secant arm; adversarial 5e-8 / 1-ulp tie bands (p4_exact_tail). A standalone harness independently confirms 186/186.Scope / risk / rollback
gvr_topk_decode.py,gvr_topk_decode_load_balance.py(a 36-line share of the pick_config/launch plumbing), and the unit-test file; no call-site, dispatch, or config-schema changes.enable_r0=Falserestores the pre-R0 kernel byte-for-byte (all R0 fields const-folded) — rollback = flip one default.r0_vseed=Falsekeeps the static ladder;p4_exact_taildefaults off for 16-bit (byte-identical there).Test plan
on SM100 (B200 / B300).
The suite is slimmed ~33% (2428s -> 1626s local B200, two consecutive full-green runs) so the CI bucket fits its default 60-minute budget without any test-db TIMEOUT bumps: input memoization, the
seqlen_sortedaxis moved from the full cross-product to a 4-cell representative test, andtest_lb_vs_referencereduced to the two SMEM-layout endpoint dtypes. All protected coverage (shortrow degrade,p4_exact_tailties, r0-equivalence incl. cs=8/big-BS,pick_config/launchpolicy locks) is unchanged; remaining runtime is dominated by first-instance cuteDSL compilation of the 133 kernel-variant signatures.