diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index e53d58693a4b..a067f2fca7a5 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7271,6 +7271,10 @@ def warmup_cute_dsl_radix_topk_decode( # ------------------------------------------------------------------ # from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import \ GvrTopKKernel as _GvrTopKKernel + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ + is_tiered_topk_supported as _is_tiered_topk_supported + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ + tiered_topk as _tiered_topk class CuteDSLGvrTopKDecodeRunner: """Runner for the GVR Top-K cuTe DSL kernel (Blackwell SM100). @@ -7292,77 +7296,37 @@ def _pick_tuning( max_seq_len: Optional[int], data_ptr: int, ) -> dict: - """Pick T / V / min_blocks_per_mp tuning knobs shared by - single-CTA / sort and LB compile paths. Returned keys match - ``_compile`` / ``_compile_lb`` param names for ``**tuning`` - spreading. + """Adapter over :meth:`GvrTopKKernel.pick_tuning` (the single + source of truth for the T / V / min_blocks_per_mp / + warp-reduce policy), shared by the single-CTA / sort and LB + compile paths. Returned keys match ``_compile`` / + ``_compile_lb`` param names for ``**tuning`` spreading. + + Intentional shell divergence from ``GvrTopKKernel.launch``: + a 32B-misaligned logits pointer is a CONTRACT VIOLATION here + (assert), while ``launch`` silently downgrades to 128-bit + loads (dev convenience for ad-hoc tensors). """ - enable_unroll_4 = True - enable_phase3_unroll = True - use_constant_hint = False - - # T=1024 needs 1 CTA/SM grid AND enough per-CTA vec work. - # Under graph capture, raise the half-prec bar so a small - # capture-N doesn't force T=1024 on small-N replays - # (~14-16% regression). - if max_seq_len is not None and torch_dtype != torch.float32: - n_thresh_t = 131072 - else: - n_thresh_t = 65536 - num_threads_per_block = (1024 if - (num_rows <= num_sms - and N_per_cta >= n_thresh_t) else 512) - # V=256-bit only helps fp32 at large N. Half-prec cvt - # doubles reg pressure (5-11% loss at K=512/1024). Caller - # must hand a contiguous (32B-aligned) tensor — torch.empty - # / row slices satisfy this; column / stride-padded layouts - # may not. - use_256bit_load = (torch_dtype == torch.float32 - and N_per_cta >= 16384) - if use_256bit_load: + cfg = _GvrTopKKernel.pick_tuning( + torch_dtype, + num_rows, + N_per_cta, + num_sms, + graph_capture=max_seq_len is not None, + ) + if cfg["use_256bit_load"]: assert data_ptr % 32 == 0, ( f"use_256bit_load=True requires 32B-aligned " f"logits.data_ptr(), got {data_ptr} % 32 = " f"{data_ptr % 32}.") - # Warp-parallel reduce only pays at 32-warp (T=1024). - enable_warp_parallel_reduce = num_threads_per_block == 1024 - - # min_blocks_per_mp: reg-vs-occupancy 3-tier. Half-prec - # prefers extra CTA/SM (cvt-ILP fits in 40 regs); fp32 - # wants mb=2 (4-LDG ILP needs ~70 regs). - vec_bits_host = 256 if use_256bit_load else 128 - vec_w_host = vec_bits_host // (32 if torch_dtype == torch.float32 - else 16) - n_vec_iters = max(1, - N_per_cta // (num_threads_per_block * vec_w_host)) - if torch_dtype == torch.float32: - if n_vec_iters < 4: - min_blocks_per_mp = 0 - elif num_rows <= num_sms: - min_blocks_per_mp = 1 - elif (num_sms * 2 < num_rows <= num_sms * 3 - and N_per_cta <= 32768): - # mb=3 packs all CTAs in 1 wave; at N>=64K kernel - # is bandwidth-bound and mb=2 wins instead. - min_blocks_per_mp = 3 - else: - min_blocks_per_mp = 2 - else: - if num_rows > num_sms: - min_blocks_per_mp = 3 - elif n_vec_iters < 4: - min_blocks_per_mp = 0 - else: - min_blocks_per_mp = 1 - return dict( - enable_unroll_4=enable_unroll_4, - enable_phase3_unroll=enable_phase3_unroll, - use_constant_hint=use_constant_hint, - num_threads_per_block=num_threads_per_block, - use_256bit_load=use_256bit_load, - enable_warp_parallel_reduce=enable_warp_parallel_reduce, - min_blocks_per_mp=min_blocks_per_mp, + enable_unroll_4=True, + enable_phase3_unroll=True, + use_constant_hint=False, + num_threads_per_block=cfg["num_threads"], + use_256bit_load=cfg["use_256bit_load"], + enable_warp_parallel_reduce=cfg["enable_warp_parallel_reduce"], + min_blocks_per_mp=cfg["min_blocks_per_mp"], ) @classmethod @@ -7569,6 +7533,23 @@ def forward( ``counters`` without ``order_row`` is rejected. """ + # Tiered-GVR fast path: fp32 / next_n >= 1 (MTP) / + # cr in {1, 4} / npad <= 262144 decode rows route to the + # direct/reg/tp CuTe DSL tiers; everything else (half-prec, LB, + # oversize npad, hw cluster cap) falls through to the in-tree + # kernel below. ``order_row`` (the LJF hint dsa.py computes for + # num_rows >= 2 * num_sms) is accepted and ignored: the GVR + # tiers launch per-row CTAs and do not consume the permutation. + # Host-only guard — no device sync. The op signature and output + # contract are unchanged (unordered int32 indices, -1 pad only + # for degenerate rows). + if _is_tiered_topk_supported(logits, pre_idx, seq_lens, + output_indices, top_k, next_n, + compress_ratio, order_row, counters): + _tiered_topk(logits, pre_idx, seq_lens, output_indices, top_k, + next_n, compress_ratio) + return + cute_dtype = _TORCH_TO_CUTLASS_DTYPE[logits.dtype] num_rows = logits.shape[0] # seq_lens is request-level, logits is row-level (next_n @@ -7621,21 +7602,8 @@ def forward( f"prepare, or use the single-CTA path.") else: if cluster_size is None: - # B200 SXM5 synth-data tuning, 2026-06-10: - # N < 64K -> 1 (sync unrecouped) - # N >= 128K, BS <= 4 -> 8 (tiny grid) - # BS * cs <= num_sms -> cs (single-wave) - # else -> 1 (multi-wave loses) - if N_row < 65536: - cluster_size = 1 - elif num_rows <= 4 and N_row >= 131072: - cluster_size = 8 - elif num_rows * 4 <= num_sms: - cluster_size = 4 - elif num_rows * 2 <= num_sms: - cluster_size = 2 - else: - cluster_size = 1 + cluster_size = _GvrTopKKernel.pick_cluster_size( + num_rows, N_row, num_sms) if cluster_size > 1: hw_max_cluster = _query_max_cluster_size() if cluster_size > hw_max_cluster: @@ -7697,7 +7665,7 @@ def forward( # ``num_rows >= 2 * num_sms``. Physical meaning: wave-2 must fit a # full SM-row's worth of CTAs so the sort has long-vs-short rows to # swap. Below that threshold the win is noise / can regress a few - # percent (B200 N∈{8K,16K,32K} sweep 2026-06-23). + # percent (measured, N in {8K,16K,32K}). @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode", mutates_args=("output_indices", ), device_types="cuda") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index f9ee2ddde4e2..68e8f7bf8a45 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -17,6 +17,10 @@ from .filtered_top_k_decode_varlen import FilteredTopKKernelVarlenDecode from .filtered_top_k_varlen_util import FilteredTopKKernelVarlen from .gvr_topk_decode import GvrParams, GvrTopKKernel +from .gvr_topk_decode_direct import DirectTopKKernel +from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk +from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel __all__ = [ @@ -25,4 +29,9 @@ "FilteredTopKKernelVarlenDecode", "GvrParams", "GvrTopKKernel", + "GvrTpKernel", + "GvrRegKernel", + "DirectTopKKernel", + "tiered_topk", + "is_tiered_topk_supported", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py index 1fe0dfa5eddc..44275ab832c3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py @@ -390,7 +390,7 @@ def __init__( self.FLT_MAX = 3.4028235e38 self.NEG_FLT_MAX = -self.FLT_MAX - # --- op#26 R0 histogram-ladder admission (default ON) --- + # --- R0 histogram-ladder admission (default ON) --- # enable_r0: replace the Phase-2 secant search with a single-pass # multi-threshold "rung ladder" admission seeded by a 256-bin # histogram over the prev-topK gathered values (P1b). @@ -398,13 +398,20 @@ def __init__( # workloads (25-cell seq-len scan) where R0 wins 24/25 vs the # secant baseline, geomean 1.33x (pro 128k 2.10x). Correctness is # value-set-exact vs torch.topk (186/186 across dtype/K/N/BS/cluster - # + tie plateaus). The secant path is retained verbatim and remains - # reachable via enable_r0=False; it is the exact fallback for the - # large-N / cold-hint (low preIdx hit-rate) regime where R0 can - # regress on the synthetic worst axis — a follow-up PR adds a - # data-driven dispatch guard to route between the two. All R0 fields - # are const-foldable, so an enable_r0=False kernel is byte-identical - # to the pre-R0 upstream base. + # + tie plateaus). + # THE SECANT PATH IS NOT DEAD CODE AND MUST NOT BE DELETED. It has + # two distinct roles: + # (a) EXACT FALLBACK, live at the default enable_r0=True: when the + # rung ladder admits no candidate (R0 miss) the row falls + # through to phase2_secant_search, so this code runs in + # production on every hint-unrepresentative row; + # (b) DIFFERENTIAL ORACLE, via enable_r0=False: it is the classic + # baseline the R0 admission is checked against + # (test_..._r0_equivalence), and the direct-drive entry used to + # bisect admission-vs-baseline regressions. + # All R0 fields are const-foldable, so an enable_r0=False kernel is + # byte-identical to the pre-R0 upstream base and the disabled branch + # costs nothing at runtime. # r0_qfracs: descending h-space quantile fractions defining the M # candidate rungs (ascending threshold values); None => no rungs. # r0_vseed: park P1's pmean (the secant init probe) as one extra @@ -423,9 +430,9 @@ def __init__( self.enable_r0 = bool(enable_r0) self.mt_unroll = int(mt_unroll) self.fb_fix = bool(fb_fix) - # C7 dispatch (op#26 host policy folded into the ctor; all gated on + # C7 dispatch (host policy folded into the ctor; all gated on # enable_r0 so an OFF kernel is byte-identical to the base): - # - qfracs default = M2D (0.85, 0.35): dispatch_r0_op26 ships M2D for + # - qfracs default = M2D (0.85, 0.35): the shipped dispatch uses M2D for # every (dtype, K, N); the M=2 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). @@ -435,7 +442,7 @@ def __init__( # single-CTA path does NOT reproduce in the cluster kernel # (latency-bound, different SMEM budget). nsys cs=4: K1024 # ~1.01x / K2048 ~1.02x / K512 wash, 0 losses, exact. Matches - # op26 dispatch_p1bc_mc (unconditional ON). + # the multi-CTA dispatch (unconditional ON). # * cs=1 (single-CTA): (dtype != fp32). The gather-cache wins # +0.8-2.8% on 16-bit (random half-prec gather is the cost) but # is flat/negative on fp32 (occupancy at kC=6144), so OFF there. @@ -445,14 +452,14 @@ def __init__( if r0_vseed is None: r0_vseed = enable_r0 if enable_r0 and r0_qfracs is None: - # Per-K default (2026-07-16 vseed full-envelope audit, 2772 + # Per-K default (full-envelope audit, 2772 # cells): with the virtual seed rung on, pmean covers q.35's # admission region for K512/K1024 (2 count columns = zero # column tax); K2048 keeps q.35 (kC/K = 2.5 makes a fat admit # costlier than a slim 2-pass miss). Without vseed, q.35 must # stay for all K (it is the only slim rung). - # K2048 low rung 0.85 -> 0.6 (2026-07-19 real-content rung - # recalibration + paired nsys cold-L2 A/B, B200): the shipped + # K2048 low rung 0.85 -> 0.6 (real-content rung recalibration, + # paired cold-L2 A/B): the shipped # 0.85 rung's admission straddles [K, kC] on real V3.2 decode # captures (bracket on 86% of steps -> one extra falsi pass); # 0.6 lands the first pass. Measured: real V3.2 geomean @@ -479,8 +486,8 @@ def __init__( kc_diet = cluster_size == 1 if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072: self.kC = 3072 - # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (2026-07-19 - # paired nsys cold-L2 A/B on B200, all cells exact). The P4 zero / + # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (paired + # cold-L2 A/B, all cells exact). The P4 zero / # atomic build / serial scan all shrink 4x; the deeper boundary-bin # recursion costs less than the saved passes at kC=6144 candidates. # Measured vs this head: real V3.2 decode captures geomean +6.1% @@ -500,7 +507,7 @@ def __init__( "r0_qfracs must be descending h (ascending threshold value)" ) self.M_thr = len(self.r0_qfracs) - # --- vseed (2026-07-16): fold P1's pmean (the secant init + # --- vseed: fold P1's pmean (the secant init # probe) into the M-ary R0 count pass as one extra "virtual rung". # Fixes the flash-1M fat-admission regression (the coarse q.85 rung # admits ~4400 candidates where pmean admits ~630 -> 7x P3/P4 cand @@ -524,14 +531,14 @@ def __init__( else 0.0 ) - # --- op#7 P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- + # --- P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- # Replaces phase4_histogram_snap's k-th-bin search + 2-pass writeback - # with a single rank-and-scatter pass (op#7 PR#15709), cutting Phase-4 + # with a single rank-and-scatter pass (PR#15709), cutting Phase-4 # barriers ~14 -> ~7. On a latency-bound kernel that is a whole-kernel # win (~1.078x, HW-invariant). enable_p4_rank_scatter_exact adds ONE # fine-histogram recursion on the straddling coarse bin so the result is # bit-exact vs torch.topk (adds a few barriers back but still < snap). - # Default ON with R0: nsys over the op22 4k-1M BS=1 best/worst envelope + # Default ON with R0: measured over the 4k-1M BS=1 best/worst envelope # gives geomean ~1.09x (K1024 1.12 / K2048 1.12 / K512 1.05) with NO # cell regressing >2%. Resolves to OFF when enable_r0 is False, so the # base kernel stays byte-identical to upstream. @@ -542,15 +549,27 @@ def __init__( self.enable_p4_rank_scatter = bool(enable_p4_rank_scatter) self.enable_p4_rank_scatter_exact = bool(enable_p4_rank_scatter_exact) # p4_exact_tail: ambiguity-gated exact tie-resolution for the fine - # straddling bin (fp32 inputs only; see phase4_rank_scatter). The - # fine recursion resolves values to range/(kNumBins*256); two fp32 - # values closer than that straddling the kK boundary inside one fine - # bin were previously picked in arrival order (observed as |miss|=1 - # with |dv| ~ 3e-6 on real Pro 512k-ISL captures). Default ON for - # fp32 rank-scatter-exact kernels; 16-bit inputs keep the arrival - # fill (their upconverted keys are already fully resolved by the - # two-level histogram, and 16-bit tie plateaus are bitwise-equal, - # where arrival order is value-exact). + # straddling bin (see phase4_rank_scatter). The fine recursion + # resolves values to range/(kNumBins*256) — WINDOW-RELATIVE, so ANY + # dtype (fp32 or upconverted 16-bit) can leave distinct values in + # one fine bin whenever the Phase-2 bracket is wide relative to the + # boundary-local ULP (e.g. fp16 1.0 vs 1.25 under a [0, 65504] + # bracket); values straddling the kK boundary inside one fine bin + # were previously picked in arrival order (observed as |miss|=1 + # with |dv| ~ 3e-6 on real fp32 captures). The tail radix re-ranks + # on the full fp32 order key — candidate keys are ALWAYS fp32 + # (16-bit inputs are upcast injectively at collect), so the repair + # is exact for every supported dtype WHEN ENABLED. Default ON for + # fp32 only: on fp32 the gate fires rarely and the fix is ~free, + # but 16-bit quantization puts value plateaus at the boundary on + # virtually every input, so the gate fires constantly — measured + # B200 envelope cost (bf16, K512/K1024 x 16k-262k x BS 1-256, + # same-process paired) is gm 1.29-1.36x, worst 2.27x, while typical + # bf16 inputs (randn and quantized-tie, 48 paired runs) are already + # value-exact without it: 16-bit misses need an adversarially wide + # Phase-2 bracket (distinct 16-bit values inside one fine bin). + # 16-bit callers that need the guarantee opt in via the knob (see + # test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit). if p4_exact_tail is None: p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32 self.p4_exact_tail = bool(p4_exact_tail) and self.enable_p4_rank_scatter_exact @@ -565,7 +584,7 @@ def __init__( # text (byte-identical PTX modulo kernel name) for A/B. # Default gate = p4_exact_tail AND top_k >= 1024: the non-firing # codegen tax concentrates at K512 cs=1 mid-N (flash 64k/128k - # -6.6/-9.1%, cross-GPU reproducible, 2026-07-20 b200-035) while the + # -6.6/-9.1%, cross-GPU reproducible) while the # fire census (pro/512k bench + 9 per-layer fixture cells) contains # NO K512 cell — so K512 keeps the original byte-identical kernel. if p4_tail_fast is None: # [p4tt] @@ -1368,7 +1387,7 @@ def block_count_ge( # path (same vec_w / 4-way-unroll / tail loops) with M static register # counters. Caches all M per-thread count columns in smem_ptcnt_multi so # the accepted rung's column seeds Phase 3 with zero rescan. This is the - # R0 admission primitive (op#18 multithresh lineage); it is only invoked + # R0 admission primitive (multi-threshold lineage); it is only invoked # from the enable_r0 path added in a later commit, so the base kernel is # unaffected. Slice + cluster form: each CTA scans [slice_start, # slice_end) and the M per-CTA totals are DSMEM all-reduced across the @@ -1617,8 +1636,18 @@ def phase2_secant_search( if nv == vlo_r or nv == vhi_r: nv = (vlo_r + vhi_r) * cutlass.Float32(0.5) if nv == vlo_r or nv == vhi_r: - thr_r = vlo_r - done_r = cutlass.Int32(2) + # ADJACENT-FLOAT bracket, same terminal as the + # leader path: a low side over the candidate + # buffer plus a high side under K means the + # boundary sits inside a bitwise-equal plateau + # wider than kC. Keep the sure-winner threshold + # and let Phase 4's plateau fill finish the row. + if clo_r > cutlass.Int32(kCC) and chi_r < cutlass.Int32(kK): + thr_r = vhi_r + done_r = cutlass.Int32(3) + else: + thr_r = vlo_r + done_r = cutlass.Int32(2) if done_r == cutlass.Int32(0): thr_r = nv par_r = par_r ^ cutlass.Int32(1) @@ -1648,6 +1677,75 @@ def phase2_secant_search( vhi_r = thr_r chi_r = cnt_r it = it + cutlass.Int32(1) + # ---- Budget-exhausted plateau collapse (mirrors the leader + # path): the refine budget can run out while the bracket is + # still wide because a tie plateau wider than kC admits no + # threshold. On exactly that signature, bisect to adjacent + # floats so the plateau terminal is exact. Every thread + # replays this from identical registers, so the branch stays + # warp-uniform and block_count_ge keeps its barrier cadence. + if ( + done_r == cutlass.Int32(0) + and clo_r > cutlass.Int32(kCC) + and chi_r >= cutlass.Int32(0) + and chi_r < cutlass.Int32(kK) + ): + itc = cutlass.Int32(0) + while itc < cutlass.Int32(64) and done_r == cutlass.Int32(0): + mid_c = (vlo_r + vhi_r) * cutlass.Float32(0.5) + if mid_c == vlo_r or mid_c == vhi_r: + thr_r = vhi_r + done_r = cutlass.Int32(3) + else: + thr_r = mid_c + par_r = par_r ^ cutlass.Int32(1) + cnt_r = self.block_count_ge( + input_row, + slice_start, + slice_end, + thr_r, + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + cutlass.Boolean(False), # do_cluster_sync (cs==1) + smem_input=smem_input, + redundant=True, + wcnt_off=par_r * cutlass.Int32(nwp2), + ) + if cnt_r >= cutlass.Int32(kK) and cnt_r <= cutlass.Int32(kCC): + done_r = cutlass.Int32(1) + elif cnt_r > cutlass.Int32(kCC): + vlo_r = thr_r + clo_r = cnt_r + else: + vhi_r = thr_r + chi_r = cnt_r + itc = itc + cutlass.Int32(1) + if done_r == cutlass.Int32(3): + # recount at the terminal threshold so Phase 3 sees + # per-thread counts for the sure-winner set. + par_r = par_r ^ cutlass.Int32(1) + cnt_r = self.block_count_ge( + input_row, + slice_start, + slice_end, + thr_r, + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + cutlass.Boolean(False), # do_cluster_sync (cs==1) + smem_input=smem_input, + redundant=True, + wcnt_off=par_r * cutlass.Int32(nwp2), + ) if done_r == cutlass.Int32(0): if clo_r <= cutlass.Int32(kCC * 2): thr_r = vlo_r @@ -1733,11 +1831,22 @@ def phase2_secant_search( nv = vhi - rng * cutlass.Float32(0.05) if nv == vlo or nv == vhi: - # Bracket exhausted — try midpoint, else give up. + # Bracket exhausted — try midpoint, else terminal. nv = (vlo + vhi) * cutlass.Float32(0.5) if nv == vlo or nv == vhi: - s_thr[0] = vlo - s_iscalars[1] = cutlass.Int32(2) # done = 2 (give up) + # ADJACENT-FLOAT bracket: every value in + # [vlo, vhi) is bitwise-equal to vlo. Low side + # overflowing the candidate buffer AND high side + # undershooting K means the boundary sits inside + # a bitwise-equal plateau wider than kC — record + # the plateau terminal (done = 3) and keep the + # sure-winner threshold vhi. + if clo > cutlass.Int32(kCC) and chi < cutlass.Int32(kK): + s_thr[0] = vhi + s_iscalars[1] = cutlass.Int32(3) # done = 3 (plateau) + else: + s_thr[0] = vlo + s_iscalars[1] = cutlass.Int32(2) # done = 2 (give up) else: s_thr[0] = nv else: @@ -1777,6 +1886,84 @@ def phase2_secant_search( cute.arch.barrier() it = it + cutlass.Int32(1) + # ---- Budget-exhausted plateau collapse ---- + # The refine budget can run out while the bracket is still wide: the + # secant step keeps making progress (the bracket shrinks every + # iteration) but a tie plateau wider than kC admits no threshold, so + # the count never lands in [kK, kCC]. In exactly that signature - + # count(>= v_lo) > kCC AND count(>= v_hi) < kK, both counts current - + # collapse the bracket by pure bisection until the ends are ADJACENT + # floats; every value in [v_lo, v_hi) is then bitwise-equal, so the + # plateau terminal (done = 3) is exact and Phase 4 completes the row + # from that tie class. A count landing in [kK, kCC] mid-collapse + # converges normally. Anything else keeps the legacy give-up below. + if ( + s_iscalars[1] == cutlass.Int32(0) + and s_iscalars[2] > cutlass.Int32(kCC) + and s_iscalars[3] >= cutlass.Int32(0) + and s_iscalars[3] < cutlass.Int32(kK) + ): + itc = cutlass.Int32(0) + while itc < cutlass.Int32(64) and s_iscalars[1] == cutlass.Int32(0): + if tidx == 0: + vlo_c = s_thr[1] + vhi_c = s_thr[2] + mid_c = (vlo_c + vhi_c) * cutlass.Float32(0.5) + if mid_c == vlo_c or mid_c == vhi_c: + s_thr[0] = vhi_c + s_iscalars[1] = cutlass.Int32(3) # plateau terminal + else: + s_thr[0] = mid_c + cute.arch.barrier() + if s_iscalars[1] == cutlass.Int32(0): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + smem_input=smem_input, + do_cluster_sync=do_cluster_sync, + ) + if tidx == 0: + c_c = s_iscalars[0] + t_c = s_thr[0] + if c_c >= cutlass.Int32(kK) and c_c <= cutlass.Int32(kCC): + s_iscalars[1] = cutlass.Int32(1) + elif c_c > cutlass.Int32(kCC): + s_thr[1] = t_c + s_iscalars[2] = c_c + else: + s_thr[2] = t_c + s_iscalars[3] = c_c + cute.arch.barrier() + itc = itc + cutlass.Int32(1) + if s_iscalars[1] == cutlass.Int32(3): + # recount at the terminal threshold so Phase 3's cached + # per-thread counts describe the sure-winner set. + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + smem_input=smem_input, + do_cluster_sync=do_cluster_sync, + ) + cute.arch.barrier() + # ---- Post-loop fallback: if still not done, force threshold ---- if tidx == 0: if s_iscalars[1] == cutlass.Int32(0): @@ -2447,9 +2634,191 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane return thr_out, sel_out # ------------------------------------------------------------------ - # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter). + # Phase 4 (alt): fused rank-and-scatter (enable_p4_rank_scatter). # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py. # ------------------------------------------------------------------ + @cute.jit + def _p4_exact_tail_radix_select( + self, + kK: cutlass.Constexpr, + kBins: cutlass.Constexpr, + num_threads: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ): + """MSB-first 4x8-bit exact radix select over the straddling + fine-bin tie set (``p4_exact_tail``) — single source shared by + the tiny-tie fast path's large-class fallback and the plain + exact-tail path (previously two verbatim copies; ``@cute.jit`` + helpers inline, so codegen is unchanged).""" + # Persistent scalars live above the 256 digit bins + # (kNumBins >= 512 always): [256] key prefix (chosen + # digits, remaining bits 0), [257] slots still to fill + # inside the current equal-prefix set, [258] ties + # strictly above the prefix (their slots precede it). + if tidx == cutlass.Int32(0): + smem_hist[256] = cutlass.Int32(0) + smem_hist[257] = need0 + smem_hist[258] = cutlass.Int32(0) + cute.arch.barrier() + for lvl in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * lvl) + iz2 = tidx + while iz2 < cutlass.Int32(256): + smem_hist[iz2] = cutlass.Int32(0) + iz2 = iz2 + cutlass.Int32(num_threads) + cute.arch.barrier() + uthr_cur = smem_hist[256] + it2 = tidx + while it2 < cand_count: + vt = smem_keys[it2] + bt = cutlass.Int32((vt - bmin_r) * inv1) + if bt < cutlass.Int32(0): + bt = cutlass.Int32(0) + if bt > cutlass.Int32(kBins - 1): + bt = cutlass.Int32(kBins - 1) + if bt == b_star: + st2 = cutlass.Int32((vt - f_lo) * finv) + if st2 < cutlass.Int32(0): + st2 = cutlass.Int32(0) + if st2 > cutlass.Int32(fbins - 1): + st2 = cutlass.Int32(fbins - 1) + if st2 == sb_star: + uk = f32_order_key(vt) + pmatch = cutlass.Int32(1) + if cutlass.const_expr(lvl > 0): + if (uk >> cutlass.Int32(shift + 8)) != ( + uthr_cur >> cutlass.Int32(shift + 8) + ): + pmatch = cutlass.Int32(0) + if pmatch == cutlass.Int32(1): + dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) + atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) + it2 = it2 + cutlass.Int32(num_threads) + cute.arch.barrier() + # Two-stage descending digit scan (mirrors the + # fine 3-step search): per-warp partial sums, + # thread0 picks the target warp, its lane0 walks + # the warp's digit range — 2*num_warps serial + # steps instead of 256. + fdw = cutlass.const_expr(256 // self.num_warps) + wsum2 = cutlass.Int32(0) + for jd in cutlass.range_constexpr(fdw): + dix = cutlass.Int32(255) - warp_id * cutlass.Int32(fdw) - cutlass.Int32(jd) + wsum2 = wsum2 + smem_hist[dix] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = wsum2 + cute.arch.barrier() + if tidx == cutlass.Int32(0): + needl = smem_hist[257] + cw = cutlass.Int32(0) + tw3 = cutlass.Int32(num_warps - 1) + f3 = cutlass.Int32(0) + for w4 in cutlass.range_constexpr(self.num_warps): + cw = cw + smem_wcnt[w4] + if cw >= needl and f3 == cutlass.Int32(0): + tw3 = cutlass.Int32(w4) + f3 = cutlass.Int32(1) + pre3 = cutlass.Int32(0) + for w5 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w5) < tw3: + pre3 = pre3 + smem_wcnt[w5] + s_iscalars[4] = pre3 # prefix above target warp + s_iscalars[0] = tw3 # target warp + cute.arch.barrier() + pre4 = s_iscalars[4] + tw4 = s_iscalars[0] + if warp_id == tw4 and lane == cutlass.Int32(0): + needl2 = smem_hist[257] + base4 = pre4 + dstar = cutlass.Int32(0) + above_d = pre4 + sd4 = cutlass.Int32(0) + for jd2 in cutlass.range_constexpr(fdw): + dix2 = cutlass.Int32(255) - tw4 * cutlass.Int32(fdw) - cutlass.Int32(jd2) + ra4 = base4 + base4 = base4 + smem_hist[dix2] + if base4 >= needl2 and sd4 == cutlass.Int32(0): + dstar = dix2 + above_d = ra4 + sd4 = cutlass.Int32(1) + smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) + smem_hist[257] = needl2 - above_d + smem_hist[258] = smem_hist[258] + above_d + cute.arch.barrier() + # Rewrite the tie slot range: ties with key > u_thr + # first (there are exactly cnt_ab of them), then the + # first need_eq bitwise-equal-to-u_thr ties in arrival + # order (value-exact by construction). Signed compare + # needs the top bit flipped (unsigned-monotonic key). + u_thr = smem_hist[256] + cnt_ab = smem_hist[258] + need_eq = smem_hist[257] + ks_thr = u_thr ^ cutlass.Int32(-2147483648) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # above-writer ctr + s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr + cute.arch.barrier() + ir2 = tidx + while ir2 < cand_count: + vr = smem_keys[ir2] + br = cutlass.Int32((vr - bmin_r) * inv1) + if br < cutlass.Int32(0): + br = cutlass.Int32(0) + if br > cutlass.Int32(kBins - 1): + br = cutlass.Int32(kBins - 1) + if br == b_star: + sr = cutlass.Int32((vr - f_lo) * finv) + if sr < cutlass.Int32(0): + sr = cutlass.Int32(0) + if sr > cutlass.Int32(fbins - 1): + sr = cutlass.Int32(fbins - 1) + if sr == sb_star: + uk2 = f32_order_key(vr) + ks2 = uk2 ^ cutlass.Int32(-2147483648) + if ks2 > ks_thr: + o2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(4), + cutlass.Int32(1), + ) + pos = rank_above_fine + o2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + elif ks2 == ks_thr: + q2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if q2 < need_eq: + pos = rank_above_fine + cnt_ab + q2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + ir2 = ir2 + cutlass.Int32(num_threads) + cute.arch.barrier() + @cute.jit def phase4_rank_scatter( self, @@ -2825,325 +3194,61 @@ def phase4_rank_scatter( tj = tj + cutlass.Int32(1) cute.arch.barrier() else: - # Persistent scalars live above the 256 digit bins - # (kNumBins >= 512 always): [256] key prefix (chosen - # digits, remaining bits 0), [257] slots still to fill - # inside the current equal-prefix set, [258] ties - # strictly above the prefix (their slots precede it). - if tidx == cutlass.Int32(0): - smem_hist[256] = cutlass.Int32(0) - smem_hist[257] = need0 - smem_hist[258] = cutlass.Int32(0) - cute.arch.barrier() - for lvl in cutlass.range_constexpr(4): - shift = cutlass.const_expr(24 - 8 * lvl) - iz2 = tidx - while iz2 < cutlass.Int32(256): - smem_hist[iz2] = cutlass.Int32(0) - iz2 = iz2 + cutlass.Int32(num_threads) - cute.arch.barrier() - uthr_cur = smem_hist[256] - it2 = tidx - while it2 < cand_count: - vt = smem_keys[it2] - bt = cutlass.Int32((vt - bmin_r) * inv1) - if bt < cutlass.Int32(0): - bt = cutlass.Int32(0) - if bt > cutlass.Int32(kBins - 1): - bt = cutlass.Int32(kBins - 1) - if bt == b_star: - st2 = cutlass.Int32((vt - f_lo) * finv) - if st2 < cutlass.Int32(0): - st2 = cutlass.Int32(0) - if st2 > cutlass.Int32(fbins - 1): - st2 = cutlass.Int32(fbins - 1) - if st2 == sb_star: - uk = f32_order_key(vt) - pmatch = cutlass.Int32(1) - if cutlass.const_expr(lvl > 0): - if (uk >> cutlass.Int32(shift + 8)) != ( - uthr_cur >> cutlass.Int32(shift + 8) - ): - pmatch = cutlass.Int32(0) - if pmatch == cutlass.Int32(1): - dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32( - 0xFF - ) - atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) - it2 = it2 + cutlass.Int32(num_threads) - cute.arch.barrier() - # Two-stage descending digit scan (mirrors the - # fine 3-step search): per-warp partial sums, - # thread0 picks the target warp, its lane0 walks - # the warp's digit range — 2*num_warps serial - # steps instead of 256. - fdw = cutlass.const_expr(256 // self.num_warps) - wsum2 = cutlass.Int32(0) - for jd in cutlass.range_constexpr(fdw): - dix = ( - cutlass.Int32(255) - - warp_id * cutlass.Int32(fdw) - - cutlass.Int32(jd) - ) - wsum2 = wsum2 + smem_hist[dix] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = wsum2 - cute.arch.barrier() - if tidx == cutlass.Int32(0): - needl = smem_hist[257] - cw = cutlass.Int32(0) - tw3 = cutlass.Int32(num_warps - 1) - f3 = cutlass.Int32(0) - for w4 in cutlass.range_constexpr(self.num_warps): - cw = cw + smem_wcnt[w4] - if cw >= needl and f3 == cutlass.Int32(0): - tw3 = cutlass.Int32(w4) - f3 = cutlass.Int32(1) - pre3 = cutlass.Int32(0) - for w5 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w5) < tw3: - pre3 = pre3 + smem_wcnt[w5] - s_iscalars[4] = pre3 # prefix above target warp - s_iscalars[0] = tw3 # target warp - cute.arch.barrier() - pre4 = s_iscalars[4] - tw4 = s_iscalars[0] - if warp_id == tw4 and lane == cutlass.Int32(0): - needl2 = smem_hist[257] - base4 = pre4 - dstar = cutlass.Int32(0) - above_d = pre4 - sd4 = cutlass.Int32(0) - for jd2 in cutlass.range_constexpr(fdw): - dix2 = ( - cutlass.Int32(255) - - tw4 * cutlass.Int32(fdw) - - cutlass.Int32(jd2) - ) - ra4 = base4 - base4 = base4 + smem_hist[dix2] - if base4 >= needl2 and sd4 == cutlass.Int32(0): - dstar = dix2 - above_d = ra4 - sd4 = cutlass.Int32(1) - smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) - smem_hist[257] = needl2 - above_d - smem_hist[258] = smem_hist[258] + above_d - cute.arch.barrier() - # Rewrite the tie slot range: ties with key > u_thr - # first (there are exactly cnt_ab of them), then the - # first need_eq bitwise-equal-to-u_thr ties in arrival - # order (value-exact by construction). Signed compare - # needs the top bit flipped (unsigned-monotonic key). - u_thr = smem_hist[256] - cnt_ab = smem_hist[258] - need_eq = smem_hist[257] - ks_thr = u_thr ^ cutlass.Int32(-2147483648) - if tidx == cutlass.Int32(0): - s_iscalars[4] = cutlass.Int32(0) # above-writer ctr - s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr - cute.arch.barrier() - ir2 = tidx - while ir2 < cand_count: - vr = smem_keys[ir2] - br = cutlass.Int32((vr - bmin_r) * inv1) - if br < cutlass.Int32(0): - br = cutlass.Int32(0) - if br > cutlass.Int32(kBins - 1): - br = cutlass.Int32(kBins - 1) - if br == b_star: - sr = cutlass.Int32((vr - f_lo) * finv) - if sr < cutlass.Int32(0): - sr = cutlass.Int32(0) - if sr > cutlass.Int32(fbins - 1): - sr = cutlass.Int32(fbins - 1) - if sr == sb_star: - uk2 = f32_order_key(vr) - ks2 = uk2 ^ cutlass.Int32(-2147483648) - if ks2 > ks_thr: - o2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(4), - cutlass.Int32(1), - ) - pos = rank_above_fine + o2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - elif ks2 == ks_thr: - q2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), - cutlass.Int32(1), - ) - if q2 < need_eq: - pos = rank_above_fine + cnt_ab + q2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr( - self.return_output_values - ): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - ir2 = ir2 + cutlass.Int32(num_threads) - cute.arch.barrier() + self._p4_exact_tail_radix_select( + kK, + kBins, + num_threads, + num_warps, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ) elif cutlass.const_expr(self.p4_exact_tail): # [p4tt] if->elif only need0 = cutlass.Int32(kK) - rank_above_fine if cnt_strad > need0 and need0 > cutlass.Int32(0): - # Persistent scalars live above the 256 digit bins - # (kNumBins >= 512 always): [256] key prefix (chosen - # digits, remaining bits 0), [257] slots still to fill - # inside the current equal-prefix set, [258] ties - # strictly above the prefix (their slots precede it). - if tidx == cutlass.Int32(0): - smem_hist[256] = cutlass.Int32(0) - smem_hist[257] = need0 - smem_hist[258] = cutlass.Int32(0) - cute.arch.barrier() - for lvl in cutlass.range_constexpr(4): - shift = cutlass.const_expr(24 - 8 * lvl) - iz2 = tidx - while iz2 < cutlass.Int32(256): - smem_hist[iz2] = cutlass.Int32(0) - iz2 = iz2 + cutlass.Int32(num_threads) - cute.arch.barrier() - uthr_cur = smem_hist[256] - it2 = tidx - while it2 < cand_count: - vt = smem_keys[it2] - bt = cutlass.Int32((vt - bmin_r) * inv1) - if bt < cutlass.Int32(0): - bt = cutlass.Int32(0) - if bt > cutlass.Int32(kBins - 1): - bt = cutlass.Int32(kBins - 1) - if bt == b_star: - st2 = cutlass.Int32((vt - f_lo) * finv) - if st2 < cutlass.Int32(0): - st2 = cutlass.Int32(0) - if st2 > cutlass.Int32(fbins - 1): - st2 = cutlass.Int32(fbins - 1) - if st2 == sb_star: - uk = f32_order_key(vt) - pmatch = cutlass.Int32(1) - if cutlass.const_expr(lvl > 0): - if (uk >> cutlass.Int32(shift + 8)) != ( - uthr_cur >> cutlass.Int32(shift + 8) - ): - pmatch = cutlass.Int32(0) - if pmatch == cutlass.Int32(1): - dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) - atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) - it2 = it2 + cutlass.Int32(num_threads) - cute.arch.barrier() - # Two-stage descending digit scan (mirrors the - # fine 3-step search): per-warp partial sums, - # thread0 picks the target warp, its lane0 walks - # the warp's digit range — 2*num_warps serial - # steps instead of 256. - fdw = cutlass.const_expr(256 // self.num_warps) - wsum2 = cutlass.Int32(0) - for jd in cutlass.range_constexpr(fdw): - dix = ( - cutlass.Int32(255) - - warp_id * cutlass.Int32(fdw) - - cutlass.Int32(jd) - ) - wsum2 = wsum2 + smem_hist[dix] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = wsum2 - cute.arch.barrier() - if tidx == cutlass.Int32(0): - needl = smem_hist[257] - cw = cutlass.Int32(0) - tw3 = cutlass.Int32(num_warps - 1) - f3 = cutlass.Int32(0) - for w4 in cutlass.range_constexpr(self.num_warps): - cw = cw + smem_wcnt[w4] - if cw >= needl and f3 == cutlass.Int32(0): - tw3 = cutlass.Int32(w4) - f3 = cutlass.Int32(1) - pre3 = cutlass.Int32(0) - for w5 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w5) < tw3: - pre3 = pre3 + smem_wcnt[w5] - s_iscalars[4] = pre3 # prefix above target warp - s_iscalars[0] = tw3 # target warp - cute.arch.barrier() - pre4 = s_iscalars[4] - tw4 = s_iscalars[0] - if warp_id == tw4 and lane == cutlass.Int32(0): - needl2 = smem_hist[257] - base4 = pre4 - dstar = cutlass.Int32(0) - above_d = pre4 - sd4 = cutlass.Int32(0) - for jd2 in cutlass.range_constexpr(fdw): - dix2 = ( - cutlass.Int32(255) - - tw4 * cutlass.Int32(fdw) - - cutlass.Int32(jd2) - ) - ra4 = base4 - base4 = base4 + smem_hist[dix2] - if base4 >= needl2 and sd4 == cutlass.Int32(0): - dstar = dix2 - above_d = ra4 - sd4 = cutlass.Int32(1) - smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) - smem_hist[257] = needl2 - above_d - smem_hist[258] = smem_hist[258] + above_d - cute.arch.barrier() - # Rewrite the tie slot range: ties with key > u_thr - # first (there are exactly cnt_ab of them), then the - # first need_eq bitwise-equal-to-u_thr ties in arrival - # order (value-exact by construction). Signed compare - # needs the top bit flipped (unsigned-monotonic key). - u_thr = smem_hist[256] - cnt_ab = smem_hist[258] - need_eq = smem_hist[257] - ks_thr = u_thr ^ cutlass.Int32(-2147483648) - if tidx == cutlass.Int32(0): - s_iscalars[4] = cutlass.Int32(0) # above-writer ctr - s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr - cute.arch.barrier() - ir2 = tidx - while ir2 < cand_count: - vr = smem_keys[ir2] - br = cutlass.Int32((vr - bmin_r) * inv1) - if br < cutlass.Int32(0): - br = cutlass.Int32(0) - if br > cutlass.Int32(kBins - 1): - br = cutlass.Int32(kBins - 1) - if br == b_star: - sr = cutlass.Int32((vr - f_lo) * finv) - if sr < cutlass.Int32(0): - sr = cutlass.Int32(0) - if sr > cutlass.Int32(fbins - 1): - sr = cutlass.Int32(fbins - 1) - if sr == sb_star: - uk2 = f32_order_key(vr) - ks2 = uk2 ^ cutlass.Int32(-2147483648) - if ks2 > ks_thr: - o2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(4), - cutlass.Int32(1), - ) - pos = rank_above_fine + o2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - elif ks2 == ks_thr: - q2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), - cutlass.Int32(1), - ) - if q2 < need_eq: - pos = rank_above_fine + cnt_ab + q2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - ir2 = ir2 + cutlass.Int32(num_threads) - cute.arch.barrier() + self._p4_exact_tail_radix_select( + kK, + kBins, + num_threads, + num_warps, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ) else: # ---- APPROX rank-and-scatter (single pass), arbitrary straddling order ---- isc = tidx @@ -3186,12 +3291,13 @@ def phase4_rank_scatter( output_values_row[i10] = self.dtype(smem_keys[i10]) output_indices_row[i10] = smem_vals[i10] i10 = i10 + cutlass.Int32(num_threads) - i11 = cand_count + tidx - while i11 < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) - output_indices_row[i11] = cutlass.Int32(-1) - i11 = i11 + cutlass.Int32(num_threads) + if s_iscalars[6] == cutlass.Int32(0): # plateau fill completes done=3 + i11 = cand_count + tidx + while i11 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[i11] = cutlass.Int32(-1) + i11 = i11 + cutlass.Int32(num_threads) # ------------------------------------------------------------------ # Phase 4: Histogram-based k-th selection + two-pass writeback @@ -3697,12 +3803,13 @@ def phase4_histogram_snap( ) output_indices_row[i10] = self._smem_ld(cutlass.Int32, vals_base, i10) i10 = i10 + cutlass.Int32(num_threads) - i11 = cand_count + tidx - while i11 < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) - output_indices_row[i11] = cutlass.Int32(-1) - i11 = i11 + cutlass.Int32(num_threads) + if s_iscalars[6] == cutlass.Int32(0): # plateau fill completes done=3 + i11 = cand_count + tidx + while i11 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[i11] = cutlass.Int32(-1) + i11 = i11 + cutlass.Int32(num_threads) # ------------------------------------------------------------------ # Main kernel — one CTA per row @@ -3920,9 +4027,13 @@ def run_one_row( # [4] out_count # [5] local cand_count (per-CTA snapshot before cluster all-reduce; # consumed by the kernel-level cluster handoff) + # [6] plateau terminal flag, captured from [1] BEFORE Phase 4 + # (Phase 4 REUSES [1] as radix scratch, so the terminal must + # never be re-read from it afterwards) + # [7] plateau fill ticket (done == 3 only) s_iscalars = smem.allocate_tensor( element_type=cutlass.Int32, - layout=cute.make_ordered_layout((6,), order=(0,)), + layout=cute.make_ordered_layout((8,), order=(0,)), byte_alignment=16, ) # Per-CTA DSMEM scratch for the cluster all-reduce of cand_count: @@ -3962,7 +4073,7 @@ def run_one_row( else: smem_input = None - # op#26 R0 admission scratch (single-CTA fast path). Allocated only + # R0 admission scratch (single-CTA fast path). Allocated only # when enable_r0; None otherwise so the base SMEM layout is byte-for- # byte unchanged and these propagate harmlessly through _run_phases' # const_expr(enable_r0)-gated branch (same idiom as s_cluster_partial @@ -4312,10 +4423,12 @@ def _run_phases( # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast # path) or the secant threshold search ---- - # enable_r0 gates to cluster_size==1 for now: op#26's R0 scans the - # full row in one CTA. The slice-parallel + cluster count-merge - # variant that lets R0 cover the cs>1 long-row branch lands in a - # later commit; until then cs>1 keeps the secant path. + # R0 covers every cluster size: at cs>1 each CTA scans its own + # slice and block_count_ge_multi cluster-merges the rung counts + # (the P1b rungs are per-CTA identical because the preIdx stats are + # full-row). The secant search below is the exact fallback taken + # when the ladder admits nothing, plus the enable_r0=False + # differential-oracle entry. if cutlass.const_expr(self.enable_r0): # P1b rung placement -> ONE M-ary R0 count pass -> accept the # tightest rung with count in [K, kC]. On a miss, fall back to @@ -4406,7 +4519,7 @@ def _run_phases( # rungs. SEED the loop with the rung bracket AND its known # counts (clo/chi) so it does log-count regula-falsi from # iter 0 with no re-measure and no separate R1 shot -> ~2-3 - # count passes (op#26 efficiency) instead of ~6. done=1 on + # count passes instead of ~6. done=1 on # accept so Phase 3 skips its retry-shrink. if bc < cutlass.Int32(0): if cutlass.const_expr(self.fb_fix): @@ -4522,28 +4635,115 @@ def _run_phases( cute.arch.barrier() rs = rs + cutlass.Int32(1) if s_iscalars[1] != cutlass.Int32(1): - # tie-plateau fail-soft: land on the measured - # undershoot side (count <= kC => no overflow). - self.block_count_ge( - input_row, - slice_start, - slice_end, - s_thr[2], - smem_ptcnt, - smem_wcnt, - s_iscalars, - s_cluster_partial, - tidx, - warp_id, - lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, - ) - cute.arch.barrier() - if tidx == cutlass.Int32(0): - s_thr[0] = s_thr[2] - s_iscalars[1] = cutlass.Int32(1) - cute.arch.barrier() + # The retry budget could not land in [K, kC]. + # ONLY the coherent undershoot-overflow corner + # (count(>= lo) > kC AND 0 <= count(>= hi) < K, + # both counts CURRENT — the retry's bracket + # widening marks a side stale with -1 and thus + # fails this guard) collapses the bracket by + # pure bisection to ADJACENT floats, where the + # plateau terminal (done = 3, threshold = hi) + # is exact: Phase 4 emits the sure winners and + # the plateau fill completes the row from the + # tie class. A mid-collapse count landing in + # [K, kC] converges normally; anything else + # (incl. an exhausted collapse budget) falls + # through to the fail-soft terminal below. + it4 = cutlass.Int32(0) + if ( + s_iscalars[2] <= cutlass.Int32(self.kC) + or s_iscalars[3] < cutlass.Int32(0) + or s_iscalars[3] >= cutlass.Int32(self.top_k) + ): + it4 = cutlass.Int32(40) # guard: skip collapse + while it4 < cutlass.Int32(40) and s_iscalars[1] == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + lo4 = s_thr[1] + hi4 = s_thr[2] + mid4 = (lo4 + hi4) * cutlass.Float32(0.5) + if mid4 == lo4 or mid4 == hi4: + s_thr[0] = hi4 + s_iscalars[1] = cutlass.Int32(3) + else: + s_thr[0] = mid4 + cute.arch.barrier() + if s_iscalars[1] == cutlass.Int32(0): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + c4 = s_iscalars[0] + t4 = s_thr[0] + if c4 >= cutlass.Int32(self.top_k) and c4 <= cutlass.Int32( + self.kC + ): + s_iscalars[1] = cutlass.Int32(1) + elif c4 > cutlass.Int32(self.kC): + s_thr[1] = t4 + s_iscalars[2] = c4 + else: + s_thr[2] = t4 + s_iscalars[3] = c4 + cute.arch.barrier() + it4 = it4 + cutlass.Int32(1) + if s_iscalars[1] == cutlass.Int32(3): + # recount at the terminal threshold so P3's + # cached per-thread counts describe the + # sure-winner set the fill completes. + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + elif s_iscalars[1] != cutlass.Int32(1): + # fail-soft (non-plateau): land on the + # measured undershoot side (count <= kC => + # no overflow; -1 pad stays the documented + # non-convergence encoding). + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[2], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_thr[0] = s_thr[2] + s_iscalars[1] = cutlass.Int32(1) + cute.arch.barrier() else: self.phase2_secant_search( input_row, @@ -4622,6 +4822,13 @@ def _run_phases( cand_count_p4 = cutlass.Int32(0) if cutlass.const_expr(cluster_size == 1): # cs=1: the single CTA per row IS the leader. + # Capture the P2 terminal BEFORE Phase 4: P4 reuses + # s_iscalars[1] as radix scratch. + if tidx == cutlass.Int32(0): + s_iscalars[6] = cutlass.Int32(0) + if s_iscalars[1] == cutlass.Int32(3): + s_iscalars[6] = cutlass.Int32(1) + cute.arch.barrier() cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) if cutlass.const_expr(self.enable_p4_rank_scatter): self.phase4_rank_scatter( @@ -4653,6 +4860,38 @@ def _run_phases( warp_id, lane, ) + # ---- plateau fill (done == 3): complete the row from the + # bitwise-equal plateau class. The terminal is only set on an + # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) + # is bitwise-equal; Phase 4 has already emitted the + # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of + # the tie class is a valid tie-aware completion. Ticket counter + # lives in the DEDICATED s_iscalars[7]. + if s_iscalars[6] == cutlass.Int32(1): + pv_lo = s_thr[1] + pv_hi = s_thr[0] + if tidx == cutlass.Int32(0): + # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] + # is radix scratch by now (same hazard as the flag). + s_iscalars[7] = cand_count_p4 + cute.arch.barrier() + ifp = tidx + while ifp < N: + vfp = cutlass.Float32(0.0) + if cutlass.const_expr(self.dtype == cutlass.Float32): + vfp = input_row[ifp] + else: + vfp = cutlass.Float32(input_row[ifp]) + if vfp >= pv_lo and vfp < pv_hi: + pfill = atomicAdd( + s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1) + ) + if pfill < cutlass.Int32(self.top_k): + if cutlass.const_expr(self.return_output_values): + output_values_row[pfill] = self.dtype(vfp) + output_indices_row[pfill] = ifp + ifp = ifp + cutlass.Int32(self.num_threads) + cute.arch.barrier() else: # cs>1: only the leader (CTA 0 in cluster) runs Phase 4. if is_leader: @@ -4699,6 +4938,13 @@ def _run_phases( # smem_keys/smem_vals (no peers to gather from). # ---- Phase 4: histogram snap + writeback ---- + # Capture the P2 terminal BEFORE Phase 4: P4 + # reuses s_iscalars[1] as radix scratch. + if tidx == cutlass.Int32(0): + s_iscalars[6] = cutlass.Int32(0) + if s_iscalars[1] == cutlass.Int32(3): + s_iscalars[6] = cutlass.Int32(1) + cute.arch.barrier() cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) if cutlass.const_expr(self.enable_p4_rank_scatter): self.phase4_rank_scatter( @@ -4731,6 +4977,39 @@ def _run_phases( lane, ) + # ---- plateau fill (done == 3): complete the row from the + # bitwise-equal plateau class. The terminal is only set on an + # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) + # is bitwise-equal; Phase 4 has already emitted the + # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of + # the tie class is a valid tie-aware completion. Ticket counter + # lives in the DEDICATED s_iscalars[7]. + if s_iscalars[6] == cutlass.Int32(1): + pv_lo = s_thr[1] + pv_hi = s_thr[0] + if tidx == cutlass.Int32(0): + # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] + # is radix scratch by now (same hazard as the flag). + s_iscalars[7] = cand_count_p4 + cute.arch.barrier() + ifp = tidx + while ifp < N: + vfp = cutlass.Float32(0.0) + if cutlass.const_expr(self.dtype == cutlass.Float32): + vfp = input_row[ifp] + else: + vfp = cutlass.Float32(input_row[ifp]) + if vfp >= pv_lo and vfp < pv_hi: + pfill = atomicAdd( + s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1) + ) + if pfill < cutlass.Int32(self.top_k): + if cutlass.const_expr(self.return_output_values): + output_values_row[pfill] = self.dtype(vfp) + output_indices_row[pfill] = ifp + ifp = ifp + cutlass.Int32(self.num_threads) + cute.arch.barrier() + # Final cluster barrier: keep peer CTAs (and their SMEM) alive # until the leader's gather + Phase 4 finish. Skipped at # do_cluster_sync=False (no peers; short-row degrade non-leaders @@ -4793,9 +5072,11 @@ def __call__( # policy as a pure function colocated with the kernel (single source of # truth), and ``launch`` is a thin variant-cache wrapper so direct-drive # users (tests, benchmarks) get the same shapes production would pick. - # The production custom op keeps its own equivalent inline policy for - # now; unifying it onto ``pick_config`` is a call-site change deferred - # to the dispatch-guard follow-up PR. + # The production custom op delegates here (``pick_cluster_size`` / + # ``pick_tuning``) — one policy, two shells. Intentional shell + # divergence: on a 32B-misaligned logits pointer the production runner + # ASSERTS (contract violation), while ``launch`` silently downgrades to + # 128-bit loads (dev convenience for ad-hoc tensors). _NUM_SMS: Optional[int] = None _LAUNCH_CACHE: dict = {} @@ -4811,69 +5092,50 @@ def _device_num_sms() -> int: return GvrTopKKernel._NUM_SMS @staticmethod - def pick_config( + def pick_cluster_size(num_rows: int, n_row: int, num_sms: int) -> int: + """Cluster-size policy: N < 64K -> 1 (sync unrecouped); tiny grid + at large N -> 8; single-wave -> 4/2; multi-wave -> 1 (row + parallelism already saturates the SMs; per-row splitting is pure + overhead past one wave).""" + if n_row < 65536: + return 1 + if num_rows <= 4 and n_row >= 131072: + return 8 + if num_rows * 4 <= num_sms: + return 4 + if num_rows * 2 <= num_sms: + return 2 + return 1 + + @staticmethod + def pick_tuning( torch_dtype, num_rows: int, - num_candidates: int, - max_seq_len: Optional[int] = None, - num_sms: Optional[int] = None, + n_per_cta: int, + num_sms: int, + graph_capture: bool, ) -> dict: - """Pick the launch-shape ctor kwargs for ``(dtype, BS, N)``. - - Mirrors the production runner policy (cluster_size auto-pick + - ``_pick_tuning``) so any caller instantiating the kernel directly - gets the same shapes the custom op would use. Rationale (B200, - nsys cold-L2, 2026-07-15 big-BS triage): a config frozen at the - BS=1 optimum (cs = N>=65536 ? 4 : 1, T=1024, mbpm=1) is geomean - 2.27x slower (max 6.0x) than the op-bench anchor at BS in - {64, 256, 1024}, while this policy is 0.95x (parity/better). - Multi-CTA splitting only pays while the grid is a single wave - (num_rows * cluster_size <= num_sms); past that, row parallelism - already saturates the SMs and per-row splitting is pure overhead. + """T / V / min_blocks_per_mp / warp-reduce policy at a given + per-CTA row width (cluster split already applied). - ``max_seq_len``: pass the peak runtime N under CUDA-graph capture - so the variant is picked for the replay shape, not the capture - shape (same contract as the custom op's ``_pick_tuning``). - - Returns kwargs for ``GvrTopKKernel(...)``: ``cluster_size``, - ``num_threads``, ``use_256bit_load``, ``min_blocks_per_mp``, - ``enable_warp_parallel_reduce``. + ``graph_capture``: raise the half-prec T=1024 bar so a small + capture-N does not pin T=1024 onto small-N replays. + Returns ``num_threads``, ``use_256bit_load``, + ``min_blocks_per_mp``, ``enable_warp_parallel_reduce``. """ import torch # local: keep the module importable without torch - if num_sms is None: - num_sms = GvrTopKKernel._device_num_sms() - n_row = max_seq_len if max_seq_len is not None else num_candidates is_fp32 = torch_dtype == torch.float32 - - # cluster_size: B200 SXM5 synth-data tuning (matches the custom - # op's auto-pick): N < 64K -> 1 (sync unrecouped); tiny grid at - # large N -> 8; single-wave -> 4/2; multi-wave -> 1. - if n_row < 65536: - cluster_size = 1 - elif num_rows <= 4 and n_row >= 131072: - cluster_size = 8 - elif num_rows * 4 <= num_sms: - cluster_size = 4 - elif num_rows * 2 <= num_sms: - cluster_size = 2 - else: - cluster_size = 1 - - # Cluster CTAs split the row, so tuning targets per-CTA work. - n_per_cta = n_row // cluster_size - # T=1024 needs 1 CTA/SM grid AND enough per-CTA vec work. Under - # graph capture, raise the half-prec bar so a small capture-N - # doesn't force T=1024 on small-N replays. - n_thresh_t = 131072 if (max_seq_len is not None and not is_fp32) else 65536 + # T=1024 needs a 1 CTA/SM grid AND enough per-CTA vec work. + n_thresh_t = 131072 if (graph_capture and not is_fp32) else 65536 num_threads = 1024 if (num_rows <= num_sms and n_per_cta >= n_thresh_t) else 512 # V=256-bit only helps fp32 at large N; half-prec cvt doubles reg - # pressure. Caller must hand a 32B-aligned contiguous tensor - # (``launch`` downgrades on misalignment). + # pressure. Requires a 32B-aligned contiguous tensor (see the + # shell-divergence note above). use_256bit_load = is_fp32 and n_per_cta >= 16384 enable_warp_parallel_reduce = num_threads == 1024 - # min_blocks_per_mp: reg-vs-occupancy 3-tier (fp32 wants ~70 regs + # min_blocks_per_mp: reg-vs-occupancy tiers (fp32 wants ~70 regs # for 4-LDG ILP -> mb<=2; half-prec fits 40 regs -> mb=3 packs # 3 CTA/SM when rows oversubscribe the device). vec_bits = 256 if use_256bit_load else 128 @@ -4897,13 +5159,44 @@ def pick_config( min_blocks_per_mp = 1 return dict( - cluster_size=cluster_size, num_threads=num_threads, use_256bit_load=use_256bit_load, min_blocks_per_mp=min_blocks_per_mp, enable_warp_parallel_reduce=enable_warp_parallel_reduce, ) + @staticmethod + def pick_config( + torch_dtype, + num_rows: int, + num_candidates: int, + max_seq_len: Optional[int] = None, + num_sms: Optional[int] = None, + ) -> dict: + """Launch-shape ctor kwargs for ``(dtype, BS, N)`` — the single + source of truth shared by the production runner + (``CuteDSLGvrTopKDecodeRunner``) and direct-drive users (tests, + benchmarks): composition of :meth:`pick_cluster_size` and + :meth:`pick_tuning`. + + ``max_seq_len``: pass the peak runtime N under CUDA-graph capture + so the variant is picked for the replay shape, not the capture + shape. + """ + if num_sms is None: + num_sms = GvrTopKKernel._device_num_sms() + n_row = max_seq_len if max_seq_len is not None else num_candidates + cluster_size = GvrTopKKernel.pick_cluster_size(num_rows, n_row, num_sms) + cfg = GvrTopKKernel.pick_tuning( + torch_dtype, + num_rows, + n_row // cluster_size, + num_sms, + graph_capture=max_seq_len is not None, + ) + cfg["cluster_size"] = cluster_size + return cfg + @classmethod def launch( cls, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py new file mode 100644 index 000000000000..5dbe90ced305 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py @@ -0,0 +1,603 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Direct (short-row) GVR top-K tier — CuTe DSL, Blackwell SM100. + +CuTe DSL translation of the original CUDA ``direct_topk_kernel``, +adapted for the production +``trtllm::cute_dsl_gvr_topk_decode`` contract (ragged N via a device +``seq_lens`` tensor + per-row degenerate identity emit; see the module +docstring of ``gvr_topk_decode_tp`` for the shared adaptation inventory). + +Exact top-K indices for short padded rows (npad <= DKCMAX = 12288), one CTA +per row, TB = 1024 threads. No threshold solve: the whole row is collected +into SMEM as packed (f2u order key << 32 | index) u64 candidates with an +in-flight 2048-bin radix histogram over the top 11 key bits, then an +11/11/10-bit radix select with whole-bin early exit and boundary-bin +compaction. Tie-aware emit (strict-greater mandatory slots + tie tickets +filling to K) makes the output value multiset equal torch.topk exactly. + +Ragged N: elements at index >= N_eff contribute key = f2u(-FLT_MAX) (the +minimum key) instead of their stale value; with the degenerate rows +(N_eff <= K) peeled off in-kernel, a masked element can never enter the +top-K, so out-of-range indices are never emitted. + +Public entry: ``direct_topk(logits, seq_lens, out, K)`` with torch tensors +(logits [BS, npad] fp32 contiguous, seq_lens [BS] int32, out [BS, K] int32). +Compiled variants are cached per (K, TB); BS and npad are dynamic (sym_int) +so one variant serves every shape. +""" + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import llvm +from cutlass.utils.smem_allocator import SmemAllocator + +FLT_MAX = 3.4028234663852886e38 +DKCMAX = 12288 # direct-path candidate capacity (mirrors the CUDA development arm) + + +@cute.jit +def _f2u(v): + """Order-preserving fp32 -> uint32 radix key (CUDA f2u). + + u ^ (sign ? 0xFFFFFFFF : 0x80000000): unsigned order of the result equals + fp32 order (NaN-free inputs). All shifts on the returned Uint32 are + logical, all compares unsigned. + """ + u = cutlass.Uint32(llvm.bitcast(cutlass.Uint32.mlir_type, v.ir_value())) + mask = (cutlass.Uint32(0) - (u >> cutlass.Uint32(31))) | cutlass.Uint32(0x80000000) + return u ^ mask + + +class DirectTopKKernel: + """One-CTA-per-row exact radix top-K for npad <= DKCMAX. + + Ctor knobs (compile-time): ``top_k`` in {512, 1024, 2048}, ``num_threads`` + (TB, default 1024), ``next_n`` / ``compress_ratio`` for the per-row + N_eff arithmetic (constexpr; the direct tier reads no hints, so MTP + support is the N_eff formula alone). SMEM layout + mirrors the CUDA development arm's DSmem with the same packed (key << 32 | idx) + u64 layout (measured: split key/idx arrays cost 2x smem + transactions in collect + emits vs CUDA's single STS.64 / LDS.64 per + candidate): + + cand[DKCMAX] u64 | side[SIDECAP] u64 (boundary-bin compaction) | + hist[2048] i32 | iwred[NWARP] i32 | sel[2] i32 | cnt[3] i32 + """ + + WARP_SIZE = 32 + + def __init__( + self, top_k: int, num_threads: int = 1024, next_n: int = 1, compress_ratio: int = 4 + ): + if top_k not in (512, 1024, 2048): + raise ValueError(f"unsupported top_k {top_k}") + if num_threads % 32 != 0: + raise ValueError("num_threads must be a multiple of 32") + self.top_k = top_k + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + # DCAP_RUNGS * TB / 2 in CUDA (side[] unions the dead ptcnt[]); the + # direct path never allocates ptcnt so we keep only side[]. + self.sidecap = 16 * num_threads // 2 + + # ------------------------------------------------------------------ + # Per-row valid length (ragged N) — mirrors the in-tree run_one_row + # arithmetic exactly (see gvr_topk_decode_tp._row_n_eff). + # ------------------------------------------------------------------ + @cute.jit + def _row_n_eff(self, seq_lens: cute.Tensor, row): + NN = cutlass.const_expr(self.next_n) + seq_len = seq_lens[row // cutlass.Int32(NN)] + actual_kv_len = seq_len - cutlass.Int32(NN) + (row % cutlass.Int32(NN)) + cutlass.Int32(1) + if cutlass.const_expr(self.compress_ratio == 1): + n_eff = actual_kv_len + else: + n_eff = actual_kv_len // cutlass.Int32(self.compress_ratio) + return n_eff + + # ------------------------------------------------------------------ + # Warp suffix-inclusive scan (higher lane = higher bins), the + # __shfl_down_sync ladder of the CUDA bin_select. + # ------------------------------------------------------------------ + @cute.jit + def _warp_suffix_scan(self, val, lane): + x = val + for d in cutlass.range_constexpr(5): + off = cutlass.const_expr(1 << d) + src = lane + cutlass.Int32(off) + if src > cutlass.Int32(31): + src = cutlass.Int32(31) + up = cute.arch.shuffle_sync(x, src) + if lane + cutlass.Int32(off) < cutlass.Int32(32): + x = x + up + return x + + # ------------------------------------------------------------------ + # bin_select: find bin b such that above(b) < want <= above(b) + # + hist[b] where above(b) = sum_{j>b} hist[j]. Writes smem_sel = + # {b, above(b)}. Trailing barrier. + # ------------------------------------------------------------------ + @cute.jit + def _bin_select( + self, + nbins: cutlass.Constexpr, + want, + smem_hist, + smem_iwred, + smem_sel, + tidx, + lane, + warp_id, + ): + per = cutlass.const_expr(nbins // self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + h = cute.make_fragment((per,), cutlass.Int32) + s_sum = cutlass.Int32(0) + for q in cutlass.range_constexpr(per): + h[q] = smem_hist[tidx * cutlass.Int32(per) + cutlass.Int32(q)] + s_sum = s_sum + h[q] + # warp suffix-inclusive over lanes (higher lane = higher bins) + x = self._warp_suffix_scan(s_sum, lane) + if lane == cutlass.Int32(0): + smem_iwred[warp_id] = x # warp total + cute.arch.barrier() + if tidx < cutlass.Int32(32): + wt = cutlass.Int32(0) + if tidx < cutlass.Int32(num_warps): + wt = smem_iwred[tidx] + wx = self._warp_suffix_scan(wt, tidx) + if tidx < cutlass.Int32(num_warps): + smem_iwred[tidx] = wx - wt # strictly-above-warp sum + cute.arch.barrier() + a_cnt = smem_iwred[warp_id] + (x - s_sum) # bins strictly above chunk + if a_cnt < want and want <= a_cnt + s_sum: + run = a_cnt + for qr in cutlass.range_constexpr(per): + q = cutlass.const_expr(per - 1 - qr) + if run < want and want <= run + h[q]: + smem_sel[0] = tidx * cutlass.Int32(per) + cutlass.Int32(q) + smem_sel[1] = run + run = run + h[q] + cute.arch.barrier() + + # ------------------------------------------------------------------ + # Warp-aggregated slot allocation: ballot + leader atomicAdd + popc. + # Whole-warp uniform trip counts are guaranteed by Cpad % 32 == 0 and + # TB % 32 == 0 (same argument as the CUDA original). + # ------------------------------------------------------------------ + @cute.jit + def _emit_prefix_ge(self, cand_cnt, shift, prefix, smem_cand, smem_cnt, out_row, tidx, lane): + """Emit indices of all candidates with (key >> shift) >= prefix; + the caller guarantees exactly K of them.""" + num_threads = cutlass.const_expr(self.num_threads) + if tidx == cutlass.Int32(0): + smem_cnt[0] = cutlass.Int32(0) + cute.arch.barrier() + cpad = (cand_cnt + cutlass.Int32(31)) & cutlass.Int32(-32) + i = tidx + while i < cpad: + valid = i < cand_cnt + kv = cutlass.Uint64(0) + if valid: + kv = smem_cand[i] + key = cutlass.Uint32(kv >> cutlass.Uint64(32)) + e = valid and ((key >> shift) >= prefix) + bal = cute.arch.vote_ballot_sync(e) + if bal != cutlass.Uint32(0): + tz = cutlass.Int32( + cute.arch.popc((bal & (cutlass.Uint32(0) - bal)) - cutlass.Uint32(1)) + ) + base = cutlass.Int32(0) + if lane == tz: + base = cute.arch.atomic_add( + smem_cnt.iterator, cutlass.Int32(cute.arch.popc(bal)), scope="cta" + ) + base = cute.arch.shuffle_sync(base, tz) + if e: + rk = cutlass.Int32( + cute.arch.popc(bal & cutlass.Uint32(cute.arch.lanemask_lt())) + ) + out_row[base + rk] = cutlass.Int32( + cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF)) + ) + i = i + cutlass.Int32(num_threads) + + @cute.jit + def _emit_final(self, cand_cnt, kth, m, nt, smem_cand, smem_cnt, out_row, tidx, lane): + """keys > kth are mandatory (slots [0, m)); keys == kth fill the + remaining nt tie tickets (slots [m, m+nt)).""" + num_threads = cutlass.const_expr(self.num_threads) + if tidx == cutlass.Int32(0): + smem_cnt[0] = cutlass.Int32(0) + smem_cnt[1] = cutlass.Int32(0) + cute.arch.barrier() + cpad = (cand_cnt + cutlass.Int32(31)) & cutlass.Int32(-32) + i = tidx + while i < cpad: + valid = i < cand_cnt + kv = cutlass.Uint64(0) + if valid: + kv = smem_cand[i] + key = cutlass.Uint32(kv >> cutlass.Uint64(32)) + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + man = valid and (key > kth) + tie = valid and (key == kth) + bm = cute.arch.vote_ballot_sync(man) + bt = cute.arch.vote_ballot_sync(tie) + if bm != cutlass.Uint32(0): + tzm = cutlass.Int32( + cute.arch.popc((bm & (cutlass.Uint32(0) - bm)) - cutlass.Uint32(1)) + ) + base_m = cutlass.Int32(0) + if lane == tzm: + base_m = cute.arch.atomic_add( + smem_cnt.iterator, cutlass.Int32(cute.arch.popc(bm)), scope="cta" + ) + base_m = cute.arch.shuffle_sync(base_m, tzm) + if man: + rk = cutlass.Int32(cute.arch.popc(bm & cutlass.Uint32(cute.arch.lanemask_lt()))) + out_row[base_m + rk] = idx + if bt != cutlass.Uint32(0): + tzt = cutlass.Int32( + cute.arch.popc((bt & (cutlass.Uint32(0) - bt)) - cutlass.Uint32(1)) + ) + base_t = cutlass.Int32(0) + if lane == tzt: + base_t = cute.arch.atomic_add( + smem_cnt.iterator + 1, cutlass.Int32(cute.arch.popc(bt)), scope="cta" + ) + base_t = cute.arch.shuffle_sync(base_t, tzt) + if tie: + p = base_t + cutlass.Int32( + cute.arch.popc(bt & cutlass.Uint32(cute.arch.lanemask_lt())) + ) + if p < nt: + out_row[m + p] = idx + i = i + cutlass.Int32(num_threads) + + # ------------------------------------------------------------------ + # Kernel + # ------------------------------------------------------------------ + @cute.kernel + def direct_topk_kernel( + self, + logits: cute.Tensor, # [BS, npad] fp32; tail beyond N_eff may be garbage + seq_lens: cute.Tensor, # [BS] int32 (uncompressed-token space) + out: cute.Tensor, # [BS, K] int32 + ): + num_threads = cutlass.const_expr(self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + kK = cutlass.const_expr(self.top_k) + sidecap = cutlass.const_expr(self.sidecap) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + lane = tidx & cutlass.Int32(31) + warp_id = tidx // cutlass.Int32(32) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[bidx, None] + out_row = out[bidx, None] + + # Ragged N: per-row valid length from seq_lens. + n_eff = self._row_n_eff(seq_lens, bidx) + + # ---- SMEM (allocated unconditionally, before the dynamic branch, + # so the DSL sizes the launch SMEM identically on every path) ---- + smem = SmemAllocator() + s_cand = smem.allocate_tensor( + cutlass.Uint64, cute.make_layout((DKCMAX,)), byte_alignment=128 + ) + s_side = smem.allocate_tensor( + cutlass.Uint64, cute.make_layout((sidecap,)), byte_alignment=128 + ) + s_hist = smem.allocate_tensor(cutlass.Int32, cute.make_layout((2048,)), byte_alignment=128) + s_iwred = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((num_warps,)), byte_alignment=128 + ) + s_sel = smem.allocate_tensor(cutlass.Int32, cute.make_layout((2,)), byte_alignment=16) + s_cnt = smem.allocate_tensor(cutlass.Int32, cute.make_layout((3,)), byte_alignment=16) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad ---- + # (mirrors the in-tree kernel; CuTe DSL has no runtime return, so the + # main body lives in the else branch). + if n_eff <= cutlass.Int32(kK): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(num_threads) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(kK): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(num_threads) + else: + # ---- Collect: vectorized load + f2u keys + in-flight 2048-bin hist + for z in cutlass.range_constexpr(2048 // num_threads): + s_hist[tidx + cutlass.Int32(z * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128, invariant=True + ) + row_addr = logits_row.iterator.toint() + vcnt = npad >> cutlass.Int32(2) # npad % 4 == 0 (host-asserted) + # npad <= DKCMAX bounds the per-thread trip count + # at U = DKCMAX/4/TB (= 3 for TB 1024), so the whole grid-stride + # loop is flattened into one predicated register batch: issue ALL + # (<=3) float4 loads back-to-back, then consume. A dynamic + # `cutlass.range(iters, unroll=4)` loop never reaches its unrolled + # body (trip <= 3) and leaves ONE load in flight per iteration — + # 1.8x the long-scoreboard stall of the CUDA arm and a stable + # ~1.03 cold-kernel gap at npad 8256. Keep the flat batch. + n_batch = cutlass.const_expr((DKCMAX // 4 + num_threads - 1) // num_threads) + frags = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(n_batch) + ] # Python-unrolled register batch + for u in cutlass.range_constexpr(n_batch): + i_vec = tidx + cutlass.Int32(u * num_threads) + if i_vec < vcnt: + src_ptr = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(i_vec) * cutlass.Int64(16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + src = cute.make_tensor(src_ptr, cute.make_layout((4,))) + cute.copy(copy_atom, src, frags[u]) + for u in cutlass.range_constexpr(n_batch): + i_vec = tidx + cutlass.Int32(u * num_threads) + if i_vec < vcnt: + gi = i_vec * cutlass.Int32(4) + for q in cutlass.range_constexpr(4): + # Ragged N: stale values beyond N_eff contribute the + # minimum key f2u(-FLT_MAX); with n_eff > K here they + # can never enter the top-K, so their (out-of-range) + # indices are never emitted. + v = frags[u][q] + if gi + cutlass.Int32(q) >= n_eff: + v = cutlass.Float32(-FLT_MAX) + key = _f2u(v) + s_cand[gi + cutlass.Int32(q)] = ( + cutlass.Uint64(key) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(gi + cutlass.Int32(q))) + cute.arch.atomic_add( + s_hist.iterator + cutlass.Int32(key >> cutlass.Uint32(21)), + cutlass.Int32(1), + scope="cta", + ) + cute.arch.barrier() + + # ---- radix_select_emit (11/11/10-bit, whole-bin early exit) ---- + if npad == cutlass.Int32(kK): + # C == K: every candidate is admitted; identity emit. (With + # seq_lens present this is unreachable — n_eff <= npad == K + # lands in the degenerate branch — but kept for structural + # parity with the CUDA source.) + io = tidx + while io < npad: + out_row[io] = cutlass.Int32( + cutlass.Uint32(s_cand[io] & cutlass.Uint64(0xFFFFFFFF)) + ) + io = io + cutlass.Int32(num_threads) + else: + want = cutlass.Int32(kK) + m = cutlass.Int32(0) + # level 0: top 11 bits (hist prebuilt during collect) + self._bin_select(2048, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b0 = s_sel[0] + a0 = s_sel[1] + h0 = s_hist[b0] + if want == a0 + h0: + # k-th boundary == bin edge: whole bin admitted + self._emit_prefix_ge( + npad, + cutlass.Uint32(21), + cutlass.Uint32(b0), + s_cand, + s_cnt, + out_row, + tidx, + lane, + ) + else: + m = m + a0 + want = want - a0 + docompact = h0 <= cutlass.Int32(sidecap) + cute.arch.barrier() # all reads of hist/sel done before reuse + if tidx == cutlass.Int32(0): + s_cnt[2] = cutlass.Int32(0) + for z1 in cutlass.range_constexpr(2048 // num_threads): + s_hist[tidx + cutlass.Int32(z1 * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + # level 1 sweep: mid 11 bits of boundary-bin members; compact + ub0 = cutlass.Uint32(b0) + i1 = tidx + while i1 < npad: + kv1 = s_cand[i1] + key1 = cutlass.Uint32(kv1 >> cutlass.Uint64(32)) + if (key1 >> cutlass.Uint32(21)) == ub0: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32( + (key1 >> cutlass.Uint32(10)) & cutlass.Uint32(0x7FF) + ), + cutlass.Int32(1), + scope="cta", + ) + if docompact: + p1 = cute.arch.atomic_add( + s_cnt.iterator + 2, cutlass.Int32(1), scope="cta" + ) + s_side[p1] = kv1 + i1 = i1 + cutlass.Int32(num_threads) + cute.arch.barrier() + self._bin_select(2048, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b1 = s_sel[0] + a1 = s_sel[1] + h1 = s_hist[b1] + p01 = (ub0 << cutlass.Uint32(11)) | cutlass.Uint32(b1) + if want == a1 + h1: + self._emit_prefix_ge( + npad, + cutlass.Uint32(10), + p01, + s_cand, + s_cnt, + out_row, + tidx, + lane, + ) + else: + m = m + a1 + want = want - a1 + cute.arch.barrier() + for z2 in cutlass.range_constexpr(1024 // num_threads): + s_hist[tidx + cutlass.Int32(z2 * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + # level 2 sweep: low 10 bits + ub1 = cutlass.Uint32(b1) + if docompact: + i2 = tidx + while i2 < h0: + key2 = cutlass.Uint32(s_side[i2] >> cutlass.Uint64(32)) + if ((key2 >> cutlass.Uint32(10)) & cutlass.Uint32(0x7FF)) == ub1: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32(key2 & cutlass.Uint32(0x3FF)), + cutlass.Int32(1), + scope="cta", + ) + i2 = i2 + cutlass.Int32(num_threads) + else: + i3 = tidx + while i3 < npad: + key3 = cutlass.Uint32(s_cand[i3] >> cutlass.Uint64(32)) + if (key3 >> cutlass.Uint32(10)) == p01: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32(key3 & cutlass.Uint32(0x3FF)), + cutlass.Int32(1), + scope="cta", + ) + i3 = i3 + cutlass.Int32(num_threads) + cute.arch.barrier() + self._bin_select(1024, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b2 = s_sel[0] + a2 = s_sel[1] + m = m + a2 + want = want - a2 + kth = (p01 << cutlass.Uint32(10)) | cutlass.Uint32(b2) + self._emit_final(npad, kth, m, want, s_cand, s_cnt, out_row, tidx, lane) + + # ------------------------------------------------------------------ + # Host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__(self, logits: cute.Tensor, seq_lens: cute.Tensor, out: cute.Tensor, stream): + num_rows = logits.shape[0] + self.direct_topk_kernel(logits, seq_lens, out).launch( + grid=(num_rows, 1, 1), + block=(self.num_threads, 1, 1), + stream=stream, + min_blocks_per_mp=1, # __launch_bounds__(TB, 1) + ) + + +# --------------------------------------------------------------------------- +# torch-facing entry with a compiled-variant cache keyed (K, TB). +# --------------------------------------------------------------------------- +_COMPILE_CACHE: dict = {} + + +def _get_compiled(top_k: int, num_threads: int = 1024, next_n: int = 1, cr: int = 4): + key = (top_k, num_threads, next_n, cr) + compiled = _COMPILE_CACHE.get(key) + if compiled is None: + from cutlass.cute import runtime as _crt + + kernel = DirectTopKKernel( + top_k=top_k, num_threads=num_threads, next_n=next_n, compress_ratio=cr + ) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, top_k), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kernel, + logits_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +_CHECKED_SIGS: set = set() + + +def _check_contract(logits, seq_lens, out, K): + assert logits.dtype == torch.float32 and logits.is_contiguous() + assert out.dtype == torch.int32 and out.is_contiguous() + assert seq_lens.dtype == torch.int32 + bs, npad = logits.shape + assert npad <= DKCMAX, f"direct path requires npad <= {DKCMAX}, got {npad}" + assert npad % 4 == 0, f"npad must be a multiple of 4, got {npad}" + assert out.shape[0] == bs and out.shape[1] == K + assert logits.data_ptr() % 16 == 0 and out.data_ptr() % 16 == 0 + + +def direct_topk( + logits: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + next_n: int = 1, + cr: int = 4, +) -> None: + """Exact top-K indices of ``logits`` rows into ``out`` (direct path). + + logits: [BS, npad] fp32 contiguous; the tail beyond each row's N_eff + may be stale garbage (masked in-kernel). + seq_lens: [BS // next_n] int32, request-level, uncompressed-token space. + out: [BS, K] int32 contiguous. + K: 512 / 1024 / 2048. + + Contract checks run once per (shape, K) signature to keep the hot + launch path at bare tvm-ffi cost. + """ + sig = (logits.shape, out.shape, K, next_n, cr) + if sig not in _CHECKED_SIGS: + _check_contract(logits, seq_lens, out, K) + _CHECKED_SIGS.add(sig) + compiled = _COMPILE_CACHE.get((K, 1024, next_n, cr)) + if compiled is None: + compiled = _get_compiled(K, next_n=next_n, cr=cr) + compiled(logits, seq_lens, out) + + +__all__ = ["DirectTopKKernel", "direct_topk", "DKCMAX"] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py new file mode 100644 index 000000000000..06bfb73c17c4 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py @@ -0,0 +1,366 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tiered-GVR top-K dispatcher — routes ``cute_dsl_gvr_topk_decode`` calls +to the CuTe DSL GVR tiers (direct / reg / tp). + +Exact transcription of the original CUDA implementation's dispatch +(``gvr_topk_launch_batched``) +(``gvr_topk_launch_batched`` + ``launch_dense`` + ``launch_tp<512>``). +The gvr streaming tier (cs=16 fallback for npad > 262144) is intentionally +NOT ported: it is unreachable inside the deployment envelope +(npad <= 262144), which :func:`is_tiered_topk_supported` enforces. + +Dispatcher guard (anything else falls back to the in-tree +``GvrTopKKernel`` path in ``CuteDSLGvrTopKDecodeRunner.forward``): + dtype == fp32, next_n >= 1 (MTP; num_rows divisible by next_n), + compress_ratio in {1, 4}, counters is None, + K in {512, 1024, 2048}, npad <= 262144, npad % 64 == 0, contiguous + 16B-aligned tensors, and the routed tier's cluster size within the + queried hardware max (never silently degrade the cluster shape). +``order_row`` is accepted and IGNORED: it is the LJF scheduling hint the +caller (``dsa.py``) computes for the in-tree persistent kernel whenever +num_rows >= 2 * num_sms, and the GVR tiers launch per-row CTAs that the +hardware schedules directly — the permutation affects neither their +correctness nor their measured performance (the tier mesh was benched in +natural row order). Rejecting it would silently turn the tiers off for every +large batch in production. ``counters`` (LB mode) still falls back: the +LB partition contract belongs to the in-tree kernel. +pre_idx / seq_lens are request-level ([num_rows // next_n, K] / +[num_rows // next_n]) — the in-tree contract. Per-row degeneracy +(N_eff <= K) and ragged N are handled INSIDE the tiers, so the guard +needs no device sync. + +Env knobs (identical semantics to the CUDA development arm's static locals; +cached at first use like those; :func:`_reset_env_cache` re-reads for tests): + TRTLLM_GVR_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); + else the bs threshold at which the tp tier takes over. + TRTLLM_GVR_DENSE_BS same, for the dense (tb=1024) reg tiers. + TRTLLM_GVR_FALLBACK_BANDS + 0 -> disable the measured fallback-band table (the + tiers serve every guarded shape); unset/other -> active. + TRTLLM_GVR_TIERS_DISABLE any value other than unset/""/"0" -> the guard + rejects everything (kill switch: every call takes the + in-tree kernel path). +Malformed numeric values fail soft: a warning is logged once and the knob +falls back to unset (-1) instead of raising on the decode path. +""" + +import os + +import torch + +from tensorrt_llm.logger import logger + +from .gvr_topk_decode_direct import DKCMAX, direct_topk +from .gvr_topk_decode_reg import reg_topk +from .gvr_topk_decode_tp import tp_cluster_size, tp_topk +from .single_pass_multi_cta_radix_topk_cluster import _query_max_cluster_size + +_BIG = 1 << 30 + +_ENV = {} # cached thresholds, mirrors the CUDA static-local caching + + +def _env_threshold(name): + t = _ENV.get(name) + if t is None: + e = os.environ.get(name) + if e is None or e.strip() == "": + t = -1 + else: + try: + t = int(e.strip()) + except ValueError: + # Tuning override with a safe baked default: fail soft + # (treat as unset) instead of killing the decode path. + logger.warning( + f"{name}={e!r} is not an integer; ignoring the override " + f"and using the baked default." + ) + t = -1 + if t == 0: + t = _BIG + _ENV[name] = t + return t + + +def _env_flag(name): + t = _ENV.get(name) + if t is None: + t = os.environ.get(name, "").strip() not in ("", "0") + _ENV[name] = t + return t + + +def _reset_env_cache(): + _ENV.clear() + _DISPATCH_CACHE.clear() # routes depend on the env thresholds + _CAP_OK_CACHE.clear() # ditto (verdict embeds the routed tier) + + +def _thresholds(npad): + tpb = _env_threshold("TRTLLM_GVR_TP_BS") + if tpb < 0: + tpb = 256 if npad <= 20480 else (128 if npad < 32768 else 16) + dnb = _env_threshold("TRTLLM_GVR_DENSE_BS") + if dnb < 0: + dnb = 8 if npad >= 163840 else (64 if npad < 32768 else _BIG) + return tpb, dnb + + +# Measured fallback band table (full-grid calibration, 2026-07-28): +# (npad, bs) buckets where the in-tree kernel is faster than every GVR +# tier by >1.10x on at least one production layer (865 real decode cells +# x 11 BS, same-rep cold-L2 nsys pairs; to recalibrate, re-run that +# paired sweep and route every (npad, bs) bucket whose per-case floor +# drops below 0.909). These are the L2-resident mid-N shapes where the +# in-tree exact-count ladder admits a leaner candidate set and its +# row-slice cluster split keeps every CTA busy through P4. Routing them +# to the in-tree kernel caps the worst case at 1.10x while keeping the +# tier win elsewhere (full-grid gm 1.40 vs the in-tree head). +# Keys are the nearest power-of-two of npad; values are inclusive bs +# ranges. TRTLLM_GVR_FALLBACK_BANDS=0 disables the table (the tiers +# serve every guarded shape). +# 2026-07-29 recalibration: the 131072 band was extended down to bs=8 for +# TRUE 128K shapes only (see _BAND_LOW_BS_NPAD_MAX). At bs=8 the reg tier +# dips to 0.82-0.91x vs the in-tree kernel on 5 production layers at +# npad=131136, which the floor guarantee cannot admit; the same bucket at +# npad~163776 wins 1.36-1.60x, and those shapes only land in this bucket +# through the nearest-power-of-two rounding, so the extension excludes them. +_FALLBACK_BANDS = { + 8192: (256, 1 << 30), + 16384: (256, 1 << 30), + 32768: (16, 1 << 30), + 65536: (16, 1 << 30), + 131072: (8, 255), + 262144: (16, 127), +} + +# The bs=8 end of the 131072 band applies only up to this npad: shapes above +# it (e.g. npad 163776) round INTO the 131072 bucket but behave like the next +# tier band, where the reg tier is far ahead. They keep the calibrated +# bs>=16 routing. +_BAND_LOW_BS_NPAD_MAX = {131072: 147456} + + +def _in_fallback_band(bs: int, npad: int) -> bool: + if _env_threshold("TRTLLM_GVR_FALLBACK_BANDS") == _BIG: # "0" -> off + return False + up = 1 << max(npad - 1, 1).bit_length() # pow2 >= npad + np2 = up if 4 * npad >= 3 * up else up >> 1 # arithmetic-midpoint nearest + band = _FALLBACK_BANDS.get(np2) + if band is None: + return False + lo, hi = band + npad_max = _BAND_LOW_BS_NPAD_MAX.get(np2) + if npad_max is not None and npad > npad_max and lo < 16: + lo = 16 + return lo <= bs <= hi + + +# tier name -> (kind, params). reg params = (cs, tb, maxv, ar). +def route(bs: int, npad: int, K: int) -> str: + """Tier the CUDA gvr_topk_launch_batched would take. Returns + 'tp' | 'direct' | 'reg(cs=C,tb=T,maxv=M,ar=A)' | 'cluster(cs=16,tb=512)' + ('cluster' is unreachable through :func:`tiered_topk` — see module docstring).""" + tpb, dnb = _thresholds(npad) + if bs >= tpb: + return "tp" + if npad > DKCMAX and bs >= dnb: + # launch_dense table + if npad <= 20480: + return "reg(cs=1,tb=1024,maxv=5,ar=8)" + if npad < 32768: + return "reg(cs=1,tb=1024,maxv=8,ar=8)" + if npad <= 65536: + return "reg(cs=2,tb=1024,maxv=8,ar=8)" + if npad <= 131072: + return "reg(cs=4,tb=1024,maxv=8,ar=8)" + if npad <= 262144: + return "reg(cs=8,tb=1024,maxv=8,ar=8)" + return "cluster(cs=16,tb=512)" + # latency ladder + if npad <= DKCMAX: + return "direct" + if npad < 16384: + return "reg(cs=1,tb=512,maxv=8,ar=8)" + if npad < 32768: + return "reg(cs=4,tb=512,maxv=4,ar=8)" + if npad <= 49152: + return "reg(cs=8,tb=512,maxv=3,ar=8)" + if npad <= 65536: + return "reg(cs=8,tb=512,maxv=4,ar=8)" + if npad <= 131072: + return "reg(cs=8,tb=512,maxv=8,ar=8)" + if npad <= 163840: + if K >= 2048: + return "reg(cs=16,tb=512,maxv=5,ar=6)" + return "reg(cs=16,tb=512,maxv=5,ar=8)" + if npad <= 262144: + if K == 2048: + return "reg(cs=16,tb=512,maxv=8,ar=8)" + return "reg(cs=16,tb=512,maxv=8,ar=6)" + return "cluster(cs=16,tb=512)" + + +def _parse_reg(tier): + body = tier[tier.index("(") + 1 : -1] + d = dict(kv.split("=") for kv in body.split(",")) + return int(d["cs"]), int(d["tb"]), int(d["maxv"]), int(d["ar"]) + + +def route_cluster_size(bs: int, npad: int, K: int) -> int: + """Cluster size the routed tier would launch with (host-only helper for + the hardware cluster-cap guard).""" + tier = route(bs, npad, K) + if tier == "tp": + return tp_cluster_size(bs, npad) + if tier == "direct": + return 1 + if tier.startswith("reg"): + return _parse_reg(tier)[0] + return 16 # cluster + + +# Per-call route()+_parse_reg() string work costs ~2.5-3us host-submit wall +# on <20us kernels (measured). The routing decision is pure in (bs, npad, K) +# [env thresholds are cached at first use, like the CUDA static locals], so +# bind it once per key to a closure. +_DISPATCH_CACHE = {} # (bs, npad, K, next_n, cr) -> callable(logits, pre, seq_lens, out) +_CAP_OK_CACHE = {} # (bs, npad, K) -> routed tier's cluster size within hw max + + +def _bind(bs, npad, K, next_n, cr): + tier = route(bs, npad, K) + if tier == "tp": + + def fn(lg, pre, sl, out): + tp_topk(lg, pre, sl, out, K, next_n, cr) + elif tier == "direct": + + def fn(lg, pre, sl, out): + direct_topk(lg, sl, out, K, next_n, cr) + elif tier.startswith("cluster"): + raise ValueError( + f"the cluster tier is not ported (npad beyond the deployment " + f"envelope); is_tiered_topk_supported must gate this out (bs={bs}, " + f"npad={npad}, K={K})" + ) + else: + cs, tb, maxv, ar = _parse_reg(tier) + + def fn(lg, pre, sl, out): + reg_topk(lg, pre, sl, out, K, cs, tb, maxv, ar, next_n, cr) + + return fn + + +def is_tiered_topk_supported( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, + next_n: int, + compress_ratio: int, + order_row, + counters, +) -> bool: + """Host-only guard for the GVR tiers (no device sync; see module + docstring). Returns False -> caller uses the in-tree kernel.""" + if _env_flag("TRTLLM_GVR_TIERS_DISABLE"): + return False + if logits.dtype != torch.float32: + return False + if next_n < 1 or compress_ratio not in (1, 4): + return False + # order_row is accepted and ignored (in-tree scheduling hint; see the + # module docstring). counters (LB mode) keeps the in-tree path. + if counters is not None: + return False + if top_k not in (512, 1024, 2048): + return False + bs, npad = logits.shape + if bs % next_n != 0: + return False + n_req = bs // next_n + if npad > 262144 or npad % 64 != 0: + return False + if _in_fallback_band(bs, npad): + return False + if not ( + logits.is_contiguous() + and pre_idx.is_contiguous() + and output_indices.is_contiguous() + and seq_lens.is_contiguous() + ): + return False + if pre_idx.shape != (n_req, top_k) or output_indices.shape != (bs, top_k): + return False + if seq_lens.shape != (n_req,) or seq_lens.dtype != torch.int32: + return False + if pre_idx.dtype != torch.int32 or output_indices.dtype != torch.int32: + return False + if ( + logits.data_ptr() % 16 != 0 + or pre_idx.data_ptr() % 16 != 0 + or output_indices.data_ptr() % 16 != 0 + ): + return False + # Cluster cap: fall back to the in-tree kernel (dispatcher-level) rather + # than silently degrading the tier's cluster shape. The verdict is pure + # in (bs, npad, K) once the env thresholds are cached, and computing it + # re-runs route() + the _parse_reg string parse (the ~2.5-3us host cost + # _DISPATCH_CACHE exists to avoid) — so memoize it the same way. + key = (bs, npad, top_k) + ok = _CAP_OK_CACHE.get(key) + if ok is None: + ok = _CAP_OK_CACHE[key] = route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() + return ok + + +def tiered_topk( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, + next_n: int = 1, + compress_ratio: int = 4, +) -> None: + """Unified GVR tier dispatch, replicating gvr_topk_launch_batched. + + logits [BS, npad] fp32 (BS = num_requests * next_n; npad multiple of + 64; per-row tail beyond N_eff may be garbage — masked in-kernel), + pre_idx [BS // next_n, K] int32 (request-level), seq_lens + [BS // next_n] int32 (request-level, uncompressed-token space), + output_indices [BS, K] int32. Caller must have passed + :func:`is_tiered_topk_supported`. + """ + bs, npad = logits.shape + key = (bs, npad, top_k, next_n, compress_ratio) + fn = _DISPATCH_CACHE.get(key) + if fn is None: + fn = _DISPATCH_CACHE[key] = _bind(bs, npad, top_k, next_n, compress_ratio) + fn(logits, pre_idx, seq_lens, output_indices) + + +__all__ = [ + "tiered_topk", + "is_tiered_topk_supported", + "route", + "route_cluster_size", + "_reset_env_cache", +] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py new file mode 100644 index 000000000000..472edfb0ef3d --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py @@ -0,0 +1,909 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Register-resident (reg) GVR Top-K tier — CuTe DSL, Blackwell SM100. + +CuTe DSL translation of the CUDA +``gvr_topk_reg`` register-resident GVR top-K kernel +(tuned CUDA head), adapted for the production +``trtllm::cute_dsl_gvr_topk_decode`` contract (see the module docstring of +``gvr_topk_decode_tp`` for the shared adaptation inventory: ragged-N +masking, pre_idx clamping, per-row degenerate identity emit). + +The row is loaded ONCE from GMEM into per-thread registers (MAXV float4s per +thread; out-of-slice lanes AND lanes beyond the row's N_eff are filled +-FLT_MAX — the ragged-N adaptation extends the existing OOR-lane idiom, so a +single mask at load time covers every downstream count / max-below / plateau +emit / collect over the register array). All subsequent passes are pure ALU — +global traffic pinned at 4*npad bytes/row. + +Phase skeleton (shared with the tp tier — helpers REUSED from +``gvr_topk_decode_tp``): + P1 : hint gather + two-stage 64-bin histogram -> rung ladder (tp.phase1) + P2 : full AR-rung register count + cluster exchange, secant refine / + max-below plateau descent (count_reg here) + P3 : rank-scatter collect (NO atomics): cached per-thread counts -> warp + inclusive scan + block scan + cluster-peer prefix -> deterministic + scatter position; keys f2u'd at push; P4's round-0 radix histogram is + built IN FLIGHT with LOCAL smem atomics, non-rank0 CTAs merge their + histograms into CTA0 with <=256 remote atomics. + P4 : CTA0-solo radix select starting directly at bin-select for round 0, + then tie-aware ticketed emit. + +``__launch_bounds__(TB, 1)`` -> ``.launch(min_blocks_per_mp=1)``. +CS=16 relies on CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, which the +CuTe DSL sets unconditionally on every kernel — the dispatcher additionally +verifies the queried hardware max cluster size and falls back to the in-tree +kernel when the tier's CS exceeds it (never degrades silently). + +CRITICAL: CuTe DSL cluster kernels need EXPLICIT trailing cluster +rendezvous where remote (DSMEM) ops can still be in flight — nvcc inserts an +implicit cluster barrier before ret in DSMEM kernels, the DSL does not +(missing it => timing-dependent CUDA 719 faults). + +Convergence notes preserved in this port (do not "simplify away"): +1. ``cutlass.utils.distributed.atomicAdd`` lowers to sys-scoped + ``atom.relaxed.SYS.shared`` which blocks ptxas' warp aggregation; the + CTA-scope relaxed ``_atomic_add_cta`` restores the fast path. +2. ``cute.arch.cluster_arrive/wait`` emit the NON-aligned barrier forms; + ``cg::cluster.sync()`` parity needs the ALIGNED pair + (``_cluster_sync_aligned``). +3. Candidates are a single packed (key<<32|idx) u64 array — one 8B DSMEM + push per candidate (halves remote-store transactions at CS>1). +""" + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import nvvm +from cutlass.cute import runtime as _crt +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from .gvr_topk_decode_tp import ( + FLT_MAX, + INF, + MAXPASS, + RUNGS, + GvrTpKernel, + _atom_shared_cluster_add_i32, + _exp2i, + _f2u_bits, + _f32_bits_u32, + _ld_shared_cluster_f32, + _ld_shared_cluster_i32, + _mapa_shared_cluster, + _shfl_down_add, + _shfl_up_add, + _st_shared_cluster_u64, +) + + +@dsl_user_op +def _atomic_add_cta(dst_ptr, val, *, loc=None, ip=None): + """CTA-scope relaxed atomicAdd on shared memory; returns the old value. + Drop-in for cutlass.utils.distributed.atomicAdd (which is sys-scoped and + blocks ptxas warp-aggregation — see module docstring note 1).""" + return cute.arch.atomic_add( + ptr=dst_ptr.llvm_ptr, val=val, sem="relaxed", scope="cta", loc=loc, ip=ip + ) + + +@dsl_user_op +def _cluster_arrive_aligned(*, loc=None, ip=None): + nvvm.cluster_arrive(aligned=True, loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_wait_aligned(*, loc=None, ip=None): + nvvm.cluster_wait(aligned=True, loc=loc, ip=ip) + + +@cute.jit +def _cluster_sync_aligned(): + _cluster_arrive_aligned() + _cluster_wait_aligned() + + +class GvrRegKernel(GvrTpKernel): + """CuTe DSL port of gvr_topk_reg (fp32, B200/B300). + + Inherits phase1 (already N_eff-aware), the DSMEM/scan idioms and + ``_row_n_eff`` from the tp port (identical CUDA source); adds + register-resident count/max-below/collect and the rank-scatter + + in-flight-histogram P3/P4. + """ + + def __init__( + self, + top_k: int, + kC: int, + cluster_size: int = 1, + ar: int = RUNGS, + maxv: int = 8, + num_threads: int = 512, + next_n: int = 1, + compress_ratio: int = 4, + ): + assert num_threads % 32 == 0 + assert ar in (6, 8) + assert cluster_size in (1, 2, 4, 8, 16) + assert 1 <= maxv <= 8 + self.top_k = top_k + self.kC = kC + self.cluster_size = cluster_size + self.ar = ar + self.uf = 4 # unused by the reg path; kept for inherited helpers + self.maxv = maxv + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + + # ------------------------------------------------------------------ + # count_reg: R-rung count over the register array. + # Per-thread counts -> s_ptcnt[r*TB + tid] (P3 rank-scatter offsets). + # Dummy (-FLT_MAX) lanes — both out-of-slice AND ragged-N-masked — + # only pass a rung == -FLT_MAX, which can never be `chosen` (its count + # > kC by construction since npad > kC). + # ------------------------------------------------------------------ + @cute.jit + def count_reg(self, R: cutlass.Constexpr, a, tidx, s_rungs, s_ptcnt): + TB = cutlass.const_expr(self.num_threads) + MAXV = cutlass.const_expr(self.maxv) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + for u in cutlass.range_constexpr(MAXV): + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # exchange_counts override: identical to the tp tier's, but the cluster + # sync is the ALIGNED barrier pair (cg::cluster.sync() parity — module + # docstring note 2). The tp tier keeps the non-aligned form it was + # validated with. + # ------------------------------------------------------------------ + @cute.jit + def exchange_counts( + self, R: cutlass.Constexpr, par, tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + cute.arch.barrier() # ptcnt final + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if wid < cutlass.Int32(R): + s = cutlass.Int32(0) + for k in cutlass.range_constexpr(TB // 32): + s = s + s_ptcnt[wid * TB + lane + cutlass.Int32(32 * k)] + s = cute.arch.warp_redux_sync(s, "add") + if cutlass.const_expr(CS == 1): + if lane == cutlass.Int32(0): + s_rcnt[wid] = s + s_rpre[wid] = cutlass.Int32(0) + else: + if lane == cutlass.Int32(0): + s_ipartial[par * cutlass.Int32(RUNGS) + wid] = s + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + else: + _cluster_sync_aligned() + if tidx < cutlass.Int32(R): + tot = cutlass.Int32(0) + pre = cutlass.Int32(0) + local_ptr = s_ipartial.iterator + (par * cutlass.Int32(RUNGS) + tidx) + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + v = _ld_shared_cluster_i32(a) + tot = tot + v + if cutlass.Int32(rr) < rank: + pre = pre + v + s_rcnt[tidx] = tot + s_rpre[tidx] = pre + cute.arch.barrier() + + # ------------------------------------------------------------------ + # max_below_reg: largest value strictly below t_hi_bound across the + # whole row (cluster-reduced). No index gate needed: dummy lanes hold + # -FLT_MAX which never raises the max. + # ------------------------------------------------------------------ + @cute.jit + def max_below_reg(self, a, t_hi_bound, par, tidx, s_fwred, s_fpartial): + CS = cutlass.const_expr(self.cluster_size) + MAXV = cutlass.const_expr(self.maxv) + NWARP = cutlass.const_expr(self.num_warps) + m = cutlass.Float32(-FLT_MAX) + for u in cutlass.range_constexpr(MAXV): + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + m = cute.arch.warp_redux_sync(m, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = m + cute.arch.barrier() + if tidx == cutlass.Int32(0): + mm = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + mm = cute.arch.fmax(mm, s_fwred[w]) + s_fpartial[par] = mm + res = cutlass.Float32(-FLT_MAX) + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + res = s_fpartial[par] + else: + _cluster_sync_aligned() + local_ptr = s_fpartial.iterator + par + for rr in cutlass.range_constexpr(CS): + ad = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + res = cute.arch.fmax(res, _ld_shared_cluster_f32(ad)) + return res + + # ------------------------------------------------------------------ + # kernel + # ------------------------------------------------------------------ + @cute.kernel + def gvr_reg_kernel( + self, logits: cute.Tensor, pre_idx: cute.Tensor, seq_lens: cute.Tensor, out_idx: cute.Tensor + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + AR = cutlass.const_expr(self.ar) + MAXV = cutlass.const_expr(self.maxv) + K = cutlass.const_expr(self.top_k) + kC = cutlass.const_expr(self.kC) + NWARP = cutlass.const_expr(self.num_warps) + + if cutlass.const_expr(CS > 1): + row = bidx // cutlass.Int32(CS) + rank = cute.arch.block_idx_in_cluster() + else: + row = bidx + rank = cutlass.Int32(0) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[row, None] + # Hint sharing: pre_idx is request-level ([num_rows // next_n, K]) — + # see the tp tier for the in-tree run_one_row parity notes. + if cutlass.const_expr(self.next_n == 1): + pre_idx_row = pre_idx[row, None] + else: + pre_idx_row = pre_idx[row // cutlass.Int32(self.next_n), None] + out_row = out_idx[row, None] + row_addr = logits_row.iterator.toint() + + # Ragged N: per-row valid length from seq_lens (inherited helper). + n_eff = self._row_n_eff(seq_lens, row) + + # ---- shared memory (order must be identical across CTAs for mapa) ---- + smem = SmemAllocator() + # Single packed (key<<32 | idx) u64 candidate array — CUDA's + # `unsigned long long cand[kC]`; one 8B DSMEM push per candidate in + # the P3 scatter (halves remote-store transactions under 8-cluster + # contention at BS>=8 — measured). + s_cand = smem.allocate_tensor( + element_type=cutlass.Uint64, + layout=cute.make_ordered_layout((kC,), order=(0,)), + byte_alignment=128, + ) + s_ptcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS * self.num_threads,), order=(0,)), + byte_alignment=128, + ) + s_hist = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((256,), order=(0,)), + byte_alignment=128, + ) + s_rungs = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rpre = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_ipartial = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((2 * RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_fpartial = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + s_fwred = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2 * self.num_warps,), order=(0,)), + byte_alignment=64, + ) + s_hminmax = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + # iscalars: [0]=sel_bin [1]=sel_above [2]=sel_count [3]=cnt_m [4]=cnt_t + s_isc = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((8,), order=(0,)), + byte_alignment=32, + ) + # P3 rank-scatter block-scan scratch (reg path only) + s_iwred = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((self.num_warps,), order=(0,)), + byte_alignment=64, + ) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad. ---- + # Cluster-uniform branch (all CTAs own the same row); no DSMEM op is + # issued on this path, so no exit rendezvous is required either. + if n_eff <= cutlass.Int32(K): + if rank == cutlass.Int32(0): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(TB) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(K): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(TB) + else: + # ---- slice + one-time register row load (launcher: vpc <= MAXV*TB) + V4 = npad >> cutlass.Int32(2) + vpc = (V4 + cutlass.Int32(CS) - cutlass.Int32(1)) // cutlass.Int32(CS) + v0 = rank * vpc + v1 = v0 + vpc + if v1 > V4: + v1 = V4 + + copy_atom = self._copy_atom() + a = cute.make_fragment((MAXV * 4,), cutlass.Float32) + frag4 = cute.make_fragment((4,), cutlass.Float32) + for u in cutlass.range_constexpr(MAXV): + i = v0 + tidx + cutlass.Int32(u * TB) + if i < v1: + self._ld_float4(copy_atom, row_addr, i, frag4) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + # Ragged N: extend the OOR-lane idiom to the valid + # length — elements at global index >= N_eff become + # -FLT_MAX dummy lanes, exactly like out-of-slice + # lanes. One mask here covers ALL downstream register + # passes (count/max-below/plateau emit/P3 collect). + val = frag4[q] + if gi + cutlass.Int32(q) >= n_eff: + val = cutlass.Float32(-FLT_MAX) + a[4 * u + q] = val + else: + for q in cutlass.range_constexpr(4): + a[4 * u + q] = cutlass.Float32(-FLT_MAX) + + xch = cutlass.Int32(0) + thr = cutlass.Float32(0.0) + chosen = cutlass.Int32(-1) + C = cutlass.Int32(0) + cbase = cutlass.Int32(0) + m_gt = cutlass.Int32(-1) + + # cr==1 hint temporal shift (in-tree pre_idx_offset parity); the + # cr>1 branch keeps the exact pre-MTP trace — see the tp tier. + if cutlass.const_expr(self.compress_ratio == 1): + hint_off = (row % cutlass.Int32(self.next_n)) + cutlass.Int32(1) + self.phase1( + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=hint_off, + ) + else: + self.phase1( + logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs + ) + # Pre-zero the P4 round-0 radix histogram now that P1 is done with + # hist. The first exchange's cluster barrier (release) publishes it + # before any CTA reaches the P3 collect that increments it in flight. + if tidx < cutlass.Int32(256): + s_hist[tidx] = cutlass.Int32(0) + self.count_reg(AR, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + + span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) + + # ---- P2: secant refine driver (redundant on every thread) ---- + t_lo = cutlass.Float32(-FLT_MAX) + t_hi = cutlass.Float32(INF) + c_hi = cutlass.Int32(0) + Rcur = cutlass.Int32(AR) + passno = cutlass.Int32(0) + running = cutlass.Int32(1) + descend_break = cutlass.Int32(0) + while running != cutlass.Int32(0): + # first rung index j with rcnt[j] >= K (rcnt ascending in j) + j = Rcur + for r_ in range(AR - 1, -1, -1): + if cutlass.Int32(r_) < Rcur: + if s_rcnt[r_] >= cutlass.Int32(K): + j = cutlass.Int32(r_) + jj = j + if jj > Rcur - cutlass.Int32(1): + jj = Rcur - cutlass.Int32(1) + cj = s_rcnt[jj] + rj = s_rungs[jj] + bj = s_rpre[jj] + found = cutlass.Int32(0) + if j < Rcur: + if cj <= cutlass.Int32(kC): + found = cutlass.Int32(1) + if found != cutlass.Int32(0): + chosen = j + thr = rj + C = cj + cbase = bj + running = cutlass.Int32(0) + else: + if j < Rcur: + if rj >= t_lo: + t_lo = rj + if j > cutlass.Int32(0): + jm = j - cutlass.Int32(1) + if s_rungs[jm] <= t_hi: + t_hi = s_rungs[jm] + c_hi = s_rcnt[jm] + descend = cutlass.Int32(0) + if passno >= cutlass.Int32(MAXPASS): + descend = cutlass.Int32(1) + e3 = passno * cutlass.Int32(3) + if e3 > cutlass.Int32(24): + e3 = cutlass.Int32(24) + step = span0 * _exp2i(e3) + dt = cutlass.Float32(0.0) + mode = cutlass.Int32(2) # 0=up-ladder, 1=down-ladder, 2=secant + if t_hi == cutlass.Float32(INF): + mode = cutlass.Int32(0) + nr0 = t_lo + step * cutlass.Float32(float(1 << (AR - 1))) + if nr0 == cutlass.Float32(INF): + descend = cutlass.Int32(1) + else: + if t_lo == cutlass.Float32(-FLT_MAX): + mode = cutlass.Int32(1) + else: + dt = (t_hi - t_lo) * cutlass.Float32(1.0 / float(AR + 1)) + nr_last = t_hi - dt * cutlass.Float32(float(AR)) + nr_first = t_hi - dt + ok = cutlass.Int32(0) + if nr_last > t_lo: + if nr_first < t_hi: + ok = cutlass.Int32(1) + if ok == cutlass.Int32(0): + descend = cutlass.Int32(1) + if descend != cutlass.Int32(0): + running = cutlass.Int32(0) + descend_break = cutlass.Int32(1) + else: + cute.arch.barrier() + if tidx == cutlass.Int32(0): + for r_ in cutlass.range_constexpr(AR): + nrv = cutlass.Float32(0.0) + if mode == cutlass.Int32(0): + nrv = t_lo + step * cutlass.Float32(float(1 << (AR - 1 - r_))) + else: + if mode == cutlass.Int32(1): + nrv = t_hi - step * cutlass.Float32(float(1 << r_)) + else: + nrv = t_hi - dt * cutlass.Float32(float(r_ + 1)) + s_rungs[r_] = nrv + cute.arch.barrier() + self.count_reg(AR, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + passno = passno + cutlass.Int32(1) + Rcur = cutlass.Int32(AR) + + # ---- plateau descent (exact max-below stepping) ---- + if descend_break != cutlass.Int32(0): + pl = cutlass.Int32(1) + while pl != cutlass.Int32(0): + vstar = self.max_below_reg( + a, t_hi, xch & cutlass.Int32(1), tidx, s_fwred, s_fpartial + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_rungs[0] = vstar + cute.arch.barrier() + self.count_reg(1, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + c = s_rcnt[0] + okc = cutlass.Int32(0) + if c >= cutlass.Int32(K): + if c <= cutlass.Int32(kC): + okc = cutlass.Int32(1) + if okc != cutlass.Int32(0): + chosen = cutlass.Int32(0) + thr = vstar + C = c + cbase = s_rpre[0] + pl = cutlass.Int32(0) + else: + if c < cutlass.Int32(K): + t_hi = vstar + c_hi = c + else: + thr = vstar + m_gt = c_hi + pl = cutlass.Int32(0) + + if m_gt >= cutlass.Int32(0): + # ---- plateau direct emit from registers ---- + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) # cnt_m + s_isc[4] = cutlass.Int32(0) # cnt_t + if cutlass.const_expr(CS > 1): + _cluster_sync_aligned() + a_m = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(0)) + a_t = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(4), cutlass.Int32(0)) + else: + cute.arch.barrier() + nt = cutlass.Int32(K) - m_gt + for u in cutlass.range_constexpr(MAXV): + i = v0 + tidx + cutlass.Int32(u * TB) + if i < v1: + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v > thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_m, cutlass.Int32(1)) + else: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1) + ) + out_row[p] = gi + cutlass.Int32(q) + else: + if v == thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_t, cutlass.Int32(1)) + else: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m_gt + p] = gi + cutlass.Int32(q) + else: + # ---- P3: rank-scatter collect from registers (NO atomics) ---- + myc = s_ptcnt[chosen * cutlass.Int32(TB) + tidx] + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + incl = myc + for o in [1, 2, 4, 8, 16]: + incl = _shfl_up_add(incl, lane, o) + if lane == cutlass.Int32(31): + s_iwred[wid] = incl + cute.arch.barrier() + if wid == cutlass.Int32(0): + v_ = cutlass.Int32(0) + if lane < cutlass.Int32(NWARP): + v_ = s_iwred[lane] + iv = v_ + for o in [x for x in [1, 2, 4, 8, 16] if x < self.num_warps]: + iv = _shfl_up_add(iv, lane, o) + if lane < cutlass.Int32(NWARP): + s_iwred[lane] = iv - v_ + cute.arch.barrier() + pos = cbase + s_iwred[wid] + (incl - myc) + + # scatter: f2u key at push + IN-FLIGHT local P4 round-0 + # histogram. thr > -FLT_MAX whenever we get here (chosen count + # <= kC < npad), so dummy lanes never pass. + if cutlass.const_expr(CS > 1): + a_cand = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + for u in cutlass.range_constexpr(MAXV): + gi = (v0 + tidx + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v >= thr: + ux = _f2u_bits(_f32_bits_u32(v)) + _atomic_add_cta( + s_hist.iterator + cutlass.Int32(ux >> cutlass.Uint32(24)), + cutlass.Int32(1), + ) + kv = (cutlass.Uint64(ux) << cutlass.Uint64(32)) | cutlass.Uint64( + cutlass.Uint32(gi + cutlass.Int32(q)) + ) + if cutlass.const_expr(CS > 1): + _st_shared_cluster_u64(a_cand + pos * cutlass.Int32(8), kv) + else: + s_cand[pos] = kv + pos = pos + cutlass.Int32(1) + + # non-rank0: merge local round-0 histogram into CTA0 + # (<=256 remote atomics) + if cutlass.const_expr(CS > 1): + if rank != cutlass.Int32(0): + cute.arch.barrier() # local hist final + a_hist = _mapa_shared_cluster(s_hist.iterator, cutlass.Int32(0)) + b = cutlass.Int32(tidx) + while b < cutlass.Int32(256): + hv = s_hist[b] + if hv != cutlass.Int32(0): + _atom_shared_cluster_add_i32(a_hist + b * cutlass.Int32(4), hv) + b = b + cutlass.Int32(TB) + _cluster_sync_aligned() + else: + cute.arch.barrier() + + # ---- P4 (CTA0 solo when CS>1) ---- + if rank == cutlass.Int32(0): + if C == cutlass.Int32(K): + ie = cutlass.Int32(tidx) + while ie < C: + out_row[ie] = cutlass.Int32( + cutlass.Uint32(s_cand[ie] & cutlass.Uint64(0xFFFFFFFF)) + ) + ie = ie + cutlass.Int32(TB) + else: + # 4x8-bit radix select; keys ALREADY f2u'd; round-0 + # histogram already built in flight -> round 0 starts + # at bin-select. + pref = cutlass.Uint32(0) + want = cutlass.Int32(K) + m = cutlass.Int32(0) + final_shift = cutlass.Uint32(0) + active = cutlass.Int32(1) + for r_ in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * r_) + if active != cutlass.Int32(0): + if cutlass.const_expr(r_ > 0): + ih = cutlass.Int32(tidx) + while ih < C: + u = cutlass.Uint32(s_cand[ih] >> cutlass.Uint64(32)) + if (u >> cutlass.Uint32(shift + 8)) == pref: + bb = cutlass.Int32( + (u >> cutlass.Uint32(shift)) & cutlass.Uint32(0xFF) + ) + _atomic_add_cta(s_hist.iterator + bb, cutlass.Int32(1)) + ih = ih + cutlass.Int32(TB) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + b8 = tidx * cutlass.Int32(8) + h = cute.make_fragment((8,), cutlass.Int32) + Ssum = cutlass.Int32(0) + for q in cutlass.range_constexpr(8): + h[q] = s_hist[b8 + cutlass.Int32(q)] + Ssum = Ssum + h[q] + s_hist[b8 + cutlass.Int32(q)] = cutlass.Int32(0) + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_down_add(x, tidx, o) + A = x - Ssum + if A < want: + if want <= A + Ssum: + run = A + for q in range(7, -1, -1): + if run < want: + if want <= run + h[q]: + s_isc[0] = b8 + cutlass.Int32(q) + s_isc[1] = run + s_isc[2] = h[q] + run = run + h[q] + cute.arch.barrier() + sel = s_isc[0] + above = s_isc[1] + m = m + above + want = want - above + pref = (pref << cutlass.Uint32(8)) | cutlass.Uint32(sel) + if s_isc[2] == want: + final_shift = cutlass.Uint32(shift) + active = cutlass.Int32(0) + kth = pref + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) + s_isc[4] = cutlass.Int32(0) + cute.arch.barrier() + nt = cutlass.Int32(K) - m + ie = cutlass.Int32(tidx) + while ie < C: + kv = s_cand[ie] + u = cutlass.Uint32(kv >> cutlass.Uint64(32)) >> final_shift + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + if u > kth: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1) + ) + out_row[p] = idx + else: + if u == kth: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m + p] = idx + ie = ie + cutlass.Int32(TB) + + # Exit rendezvous — PLATEAU PATH ONLY (wave-overlap fix). + # On the P3/P4 path the post-merge _cluster_sync_aligned() above is + # the LAST remote-op rendezvous (P4 is CTA0-local smem + global + # stores only), so non-rank0 CTAs exit right after it — mirroring + # CUDA's `if (rank != 0) return;` and freeing 15/16 SMs one + # P4-duration earlier (~+5us/wave at BS>=8 otherwise). The + # plateau-emit path still has remote atomics in flight up to its + # end (timing-dependent CUDA 719 territory without this barrier — + # the DSL emits NO implicit pre-ret cluster barrier) and therefore + # KEEPS the trailing rendezvous. m_gt is cluster-uniform + # (redundantly computed driver), so this is a uniform branch + # around the cluster barrier. Degenerate rows (outer branch) + # issue no remote ops, so skipping it there is safe too. + if cutlass.const_expr(CS > 1): + if m_gt >= cutlass.Int32(0): + _cluster_sync_aligned() + + # ------------------------------------------------------------------ + # host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + seq_lens: cute.Tensor, + out_idx: cute.Tensor, + stream, + ): + num_rows = logits.shape[0] + CS = cutlass.const_expr(self.cluster_size) + self.gvr_reg_kernel(logits, pre_idx, seq_lens, out_idx).launch( + grid=(num_rows * CS, 1, 1), + block=(self.num_threads, 1, 1), + cluster=(CS, 1, 1) if cutlass.const_expr(CS > 1) else None, + stream=stream, + min_blocks_per_mp=1, + ) + + +# --------------------------------------------------------------------------- +# compile cache + public entry (explicit variant params; dispatcher selects) +# --------------------------------------------------------------------------- +_LAUNCH_CACHE = {} + + +def _get_compiled(K: int, CS: int, TB: int, MAXV: int, AR: int, next_n: int = 1, cr: int = 4): + key = (K, CS, TB, MAXV, AR, next_n, cr) + compiled = _LAUNCH_CACHE.get(key) + if compiled is None: + kC = 8192 if K >= 2048 else 6144 + kern = GvrRegKernel( + top_k=K, + kC=kC, + cluster_size=CS, + ar=AR, + maxv=MAXV, + num_threads=TB, + next_n=next_n, + compress_ratio=cr, + ) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + # pre_idx is request-level (num_rows // next_n rows); keep the shared + # n_rows sym at next_n == 1 (identical compiled artifact to v1). + n_pre = n_rows if next_n == 1 else cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + pre_idx_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_pre, K), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kern, + logits_fake, + pre_idx_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _LAUNCH_CACHE[key] = compiled + return compiled + + +_FAST = {} # (K, cs, tb, maxv, ar, npad) -> compiled (contract checked once) + + +def reg_topk( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + cs: int, + tb: int, + maxv: int, + ar: int, + next_n: int = 1, + cr: int = 4, +) -> None: + """CuTe DSL gvr_topk_reg. logits [BS, npad] fp32 (npad + mult of 64; tail beyond each row's N_eff may be garbage — masked at the + register load), pre_idx [BS // next_n, K] int32 hint (request-level), + seq_lens [BS // next_n] int32 (uncompressed-token space), out [BS, K] + int32. Variant params are explicit — the dispatcher selects; compile + cache keyed (K, cs, tb, maxv, ar, next_n, cr).""" + npad = logits.shape[1] + key = (K, cs, tb, maxv, ar, npad, next_n, cr) + fn = _FAST.get(key) + if fn is None: + assert npad % 64 == 0, f"npad {npad} not a multiple of 64" + vpc = (npad // 4 + cs - 1) // cs + assert vpc <= maxv * tb, ( + f"slice {vpc} float4s exceeds MAXV*TB={maxv * tb} (npad={npad}, cs={cs})" + ) + kc = 8192 if K >= 2048 else 6144 + assert npad > kc, f"npad {npad} <= kC {kc}: reg path has no trivial branch" + fn = _get_compiled(K, cs, tb, maxv, ar, next_n, cr) + _FAST[key] = fn + fn(logits, pre_idx, seq_lens, out) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py new file mode 100644 index 000000000000..9c194a556546 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py @@ -0,0 +1,1949 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Throughput (tp) GVR Top-K tier — CuTe DSL, Blackwell SM100. + +CuTe DSL translation of the CUDA ``gvr_topk_tp`` +throughput GVR top-K kernel (tuned CUDA head), +adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract: + +* RAGGED N: production logits tails beyond the per-row valid length are + stale garbage (NOT -FLT_MAX pad, unlike the standalone bench harness). A device + ``seq_lens`` tensor is threaded into every tier and every global row read + is predicated at the LANE level: element index >= N_eff substitutes + -FLT_MAX AFTER the (always in-bounds, buffer width = logits.shape[1]) + float4 load. N_eff mirrors the in-tree ``run_one_row`` arithmetic: + ``(seq_lens[req] - next_n + row % next_n + 1) // compress_ratio``. +* pre_idx hardening: the P1 hint gather clamps out-of-range hint indices + into [0, N_eff-1] (production cold-start is all zeros; arbitrary garbage + must neither fault nor corrupt the rung ladder). +* Degenerate rows (N_eff <= K): per-row in-kernel identity emit + [0..N_eff-1] plus -1 padding to K, matching the in-tree kernel. + +Faithful phase-by-phase port otherwise: + P1 : hint gather + minmax + two-stage 64-bin histogram -> CCDF rung ladder + P2a : uniform every-32nd-float4 sampled multi-rung count (+ per-rung + float4 occupancy for a clustering-aware sigma) -> 4-stage pivot: + stage 0 is the hint-ladder ADMISSION (in-tree R0 parity) — + tightest rung whose sampled CI sits inside [K, 0.6*kC] (the + legacy pivot-band hi) — stage 0b prefers the legacy band pick + when it is strictly leaner and safe by its clustering-aware + 1.5-sigma lower CI, and the legacy 3-stage pick handles rows + where no rung qualifies + P2b : ONE fused streaming pass: exact counts at {pivot, rescue rung} + + optimistic collect of packed (key<<32|idx) u64 candidates + >= pivot into CTA0 smem (capped kC). Pivot count in [K, kC] => + the candidates are reused as-is (1-pass admission); pivot under K + but rescue in-window => ONE collect re-stream + P2c : multi-rung secant refine / max-below plateau descent when the pivot + count misses [K, kC] + P3 : candidate reuse (thr == tpush, no overflow) or one re-stream collect + P4 : CTA0-solo 4x8-bit radix select + tie-aware ticketed emit + (plus the trivial npad <= kC path and the plateau direct-emit path) + +Templates -> compile-time ctor knobs; compile cache keyed (K, CS, AR, UF, TB). +``__launch_bounds__(TB, 2)`` -> ``.launch(min_blocks_per_mp=2)``. +grid dim3(CS, BS) -> 1-D grid BS*CS with cluster=(CS,1,1); row = bidx // CS. + +DSMEM (CS > 1) via inline PTX: mapa.shared::cluster + ld/st.shared::cluster ++ atom.relaxed.cluster.shared::cluster.add.u32. Writer-side cluster syncs use +FULL cluster_arrive (release) — never relaxed (a relaxed arrive has no +release semantics, so peer CTAs could observe stale DSMEM: see +``cluster_arrive_relaxed`` DSMEM race observed in development). +""" + +import math + +import cutlass +import cutlass.cute as cute +import cutlass.cute.math as cmath +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute import runtime as _crt +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils.distributed import atomicAdd +from cutlass.utils.smem_allocator import SmemAllocator + +RUNGS = 8 +MAXPASS = 8 +SS = 32 # P2a sample stride (float4s) + +# DATA-ADAPTIVE admission (measured full-grid verdict). +# This is the baseline admission machinery +# BYTE-PARITY on every streaming pass, the fused R=2 count-collect, the +# reuse rule, the P4 select and the P2c driver (a dedicated-window-rung variant +# rungs and a later R=4 window variant both taxed the whole kernel 4-15% +# globally — even trivial-branch cells — via register pressure/extra +# compares/P4 prefilter, and the former additionally carried an adversarial +# under-emit bug). The adaptivity is a PICK-ONLY delta: +# * stage 0c LEAN-PIVOT OVERRIDE: when the stage-0 CI admission FAILED +# (the pick in hand is a band/2-sigma/overshoot fallback: fat or +# undershoot-prone), npad <= 262144, and a strictly LEANER ladder +# rung's sampled count lands in [K, kC] by its clustering-aware +# 1.5-sigma CI on both sides, that rung becomes the pivot (the +# rescue is recomputed from it by the standard next-fatter rule). +# A fat band pick (est ~3-4x K, the P3-push/P4-radix fatness the +# mechanism probe identified as the whole residual-band gap) is +# replaced by a lean CI-backed pivot; an undershoot costs ONE rescue +# re-stream — the baseline's own economics. +# * kC pinned to the flat budget 8192 (K>=2048) / 6144 +# (the K-scaled diet was falsified as pure harm). +# * ladder quantiles pinned WIDE (baseline): ablation showed the re-placed +# spread under this admission is a net residual harm. +# * the C4 occupancy CS cut keeps the fix-head default (neutral). +FLT_MAX = 3.4028234663852886e38 +INF = float("inf") + + +# --------------------------------------------------------------------------- +# DSMEM primitives (inline PTX) +# --------------------------------------------------------------------------- +@dsl_user_op +def _mapa_shared_cluster(smem_ptr, peer_rank, *, loc=None, ip=None): + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [smem_ptr_i32, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_i32(mapped_addr, *, loc=None, ip=None): + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.u32 $0, [$1];", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_f32(mapped_addr, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_shared_cluster_i32(mapped_addr, val, *, loc=None, ip=None): + llvm.inline_asm( + res=None, + operands_=[mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + asm_string="st.shared::cluster.u32 [$0], $1;", + constraints="r,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_shared_cluster_u64(mapped_addr, val, *, loc=None, ip=None): + """One 8B DSMEM candidate push (CUDA: dst[pos] = (u64)key<<32 | idx). + + Measured: a single packed u64 store halves remote-store + transactions vs two 4B pushes under CS>1 cluster contention (matches + the CUDA arm's ``unsigned long long cand[kC]``). Do NOT split back + into key/idx 4B stores — the split form measurably regresses the + K-scaled CS>1 band (pro/v32 BS16-64).""" + llvm.inline_asm( + res=None, + operands_=[mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + asm_string="st.shared::cluster.u64 [$0], $1;", + constraints="r,l", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_shared_cluster_add_i32(mapped_addr, val, *, loc=None, ip=None): + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.relaxed.cluster.shared::cluster.add.u32 $0, [$1], $2;", + "=r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _f32_bits_u32(float_val): + """Raw fp32 bits as Uint32 (bit-cast).""" + return cutlass.Uint32(llvm.bitcast(cutlass.Uint32.mlir_type, float_val.ir_value())) + + +def _i32_bits_f32(int_val): + return cutlass.Float32(llvm.bitcast(cutlass.Float32.mlir_type, int_val.ir_value())) + + +@dsl_user_op +def _fmin_f32(a, b, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lg2_f32(a, *, loc=None, ip=None): + """lg2.approx.f32 — thread-0 serial pick section only (lean interp).""" + return cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [a.ir_value(loc=loc, ip=ip)], + "lg2.approx.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _f2u_bits(u): + """Order-preserving fp32-bits -> Uint32 key: u ^ (sign ? 0xFFFFFFFF : 0x80000000).""" + neg = cutlass.Uint32(0) - (u >> cutlass.Uint32(31)) # 0 or 0xFFFFFFFF + return u ^ (neg | cutlass.Uint32(0x80000000)) + + +@cute.jit +def _exp2i(e): + """2.0**e for dynamic int e in [0, 127] via exponent-bit construction.""" + bits = (e + cutlass.Int32(127)) << cutlass.Int32(23) + return _i32_bits_f32(bits) + + +@cute.jit +def _shfl_up_add(val, lane, offset: cutlass.Constexpr): + """Inclusive-scan step: val += shfl_up(val, offset) gated by lane >= offset.""" + src = lane - cutlass.Int32(offset) + if src < 0: + src = cutlass.Int32(0) + other = cute.arch.shuffle_sync(val, src) + if lane >= cutlass.Int32(offset): + val = val + other + return val + + +@cute.jit +def _shfl_down_add(val, lane, offset: cutlass.Constexpr): + """Suffix-scan step: val += shfl_down(val, offset) gated by lane + offset < 32.""" + src = lane + cutlass.Int32(offset) + if src > 31: + src = cutlass.Int32(31) + other = cute.arch.shuffle_sync(val, src) + if lane + cutlass.Int32(offset) < cutlass.Int32(32): + val = val + other + return val + + +@cute.jit +def _mask_tail(v, gidx, n_eff): + """Ragged-N lane predicate: value at global element index ``gidx`` is + replaced by -FLT_MAX when ``gidx >= n_eff``. + + The float4 LOAD itself is always in-bounds (the logits buffer is + allocated to npad = logits.shape[1]); only the stale VALUE beyond the + row's valid length must be masked so garbage can never enter counts, + max-below, candidate pushes, or emits.""" + r = v + if gidx >= n_eff: + r = cutlass.Float32(-FLT_MAX) + return r + + +class GvrTpKernel: + """CuTe DSL port of gvr_topk_tp (fp32, B200/B300). + + Ctor knobs mirror the CUDA template params plus the production + ``next_n`` / ``compress_ratio`` contract (compile-time constexpr: + N_eff arithmetic, request-level hint-row sharing and the cr==1 hint + temporal shift all mirror the in-tree ``run_one_row`` / + ``phase1_preidx_stats`` formulas exactly; the next_n==1 / cr==4 hot + path traces identically to the v1 port). + """ + + WARP_SIZE = 32 + + def __init__( + self, + top_k: int, + kC: int, + cluster_size: int = 1, + ar: int = RUNGS, + uf: int = 4, + num_threads: int = 512, + next_n: int = 1, + compress_ratio: int = 4, + ): + assert num_threads % 32 == 0 + assert ar in (6, 8) + assert cluster_size in (1, 2, 4, 8) + # Lean-pivot interpolation target (per-K: K512 rows carry + # 5-20x sparser P2a samples, so 2.0K keeps interpolation-error + # margin; K>=1024 takes the measured-win 1.5K). + self.lean_tgt = (2 * top_k) if top_k < 1024 else (3 * top_k) // 2 + self.top_k = top_k + self.kC = kC + self.cluster_size = cluster_size + self.ar = ar + self.uf = uf + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + + # ------------------------------------------------------------------ + # float4 (128-bit) strided streaming load helper pieces + # ------------------------------------------------------------------ + def _copy_atom(self): + return cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128 + ) + + @cute.jit + def _ld_float4(self, copy_atom, row_addr, v_idx, frag): + """Load float4 #v_idx of the row into frag[0..3].""" + p = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(v_idx) * cutlass.Int64(16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + src = cute.make_tensor(p, cute.make_layout((4,))) + cute.copy(copy_atom, src, frag) + + # ------------------------------------------------------------------ + # Per-row valid length (ragged N). Mirrors the in-tree run_one_row + # arithmetic EXACTLY (gvr_topk_decode.py): seq_lens is request-level + # and in uncompressed-token space; logits live in compressed space + # when cr > 1. + # ------------------------------------------------------------------ + @cute.jit + def _row_n_eff(self, seq_lens: cute.Tensor, row): + NN = cutlass.const_expr(self.next_n) + seq_len = seq_lens[row // cutlass.Int32(NN)] + actual_kv_len = seq_len - cutlass.Int32(NN) + (row % cutlass.Int32(NN)) + cutlass.Int32(1) + if cutlass.const_expr(self.compress_ratio == 1): + n_eff = actual_kv_len + else: + n_eff = actual_kv_len // cutlass.Int32(self.compress_ratio) + return n_eff + + # ------------------------------------------------------------------ + # count_pass: R-rung exact count over [v0, v1) float4s, + # per-thread counts -> s_ptcnt[r*TB + tid]. + # ------------------------------------------------------------------ + @cute.jit + def count_pass( + self, + R: cutlass.Constexpr, + U: cutlass.Constexpr, + row_addr, + v0, + v1, + n_eff, + tidx, + s_rungs, + s_ptcnt, + ): + # Explicit U-batched loads (CUDA `float4 a[U]` idiom) — a + # `cutlass.range(unroll=U)` loop leaves ONE load in flight per iter + # and costs ~11% kernel time in the DRAM-bound regimes. Keep the + # Python-unrolled register batch. + TB = cutlass.const_expr(self.num_threads) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + frags = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(U) + ] # Python-unrolled register batch + # Mask hoist: float4s below vmain = min(v1, n_eff >> 2) are fully + # valid (4i+3 < n_eff), so the main loop runs mask-free (the + # per-element gi compare + select was ~10% of whole-kernel + # instructions); only [vmain, v1) — the n_eff boundary float4 and + # the pad tail — takes the masked epilogue. Bit-exact: the mask + # only changes values at gi >= n_eff, all of which live in + # [vmain, v1). + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull + i = v0 + tidx + while i + cutlass.Int32((U - 1) * TB) < vmain: + for u in cutlass.range_constexpr(U): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(U): + for q in cutlass.range_constexpr(4): + v = frags[u][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(U * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # sample_count: uniform every-SS-th float4 over the slice. + # Each per-thread accumulator PACKS the per-rung hit count (low 16 + # bits) with the per-rung float4 OCCUPANCY (high 16 bits) — the number + # of sampled float4s with at least one hit. On spatially-clustered + # real rows the 4 values of one float4 are strongly correlated, so the + # effective independent sample count is the occupancy, not the hit + # count: the admission stage sizes its confidence interval with the + # compound-Poisson sigma cnt/sqrt(occ) (equal to the classic + # sqrt(cnt) Poisson sigma on IID data where occ == cnt). Packing + # keeps registers, SMEM and the exchange at their pre-admission + # sizes; no field overflow: the tier envelope pins npad <= 262144, so + # cluster-total cnt <= npad/32 = 8192 < 2^16 and occ <= npad/128. + # ------------------------------------------------------------------ + @cute.jit + def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt): + TB = cutlass.const_expr(self.num_threads) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + frag = cute.make_fragment((4,), cutlass.Float32) + j = cutlass.Int32(tidx) + while v0 + j * cutlass.Int32(SS) < v1: + self._ld_float4(copy_atom, row_addr, v0 + j * cutlass.Int32(SS), frag) + gi = (v0 + j * cutlass.Int32(SS)) << cutlass.Int32(2) + v0m = _mask_tail(frag[0], gi, n_eff) + v1m = _mask_tail(frag[1], gi + cutlass.Int32(1), n_eff) + v2m = _mask_tail(frag[2], gi + cutlass.Int32(2), n_eff) + v3m = _mask_tail(frag[3], gi + cutlass.Int32(3), n_eff) + for r in cutlass.range_constexpr(R): + c4 = ( + cutlass.Int32(v0m >= tr[r]) + + cutlass.Int32(v1m >= tr[r]) + + cutlass.Int32(v2m >= tr[r]) + + cutlass.Int32(v3m >= tr[r]) + ) + cnt[r] = cnt[r] + c4 + (cutlass.Int32(c4 > cutlass.Int32(0)) << cutlass.Int32(16)) + j = j + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # exchange_counts: warp-fused CTA reduce -> cluster sum + prefix. + # rcnt[r] = cluster-wide count, rpre[r] = exclusive prefix over lower ranks. + # P2a rows carry packed (occ << 16 | cnt) accumulators; the integer sum + # distributes over the packed fields (no overflow: see sample_count). + # ------------------------------------------------------------------ + @cute.jit + def exchange_counts( + self, R: cutlass.Constexpr, par, tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + cute.arch.barrier() # ptcnt final + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if wid < cutlass.Int32(R): + s = cutlass.Int32(0) + for k in cutlass.range_constexpr(TB // 32): + s = s + s_ptcnt[wid * TB + lane + cutlass.Int32(32 * k)] + s = cute.arch.warp_redux_sync(s, "add") + if cutlass.const_expr(CS == 1): + if lane == cutlass.Int32(0): + s_rcnt[wid] = s + s_rpre[wid] = cutlass.Int32(0) + else: + if lane == cutlass.Int32(0): + s_ipartial[par * cutlass.Int32(RUNGS) + wid] = s + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + else: + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if tidx < cutlass.Int32(R): + tot = cutlass.Int32(0) + pre = cutlass.Int32(0) + local_ptr = s_ipartial.iterator + (par * cutlass.Int32(RUNGS) + tidx) + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + v = _ld_shared_cluster_i32(a) + tot = tot + v + if cutlass.Int32(rr) < rank: + pre = pre + v + s_rcnt[tidx] = tot + s_rpre[tidx] = pre + cute.arch.barrier() + + # ------------------------------------------------------------------ + # max_below_pass: largest value strictly below t_hi_bound (cluster-reduced). + # ------------------------------------------------------------------ + @cute.jit + def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred, s_fpartial): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + NWARP = cutlass.const_expr(self.num_warps) + copy_atom = self._copy_atom() + m = cutlass.Float32(-FLT_MAX) + # Explicit U=4 batched loads (CUDA `float4 a[4]` idiom). + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull + i = v0 + tidx + while i + cutlass.Int32(3 * TB) < vmain: + for u in cutlass.range_constexpr(4): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(4): + for q in cutlass.range_constexpr(4): + v = frags[u][q] + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(4 * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(TB) + m = cute.arch.warp_redux_sync(m, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = m + cute.arch.barrier() + if tidx == cutlass.Int32(0): + mm = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + mm = cute.arch.fmax(mm, s_fwred[w]) + s_fpartial[par] = mm + res = cutlass.Float32(-FLT_MAX) + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + res = s_fpartial[par] + else: + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + local_ptr = s_fpartial.iterator + par + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + res = cute.arch.fmax(res, _ld_shared_cluster_f32(a)) + return res + + # ------------------------------------------------------------------ + # phase1: hint gather + stats + rung ladder from hint-value CCDF. + # ------------------------------------------------------------------ + @cute.jit + def phase1( + self, + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=None, # cr==1 temporal shift ((row % next_n) + 1); None => 0 + ): + TB = cutlass.const_expr(self.num_threads) + AR = cutlass.const_expr(self.ar) + NWARP = cutlass.const_expr(self.num_warps) + K = cutlass.const_expr(self.top_k) + if tidx < cutlass.Int32(64): + s_hist[tidx] = cutlass.Int32(0) + hv = cute.make_fragment((4,), cutlass.Float32) + hok = cute.make_fragment((4,), cutlass.Int32) + mn = cutlass.Float32(FLT_MAX) + mx = cutlass.Float32(-FLT_MAX) + for jj in cutlass.range_constexpr(4): + j = tidx + cutlass.Int32(jj * TB) + hok[jj] = cutlass.Int32(0) + hv[jj] = cutlass.Float32(0.0) + if j < cutlass.Int32(K): + # pre_idx hardening: clamp hints into [0, N_eff-1] (not + # [0, npad-1] as in the standalone bench port). Production + # cold-start is all zeros and arbitrary garbage must not + # crash or corrupt; clamping (rather than skipping, which + # the in-tree phase1_preidx_stats does) also keeps the CCDF + # ladder seeded with real in-range values so an all-OOR + # hint cannot collapse the ladder onto FLT_MAX and force a + # pathological one-value-per-pass plateau descent. + hidx = pre_idx_row[j] + # cr==1 temporal shift, mirroring the in-tree run_one_row + # pre_idx_offset ((row % next_n) + 1); shifted-out-of-range + # hints fall into the same clamp below. const_expr'd out on + # cr>1 builds (hint_off is None => identical trace). + if cutlass.const_expr(hint_off is not None): + hidx = hidx + hint_off + if hidx < cutlass.Int32(0): + hidx = cutlass.Int32(0) + if hidx > n_eff - cutlass.Int32(1): + hidx = n_eff - cutlass.Int32(1) + v = logits_row[hidx] + hv[jj] = v + hok[jj] = cutlass.Int32(1) + mn = _fmin_f32(mn, v) + mx = cute.arch.fmax(mx, v) + mn = cute.arch.warp_redux_sync(mn, "fmin") + mx = cute.arch.warp_redux_sync(mx, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = mn + s_fwred[cutlass.Int32(NWARP) + wid] = mx + cute.arch.barrier() + if tidx == cutlass.Int32(0): + a = cutlass.Float32(FLT_MAX) + b = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + a = _fmin_f32(a, s_fwred[w]) + b = cute.arch.fmax(b, s_fwred[NWARP + w]) + s_hminmax[0] = a + s_hminmax[1] = b + cute.arch.barrier() + hmin = s_hminmax[0] + hmax = s_hminmax[1] + if hmax - hmin > cutlass.Float32(0.0): + # Stage 1: coarse 64-bin hist over [hmin, hmax] -> 97% trim point. + scale = cutlass.Float32(64.0) / (hmax - hmin) + for jj in cutlass.range_constexpr(4): + if hok[jj] != cutlass.Int32(0): + b1 = cutlass.Int32((hv[jj] - hmin) * scale) + if b1 < cutlass.Int32(0): + b1 = cutlass.Int32(0) + if b1 > cutlass.Int32(63): + b1 = cutlass.Int32(63) + atomicAdd(s_hist.iterator + b1, cutlass.Int32(1)) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + h0 = s_hist[cutlass.Int32(62) - cutlass.Int32(2) * tidx] + h1 = s_hist[cutlass.Int32(63) - cutlass.Int32(2) * tidx] + Ssum = h0 + h1 + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_up_add(x, tidx, o) + A = x - Ssum + binw = (hmax - hmin) * cutlass.Float32(1.0 / 64.0) + qtrim = cutlass.const_expr((K * 97) // 100) + if A < cutlass.Int32(qtrim): + if x >= cutlass.Int32(qtrim): + b2 = cutlass.Int32(62) - cutlass.Int32(2) * tidx + if A + h1 >= cutlass.Int32(qtrim): + b2 = cutlass.Int32(63) - cutlass.Int32(2) * tidx + s_hminmax[0] = hmin + binw * cutlass.Float32(b2) + if tidx == cutlass.Int32(31): + if x < cutlass.Int32(qtrim): + s_hminmax[0] = hmin + if tidx == cutlass.Int32(0): + s_rungs[AR - 1] = hmin + # barrier fold: re-zero both hist halves for stage 2 + s_hist[tidx] = cutlass.Int32(0) + s_hist[tidx + cutlass.Int32(32)] = cutlass.Int32(0) + cute.arch.barrier() + tlow = s_hminmax[0] + if hmax - tlow > cutlass.Float32(0.0): + # Stage 2: fine 64-bin hist over [tlow, hmax] -> rung quantiles. + scale2 = cutlass.Float32(64.0) / (hmax - tlow) + for jj in cutlass.range_constexpr(4): + if hok[jj] != cutlass.Int32(0): + if hv[jj] >= tlow: + b3 = cutlass.Int32((hv[jj] - tlow) * scale2) + if b3 < cutlass.Int32(0): + b3 = cutlass.Int32(0) + if b3 > cutlass.Int32(63): + b3 = cutlass.Int32(63) + atomicAdd(s_hist.iterator + b3, cutlass.Int32(1)) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + h0 = s_hist[cutlass.Int32(62) - cutlass.Int32(2) * tidx] + h1 = s_hist[cutlass.Int32(63) - cutlass.Int32(2) * tidx] + Ssum = h0 + h1 + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_up_add(x, tidx, o) + A = x - Ssum + binw2 = (hmax - tlow) * cutlass.Float32(1.0 / 64.0) + # WIDE (baseline) quantile spread, pinned — the + # P2A path must stay baseline-parity (ablation: the re-placed + # spread under this admission is a net residual harm). + if cutlass.const_expr(AR == 6): + qt = ( + (K * 15) // 100, + (K * 40) // 100, + (K * 70) // 100, + (K * 92) // 100, + ) + else: + qt = ( + (K * 10) // 100, + (K * 25) // 100, + (K * 45) // 100, + (K * 65) // 100, + (K * 82) // 100, + (K * 94) // 100, + ) + tot = cute.arch.shuffle_sync(x, cutlass.Int32(31)) + for r in cutlass.range_constexpr(AR - 2): + qtr = cutlass.const_expr(qt[r]) + if A < cutlass.Int32(qtr): + if x >= cutlass.Int32(qtr): + b4 = cutlass.Int32(62) - cutlass.Int32(2) * tidx + if A + h1 >= cutlass.Int32(qtr): + b4 = cutlass.Int32(63) - cutlass.Int32(2) * tidx + s_rungs[r + 1] = tlow + binw2 * cutlass.Float32(b4) + if tidx == cutlass.Int32(31): + if tot < cutlass.Int32(qtr): + s_rungs[r + 1] = tlow + if tidx == cutlass.Int32(0): + s_rungs[0] = hmax + (hmax - tlow) + cute.arch.barrier() + else: + # degenerate trim: all rungs above the floor = hmax + if tidx < cutlass.Int32(AR - 1): + s_rungs[tidx] = hmax + cute.arch.barrier() + else: + # degenerate: all hint values equal + if tidx < cutlass.Int32(AR): + s_rungs[tidx] = hmin + cute.arch.barrier() + + # ------------------------------------------------------------------ + # fused_count_collect: exact counts at {rungs[0]=pivot, + # rungs[1]=rescue} + push (key,idx) >= tpush into CTA0 cand, capped + # kcap. The rescue rung (next fatter ladder rung) replaces HEAD's hmin + # column at identical register cost: on a pivot undershoot the driver + # accepts the rescue with ONE collect re-stream instead of the + # multi-pass secant loop (hmin was a bracket-only data point; the + # secant fallback re-derives its bracket from the rescue instead). + # Keys are RAW fp32 bits (f2u happens in P4 round 0), matching CUDA. + # Candidates are a single packed (key<<32 | idx) u64 array — CUDA's + # `unsigned long long cand[kC]` — one 8B push per candidate. + # ------------------------------------------------------------------ + @cute.jit + def _push_cand( + self, + a_cnt, + a_st, + s_cand, + s_isc, + kcap: cutlass.Constexpr, + capped: cutlass.Constexpr, + v, + gidx, + ): + """One candidate push: kv = (raw fp32 bits << 32) | index, single 8B + store (st.shared::cluster.u64 remote at CS>1, local u64 store at CS1). + a_cnt/a_st are the pre-mapa'd CTA0 addresses of cnt_c / cand[0].""" + CS = cutlass.const_expr(self.cluster_size) + kv = (cutlass.Uint64(_f32_bits_u32(v)) << cutlass.Uint64(32)) | cutlass.Uint64( + cutlass.Uint32(gidx) + ) + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_cnt, cutlass.Int32(1)) + if cutlass.const_expr(capped): + if p < cutlass.Int32(kcap): + _st_shared_cluster_u64(a_st + p * cutlass.Int32(8), kv) + else: + _st_shared_cluster_u64(a_st + p * cutlass.Int32(8), kv) + else: + p = atomicAdd(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(1)) + if cutlass.const_expr(capped): + if p < cutlass.Int32(kcap): + s_cand[p] = kv + else: + s_cand[p] = kv + + @cute.jit + def fused_count_collect( + self, + U: cutlass.Constexpr, + row_addr, + v0, + v1, + n_eff, + tpush, + tidx, + s_rungs, + s_ptcnt, + s_cand, + s_isc, + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + kcap = cutlass.const_expr(self.kC) + R = cutlass.const_expr(2) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + a_cnt = cutlass.Int32(0) + a_st = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) + a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + # Explicit U-batched loads (CUDA `float4 a[U]` idiom), main + + # vec-tail while-loops (same fix as count_pass). Mask + # hoist as in count_pass: [v0, vmain) mask-free (gi computed only + # inside the rare push branch), [vmain, v1) masked epilogue. + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(U)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull + i = v0 + tidx + while i + cutlass.Int32((U - 1) * TB) < vmain: + for u in cutlass.range_constexpr(U): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(U): + for q in cutlass.range_constexpr(4): + v = frags[u][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + True, + v, + ((i + cutlass.Int32(u * TB)) << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(U * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + True, + v, + (i << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, True, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # collect_at: plain streaming collect at thr (uncapped; caller + # guarantees count(thr) <= kC). + # ------------------------------------------------------------------ + @cute.jit + def collect_at(self, row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + kcap = cutlass.const_expr(self.kC) + copy_atom = self._copy_atom() + a_cnt = cutlass.Int32(0) + a_st = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) + a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + # Explicit U=4 batched loads. nvcc auto-unrolls the + # CUDA collect_at loop 4x with 4 LDG.E.128 issued back-to-back; a + # 1-deep loop was the dominant stall site (36% of all warp-stall + # samples) on cells whose reuse check fails at big npad x big BS. + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull + i = v0 + tidx + while i + cutlass.Int32(3 * TB) < vmain: + for u in cutlass.range_constexpr(4): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(4): + for q in cutlass.range_constexpr(4): + v = frags[u][q] + if v >= thr: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + False, + v, + ((i + cutlass.Int32(u * TB)) << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(4 * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + if v >= thr: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + False, + v, + (i << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + if v >= thr: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, False, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(TB) + + # ------------------------------------------------------------------ + # kernel + # ------------------------------------------------------------------ + @cute.kernel + def gvr_tp_kernel( + self, logits: cute.Tensor, pre_idx: cute.Tensor, seq_lens: cute.Tensor, out_idx: cute.Tensor + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + AR = cutlass.const_expr(self.ar) + UF = cutlass.const_expr(self.uf) + K = cutlass.const_expr(self.top_k) + kC = cutlass.const_expr(self.kC) + + if cutlass.const_expr(CS > 1): + row = bidx // cutlass.Int32(CS) + rank = cute.arch.block_idx_in_cluster() + else: + row = bidx + rank = cutlass.Int32(0) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[row, None] + # Hint sharing: pre_idx is request-level ([num_rows // next_n, K]); + # the next_n MTP rows of one request share the same hint row — + # mirrors the in-tree run_one_row's pre_idx_row_idx = row // next_n. + if cutlass.const_expr(self.next_n == 1): + pre_idx_row = pre_idx[row, None] + else: + pre_idx_row = pre_idx[row // cutlass.Int32(self.next_n), None] + out_row = out_idx[row, None] + row_addr = logits_row.iterator.toint() + + # Ragged N: per-row valid length from seq_lens (see _row_n_eff). + n_eff = self._row_n_eff(seq_lens, row) + + # ---- shared memory (order must be identical across CTAs for mapa) ---- + smem = SmemAllocator() + # Packed (key<<32 | idx) u64 candidates = CUDA's + # `unsigned long long cand[kC]` (one 8B push per candidate). + s_cand = smem.allocate_tensor( + element_type=cutlass.Uint64, + layout=cute.make_ordered_layout((kC,), order=(0,)), + byte_alignment=128, + ) + s_ptcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS * self.num_threads,), order=(0,)), + byte_alignment=128, + ) + s_hist = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((256,), order=(0,)), + byte_alignment=128, + ) + s_rungs = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rpre = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_ipartial = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((2 * RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_fpartial = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + s_fwred = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2 * self.num_warps,), order=(0,)), + byte_alignment=64, + ) + s_hminmax = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + # iscalars: [0]=sel_bin [1]=sel_above [2]=sel_count [3]=cnt_m [4]=cnt_t [5]=cnt_c + s_isc = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((8,), order=(0,)), + byte_alignment=32, + ) + + # slice in 64-float units (npad multiple of 64); v0/v1 are float4 indices + units = npad >> cutlass.Int32(6) + u0 = (units * rank) // cutlass.Int32(CS) + u1 = (units * (rank + cutlass.Int32(1))) // cutlass.Int32(CS) + v0 = u0 << cutlass.Int32(4) + v1 = u1 << cutlass.Int32(4) + + xch = cutlass.Int32(0) + thr = cutlass.Float32(0.0) + tpush = cutlass.Float32(0.0) + C = cutlass.Int32(0) + m_gt = cutlass.Int32(-1) + span0 = cutlass.Float32(1e-3) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad. ---- + # Mirrors the in-tree kernel's degenerate branch. n_eff is uniform + # across the cluster (all CTAs own the same row), so this is a + # cluster-uniform branch; non-leader CTAs fall through to the exit + # rendezvous (CuTe DSL has no runtime return). + if n_eff <= cutlass.Int32(K): + if rank == cutlass.Int32(0): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(TB) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(K): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(TB) + else: + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[5] = cutlass.Int32(0) # cnt_c + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + if npad <= cutlass.Int32(kC): + # trivial: whole row fits the candidate buffer. Masked tail + # lanes push -FLT_MAX candidates; they can never displace a + # real value because n_eff > K here (kth value > -FLT_MAX). + if tidx < cutlass.Int32(RUNGS): + s_rungs[tidx] = cutlass.Float32(-FLT_MAX) + cute.arch.barrier() + self.count_pass(1, 8, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, cutlass.Int32(0), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + thr = cutlass.Float32(-FLT_MAX) + C = s_rcnt[0] + self.collect_at(row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc) + cute.arch.barrier() + else: + # cr==1 hint temporal shift (in-tree pre_idx_offset parity: + # (row % next_n) + 1 maps prev-step indices into this step's + # KV space); cr>1 keeps the exact pre-MTP trace (no offset + # value is even computed). + if cutlass.const_expr(self.compress_ratio == 1): + hint_off = (row % cutlass.Int32(self.next_n)) + cutlass.Int32(1) + self.phase1( + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=hint_off, + ) + else: + self.phase1( + logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs + ) + hmin_floor = s_rungs[AR - 1] + span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) + # P2a: sampled ladder count -> pivot pick. The exchanged rows + # carry packed (occ << 16 | cnt) values (admission sigma + # input); every consumer below unpacks with & 0xFFFF / >> 16. + self.sample_count(AR, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + # 4-stage pivot pick (R0-admission tightest-in-window -> + # band target -> 2-sigma bound -> gamble fallback) + lo = cutlass.const_expr((3 * K) // 2) + hi = cutlass.const_expr((6 * kC) // 10) + tgt_py = min(max(3 * K, (3 * K) // 2), (6 * kC) // 10) + tgt = cutlass.Int32(tgt_py) + best = cutlass.Int32(AR - 1) + bestd = cutlass.Int32(0x7FFFFFFF) + # Stage 0 — hint-ladder ADMISSION (in-tree R0 + # parity): accept the TIGHTEST ladder rung + # (highest threshold = smallest admitted-candidate set) + # whose sampled-count confidence interval lies inside + # the [K, kC] acceptance window, mirroring the R0 rule + # "smallest exact count in [K, kC]". The old stage-1 + # band ([1.5K, 0.6kC], target ~3K) systematically picks + # a FAT rung on real high-hit-rate rows (2-4x more P3 + # pushes + P4 candidates than needed — the v32/pro + # BS>=16 losses vs the in-tree R0 kernel) or, when + # spatially-clustered real data inflates a sampled + # estimate into the band while the true count is < K, + # an undershooting one (whole-row secant + re-stream: + # the flash_512k 1.6-1.8x losses). + # Sigma is clustering-aware: est = SS*cnt, sigma = + # SS*cnt/sqrt(occ) (compound Poisson over occupied + # float4 clumps; equals the classic sqrt(SS*est) when + # occ == cnt, i.e. IID rows). + # Fix round (pr2 full-grid regressions), all changes on + # the PICK only — the fused pass / rescue / exactness + # machinery is byte-identical to the admission commit: + # * upper acceptance bound tightened from kC to the + # legacy pivot-band hi (0.6*kC): an in-window-but- + # fat admitted rung inflates P3 pushes + P4 + # candidates 2-4x over the ~3K legacy target; + # * K2048 lower margin raised 1.5 -> 1.85 sigma: at + # 1.5 the admitted rung's true count lands under K + # on clustered rows (v32_32k_L23 rung est 2720, + # sigma 374, true 1959 < K) and the rescue + # re-stream costs more than the fat legacy pick it + # displaced; 1.85 keeps the genuine tight admits + # (v32_32k_L50: margin +1.92 sigma, true 2917); + # * stage 0b below: when the legacy band pick is + # strictly LEANER than the admitted rung and safe + # by its own clustering-aware 1.5-sigma lower CI, + # prefer it (mlo=2.0 rows whose next-leaner rung + # sits at 1.8-2.0 sigma otherwise admit a 2.5x + # fatter set: pro_64k_L06 4134 vs legacy 1671). + # Rungs are descending in j, so the first passing j is + # the tightest. + # All margin tests below are sqrt/div-free: with + # sigma = est/sqrt(occ), "est - m*sigma >= K" is + # "(est-K)^2 * occ >= m^2 * est^2" (est >= K), and + # "est + 2*sigma <= U" is "4*est^2 <= (U-est)^2 * occ" + # (est <= U). The pick is a THREAD-0 serial section + # between two CTA barriers; the sqrt+fdiv chain of the + # first admission cut measured ~2-5% whole-kernel on + # L2-resident accept-path rows. + mlo2 = cutlass.const_expr(1.85 * 1.85 if K >= 2048 else 4.0) + ubnd = cutlass.const_expr(float((6 * kC) // 10)) + adm_est = cutlass.Int32(0x7FFFFFFF) + for j in cutlass.range_constexpr(AR): + cnt_j = s_rcnt[j] & cutlass.Int32(0xFFFF) + if cnt_j > cutlass.Int32(0): + if bestd == cutlass.Int32(0x7FFFFFFF): + occ_j = s_rcnt[j] >> cutlass.Int32(16) + if occ_j < cutlass.Int32(1): + occ_j = cutlass.Int32(1) + focc = cutlass.Float32(occ_j) + fest = cutlass.Float32(cnt_j * cutlass.Int32(SS)) + fe2 = fest * fest + a_lo = fest - cutlass.Float32(float(K)) + bhi = cutlass.Float32(ubnd) - fest + if a_lo >= cutlass.Float32(0.0): + if a_lo * a_lo * focc >= cutlass.Float32(mlo2) * fe2: + if bhi >= cutlass.Float32(0.0): + if cutlass.Float32(4.0) * fe2 <= bhi * bhi * focc: + bestd = cutlass.Int32(0) + best = cutlass.Int32(j) + adm_est = cnt_j * cutlass.Int32(SS) + # Record the stage-0 CI-admission outcome + # BEFORE stages 0b/1/2 mutate bestd; a stage-0 admit + # is the "confident P2A" signal — those rows never + # take the stage-0c override. + s0_ok = cutlass.Int32(0) + if bestd == cutlass.Int32(0): + s0_ok = cutlass.Int32(1) + # Stage 0b — legacy band pick (min |est - tgt| in + # [1.5K, 0.6kC], EXACTLY the pr1 stage-1 rule), then + # combine: with no admission it is taken as-is (pr1 + # parity); with an admitted rung it OVERRIDES only + # when strictly leaner AND safe by its own + # clustering-aware 1.5-sigma lower CI. + jb = cutlass.Int32(AR - 1) + jbd = cutlass.Int32(0x7FFFFFFF) + for j in cutlass.range_constexpr(AR): + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if est >= cutlass.Int32(lo): + if est <= cutlass.Int32(hi): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < jbd: + jbd = dd + jb = cutlass.Int32(j) + if jbd != cutlass.Int32(0x7FFFFFFF): + if bestd == cutlass.Int32(0x7FFFFFFF): + bestd = jbd + best = jb + else: + cnt_b = s_rcnt[jb] & cutlass.Int32(0xFFFF) + est_b = cnt_b * cutlass.Int32(SS) + if est_b < adm_est: + occ_b = s_rcnt[jb] >> cutlass.Int32(16) + if occ_b < cutlass.Int32(1): + occ_b = cutlass.Int32(1) + fb = cutlass.Float32(est_b) + ab = fb - cutlass.Float32(float(K)) + if ab >= cutlass.Float32(0.0): + if ( + ab * ab * cutlass.Float32(occ_b) + >= cutlass.Float32(2.25) * fb * fb + ): + bestd = jbd + best = jb + if bestd == cutlass.Int32(0x7FFFFFFF): + for j in cutlass.range_constexpr(AR): + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if est > cutlass.Int32(0): + g = cutlass.Float32(2.0) * cmath.sqrt( + cutlass.Float32(cutlass.Int32(SS) * est) + ) + fest = cutlass.Float32(est) + if fest - g >= cutlass.Float32(float(K)): + if fest + g <= cutlass.Float32(float(kC)): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < bestd: + bestd = dd + best = cutlass.Int32(j) + if bestd == cutlass.Int32(0x7FFFFFFF): + if (s_rcnt[AR - 1] & cutlass.Int32(0xFFFF)) * cutlass.Int32( + SS + ) > cutlass.Int32(kC): + for j in cutlass.range_constexpr(AR): + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if est > cutlass.Int32(0): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < bestd: + bestd = dd + best = cutlass.Int32(j) + # ---- stage 0c: lean-pivot override ---- + # When the stage-0 CI admission FAILED, npad <= + # 262144, and some ladder rung is (i) strictly LEANER + # than the pick in hand and (ii) lands in [K, kC] by + # its own clustering-aware 1.5-sigma CI on BOTH sides + # (a looser window than stage 0's [K, 0.6kC] at + # 1.85/2.0-sigma — rows passing THAT never get here), + # make IT the pivot. Rungs are descending in j: scan + # ascending and keep the FIRST (tightest) qualifier. + # Hint-unrepresentative rows (the measured harm + # band) have no CI-qualifying rung and keep the baseline + # pick. Same sqrt/div-free margin algebra as stage 0; + # everything downstream (fused R=2, rescue, reuse, + # P4, secant) is byte-parity with the baseline. + # npad floor 16384: below it the stride-32 sample sees + # <= ~128 float4s, the occ-aware CI fires on noise and + # misfire rescues cost 3-5% (measured, npad~8K probe); + # the absolute win-room there is small anyway. + gate_ok = cutlass.Int32(0) + if npad >= cutlass.Int32(16384): + if npad <= cutlass.Int32(262144): + gate_ok = cutlass.Int32(1) + if gate_ok != cutlass.Int32(0): + if s0_ok == cutlass.Int32(0): + est_pick = cutlass.Int32(0x7FFFFFFF) + if bestd != cutlass.Int32(0x7FFFFFFF): + est_pick = (s_rcnt[best] & cutlass.Int32(0xFFFF)) * cutlass.Int32( + SS + ) + jw = cutlass.Int32(AR) + for w_ in cutlass.range_constexpr(AR): + if jw == cutlass.Int32(AR): + cnt_w = s_rcnt[w_] & cutlass.Int32(0xFFFF) + if cnt_w > cutlass.Int32(0): + estw = cnt_w * cutlass.Int32(SS) + if estw < est_pick: + occ_w = s_rcnt[w_] >> cutlass.Int32(16) + if occ_w < cutlass.Int32(1): + occ_w = cutlass.Int32(1) + foccw = cutlass.Float32(occ_w) + festw = cutlass.Float32(estw) + few2 = festw * festw + a_low = festw - cutlass.Float32(float(K)) + bhiw = cutlass.Float32(float(kC)) - festw + if a_low >= cutlass.Float32(0.0): + if ( + a_low * a_low * foccw + >= cutlass.Float32(2.25) * few2 + ): + if bhiw >= cutlass.Float32(0.0): + if ( + cutlass.Float32(2.25) * few2 + <= bhiw * bhiw * foccw + ): + jw = cutlass.Int32(w_) + if jw < cutlass.Int32(AR): + best = jw + bestd = cutlass.Int32(0) + tp_ = s_rungs[best] + # Rescue rung: the next FATTER ladder rung below the + # pivot. If the pivot's exact count lands under K + # (sampling error on a clustered row), the driver can + # accept the rescue with ONE collect re-stream instead + # of the multi-pass secant loop. Read before the ladder + # slots are overwritten. + resc_ = hmin_floor + if best < cutlass.Int32(AR - 1): + resc_ = s_rungs[best + cutlass.Int32(1)] + # ---- lean pivot (per-K interpolated push threshold, + # composed AFTER stage 0c): the pick in hand — band, + # 2-sigma, overshoot fallback, or a stage-0c lean rung + # that is still fat — is structurally FAT on the + # 32k-128k residual band (exact Cp 2.4-4.7x K; P3 + # pushes + P4 radix pay for it). When the picked + # rung's sampled est >= 1.8K, interpolate tpush + # between the pick and the next tighter rung with + # est < lean_tgt, targeting count ~= lean_tgt + # (log2-count interpolation: tails are ~exponential, + # the linear form undershoots). Rescue = the ORIGINAL + # pick, whose exact count the fused pass computes + # anyway: an undershoot costs one collect re-stream, + # cheap under the npad <= 98304 gate; the 256k-1024k + # guard band keeps the stock pick bit-for-bit. + if npad <= cutlass.Int32(98304): + estp = (s_rcnt[best] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if estp >= cutlass.Int32((9 * K) // 5): + jm = cutlass.Int32(-1) + estm = cutlass.Int32(1) + for j in cutlass.range_constexpr(AR): + if cutlass.Int32(j) < best: + cj = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if cj < cutlass.Int32(self.lean_tgt): + jm = cutlass.Int32(j) + estm = cj + if jm >= cutlass.Int32(0): + if estm < cutlass.Int32(1): + estm = cutlass.Int32(1) + fp = _lg2_f32(cutlass.Float32(estp)) + fm = _lg2_f32(cutlass.Float32(estm)) + den = fp - fm + if den < cutlass.Float32(1e-6): + den = cutlass.Float32(1e-6) + frac = (fp - cutlass.Float32(math.log2(self.lean_tgt))) / den + resc_ = tp_ + tp_ = tp_ + (s_rungs[jm] - tp_) * frac + s_rungs[0] = tp_ + s_rungs[1] = resc_ + cute.arch.barrier() + tpush = s_rungs[0] + self.fused_count_collect( + UF, row_addr, v0, v1, n_eff, tpush, tidx, s_rungs, s_ptcnt, s_cand, s_isc + ) + self.exchange_counts( + 2, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + + # ---- P2c: secant refine driver (redundant on every thread) ---- + t_lo = cutlass.Float32(-FLT_MAX) + t_hi = cutlass.Float32(INF) + c_hi = cutlass.Int32(0) + Rcur = cutlass.Int32(2) + passno = cutlass.Int32(0) + running = cutlass.Int32(1) + descend_break = cutlass.Int32(0) + while running != cutlass.Int32(0): + # first rung index j with rcnt[j] >= K (rcnt ascending in j) + j = Rcur + for r_ in range(RUNGS - 1, -1, -1): + if cutlass.Int32(r_) < Rcur: + if s_rcnt[r_] >= cutlass.Int32(K): + j = cutlass.Int32(r_) + jj = j + if jj > Rcur - cutlass.Int32(1): + jj = Rcur - cutlass.Int32(1) + cj = s_rcnt[jj] + rj = s_rungs[jj] + found = cutlass.Int32(0) + if j < Rcur: + if cj <= cutlass.Int32(kC): + found = cutlass.Int32(1) + if found != cutlass.Int32(0): + thr = rj + C = cj + running = cutlass.Int32(0) + else: + if j < Rcur: + if rj >= t_lo: + t_lo = rj + if j > cutlass.Int32(0): + jm = j - cutlass.Int32(1) + if s_rungs[jm] <= t_hi: + t_hi = s_rungs[jm] + c_hi = s_rcnt[jm] + descend = cutlass.Int32(0) + if passno >= cutlass.Int32(MAXPASS): + descend = cutlass.Int32(1) + # ladder params (uniform; recomputed per thread) + e3 = passno * cutlass.Int32(3) + if e3 > cutlass.Int32(24): + e3 = cutlass.Int32(24) + step = span0 * _exp2i(e3) + dt = cutlass.Float32(0.0) + mode = cutlass.Int32(2) # 0=up-ladder,1=down-ladder,2=secant + if t_hi == cutlass.Float32(INF): + mode = cutlass.Int32(0) + nr0 = t_lo + step * cutlass.Float32(float(1 << (AR - 1))) + if nr0 == cutlass.Float32(INF): + descend = cutlass.Int32(1) + else: + if t_lo == cutlass.Float32(-FLT_MAX): + mode = cutlass.Int32(1) + else: + dt = (t_hi - t_lo) * cutlass.Float32(1.0 / float(AR + 1)) + nr_last = t_hi - dt * cutlass.Float32(float(AR)) + nr_first = t_hi - dt + ok = cutlass.Int32(0) + if nr_last > t_lo: + if nr_first < t_hi: + ok = cutlass.Int32(1) + if ok == cutlass.Int32(0): + descend = cutlass.Int32(1) + if descend != cutlass.Int32(0): + running = cutlass.Int32(0) + descend_break = cutlass.Int32(1) + else: + cute.arch.barrier() + if tidx == cutlass.Int32(0): + for r_ in cutlass.range_constexpr(AR): + nrv = cutlass.Float32(0.0) + if mode == cutlass.Int32(0): + nrv = t_lo + step * cutlass.Float32( + float(1 << (AR - 1 - r_)) + ) + else: + if mode == cutlass.Int32(1): + nrv = t_hi - step * cutlass.Float32(float(1 << r_)) + else: + nrv = t_hi - dt * cutlass.Float32(float(r_ + 1)) + s_rungs[r_] = nrv + cute.arch.barrier() + self.count_pass(AR, UF, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + passno = passno + cutlass.Int32(1) + Rcur = cutlass.Int32(AR) + + # ---- plateau descent (exact max-below stepping) ---- + if descend_break != cutlass.Int32(0): + pl = cutlass.Int32(1) + while pl != cutlass.Int32(0): + vstar = self.max_below_pass( + row_addr, + v0, + v1, + n_eff, + t_hi, + xch & cutlass.Int32(1), + tidx, + s_fwred, + s_fpartial, + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_rungs[0] = vstar + cute.arch.barrier() + self.count_pass(1, 8, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + c = s_rcnt[0] + okc = cutlass.Int32(0) + if c >= cutlass.Int32(K): + if c <= cutlass.Int32(kC): + okc = cutlass.Int32(1) + if okc != cutlass.Int32(0): + thr = vstar + C = c + pl = cutlass.Int32(0) + else: + if c < cutlass.Int32(K): + t_hi = vstar + c_hi = c + else: + thr = vstar + m_gt = c_hi + pl = cutlass.Int32(0) + + # ---- candidate reuse check / re-stream collect ---- + if m_gt < cutlass.Int32(0): + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + if cutlass.const_expr(CS > 1): + a_cnt0 = _mapa_shared_cluster( + s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0) + ) + dcnt = _ld_shared_cluster_i32(a_cnt0) + else: + dcnt = s_isc[5] + reuse = cutlass.Int32(0) + if thr == tpush: + if dcnt == C: + reuse = cutlass.Int32(1) + if reuse == cutlass.Int32(0): + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[5] = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + self.collect_at(row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc) + cute.arch.barrier() + + # ---- plateau direct emit (all CTAs stream their slice) ---- + if m_gt >= cutlass.Int32(0): + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) # cnt_m + s_isc[4] = cutlass.Int32(0) # cnt_t + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + a_m = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(0)) + a_t = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(4), cutlass.Int32(0)) + else: + cute.arch.barrier() + nt = cutlass.Int32(K) - m_gt + copy_atom = self._copy_atom() + frag = cute.make_fragment((4,), cutlass.Float32) + ii = v0 + tidx + while ii < v1: + self._ld_float4(copy_atom, row_addr, ii, frag) + gi = ii << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frag[q], gi + cutlass.Int32(q), n_eff) + if v > thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_m, cutlass.Int32(1)) + else: + p = atomicAdd(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1)) + out_row[p] = gi + cutlass.Int32(q) + else: + if v == thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_t, cutlass.Int32(1)) + else: + p = atomicAdd( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m_gt + p] = gi + cutlass.Int32(q) + ii = ii + cutlass.Int32(TB) + else: + # ---- P4 (CTA0 solo when CS>1) ---- + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if rank == cutlass.Int32(0): + if C == cutlass.Int32(K): + ie = cutlass.Int32(tidx) + while ie < C: + out_row[ie] = cutlass.Int32( + cutlass.Uint32(s_cand[ie] & cutlass.Uint64(0xFFFFFFFF)) + ) + ie = ie + cutlass.Int32(TB) + else: + # 4x8-bit radix select over cand keys (f2u in round 0) + if tidx < cutlass.Int32(256): + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() + pref = cutlass.Uint32(0) + want = cutlass.Int32(K) + m = cutlass.Int32(0) + final_shift = cutlass.Uint32(0) + active = cutlass.Int32(1) + for r_ in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * r_) + if active != cutlass.Int32(0): + ih = cutlass.Int32(tidx) + while ih < C: + if cutlass.const_expr(r_ == 0): + kv = s_cand[ih] + raw = cutlass.Uint32(kv >> cutlass.Uint64(32)) + u = _f2u_bits(raw) + s_cand[ih] = (cutlass.Uint64(u) << cutlass.Uint64(32)) | ( + kv & cutlass.Uint64(0xFFFFFFFF) + ) + atomicAdd( + s_hist.iterator + + cutlass.Int32(u >> cutlass.Uint32(24)), + cutlass.Int32(1), + ) + else: + u = cutlass.Uint32(s_cand[ih] >> cutlass.Uint64(32)) + if (u >> cutlass.Uint32(shift + 8)) == pref: + b = cutlass.Int32( + (u >> cutlass.Uint32(shift)) & cutlass.Uint32(0xFF) + ) + atomicAdd(s_hist.iterator + b, cutlass.Int32(1)) + ih = ih + cutlass.Int32(TB) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + b8 = tidx * cutlass.Int32(8) + h = cute.make_fragment((8,), cutlass.Int32) + Ssum = cutlass.Int32(0) + for q in cutlass.range_constexpr(8): + h[q] = s_hist[b8 + cutlass.Int32(q)] + Ssum = Ssum + h[q] + s_hist[b8 + cutlass.Int32(q)] = cutlass.Int32(0) + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_down_add(x, tidx, o) + A = x - Ssum + if A < want: + if want <= A + Ssum: + run = A + for q in range(7, -1, -1): + if run < want: + if want <= run + h[q]: + s_isc[0] = b8 + cutlass.Int32(q) + s_isc[1] = run + s_isc[2] = h[q] + run = run + h[q] + cute.arch.barrier() + sel = s_isc[0] + above = s_isc[1] + m = m + above + want = want - above + pref = (pref << cutlass.Uint32(8)) | cutlass.Uint32(sel) + if s_isc[2] == want: + final_shift = cutlass.Uint32(shift) + active = cutlass.Int32(0) + kth = pref + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) + s_isc[4] = cutlass.Int32(0) + cute.arch.barrier() + nt = cutlass.Int32(K) - m + ie = cutlass.Int32(tidx) + while ie < C: + kv = s_cand[ie] + u = cutlass.Uint32(kv >> cutlass.Uint64(32)) >> final_shift + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + if u > kth: + p = atomicAdd(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1)) + out_row[p] = idx + else: + if u == kth: + p = atomicAdd( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m + p] = idx + ie = ie + cutlass.Int32(TB) + + # Exit rendezvous: DSMEM (mapa'd peer smem) must stay valid until every + # CTA of the cluster is done issuing remote ops (plateau emit's remote + # cnt_m/cnt_t atomics can land after CTA0 finishes its own slice). nvcc + # inserts an implicit cluster barrier before ret for DSMEM kernels; CuTe + # DSL does not — without this, plateau rows fault with a + # timing-dependent cudaErrorLaunchFailure (CUDA 719, hidden under + # compute-sanitizer synccheck). DO NOT remove or make conditional. + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # ------------------------------------------------------------------ + # host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + seq_lens: cute.Tensor, + out_idx: cute.Tensor, + stream, + ): + num_rows = logits.shape[0] + CS = cutlass.const_expr(self.cluster_size) + self.gvr_tp_kernel(logits, pre_idx, seq_lens, out_idx).launch( + grid=(num_rows * CS, 1, 1), + block=(self.num_threads, 1, 1), + cluster=(CS, 1, 1) if cutlass.const_expr(CS > 1) else None, + stream=stream, + min_blocks_per_mp=2, + ) + + +# --------------------------------------------------------------------------- +# compile cache + public entry (mirrors launch_tp<512> selection exactly) +# --------------------------------------------------------------------------- +_LAUNCH_CACHE = {} + + +def _p2floor(x: int) -> int: + p = 1 + while p * 2 <= x: + p *= 2 + return p + + +def _two_waves_rows() -> int: + """2 * SM count of the current device (cached): the co-residency row + budget behind the CS selection below (min_blocks_per_mp=2). Mirrors + ``_get_num_sms`` in ``cute_dsl_custom_ops.py`` (not importable from + here — the import runs the other way). NOTE: only this term scales + with the device; every other constant in ``tp_cluster_size`` and the + dispatch band tables is frozen B200 calibration (the kernel family + measured HW-invariant in earlier B200/B300 cross-arch A/Bs, so the + bands are kept device-independent on purpose). + """ + if not hasattr(_two_waves_rows, "_value"): + _two_waves_rows._value = 2 * torch.cuda.get_device_properties().multi_processor_count + return _two_waves_rows._value + + +def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, cr: int = 4): + key = (K, CS, AR, UF, TB, next_n, cr) + compiled = _LAUNCH_CACHE.get(key) + if compiled is None: + # Flat candidate budget, never the K-scaled + # diet (falsified as pure harm). + kC = 8192 if K >= 2048 else 6144 + kern = GvrTpKernel( + top_k=K, + kC=kC, + cluster_size=CS, + ar=AR, + uf=UF, + num_threads=TB, + next_n=next_n, + compress_ratio=cr, + ) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + # pre_idx is request-level: num_rows // next_n rows. At next_n == 1 + # keep the shared n_rows sym (identical compiled artifact to the v1 + # port); at next_n > 1 the row counts differ, so use a distinct sym. + n_pre = n_rows if next_n == 1 else cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + pre_idx_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_pre, K), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kern, + logits_fake, + pre_idx_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _LAUNCH_CACHE[key] = compiled + return compiled + + +_FAST = {} # (K, bs, npad) -> compiled variant (hot-path single dict hit) + + +def tp_cluster_size(bs: int, npad: int) -> int: + """CS selection of launch_tp<512>: co-residency + slice floor. + Occupancy cut for mid-N large-BS (measured neutral-to-positive).""" + # C4 occupancy cut kept at fix-head default (ablation: neutral). + if npad < 65536 and bs >= 64: + return 1 + cs = 1 + if bs < 128: + two_waves = _two_waves_rows() # 296 on B200 (2 x 148 SMs) + cs = _p2floor(two_waves // bs) if bs <= two_waves else 1 + capn = _p2floor(npad // 8192 if npad // 8192 > 0 else 1) + if cs > capn: + cs = capn + if cs > 8: + cs = 8 + return cs + + +def _dispatch(K: int, bs: int, npad: int, next_n: int = 1, cr: int = 4): + """launch_tp<512> selection: CS by co-residency + slice floor, UF by depth.""" + assert npad % 64 == 0 + TB = 512 + AR = RUNGS + cs = tp_cluster_size(bs, npad) + if cs > 1: + uf = 4 + else: + uf = 8 if npad >= 16384 else 4 + return _get_compiled(K, cs, AR, uf, TB, next_n, cr) + + +def tp_topk( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + next_n: int = 1, + cr: int = 4, +) -> None: + """CuTe DSL gvr_topk_tp. logits [BS, npad] fp32 (npad mult of 64; + tail beyond each row's N_eff may be garbage — masked in-kernel), + pre_idx [BS // next_n, K] int32 hint (request-level), seq_lens + [BS // next_n] int32 (uncompressed-token space), out [BS, K] int32. + Launch-time CS/UF selection replicates launch_tp<512>.""" + sh = logits.shape + key = (K, sh[0], sh[1], next_n, cr) + fn = _FAST.get(key) + if fn is None: + fn = _dispatch(K, sh[0], sh[1], next_n, cr) + _FAST[key] = fn + fn(logits, pre_idx, seq_lens, out) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 95938a31cd98..747679751e5b 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -16,10 +16,11 @@ l0_b300: backend: pytorch tests: # ------------- PyTorch tests --------------- - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/thop/serial - unittest/_torch/executor # 250s diff --git a/tests/unittest/_torch/attention/sparse/conftest.py b/tests/unittest/_torch/attention/sparse/conftest.py new file mode 100644 index 000000000000..02af87693b6f --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/conftest.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared checkers for the top-K decode test modules in this directory. + +The tie-aware checker lives here (single definition) and is handed to tests +via the ``tie_aware_check`` fixture: conftest is collected by pytest +regardless of rootdir/package-resolution style, so this sidesteps the +cross-module import problem that previously forced a keep-in-sync duplicate +of the checker in ``test_cute_dsl_gvr_topk_tiers.py``. +""" + +import pytest +import torch + + +def _tie_aware_check_impl( + out_indices: torch.Tensor, + logits: torch.Tensor, + seq_lens: torch.Tensor, + top_k: int, + next_n: int, + compress_ratio: int = 1, + ref_vals_cache: dict | None = None, +) -> None: + """Vectorized multi-row tie-aware correctness check with strict sort+allclose. + + Per row r: scan range is ``logits[r, :N_eff(r)]`` where N_eff mirrors + the kernel's exact formula (see ``GvrTopKKernel.gvr_topk_kernel``): + + actual_kv_len = seq_lens[r // next_n] - next_n + (r % next_n) + 1 + N_eff = actual_kv_len // compress_ratio # cr=1 is identity + + Reference ``torch.topk`` is masked to this range so the reference and + kernel scan exactly the same columns under any (next_n, cr) combo + (including next_n>=3 + cr>=2 where the floor-division makes per-row + N_eff vary within a group). + + All checks (out-of-range, duplicates, n_below, sort+allclose) run as + batched GPU ops; only assertion-failure diagnostics fall back to host. + + ``ref_vals_cache``: optional id-keyed memo dict for the reference top-K + values. Pass it ONLY when (logits, seq_lens) are pinned alive for the + process lifetime by the caller (id-keying on transient tensors would + alias after GC); the caller owns that guarantee. + """ + num_rows, top_k_out = out_indices.shape + assert top_k_out == top_k + device = logits.device + logits_f32 = logits.to(torch.float32) + N = logits.shape[1] + + # Per-row N_eff mirroring the kernel formula. seq_lens is per-group + # (length num_rows // next_n); broadcast across the next_n rows of + # each group before computing actual_kv_len // cr. + row_idx = torch.arange(num_rows, device=device) + group_idx = row_idx // next_n + ofs = row_idx % next_n + seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] + actual_kv_len = seq_lens_per_row - next_n + ofs + 1 + N_eff = actual_kv_len // compress_ratio # [num_rows] + + # Reference per-row top-K, sorted descending, over logits masked beyond + # per-row N_eff. + ref_key = None + if ref_vals_cache is not None: + ref_key = (id(logits), id(seq_lens), top_k, next_n, compress_ratio) + if ref_key is not None and ref_key in ref_vals_cache: + ref_vals = ref_vals_cache[ref_key] + else: + col_idx = torch.arange(N, device=device) + in_range_mask = col_idx[None, :] < N_eff[:, None] # [num_rows, N] + masked_logits = torch.where(in_range_mask, logits_f32, float("-inf")) + ref_vals, _ = torch.topk(masked_logits, k=top_k, largest=True, sorted=True, dim=-1) + if ref_key is not None: + ref_vals_cache[ref_key] = ref_vals + + # ---- 1. Out-of-range / -1 placeholder check (single fused mask) ---- + out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) + if bool(out_of_range.any().item()): + bad_row = int(out_of_range.any(dim=1).int().argmax().item()) + bad_indices = out_indices[bad_row].cpu().tolist() + raise AssertionError( + f"row={bad_row}: kernel returned out-of-range index " + f"(N_eff={int(N_eff[bad_row].item())}, indices={bad_indices})" + ) + + # ---- 2. Duplicate-index check (sort each row, scan consecutive eq) ---- + sorted_idx, _ = out_indices.sort(dim=-1) + has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) + if bool(has_dup.any().item()): + bad_row = int(has_dup.int().argmax().item()) + raise AssertionError( + f"row={bad_row}: kernel returned duplicate indices: " + f"{out_indices[bad_row].cpu().tolist()}" + ) + + # ---- 3. Gather selected values (safe — already in range) ---- + sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) + + # ---- 4. n_below check vs per-row K-th value ---- + kth_vals = ref_vals[:, -1:] # [num_rows, 1] + n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) + if bool((n_below_per_row > 0).any().item()): + bad_row = int(n_below_per_row.argmax().item()) + n_below = int(n_below_per_row[bad_row].item()) + kth = float(kth_vals[bad_row, 0].item()) + raise AssertionError( + f"row={bad_row}: {n_below} selected values < Kth-rank value ({kth:.6f})" + ) + + # ---- 5. Strict: sorted-value multiset == torch.topk reference ---- + sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) + diff = (sel_sorted - ref_vals).abs() + if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): + per_row_max = diff.max(dim=-1).values + bad_row = int(per_row_max.argmax().item()) + max_diff = float(per_row_max[bad_row].item()) + raise AssertionError(f"row={bad_row}: sorted-value mismatch — max diff {max_diff:.4e}") + + +@pytest.fixture +def tie_aware_check(): + """The shared tie-aware top-K correctness checker (plain function).""" + return _tie_aware_check_impl diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py index f8f7215eb6ee..789749bff5b4 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py @@ -22,6 +22,9 @@ from cutlass.cute import runtime as _crt import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_dispatch as _tier_dispatch, +) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import ( GvrTopKKernel as _GvrTopKKernel, ) @@ -33,6 +36,22 @@ ) +@pytest.fixture(autouse=True) +def _tiers_off(monkeypatch): + """This module tests the IN-TREE kernel contract. Its op-level cases + (fp32 shapes within the tier envelope) would otherwise be routed to the + GVR tiers by the dispatcher — which also ignores the explicit + ``cluster_size`` these tests parametrize over — so pin the op to the + in-tree path. The GVR tiers are covered by + ``test_cute_dsl_gvr_topk_tiers.py`` (which flips the routing the other + way: fallback bands off so every test reaches a GVR tier).""" + monkeypatch.setenv("TRTLLM_GVR_TIERS_DISABLE", "1") + _tier_dispatch._reset_env_cache() + yield + monkeypatch.delenv("TRTLLM_GVR_TIERS_DISABLE", raising=False) + _tier_dispatch._reset_env_cache() + + def _make_inputs_impl( num_rows: int, N: int, @@ -133,6 +152,19 @@ def _make_inputs_impl( _ref_vals_cache: dict = {} +def _gvr_check(check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1): + """Delegate to the shared conftest checker (``tie_aware_check`` fixture), + attaching this module's reference-values memo only when (logits, + seq_lens) are pinned for the process lifetime by ``_inputs_cache`` + (id-keyed caching on transient tensors would alias after GC).""" + ref = ( + _ref_vals_cache + if any(logits is v[0] and seq_lens is v[2] for v in _inputs_cache.values()) + else None + ) + check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio, ref_vals_cache=ref) + + def _make_inputs( num_rows: int, N: int, @@ -179,108 +211,6 @@ def _make_inputs( return _inputs_cache[key] -def _tie_aware_check( - out_indices: torch.Tensor, - logits: torch.Tensor, - seq_lens: torch.Tensor, - top_k: int, - next_n: int, - compress_ratio: int = 1, -) -> None: - """Vectorized multi-row tie-aware correctness check with strict sort+allclose. - - Per row r: scan range is ``logits[r, :N_eff(r)]`` where N_eff mirrors - the kernel's exact formula (see ``GvrTopKKernel.gvr_topk_kernel``): - - actual_kv_len = seq_lens[r // next_n] - next_n + (r % next_n) + 1 - N_eff = actual_kv_len // compress_ratio # cr=1 is identity - - Reference ``torch.topk`` is masked to this range so the reference and - kernel scan exactly the same columns under any (next_n, cr) combo - (including next_n>=3 + cr>=2 where the floor-division makes per-row - N_eff vary within a group). - - All checks (out-of-range, duplicates, n_below, sort+allclose) run as - batched GPU ops; only assertion-failure diagnostics fall back to host. - """ - num_rows, top_k_out = out_indices.shape - assert top_k_out == top_k - device = logits.device - logits_f32 = logits.to(torch.float32) - N = logits.shape[1] - - # Per-row N_eff mirroring the kernel formula. seq_lens is per-group - # (length num_rows // next_n); broadcast across the next_n rows of - # each group before computing actual_kv_len // cr. - row_idx = torch.arange(num_rows, device=device) - group_idx = row_idx // next_n - ofs = row_idx % next_n - seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] - actual_kv_len = seq_lens_per_row - next_n + ofs + 1 - N_eff = actual_kv_len // compress_ratio # [num_rows] - - # Reference per-row top-K, sorted descending, over logits masked beyond - # per-row N_eff. Memoized when (logits, seq_lens) come from the pinned - # ``_inputs_cache`` (identity match), since the reference only depends on - # (logits, seq_lens, top_k, next_n, compress_ratio) — not on the launch - # variant (cluster_size / order_row / ...) the test is exercising. - ref_key = None - if any(logits is v[0] and seq_lens is v[2] for v in _inputs_cache.values()): - ref_key = (id(logits), id(seq_lens), top_k, next_n, compress_ratio) - if ref_key is not None and ref_key in _ref_vals_cache: - ref_vals = _ref_vals_cache[ref_key] - else: - col_idx = torch.arange(N, device=device) - in_range_mask = col_idx[None, :] < N_eff[:, None] # [num_rows, N] - masked_logits = torch.where(in_range_mask, logits_f32, float("-inf")) - ref_vals, _ = torch.topk(masked_logits, k=top_k, largest=True, sorted=True, dim=-1) - if ref_key is not None: - _ref_vals_cache[ref_key] = ref_vals - - # ---- 1. Out-of-range / -1 placeholder check (single fused mask) ---- - out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) - if bool(out_of_range.any().item()): - bad_row = int(out_of_range.any(dim=1).int().argmax().item()) - bad_indices = out_indices[bad_row].cpu().tolist() - raise AssertionError( - f"row={bad_row}: kernel returned out-of-range index " - f"(N_eff={int(N_eff[bad_row].item())}, indices={bad_indices})" - ) - - # ---- 2. Duplicate-index check (sort each row, scan consecutive eq) ---- - sorted_idx, _ = out_indices.sort(dim=-1) - has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) - if bool(has_dup.any().item()): - bad_row = int(has_dup.int().argmax().item()) - raise AssertionError( - f"row={bad_row}: kernel returned duplicate indices: " - f"{out_indices[bad_row].cpu().tolist()}" - ) - - # ---- 3. Gather selected values (safe — already in range) ---- - sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) - - # ---- 4. n_below check vs per-row K-th value ---- - kth_vals = ref_vals[:, -1:] # [num_rows, 1] - n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) - if bool((n_below_per_row > 0).any().item()): - bad_row = int(n_below_per_row.argmax().item()) - n_below = int(n_below_per_row[bad_row].item()) - kth = float(kth_vals[bad_row, 0].item()) - raise AssertionError( - f"row={bad_row}: {n_below} selected values < Kth-rank value ({kth:.6f})" - ) - - # ---- 5. Strict: sorted-value multiset == torch.topk reference ---- - sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) - diff = (sel_sorted - ref_vals).abs() - if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): - per_row_max = diff.max(dim=-1).values - bad_row = int(per_row_max.argmax().item()) - max_diff = float(per_row_max[bad_row].item()) - raise AssertionError(f"row={bad_row}: sorted-value mismatch — max diff {max_diff:.4e}") - - @skip_not_sm100 @pytest.mark.parametrize( "dtype,top_k", @@ -312,6 +242,7 @@ def test_cute_dsl_gvr_topk_decode( compress_ratio, preidx_hit_rate, cluster_size, + tie_aware_check, ): """Compare custom op output against torch.topk reference (tie-aware). @@ -359,7 +290,9 @@ def test_cute_dsl_gvr_topk_decode( ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) @skip_not_sm100 @@ -380,7 +313,7 @@ def test_cute_dsl_gvr_topk_decode( ], ) def test_cute_dsl_gvr_topk_decode_seqlen_sorted( - dtype, top_k, N, batch_size, varlen, next_n, compress_ratio, cluster_size + dtype, top_k, N, batch_size, varlen, next_n, compress_ratio, cluster_size, tie_aware_check ): """LJF host-side dispatch order: ``order_row`` = descending argsort of ``seq_lens`` passed through the custom op. @@ -421,7 +354,9 @@ def test_cute_dsl_gvr_topk_decode_seqlen_sorted( ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) # =========================================================================== @@ -441,7 +376,9 @@ def test_cute_dsl_gvr_topk_decode_seqlen_sorted( "dtype,top_k", [(torch.bfloat16, 512), (torch.float32, 2048)], ) -def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, cluster_size): +def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary( + dtype, top_k, cluster_size, tie_aware_check +): """GVR top-K multi-CTA short-row degrade: correctness at the cluster transition boundary. When ``cluster_size > 1``, each row is dispatched to a cluster of @@ -512,8 +449,14 @@ def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, clu cluster_size=cluster_size, ) torch.cuda.synchronize() - _tie_aware_check( - out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + _gvr_check( + tie_aware_check, + out_indices, + logits, + seq_lens, + top_k, + next_n, + compress_ratio=compress_ratio, ) # Mixed batch: alternating degrade (even rows) and co-op (odd rows). @@ -548,7 +491,9 @@ def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, clu cluster_size=cluster_size, ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) # =========================================================================== @@ -638,7 +583,9 @@ def test_lb_prepare_partition(B, ratio): ) @pytest.mark.parametrize("batch_size", [4, 32]) @pytest.mark.parametrize("next_n", [1, 2]) -def test_lb_main_branches(dtype, top_k, scenario, N, seq_lens_mode, batch_size, next_n): +def test_lb_main_branches( + dtype, top_k, scenario, N, seq_lens_mode, batch_size, next_n, tie_aware_check +): """Each LB branch (all_long / all_short / mixed) produces correct top-K. For ``mixed_half`` half the rows are forced to be short (seq_len < threshold) @@ -703,7 +650,7 @@ def test_lb_main_branches(dtype, top_k, scenario, N, seq_lens_mode, batch_size, elif scenario == "mixed_half": assert n_long == batch_size - batch_size // 2 - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1) + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1) @skip_not_sm100 @@ -733,6 +680,7 @@ def test_lb_vs_reference( batch_size, compress_ratio, preidx_hit_rate, + tie_aware_check, ): """LB kernel output matches torch.topk tie-aware reference across the same param sweep used by the single-CTA UT.""" @@ -779,7 +727,8 @@ def test_lb_vs_reference( max_batch_size=max_batch_size, ) torch.cuda.synchronize() - _tie_aware_check( + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, @@ -934,7 +883,9 @@ def _make_r0_pre_idx(logits, top_k, hint, seed): @pytest.mark.parametrize("batch_size", [1, 16]) @pytest.mark.parametrize("hint", ["real", "rand"]) @pytest.mark.parametrize("cluster_size", [1, 4, 8]) -def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hint, cluster_size): +def test_cute_dsl_gvr_topk_decode_r0_equivalence( + dtype, top_k, N, batch_size, hint, cluster_size, tie_aware_check +): """R0 admission (``enable_r0=True``, the new default) selects the same top-K as the secant baseline (``enable_r0=False``), by index set. @@ -965,8 +916,8 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hi ) # 1. Both arms independently produce a valid top-K (tie-aware value set). - _tie_aware_check(out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) - _tie_aware_check(out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) # 2. Equivalence. fp32 logits are ALMOST tie-free, so R0 and base must # return the identical index set (order-independent) — but randn @@ -997,7 +948,7 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hi ) @pytest.mark.parametrize("hint", ["real", "rand"]) def test_cute_dsl_gvr_topk_decode_r0_equivalence_bigbs( - dtype, top_k, N, batch_size, hint, cluster_size + dtype, top_k, N, batch_size, hint, cluster_size, tie_aware_check ): """Big-batch R0-vs-secant equivalence: multi-wave grids only. @@ -1017,8 +968,8 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence_bigbs( out_r0 = _run_gvr_direct( logits, pre_idx, seq_lens, top_k, enable_r0=True, cluster_size=cluster_size ) - _tie_aware_check(out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) - _tie_aware_check(out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) if dtype == torch.float32: _assert_index_sets_equal_tie_aware(out_base, out_r0, logits) @@ -1111,7 +1062,7 @@ def test_cute_dsl_gvr_topk_decode_pick_config_policy(): (torch.bfloat16, 1024, 65536, 256), # cs=1 multi-wave big-BS ], ) -def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size): +def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size, tie_aware_check): """``GvrTopKKernel.launch`` (pick_config + variant cache) produces a valid top-K at every launch-shape regime the policy can pick, including cluster_size=8. Direct-drive users get production-equivalent shapes.""" @@ -1125,13 +1076,149 @@ def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size) _GvrTopKKernel.launch(logits, pre_idx, seq_lens, out, top_k) torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out, logits, seq_lens, top_k, next_n=1, compress_ratio=1) # Override path: forcing the secant arm through launch() must also be a # valid top-K and (fp32, tie-free) the identical index set. out_sec = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") _GvrTopKKernel.launch(logits, pre_idx, seq_lens, out_sec, top_k, enable_r0=False) torch.cuda.synchronize() - _tie_aware_check(out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1) if dtype == torch.float32: assert torch.equal(out.sort(dim=-1).values, out_sec.sort(dim=-1).values) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): + """16-bit exact-tail adversarial: two DISTINCT half-precision values in + ONE fine bin straddling the K boundary (window-relative binning cannot + separate them under a wide Phase-2 bracket — e.g. fp16 1.0 vs 1.25 + under a [0, 65504]-scale bracket). The 16-bit default stays OFF (the + ambiguity gate fires on virtually every 16-bit input, measured gm + 1.29-1.36x envelope cost, while typical 16-bit inputs are value-exact + without the tail); this covers the explicit OPT-IN: with + ``p4_exact_tail=True`` the tail radix must keep every 1.25 above + every 1.0.""" + torch.manual_seed(11) + top_k, n = 1024, 32768 + bs = 2 + lo = torch.full((bs, n), -1.0, dtype=dtype, device="cuda") + # wide bracket anchors (force a wide Phase-2 window) + lo[:, :8] = torch.tensor( + [60000.0, 40000.0, 20000.0, 10000.0, 5000.0, 2500.0, 1200.0, 600.0], device="cuda" + ).to(dtype) + # boundary: (top_k - 8 - 512) high-tie values 1.25 and 1024 low-tie 1.0 + n_hi, n_lo = top_k - 8 - 512, 1024 + perm = torch.randperm(n - 8, device="cuda") + 8 + hi_pos, lo_pos = perm[:n_hi], perm[n_hi : n_hi + n_lo] + for r in range(bs): + lo[r, hi_pos] = torch.tensor(1.25, dtype=dtype, device="cuda") + lo[r, lo_pos] = torch.tensor(1.0, dtype=dtype, device="cuda") + seq_lens = torch.full((bs,), n, dtype=torch.int32, device="cuda") + pre = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre[:, 0] = lo.float().argmax(dim=-1).int() + pre[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1, p4_exact_tail=True) + torch.cuda.synchronize() + sel = torch.gather(lo.float(), -1, out.long()) + # every 1.25 must be selected before any 1.0 fills the remainder + assert int((sel == 1.25).sum()) == n_hi * bs, ( + f"exact-tail 16-bit: {(sel == 1.25).sum()} of {n_hi * bs} high-tie values selected" + ) + ref = torch.topk(lo.float(), top_k, dim=-1).values.sort(-1).values + got = sel.sort(-1).values + assert torch.equal(ref, got) + + +def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): + """The production runner's tuning adapter must agree with the kernel's + pick_cluster_size/pick_tuning single source across a shape sweep + (guards the de-duplicated launch-shape policy against drift).""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLGvrTopKDecodeRunner as R + + num_sms = 148 + for torch_dtype in (torch.float32, torch.bfloat16, torch.float16): + for num_rows in (1, 4, 32, 148, 300, 512): + for n in (4096, 16384, 65536, 131072, 262144): + for msl in (None, 262144): + cs = _GvrTopKKernel.pick_cluster_size( + num_rows, msl if msl is not None else n, num_sms + ) + cfg = _GvrTopKKernel.pick_config( + torch_dtype, num_rows, n, max_seq_len=msl, num_sms=num_sms + ) + assert cfg["cluster_size"] == cs + tuning = R._pick_tuning( + torch_dtype, + num_rows, + (msl if msl is not None else n) // cs, + num_sms, + msl, + 0, + ) + assert tuning["num_threads_per_block"] == cfg["num_threads"] + assert tuning["use_256bit_load"] == cfg["use_256bit_load"] + assert tuning["min_blocks_per_mp"] == cfg["min_blocks_per_mp"] + assert ( + tuning["enable_warp_parallel_reduce"] == cfg["enable_warp_parallel_reduce"] + ) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize( + "variant", + # base_r0off exercises the classic secant admission (the exact fallback the + # production path takes when the R0 ladder misses): its refine budget can + # run out while the bracket is still wide, which is a different route into + # the plateau terminal than the R0 ladder's. + ["rank_scatter_cs1", "rank_scatter_cs4", "snap_cs1", "base_r0off", "base_r0off_cs4"], +) +def test_cute_dsl_gvr_topk_decode_plateau_terminal(dtype, variant): + """Adversarial plateau terminal (done == 3): a bitwise-equal plateau + WIDER than the candidate buffer (kC) straddling the K boundary. Any + threshold either overflows the buffer (>= plateau) or undershoots K + (> plateau), so the Phase-2 bracket collapses to adjacent floats. + The plateau terminal must emit the sure winners plus plateau members + (exact tie-aware) instead of the old -1 pad / unfilled tail.""" + torch.manual_seed(23) + top_k, n, bs = 1024, 32768, 2 + n_hi, n_plateau = 512, 8500 # kC = 6144 for K=1024; 8500 > kC + lo = torch.full((bs, n), -1.0, dtype=dtype, device="cuda") + for r in range(bs): + perm = torch.randperm(n, device="cuda") + hi = perm[:n_hi] + pl = perm[n_hi : n_hi + n_plateau] + lo[r, hi] = (5.0 + torch.arange(n_hi, device="cuda").float() * 0.01).to(dtype) + lo[r, pl] = torch.tensor(1.0, dtype=dtype, device="cuda") + seq_lens = torch.full((bs,), n, dtype=torch.int32, device="cuda") + pre = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre[:, 0] = lo.float().argmax(dim=-1).int() + pre[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + overrides = {} + if variant == "rank_scatter_cs4": + overrides["cluster_size"] = 4 + elif variant == "snap_cs1": + overrides["enable_p4_rank_scatter"] = False + elif variant == "base_r0off": + overrides["enable_r0"] = False + elif variant == "base_r0off_cs4": + overrides["enable_r0"] = False + overrides["cluster_size"] = 4 + _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1, **overrides) + torch.cuda.synchronize() + assert int((out < 0).sum()) == 0, ( + f"{int((out < 0).sum())} unfilled/-1 slots in the plateau terminal" + ) + # in-range, unique + assert out.max() < n and out.min() >= 0 + for r in range(bs): + assert out[r].unique().numel() == top_k + sel = torch.gather(lo.float(), -1, out.long()) + assert int((sel >= 5.0).sum()) == n_hi * bs, "missing sure winners" + assert int((sel == 1.0).sum()) == (top_k - n_hi) * bs, "remaining slots must be plateau members" + ref = torch.topk(lo.float(), top_k, dim=-1).values.sort(-1).values + assert torch.equal(sel.sort(-1).values, ref) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py new file mode 100644 index 000000000000..a49be7eb0ddc --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py @@ -0,0 +1,976 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tiered GVR (direct/reg/tp CuTe DSL tiers) top-K decode tests. + +CI-sized exactness grid for the guarded fp32 fast path inside +``trtllm::cute_dsl_gvr_topk_decode`` (next_n >= 1, cr in {1, 4}): every reg +launch-table instance once, the direct and tp tiers at a few npad each, +ragged (varlen) rows with POISONED tails (stale-garbage simulation: +1e30 +beyond N_eff, which would dominate the top-K if the ragged-N masking were +broken), quantized-tie inputs, degenerate rows, pre_idx hardening, +host-only route-table asserts, a dispatcher-fallback check (bf16 routes to +the in-tree kernel), and the MTP axis: next_n in {2, 3, 4} x cr in {1, 4} +x all three tiers, each case checked against BOTH the torch.topk host +N_eff/offset simulation (the shared ``tie_aware_check`` conftest fixture) +and a differential in-tree arm (same inputs through the in-tree kernel via +the ``TRTLLM_GVR_TIERS_DISABLE`` kill switch, per-row value-multiset equality). +""" + +import os + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_dispatch as tier_dispatch, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.single_pass_multi_cta_radix_topk_cluster import ( # noqa: E501 + _query_max_cluster_size, +) +from tensorrt_llm._utils import get_sm_version + +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason=f"CuTe DSL GVR tiers only support SM 100/103, got SM {get_sm_version()}", +) + +CR = 4 # default axis of the legacy (pre-MTP) cases: compress_ratio == 4 (DSv4) +NEXT_N = 1 # default axis of the legacy (pre-MTP) cases: next_n == 1 + + +@pytest.fixture(autouse=True) +def _tiers_bands_off(monkeypatch): + """Kernel-contract tests must reach the GVR tiers: disable the measured + fallback band table (it legitimately routes several tested (npad, bs) + shapes to the in-tree kernel in production). The table itself is covered + by ``test_tiers_fallback_band_table``.""" + monkeypatch.setenv("TRTLLM_GVR_FALLBACK_BANDS", "0") + tier_dispatch._reset_env_cache() + yield + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() + + +@skip_not_sm100 +@pytest.mark.parametrize("bands", ["on", "off"]) +def test_tiers_cuda_graph_capture_replay(monkeypatch, bands, tie_aware_check): + """The op must be CUDA-graph capturable and replay-consistent on both + sides of the fallback band table. Shape (npad=32768, bs=64) sits inside + a fallback bucket: bands=on captures the in-tree branch, bands=off the + tp tier branch — the dispatch decision is host-side per shape, so it + bakes into the graph at capture time and must hold across replays, + including replays over in-place rewritten inputs.""" + if bands == "on": + # the autouse fixture pinned the table off; restore the default + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() + bs, npad, top_k = 64, 32768, 2048 + logits, pre_idx, seq_lens = _make_tier_inputs(bs, npad, top_k, seed=1234) + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + def call(): + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=NEXT_N, compress_ratio=CR + ) + + # warmup outside capture: JIT compile + any lazy init, on a side stream + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(2): + call() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + call() + g.replay() + torch.cuda.synchronize() + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + # replay over in-place rewritten inputs (fresh logits + fresh hints) + logits2, pre_idx2, seq_lens2 = _make_tier_inputs(bs, npad, top_k, seed=4321) + logits.copy_(logits2) + pre_idx.copy_(pre_idx2) + seq_lens.copy_(seq_lens2) + g.replay() + torch.cuda.synchronize() + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +def test_tiers_fallback_band_table(monkeypatch): + """The measured (npad, bs) fallback buckets route to the in-tree kernel + by default; the kill-switch restores tier service; neighbours are not + over-routed. Bands: full-grid calibration (2026-07-28), 131072 + lower bound recalibrated to 8 (2026-07-29).""" + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() + inb = tier_dispatch._in_fallback_band + # routed buckets (bucket floor <0.909 vs the in-tree head) + assert inb(256, 8192) and inb(1024, 8192) + assert inb(16, 32768) and inb(1024, 65536) + assert inb(128, 131072) and inb(64, 262144) + assert inb(8, 131072) and inb(8, 131136) # 128K x BS8 (reg-tier floor) + # ... but shapes that only ROUND into the 131072 bucket keep the + # calibrated bs>=16 routing: at bs=8 the reg tier is 1.36-1.60x ahead + # there, so the low-bs extension must not capture them. + assert not inb(8, 163776) + assert inb(16, 163776) and inb(255, 163776) and not inb(256, 163776) + assert inb(48, 40960) # off-grid npad resolves to the nearest pow2 + # neighbours stay on the tiers + assert not inb(128, 8192) # direct/tp win band + assert not inb(8, 32768) # latency reg band + assert not inb(256, 131072) and not inb(128, 262144) # large-N tp win + assert not inb(1024, 4096) # small-npad tp win + # kill-switch + monkeypatch.setenv("TRTLLM_GVR_FALLBACK_BANDS", "0") + tier_dispatch._reset_env_cache() + assert not inb(64, 32768) + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() + + +def _make_tier_inputs( + bs: int, + npad: int, + top_k: int, + seed: int, + kind: str = "randn", + varlen: bool = True, + preidx: str = "mixed", + hit_rate: float = 0.5, +): + """Build (logits fp32, pre_idx int32, seq_lens int32) for the tier path. + + ``kind='ties'`` quantizes the logits to 0.25 steps so massive value + plateaus straddle the K-th boundary (exercises the tie-ticket emits and + the max-below plateau-descent path). Ragged rows: the tail beyond each + row's N_eff is POISONED with +1e30 — stale garbage that would win the + top-K if the lane masking were broken (production tails are stale + values, not -FLT_MAX pad). + + ``preidx``: 'mixed' = argmax slot 0 + ``hit_rate`` real topk hints + (default ~50%); 'zeros' = all-zero cold start; 'oor' = out-of-range + garbage (negative / >= npad) that the kernels must clamp harmlessly. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + logits = torch.randn(bs, npad, dtype=torch.float32, device="cuda") * 2.0 + if kind == "ties": + logits = (logits * 4.0).round() * 0.25 + + if varlen: + lo = (top_k + 1) * CR # keep every row non-degenerate (N_eff > K) + seq_lens = torch.randint(lo, npad * CR + 1, (bs,), dtype=torch.int32, device="cuda") + else: + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + + n_eff = (seq_lens.long() // CR).clamp(max=npad) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + valid_logits = torch.where(tail, torch.full_like(logits, float("-inf")), logits) + argmax_idx = valid_logits.argmax(dim=-1).int() + if preidx == "zeros": + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + elif preidx == "oor": + pre_idx = torch.randint(-npad, 4 * npad, (bs, top_k), dtype=torch.int32, device="cuda") + else: + ref_topk = valid_logits.topk(top_k, dim=-1).indices.int() + keep = torch.rand(ref_topk.shape, device="cuda") < hit_rate + junk = torch.arange(top_k, dtype=torch.int32, device="cuda").expand(bs, -1) + # junk arange is in-range for every row: N_eff > top_k here. + pre_idx = torch.where(keep, ref_topk, junk).contiguous() + pre_idx[:, 0] = argmax_idx + return logits, pre_idx, seq_lens + + +def _run_op(logits, pre_idx, seq_lens, top_k): + out = torch.empty(logits.shape[0], top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out, + top_k=top_k, + next_n=NEXT_N, + compress_ratio=CR, + ) + torch.cuda.synchronize() + return out + + +def _assert_tier_routes(logits, pre_idx, seq_lens, top_k, expected_tier): + """Host-only: the dispatcher must accept this call and route it to + ``expected_tier``.""" + out = torch.empty(logits.shape[0], top_k, dtype=torch.int32, device="cuda") + assert tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ), "expected the tiered fast path to accept this call" + bs, npad = logits.shape + tier = tier_dispatch.route(bs, npad, top_k) + assert tier == expected_tier, f"route({bs}, {npad}, {top_k}) = {tier} != {expected_tier}" + + +def _skip_if_cluster_capped(bs, npad, top_k): + cs = tier_dispatch.route_cluster_size(bs, npad, top_k) + torch.zeros(1, device="cuda") # the driver-API query needs a live context + hw_max = _query_max_cluster_size() + if cs > hw_max: + pytest.skip(f"tier cluster size {cs} exceeds device max {hw_max}") + + +# --------------------------------------------------------------------------- +# Host-only route-table asserts (mirror of the original CUDA dispatch). +# --------------------------------------------------------------------------- +def test_tiers_route_table(): + r = tier_dispatch.route + # latency ladder + assert r(1, 4096, 512) == "direct" + assert r(1, 12288, 2048) == "direct" + assert r(1, 14336, 512) == "reg(cs=1,tb=512,maxv=8,ar=8)" + assert r(1, 24576, 512) == "reg(cs=4,tb=512,maxv=4,ar=8)" + assert r(1, 49152, 512) == "reg(cs=8,tb=512,maxv=3,ar=8)" + assert r(1, 65536, 1024) == "reg(cs=8,tb=512,maxv=4,ar=8)" + assert r(8, 131072, 512) == "reg(cs=8,tb=512,maxv=8,ar=8)" + assert r(1, 163840, 2048) == "reg(cs=16,tb=512,maxv=5,ar=6)" + assert r(4, 163840, 512) == "reg(cs=16,tb=512,maxv=5,ar=8)" + assert r(1, 262144, 2048) == "reg(cs=16,tb=512,maxv=8,ar=8)" + assert r(4, 262144, 1024) == "reg(cs=16,tb=512,maxv=8,ar=6)" + # dense table + assert r(64, 20480, 512) == "reg(cs=1,tb=1024,maxv=5,ar=8)" + assert r(64, 28672, 512) == "reg(cs=1,tb=1024,maxv=8,ar=8)" + assert r(8, 262144, 2048) == "reg(cs=8,tb=1024,maxv=8,ar=8)" + # tp takeover + assert r(256, 20480, 512) == "tp" + assert r(128, 28672, 512) == "tp" + assert r(16, 65536, 1024) == "tp" + assert r(128, 262144, 2048) == "tp" + # gvr tier only beyond the deployment envelope (guarded out upstream) + assert r(1, 262208, 512).startswith("cluster") + + +def test_tiers_route_env_knobs(monkeypatch): + """TRTLLM_GVR_TP_BS / TRTLLM_GVR_DENSE_BS keep the CUDA development arm's semantics: + unset/-1 -> baked bands, 0 -> disable, else explicit bs threshold.""" + r = tier_dispatch.route + try: + monkeypatch.setenv("TRTLLM_GVR_DENSE_BS", "8") + tier_dispatch._reset_env_cache() + assert r(8, 65536, 512) == "reg(cs=2,tb=1024,maxv=8,ar=8)" + assert r(8, 131072, 1024) == "reg(cs=4,tb=1024,maxv=8,ar=8)" + + monkeypatch.setenv("TRTLLM_GVR_TP_BS", "0") # 0 -> disable tp + tier_dispatch._reset_env_cache() + assert r(1024, 65536, 512) != "tp" + + monkeypatch.setenv("TRTLLM_GVR_TP_BS", "4") + tier_dispatch._reset_env_cache() + assert r(4, 65536, 512) == "tp" + finally: + monkeypatch.delenv("TRTLLM_GVR_TP_BS", raising=False) + monkeypatch.delenv("TRTLLM_GVR_DENSE_BS", raising=False) + tier_dispatch._reset_env_cache() + + +@skip_not_sm100 +def test_tiers_cluster_cap_verdict_memoized(): + """The guard's cluster-cap verdict is memoized per (bs, npad, K) — + the per-call route()+_parse_reg host cost the _DISPATCH_CACHE comment + motivates must not be re-paid by is_tiered_topk_supported on every eager + forward — and _reset_env_cache clears it (routes depend on the env + thresholds).""" + bs, npad, top_k = 4, 65536, 512 + logits, pre_idx, seq_lens = _make_tier_inputs(bs, npad, top_k, seed=0) + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + tier_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in tier_dispatch._CAP_OK_CACHE + ok = tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ) + expected = tier_dispatch.route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() + assert ok == expected + assert tier_dispatch._CAP_OK_CACHE[(bs, npad, top_k)] == expected + # Second call must hit the memo (same verdict, no recompute observable + # beyond the cache entry staying put). + assert ( + tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ) + == expected + ) + tier_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in tier_dispatch._CAP_OK_CACHE + + +# --------------------------------------------------------------------------- +# reg tier: every launch-table instance's ROUTE is asserted (host-only, +# free); a LAUNCH (JIT compile ~5-7s each) runs only for the 6 instances +# that together cover every codegen axis value (cs {1,2,8,16} x tb +# {512,1024} x ar {6,8} x maxv {5,8} + the dense-knob-only path) — the +# dropped launches differ only in axis combinations, not in code paths. +# The cs=2 dense instance is reachable only through the TRTLLM_GVR_DENSE_BS +# knob (default bands route around it), doubling as the env-knob dispatch +# test on a live launch. +# --------------------------------------------------------------------------- +_REG_INSTANCES = [ + # (npad, bs, K, dense_env, expected tier, launch) + (14336, 1, 512, None, "reg(cs=1,tb=512,maxv=8,ar=8)", True), + (24576, 1, 512, None, "reg(cs=4,tb=512,maxv=4,ar=8)", False), + (49152, 1, 512, None, "reg(cs=8,tb=512,maxv=3,ar=8)", False), + (65536, 1, 1024, None, "reg(cs=8,tb=512,maxv=4,ar=8)", False), + (131072, 8, 512, None, "reg(cs=8,tb=512,maxv=8,ar=8)", True), + (163840, 1, 2048, None, "reg(cs=16,tb=512,maxv=5,ar=6)", True), + (163840, 4, 512, None, "reg(cs=16,tb=512,maxv=5,ar=8)", False), + (262144, 1, 2048, None, "reg(cs=16,tb=512,maxv=8,ar=8)", False), + (262144, 4, 1024, None, "reg(cs=16,tb=512,maxv=8,ar=6)", False), + (20480, 64, 512, None, "reg(cs=1,tb=1024,maxv=5,ar=8)", False), + (28672, 64, 512, None, "reg(cs=1,tb=1024,maxv=8,ar=8)", True), + (65536, 8, 512, "8", "reg(cs=2,tb=1024,maxv=8,ar=8)", True), + (131072, 8, 1024, "8", "reg(cs=4,tb=1024,maxv=8,ar=8)", False), + (262144, 8, 2048, None, "reg(cs=8,tb=1024,maxv=8,ar=8)", True), +] + + +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,dense_env,expected,launch", + _REG_INSTANCES, + ids=[t[4] + f"_n{t[0]}_bs{t[1]}_k{t[2]}" for t in _REG_INSTANCES], +) +@pytest.mark.parametrize("kind", ["randn", "ties"]) +def test_tiers_reg_launch_table( + npad, bs, top_k, dense_env, expected, launch, kind, monkeypatch, tie_aware_check +): + if not launch and kind == "ties": + pytest.skip("route-assert-only instance (single kind suffices)") + _skip_if_cluster_capped(bs, npad, top_k) + try: + if dense_env is not None: + monkeypatch.setenv("TRTLLM_GVR_DENSE_BS", dense_env) + tier_dispatch._reset_env_cache() + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=npad + bs + top_k, kind=kind, varlen=True + ) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, expected) + if launch: + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + finally: + if dense_env is not None: + monkeypatch.delenv("TRTLLM_GVR_DENSE_BS", raising=False) + tier_dispatch._reset_env_cache() + + +# --------------------------------------------------------------------------- +# direct tier (npad <= DKCMAX): BS 1/8/64, randn + quantized ties, ragged + +# uniform seq_lens. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,kind,varlen", + [ + (4096, 1, 512, "randn", True), + (8256, 8, 512, "ties", True), + (12288, 64, 2048, "randn", True), + (4096, 8, 512, "randn", False), # uniform N_eff == npad (no tail) + ], +) +def test_tiers_direct(npad, bs, top_k, kind, varlen, tie_aware_check): + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=npad * 3 + bs, kind=kind, varlen=varlen + ) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "direct") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +# --------------------------------------------------------------------------- +# tp tier (bs >= tp threshold): trivial npad<=kC path, cs>1 cluster path and +# cs=1 big-npad path. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,kind", + [ + (4096, 256, 512, "randn"), # cs=1, trivial npad <= kC path + (65536, 16, 1024, "ties"), # cs=8 cluster path + tie plateaus + (262144, 128, 2048, "randn"), # cs=1, uf=8 streaming path + ], +) +def test_tiers_tp(npad, bs, top_k, kind, tie_aware_check): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=npad + 7 * bs, kind=kind, varlen=True + ) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +# --------------------------------------------------------------------------- +# Degenerate rows (N_eff <= K): in-kernel identity emit [0..N_eff-1] + -1 +# pad, mixed with normal rows in the same batch, on all three tiers. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,expected_kind", + [ + (4096, 8, 2048, "direct"), + (65536, 8, 512, "reg"), + (65536, 16, 1024, "tp"), + ], +) +def test_tiers_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=11, kind="randn", varlen=True + ) + # Rows 0..3: degenerate (N_eff = 0 / 1 / K-1 / K); the rest keep their + # non-degenerate varlen lengths. Poison the tails accordingly. + for i, sl in enumerate([NEXT_N, CR + NEXT_N, (top_k - 1) * CR, top_k * CR]): + seq_lens[i] = sl + n_eff = (seq_lens.long() // CR).clamp(max=npad) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + tier = tier_dispatch.route(bs, npad, top_k) + assert tier.startswith(expected_kind) or tier == expected_kind + out = _run_op(logits, pre_idx, seq_lens, top_k) + + for i in range(bs): + ne = int(n_eff[i].item()) + if ne <= top_k: + expect = torch.full((top_k,), -1, dtype=torch.int32, device="cuda") + expect[:ne] = torch.arange(ne, dtype=torch.int32, device="cuda") + assert torch.equal(out[i], expect), ( + f"row {i} (N_eff={ne}): degenerate identity emit mismatch: " + f"{out[i, : max(ne + 2, 8)].cpu().tolist()}..." + ) + else: + tie_aware_check( + out[i : i + 1], + logits[i : i + 1], + seq_lens[i : i + 1], + top_k, + next_n=NEXT_N, + compress_ratio=CR, + ) + + +# --------------------------------------------------------------------------- +# pre_idx hardening: all-zero (production cold start) and out-of-range +# garbage hints must neither fault nor corrupt the output. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize("preidx", ["zeros", "oor"]) +@pytest.mark.parametrize( + "npad,bs,top_k", + [ + (65536, 4, 512), # reg tier + (65536, 16, 1024), # tp tier + ], +) +def test_tiers_preidx_hardening(npad, bs, top_k, preidx, tie_aware_check): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=13, kind="randn", varlen=True, preidx=preidx + ) + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +# --------------------------------------------------------------------------- +# tp-tier hint-ladder ADMISSION fast path (R0 parity): the P2a stage-0 pick +# pushes candidates at the tightest ladder rung whose sampled CI lies in +# [K, 0.6*kC] (the legacy pivot-band hi; a leaner SAFE legacy band pick +# overrides it — stage 0b), so high-hit-rate rows finish in one fused +# streaming pass. These +# cases pin the admission decision surface: hit-rate extremes, tie plateaus +# AT the admission threshold, forced count > kC overflow fallback, and a +# batch mixing admitted and fallback rows (per-row escape independence). +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize("hit_rate", [0.85, 0.15]) +@pytest.mark.parametrize( + "npad,bs,top_k", + [ + (65536, 16, 1024), # cs=8 cluster tp (pro-like) + # cs=8 cluster tp, K=2048 (v32-like); npad pinned to 65536 so the + # JIT variant is shared with the overflow test (route/cs and the + # compiled kernel key on npad-derived cs, not on npad itself) + (65536, 16, 2048), + (131136, 128, 512), # cs=1 streaming tp (flash 512k production shape) + ], +) +def test_tiers_tp_admission_hitrate(npad, bs, top_k, hit_rate, tie_aware_check): + """High-hr rows must admit (1-pass) and low-hr rows must stay exact via + the pivot/secant fallback; both paths must produce exact top-K.""" + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=npad + bs + int(hit_rate * 100), varlen=True, hit_rate=hit_rate + ) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +@skip_not_sm100 +def test_tiers_tp_admission_tie_plateau(tie_aware_check): + """Tie plateau AT the admission threshold: K/2 distinct high values + + a 3K-wide exact-tie plateau straddling the K-th rank. The hint set is + the true top-K, so the ladder rungs land ON the plateau value and the + admitted candidate set is tie-degenerate; the P4 tie-ticket emit must + still return an exact value multiset.""" + bs, npad, top_k = 16, 65536, 1024 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + torch.manual_seed(31) + torch.cuda.manual_seed(31) + logits = -torch.rand(bs, npad, dtype=torch.float32, device="cuda") - 1.0 + for r in range(bs): + perm = torch.randperm(npad, device="cuda") + hi = perm[: top_k // 2] + plateau = perm[top_k // 2 : top_k // 2 + 3 * top_k] + logits[r, hi] = 5.0 + torch.arange(top_k // 2, device="cuda").float() * 0.01 + logits[r, plateau] = 1.0 # exact ties straddling rank K + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +@skip_not_sm100 +@pytest.mark.parametrize("top_k,n_plateau", [(1024, 8000), (2048, 12000)]) +def test_tiers_tp_admission_overflow_fallback(top_k, n_plateau, tie_aware_check): + """count(>= any admissible threshold) > kC: K-1 distinct high values, + then an n_plateau-wide exact-tie plateau (n_plateau > kC = 6144/8192). + Every threshold at or below the plateau overflows the candidate buffer + (the uncapped push counter detects it) and every threshold above it + undershoots (K-1 < K), so admission can never accept and the kernel + must fall through to the max-below plateau-descent path and its direct + emit. Exact answer: the K-1 highs plus exactly one plateau member.""" + bs, npad = 16, 65536 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + torch.manual_seed(37) + torch.cuda.manual_seed(37) + logits = -torch.rand(bs, npad, dtype=torch.float32, device="cuda") - 1.0 + for r in range(bs): + perm = torch.randperm(npad, device="cuda") + hi = perm[: top_k - 1] + plateau = perm[top_k - 1 : top_k - 1 + n_plateau] + logits[r, hi] = 5.0 + torch.arange(top_k - 1, device="cuda").float() * 0.01 + logits[r, plateau] = 1.0 + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +@skip_not_sm100 +def test_tiers_tp_admission_mixed_batch(tie_aware_check): + """Ragged batch mixing rows that admit (true-top-K hints) with rows + that must fall back (all-zero cold-start hints): per-row admission + escape is independent, so every row must stay exact regardless of + which path its cluster takes.""" + bs, npad, top_k = 16, 65536, 1024 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=41, varlen=True, hit_rate=0.9 + ) + pre_idx[1::2] = 0 # odd rows: cold-start (degenerate ladder -> fallback) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + + +# --------------------------------------------------------------------------- +# Dispatcher fallback: unsupported inputs must route to the in-tree kernel +# and still produce its results (op contract unchanged). +# --------------------------------------------------------------------------- +@skip_not_sm100 +def test_tiers_dispatcher_fallback_bf16(): + bs, npad, top_k = 4, 65536, 512 + torch.manual_seed(21) + torch.cuda.manual_seed(21) + logits = (torch.randn(bs, npad, device="cuda") * 2.0).to(torch.bfloat16) + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre_idx[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + pre_idx[:, 0] = logits.float().argmax(dim=-1).int() + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + assert not tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ), "bf16 must NOT take the tiered fast path" + # Guard-only on purpose (host, no JIT): the bf16 in-tree execution the + # op falls back to is exhaustively covered by + # test_cute_dsl_gvr_topk_decode.py; compiling that variant here again + # costs ~10s of CI for no added coverage. + + +@skip_not_sm100 +def test_tiers_dispatcher_fallback_bad_shapes(): + """Host-only: MTP contract violations must fall back to the in-tree + kernel (which asserts on them) rather than mis-launch a tier.""" + bs, npad, top_k = 4, 65536, 512 + logits = torch.randn(bs * 2, npad, dtype=torch.float32, device="cuda") + seq_lens = torch.full((bs,), npad * 4, dtype=torch.int32, device="cuda") + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs * 2, top_k, dtype=torch.int32, device="cuda") + ok = tier_dispatch.is_tiered_topk_supported + # next_n=2 with request-level pre_idx/seq_lens: accepted. + assert ok(logits, pre_idx, seq_lens, out, top_k, 2, 4, None, None) + # cr outside {1, 4}: rejected. + assert not ok(logits, pre_idx, seq_lens, out, top_k, 2, 2, None, None) + # num_rows not divisible by next_n: rejected. + assert not ok(logits, pre_idx, seq_lens, out, top_k, 3, 4, None, None) + # row-level (non-request-level) pre_idx under next_n=2: rejected. + pre_row = torch.zeros(bs * 2, top_k, dtype=torch.int32, device="cuda") + assert not ok(logits, pre_row, seq_lens, out, top_k, 2, 4, None, None) + # row-level seq_lens under next_n=2: rejected. + sl_row = torch.full((bs * 2,), npad * 4, dtype=torch.int32, device="cuda") + assert not ok(logits, pre_idx, sl_row, out, top_k, 2, 4, None, None) + + +@skip_not_sm100 +def test_tiers_accepts_order_row(tie_aware_check): + """``order_row`` (the LJF hint dsa.py computes for every batch with + num_rows >= 2 * num_sms) must NOT turn the tiers off: the guard accepts it and + the tiers ignore it, so the per-row index SET is identical with and + without the permutation (emission order is unordered by contract). + Reuses the (4096, 256, 512) tp cell compiled by ``test_tiers_tp`` (no + extra JIT).""" + npad, bs, top_k = 4096, 256, 512 + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_tier_inputs( + bs, npad, top_k, seed=npad + 7 * bs, kind="randn", varlen=True + ) + order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() + out_dummy = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + assert tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out_dummy, top_k, NEXT_N, CR, order_row, None + ), "guard must accept order_row (scheduling hint, ignored by the tiers)" + out_plain = _run_op(logits, pre_idx, seq_lens, top_k) + out_ordered = torch.empty_like(out_plain) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_ordered, + top_k=top_k, + next_n=NEXT_N, + compress_ratio=CR, + order_row=order_row, + ) + torch.cuda.synchronize() + assert torch.equal(out_plain.sort(-1).values, out_ordered.sort(-1).values), ( + "tier per-row index set must be independent of the ignored order_row permutation" + ) + tie_aware_check(out_ordered, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + # Production shape sanity (guard-only, no launch): num_rows >= 2*num_sms + # — the batch band where dsa.py always sends order_row — stays accepted. + big_bs = 2 * torch.cuda.get_device_properties().multi_processor_count + 8 + logits_b = torch.randn(big_bs, npad, dtype=torch.float32, device="cuda") + pre_b = torch.zeros(big_bs, top_k, dtype=torch.int32, device="cuda") + sl_b = torch.full((big_bs,), npad * CR, dtype=torch.int32, device="cuda") + out_b = torch.empty(big_bs, top_k, dtype=torch.int32, device="cuda") + order_b = torch.argsort(sl_b.long(), descending=True).int().contiguous() + assert tier_dispatch.is_tiered_topk_supported( + logits_b, pre_b, sl_b, out_b, top_k, NEXT_N, CR, order_b, None + ), "large-batch (num_rows >= 2*num_sms) calls with order_row must stay on the tiers" + + +@skip_not_sm100 +def test_tiers_disable_kill_switch(monkeypatch): + """TRTLLM_GVR_TIERS_DISABLE rejects everything at the guard (host-only).""" + bs, npad, top_k = 4, 65536, 512 + logits = torch.randn(bs, npad, dtype=torch.float32, device="cuda") + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + ok = tier_dispatch.is_tiered_topk_supported + assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + monkeypatch.setenv("TRTLLM_GVR_TIERS_DISABLE", "1") + tier_dispatch._reset_env_cache() + assert not ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + monkeypatch.delenv("TRTLLM_GVR_TIERS_DISABLE", raising=False) + tier_dispatch._reset_env_cache() + assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + + +def test_tiers_env_malformed_soft_fail(monkeypatch): + """Malformed tuning-knob values fail soft (warn + baked default), never + raise on the decode path.""" + baseline = tier_dispatch.route(256, 20480, 512) + for bad in ("true", " 16x", "8.5"): + monkeypatch.setenv("TRTLLM_GVR_TP_BS", bad) + tier_dispatch._reset_env_cache() + assert tier_dispatch.route(256, 20480, 512) == baseline + # whitespace-padded but valid values still parse. + monkeypatch.setenv("TRTLLM_GVR_TP_BS", " 0 ") + tier_dispatch._reset_env_cache() + assert tier_dispatch.route(256, 20480, 512).startswith(("reg", "direct")) + monkeypatch.delenv("TRTLLM_GVR_TP_BS", raising=False) + tier_dispatch._reset_env_cache() + + +# --------------------------------------------------------------------------- +# MTP (next_n > 1) + cr in {1, 4}: exactness on all three tiers, checked +# against BOTH the torch.topk host N_eff/offset simulation +# (``tie_aware_check``) and a differential in-tree arm — the SAME inputs +# through the in-tree kernel (forced via ``TRTLLM_GVR_TIERS_DISABLE``; the guard +# accepts+ignores order_row), per-row value-multiset equality. Ragged: +# seq_lens make N_eff differ across requests AND (via row % next_n) across +# the MTP rows of one request; tails are poisoned with +1e30. +# --------------------------------------------------------------------------- +def _n_eff_rows(seq_lens, num_rows, next_n, cr): + r = torch.arange(num_rows, device="cuda") + sl = seq_lens.to(device="cuda", dtype=torch.long)[r // next_n] + return (sl - next_n + (r % next_n) + 1) // cr + + +def _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed, kind, preidx, hit_rate=0.5, argmax_slot0=True +): + """(logits [bs_req*next_n, npad] fp32, pre_idx [bs_req, K] int32, + seq_lens [bs_req] int32) with poisoned per-row tails. ``preidx``: + 'noised' = per-request true-top-K hints (host-simulated cr==1 temporal + offset: hint = ref_idx - 1) mixed with junk at ``hit_rate``; 'random' = + out-of-range garbage the kernels must clamp; 'zeros' = cold start. + + ``argmax_slot0=True`` (default) enforces the op contract + ``pre_idx[..., 0] = per-group argmax`` (over the min-N_eff window, + mirroring the in-tree test suite). The in-tree kernel REQUIRES this + invariant for exactness — the differential arm is only valid with it. + The GVR tiers do not require it (clamp hardening); pass False for + tier-only robustness cases.""" + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + num_rows = bs_req * next_n + logits = torch.randn(num_rows, npad, dtype=torch.float32, device="cuda") * 2.0 + if kind == "ties": + logits = (logits * 4.0).round() * 0.25 + lo = (top_k + 1) * cr + next_n # every row non-degenerate (N_eff > K) + seq_lens = torch.randint(lo, npad * cr + 1, (bs_req,), dtype=torch.int32, device="cuda") + n_eff = _n_eff_rows(seq_lens, num_rows, next_n, cr) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + if preidx == "zeros": + pre_idx = torch.zeros(bs_req, top_k, dtype=torch.int32, device="cuda") + elif preidx == "random": + pre_idx = torch.randint(-npad, 4 * npad, (bs_req, top_k), dtype=torch.int32, device="cuda") + else: # noised + valid = torch.where(tail, torch.full_like(logits, float("-inf")), logits) + ref = valid[::next_n].topk(top_k, dim=-1).indices.int() + # host-side temporal-offset simulation: the cr==1 kernels read + # logits[hint + (row % next_n) + 1], so a prev-step hint for a + # top value at index i is (i - 1) for the first MTP row. + off = 1 if cr == 1 else 0 + good = (ref - off).clamp(min=0) + keep = torch.rand(good.shape, device="cuda") < hit_rate + junk = torch.arange(top_k, dtype=torch.int32, device="cuda").expand(bs_req, -1) + pre_idx = torch.where(keep, good, junk).contiguous() + if argmax_slot0: + min_ne = int(n_eff.min().item()) + pre_idx[:, 0] = logits[::next_n, :min_ne].argmax(dim=-1).int() + return logits, pre_idx, seq_lens + + +def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): + """tier arm + in-tree differential arm (the ``TRTLLM_GVR_TIERS_DISABLE`` kill + switch forces the in-tree path; ``order_row`` no longer does — the guard + accepts and ignores it). The ref call still passes ``order_row`` so the + in-tree sort-indirect path stays exercised. Returns (out_tier, out_ref).""" + num_rows = logits.shape[0] + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + assert tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None + ), "expected the tiered fast path to accept this MTP call" + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr + ) + order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() + out_ref = torch.empty_like(out) + os.environ["TRTLLM_GVR_TIERS_DISABLE"] = "1" + tier_dispatch._reset_env_cache() + try: + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_ref, + top_k=top_k, + next_n=next_n, + compress_ratio=cr, + order_row=order_row, + ) + finally: + del os.environ["TRTLLM_GVR_TIERS_DISABLE"] + tier_dispatch._reset_env_cache() + torch.cuda.synchronize() + return out, out_ref + + +_MTP_TIER_CELLS = { + # tier -> (npad, bs_req_base, top_k) — bs_req is scaled so num_rows = + # bs_req * next_n stays in the tier's route band for every next_n. + "direct": (4096, 4, 512), + "reg": (24576, 1, 2048), # reg(cs=4,tb=512,maxv=4,ar=8) at num_rows < 16 + "tp": (65536, 8, 1024), # tp takeover at num_rows >= 16 (npad >= 32768) +} + +# (tier, next_n, cr) — every (next_n, cr) constexpr pair is a separate JIT +# compile of BOTH arms (~15-34s each), so the full 3-tier x {2,3,4} x {1,4} +# cross is prohibitively slow in CI. Coverage kept: the tp tier (the most +# complex MTP arithmetic: cluster exchange + admission escape) runs the +# full {2,3} x {1,4} cross; direct/reg run complementary (next_n, cr) +# diagonals so each tier still sees odd/even next_n and both cr values. +# next_n=4 is dropped: 2 covers the even/row-sharing arithmetic, 3 covers +# odd division (and is the production MTP depth). +_MTP_COMBOS = [ + ("tp", 2, 1), + ("tp", 2, 4), + ("tp", 3, 1), + ("tp", 3, 4), + ("direct", 2, 4), + ("direct", 3, 1), + ("reg", 2, 1), + ("reg", 3, 4), +] + +_MTP_KINDS = [ + ("randn", "random"), # random logits + OOR-garbage hints (clamp hardening) + ("randn", "noised"), # realistic noised hints (+ host offset simulation) + ("randn", "zeros"), # all-zero cold-start hints + ("ties", "noised"), # quantized tie plateaus +] + + +@skip_not_sm100 +@pytest.mark.parametrize( + "tier,next_n,cr", _MTP_COMBOS, ids=[f"{t}-{n}-{c}" for t, n, c in _MTP_COMBOS] +) +@pytest.mark.parametrize("kind,preidx", _MTP_KINDS, ids=[f"{k}-{p}" for k, p in _MTP_KINDS]) +def test_tiers_mtp_exactness(tier, next_n, cr, kind, preidx, tie_aware_check): + npad, bs_req, top_k = _MTP_TIER_CELLS[tier] + num_rows = bs_req * next_n + if tier == "tp" and num_rows < 16: + bs_req = (16 + next_n - 1) // next_n + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + assert tier_dispatch.route(num_rows, npad, top_k).startswith(tier) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed=npad + 13 * next_n + cr, kind=kind, preidx=preidx + ) + out, out_ref = _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr) + # Independent torch.topk + host N_eff/offset simulation reference. + tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + # Differential oracle: per-row value multiset equal to the in-tree arm. + sel = torch.gather(logits, -1, out.long()).sort(-1, descending=True).values + sel_ref = torch.gather(logits, -1, out_ref.long()).sort(-1, descending=True).values + assert torch.equal(sel, sel_ref), ( + f"tier vs in-tree value-multiset mismatch (next_n={next_n}, cr={cr}, tier={tier})" + ) + + +@skip_not_sm100 +@pytest.mark.parametrize("next_n,cr", [(2, 1), (3, 4)]) +@pytest.mark.parametrize("preidx", ["zeros", "random"]) +def test_tiers_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): + """Tier-only MTP robustness: hint sets that VIOLATE the op's argmax- + slot-0 contract (pure zeros / out-of-range garbage) must neither fault + nor break exactness on the GVR tiers. No differential arm here: the + in-tree kernel requires the argmax invariant for exactness, the GVR + tiers deliberately do not (clamp hardening). Cell and (next_n, cr) + pairs match the reg combos of ``test_tiers_mtp_exactness`` so the JIT + variants are reused (zero extra compiles).""" + bs_req, npad, top_k = 1, 24576, 2048 # reg tier at num_rows 2/3 + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, + next_n, + cr, + npad, + top_k, + seed=19 + next_n + cr, + kind="randn", + preidx=preidx, + argmax_slot0=False, + ) + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + assert tier_dispatch.is_tiered_topk_supported( + logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None + ) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr + ) + torch.cuda.synchronize() + tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + + +@skip_not_sm100 +@pytest.mark.parametrize("next_n", [2, 3]) +@pytest.mark.parametrize("cr", [1, 4]) +def test_tiers_mtp_degenerate_rows(next_n, cr, tie_aware_check): + """Degenerate MTP requests (N_eff <= K for some or all of the next_n + rows): identity emit [0..N_eff-1] + -1 pad must hold per ROW, with the + per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr.""" + bs_req, npad, top_k = 8, 65536, 1024 + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed=17, kind="randn", preidx="noised" + ) + # Requests 0..3: degenerate / boundary — N_eff spans 0, 1, K-1, K, K+1 + # across their MTP rows. + for i, sl in enumerate([next_n, cr + next_n, (top_k - 1) * cr + next_n - 1, top_k * cr]): + seq_lens[i] = sl + n_eff = _n_eff_rows(seq_lens, num_rows, next_n, cr) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff.clamp(min=0)[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + out, out_ref = _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr) + for r in range(num_rows): + ne = int(n_eff[r].item()) + if ne <= top_k: + expect = torch.full((top_k,), -1, dtype=torch.int32, device="cuda") + if ne > 0: + expect[:ne] = torch.arange(ne, dtype=torch.int32, device="cuda") + assert torch.equal(out[r], expect), ( + f"row {r} (N_eff={ne}): degenerate identity emit mismatch" + ) + else: + # Per-row torch reference: collapse the row's MTP arithmetic + # into an equivalent next_n=1 seq_lens so the shared checker + # scans exactly N_eff(r) columns. + sl_row = seq_lens[r // next_n : r // next_n + 1] - next_n + (r % next_n) + 1 + tie_aware_check( + out[r : r + 1], logits[r : r + 1], sl_row, top_k, next_n=1, compress_ratio=cr + ) + nd = (n_eff > top_k).nonzero().flatten() + if nd.numel(): + sel = torch.gather(logits[nd], -1, out[nd].long()).sort(-1, descending=True).values + sel_ref = torch.gather(logits[nd], -1, out_ref[nd].long()).sort(-1, descending=True).values + assert torch.equal(sel, sel_ref)