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 c879c57c36cf..1fe0dfa5eddc 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 @@ -26,11 +26,13 @@ cluster_size: 1 (default), 2, 4 (B200 GPC limit caps at ~16). """ +import math from dataclasses import dataclass from typing import Optional import cutlass import cutlass.cute as cute +import cutlass.cute.math as cmath from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op from cutlass.utils.distributed import atomicAdd @@ -127,6 +129,25 @@ def float_as_uint32(float_val): return llvm.bitcast(cutlass.Uint32.mlir_type, float_val.ir_value()) +def float_as_int32(float_val): + """Interpret FP32 value as int32 bit pattern (cuTe DSL bit-cast).""" + return cutlass.Int32(llvm.bitcast(cutlass.Int32.mlir_type, float_val.ir_value())) + + +def f32_order_key(float_val): + """Order-preserving fp32 -> int32 key (unsigned-monotonic bit pattern). + + ``s ^ ((s >> 31) | 0x80000000)``: positive floats map to + ``bits | 0x80000000``, negative floats to ``~bits`` — the standard radix + transform whose UNSIGNED order equals fp32 order (NaN-free inputs). The + returned Int32 must only be consumed digit-wise (``(k >> s) & 0xFF``) or + via equality / prefix-equality; for a full ordered compare, flip the top + bit first (``k ^ 0x80000000`` is signed-monotonic). + """ + s = float_as_int32(float_val) + return s ^ ((s >> cutlass.Int32(31)) | cutlass.Int32(-2147483648)) + + def _fmin_f32_inline(a, b): """Single PTX ``min.f32`` → one SASS FMNMX. @@ -204,9 +225,13 @@ class GvrTopKKernel: Algorithm phases: P1: preIdx Min/Max/Mean → initial threshold - P2: Secant threshold search loop (count-only) + P1b: 256-bin histogram over prev-topK gathered values → M rung + thresholds (enable_r0 only) + P2: threshold admission — default (enable_r0=True) is a single-pass + multi-threshold rung-ladder; enable_r0=False keeps the classic + secant threshold search loop (count-only), also the R0-miss fallback P3: Ballot-free candidate collect into smem keys[]/vals[] - P4: Histogram snap (cand → exact top-K) + writeback + P4: rank-and-scatter (enable_r0) / histogram snap → exact top-K + writeback For different compress_ratio: cr = 1: preIdxOffset = (row_idx % next_n) + 1. V3.2 decode +1 temporal shift. @@ -231,6 +256,18 @@ def __init__( enable_smem_cache: bool = False, smem_cache_elems: int = 32768, seqlen_sorted: bool = False, + kc_diet: Optional[bool] = None, + enable_r0: bool = True, + r0_qfracs: Optional[tuple] = None, + mt_unroll: int = 4, + p1b_cache: Optional[bool] = None, + fb_fix: bool = True, + fb_alpha: float = 0.2, + r0_vseed: Optional[bool] = None, + enable_p4_rank_scatter: Optional[bool] = None, + enable_p4_rank_scatter_exact: Optional[bool] = None, + p4_exact_tail: Optional[bool] = None, + p4_tail_fast: Optional[bool] = None, # [p4tt] p4_warp_redundant: bool = True, p2_warp_redundant: bool = True, ): @@ -353,6 +390,188 @@ def __init__( self.FLT_MAX = 3.4028235e38 self.NEG_FLT_MAX = -self.FLT_MAX + # --- op#26 R0 histogram-ladder admission (default ON) --- + # 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). + # DEFAULT True: validated on real DSv4/V3.2 decode-capture + # 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. + # 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 + # "virtual seed" rung column in the M-ary count pass (no extra + # memory traffic or sync; the column reuses the secant per-thread + # count buffer, so SMEM does not grow). Adapts the admission + # ladder to the row's value distribution: fixes the fat-admission + # regime (a coarse quantile rung admitting ~kC candidates where + # pmean admits ~K) and donates a measured interior bracket point + # to the fallback refine on a full miss. None => enable_r0. + # mt_unroll: 4-way unroll factor for block_count_ge_multi. + # p1b_cache: stash the K gathered preIdx values in SMEM so P1b skips + # a second GMEM random gather (dtype-gated in a later commit). + # fb_fix: R0-miss fallback re-measures the rung bracket ends before + # refining (excludes the R2-class unmeasured-seed failure mode). + 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 + # 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 + # 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). + # - p1b_cache default is cs-aware: + # * cs>1 (cluster): ON for ALL dtypes. The SMEM gather-cache win + # holds and the fp32 occupancy regression that hurts the + # 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). + # * 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. + # - kC-diet: K512 single-CTA -> kC=3072 (saves 16KB SMEM; 16-bit win, + # fp32 neutral). kC>=2560 is the K512 16-bit tie-safety contract so + # 3072 is safe; the cluster port and K1024/K2048 stay stock. + 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 + # 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 + # 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 + # +2.2-2.8% across fp32/bf16/fp16 and the full BS grid (8K + # rung +10-13% at every BS, no loser cell), favorable + # synthetic +9-11%, adverse synthetic wash, exact everywhere. + # K512/K1024 unchanged: moving or widening their ladder + # measured wash-to-loss (the extra count column costs 3-7%). + if top_k == 2048: + r0_qfracs = (0.6, 0.35) if r0_vseed else (0.85, 0.35) + else: + r0_qfracs = (0.85,) if r0_vseed else (0.85, 0.35) + if enable_r0 and p1b_cache is None: + if cluster_size > 1: + p1b_cache = True + else: + p1b_cache = dtype != cutlass.Float32 + self.p1b_cache = bool(p1b_cache) + # kc_diet: None → diet iff single-CTA (tuned default). The LB hybrid + # kernel passes False for BOTH member instances so their SMEM layouts + # stay byte-identical (the DSL sizes the launch from the last-traced + # SmemAllocator only; see GvrTopKLBKernel). + if kc_diet is None: + 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 / + # 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% + # (fp32) / +10.9% (bf16) / +6.3% (fp16); favorable synthetic + # +5.2-11.0%, adverse synthetic +5.1-10.6%; no losing cell + # (fp32 min 0.994, bf16 min 1.035, fp16 min 0.999). Gated on + # enable_r0 so the retained secant path (which shares GvrParams + # and its own P4 histogram) stays byte-identical. P1b reuses this + # buffer and needs >= 256 bins, so 512 is safe. K512/K1024 + # measured as a wash under the same protocol and stay stock. + if enable_r0 and top_k == 2048 and self.kNumBins > 512: + self.kNumBins = 512 + self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else () + if self.r0_qfracs: + assert all(0.0 < q < 1.0 for q in self.r0_qfracs), self.r0_qfracs + assert list(self.r0_qfracs) == sorted(self.r0_qfracs, reverse=True), ( + "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 + # 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 + # cost) and, on a true miss, donates a measured interior bracket + # point to the fallback refine. Const-folded: r0_vseed=False kernels + # are byte-identical to before. M_qf = rungs P1b places from qneeds; + # M_thr = total columns counted/admitted (M_qf + 1 when vseed). + self.r0_vseed = bool(r0_vseed) and bool(enable_r0) and self.M_thr > 0 + self.M_qf = self.M_thr + if self.r0_vseed: + self.M_thr = self.M_qf + 1 + # need[m] = ceil(q_m * K) prev-topK values >= rung m. + self.qneeds = tuple(max(1, int(math.ceil(q * self.top_k))) for q in self.r0_qfracs) + # R1 inline shot aim in log2-count space: geometric center of the + # [K, kC] acceptance window. + self.log2_r1aim = math.log2(math.sqrt(self.top_k * self.kC)) if self.r0_qfracs else 0.0 + # fb_fix interior aim (HLS grid optimum): log2(K * (kC/K)**fb_alpha). + self.log2_mstar = ( + math.log2(self.top_k * (self.kC / self.top_k) ** float(fb_alpha)) + if self.r0_qfracs + else 0.0 + ) + + # --- op#7 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 + # 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 + # 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. + if enable_p4_rank_scatter is None: + enable_p4_rank_scatter = bool(enable_r0) + if enable_p4_rank_scatter_exact is None: + enable_p4_rank_scatter_exact = bool(enable_p4_rank_scatter) + 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). + 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 + # [p4tt] p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the + # exact-tail fire branch. When the (b*, sb*) tie class holds <= 128 + # entries (the real firing cells have 2), ONE candidate pass collects + # (value_bits, cand_idx) pairs into SMEM and thread0 selects the + # top-need exactly, replacing the 4 unconditional radix passes + # (~5.3us -> ~1 pass on pro/512k). Larger tie classes fall through to + # the existing radix select. Pure optimization (the radix backstop + # keeps exactness identical either way); False compiles the original + # 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 + # 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] + p4_tail_fast = self.p4_exact_tail and top_k >= 1024 + self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail # [p4tt] + # ------------------------------------------------------------------ # SMEM slice cache loader. Streams this CTA's slice GMEM → SMEM via # LDG → STS so Phase 2/3 can read LDS instead of re-streaming GMEM. @@ -555,6 +774,10 @@ def phase1_preidx_stats( tidx, warp_id, lane, + smem_gath=None, # cute.Tensor [top_k] f32 or None (p1b_cache): stash + # the gathered value per preIdx slot so P1b skips a 2nd GMEM gather. + s_mt_thr=None, # r0_vseed: P1 also parks pmean in the last rung + # column (visibility via P1's own trailing barrier -> zero extra sync). ): """preIdx scan + warp reduce + block aggregate + initial threshold. @@ -579,8 +802,12 @@ def phase1_preidx_stats( i = tidx + cutlass.Int32(u * self.num_threads) raw = pre_idx_row[i] idx = raw + pre_idx_offset + if cutlass.const_expr(smem_gath is not None): + smem_gath[i] = cutlass.Float32(self.NEG_FLT_MAX) if idx >= 0 and idx < N: v = self._load_fp32(input_row, idx) + if cutlass.const_expr(smem_gath is not None): + smem_gath[i] = v local_max = cute.arch.fmax(local_max, v) local_min = _fmin_f32_inline(local_min, v) local_sum = local_sum + v @@ -594,8 +821,12 @@ def phase1_preidx_stats( idx = cutlass.Int32(-1) if tidx < cutlass.Int32(pre_idx_count): idx = pre_idx_row[tidx] + pre_idx_offset + if cutlass.const_expr(smem_gath is not None): + smem_gath[tidx] = cutlass.Float32(self.NEG_FLT_MAX) if idx >= 0 and idx < N: v = self._load_fp32(input_row, idx) + if cutlass.const_expr(smem_gath is not None): + smem_gath[tidx] = v local_max = cute.arch.fmax(local_max, v) local_min = _fmin_f32_inline(local_min, v) local_sum = local_sum + v @@ -660,6 +891,8 @@ def phase1_preidx_stats( pmean = (pmin + pmax) * cutlass.Float32(0.5) cnt_lo_seed = pre_idx_count + (pre_idx_count >> 2) s_thr[0] = pmean + if cutlass.const_expr(self.r0_vseed): + s_mt_thr[self.M_thr - 1] = pmean s_thr[1] = pmin s_thr[2] = pmax s_iscalars[0] = cutlass.Int32(0) # cand_count @@ -695,6 +928,8 @@ def phase1_preidx_stats( cnt_lo_seed = pre_idx_count + (pre_idx_count >> 2) s_thr[0] = pmean + if cutlass.const_expr(self.r0_vseed): + s_mt_thr[self.M_thr - 1] = pmean s_thr[1] = pmin s_thr[2] = pmax s_iscalars[0] = cutlass.Int32(0) @@ -704,6 +939,173 @@ def phase1_preidx_stats( s_iscalars[4] = cutlass.Int32(0) cute.arch.barrier() + # ------------------------------------------------------------------ + # P1b — 256-bin SMEM histogram over the prev-topK gathered values + # (band [v_lo, v_hi] = P1's pmin/pmax = s_thr[1]/s_thr[2]), then M + # h-space quantile rungs into s_mt_thr (ascending value order). Reuses + # the Phase-4 smem_hist buffer (kNumBins >= 512 >= 256 in every spec; + # Phase 4 re-zeroes it later). Provides the R0 admission placement; it + # is only invoked from the enable_r0 path (added in a follow-up commit), + # so the base kernel is unaffected. + # ------------------------------------------------------------------ + @cute.jit + def phase1b_hspace_rungs( + self, + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ): + M = cutlass.const_expr(self.M_qf) + NB = cutlass.const_expr(256) + SEG = cutlass.const_expr(8) # NB / WARP_SIZE bins per lane + num_threads = cutlass.const_expr(self.num_threads) + + jz = tidx + while jz < cutlass.Int32(NB): + smem_hist[jz] = cutlass.Int32(0) + jz = jz + cutlass.Int32(num_threads) + cute.arch.barrier() + + v_lo = s_thr[1] + v_hi = s_thr[2] + width = (v_hi - v_lo) / cutlass.Float32(NB) # caller guards v_hi > v_lo + inv_w = cutlass.Float32(1.0) / width + + ig = tidx + while ig < cutlass.Int32(pre_idx_count): + idx = pre_idx_row[ig] + pre_idx_offset + if idx >= cutlass.Int32(0) and idx < N: + v = cutlass.Float32(input_row[idx]) + bf = (v - v_lo) * inv_w + b = cutlass.Int32(bf) + if b < cutlass.Int32(0): + b = cutlass.Int32(0) + if b > cutlass.Int32(NB - 1): + b = cutlass.Int32(NB - 1) + atomicAdd(smem_hist.iterator + b, cutlass.Int32(1)) + ig = ig + cutlass.Int32(num_threads) + cute.arch.barrier() + + # Warp-0-parallel rung extraction (a tid0 256-bin serial walk is a + # ~10-15us per-CTA dependency chain). Lane l owns the SEG consecutive + # bins descending from bin NB-1-l*SEG; segment sums -> 5-step shfl_up + # inclusive scan gives each lane the cumulative count of all + # higher-value bins; each lane then walks its SEG bins once and fires + # rung m at the unique crossing bin (cum_before < qneeds[m] <= + # cum_at). qfracs descending in h => thresholds ascending in m. + if warp_id == cutlass.Int32(0): + top = cutlass.Int32(NB - 1) - lane * cutlass.Int32(SEG) + seg_frag = cute.make_fragment((SEG,), cutlass.Int32) + part = cutlass.Int32(0) + for j in cutlass.range_constexpr(SEG): + v8 = smem_hist[top - cutlass.Int32(j)] + seg_frag[j] = v8 + part = part + v8 + tp = part + for off_i in cutlass.range_constexpr(5): + off_v = cutlass.const_expr(1 << off_i) + other = cute.arch.shuffle_sync_up(tp, off_v, mask_and_clamp=0) + if lane >= cutlass.Int32(off_v): + tp = tp + other + excl = tp - part # cum of all bins above my segment + total = cute.arch.shuffle_sync(tp, cutlass.Int32(self.WARP_SIZE - 1)) + run = cutlass.Int32(0) + for j in cutlass.range_constexpr(SEG): + run = run + seg_frag[j] + cum_at = excl + run + cum_before = cum_at - seg_frag[j] + for m in cutlass.range_constexpr(M): + if cum_at >= cutlass.Int32(self.qneeds[m]) and cum_before < cutlass.Int32( + self.qneeds[m] + ): + s_mt_thr[m] = v_lo + cutlass.Float32(top - cutlass.Int32(j)) * width + # unfired rungs (heavy invalid-preIdx rows: total < need): v_lo + if lane == 0: + for m in cutlass.range_constexpr(M): + if total < cutlass.Int32(self.qneeds[m]): + s_mt_thr[m] = v_lo + cute.arch.barrier() + + # ------------------------------------------------------------------ + # P1b (p1b_cache variant) — build the rung histogram from the SMEM + # gathered values that P1 stashed (smem_gath), skipping P1b's second + # GMEM random gather. Sentinel NEG_FLT_MAX marks invalid/out-of-range + # preIdx slots. Rung extraction is identical to phase1b_hspace_rungs. + # ------------------------------------------------------------------ + @cute.jit + def phase1b_hspace_rungs_cached( + self, pre_idx_count, smem_gath, smem_hist, s_thr, s_mt_thr, tidx, warp_id, lane + ): + M = cutlass.const_expr(self.M_qf) + NB = cutlass.const_expr(256) + SEG = cutlass.const_expr(8) + num_threads = cutlass.const_expr(self.num_threads) + + jz = tidx + while jz < cutlass.Int32(NB): + smem_hist[jz] = cutlass.Int32(0) + jz = jz + cutlass.Int32(num_threads) + cute.arch.barrier() + + v_lo = s_thr[1] + v_hi = s_thr[2] + width = (v_hi - v_lo) / cutlass.Float32(NB) + inv_w = cutlass.Float32(1.0) / width + + ig = tidx + while ig < cutlass.Int32(pre_idx_count): + v = smem_gath[ig] + if v > cutlass.Float32(self.NEG_FLT_MAX): + bf = (v - v_lo) * inv_w + b = cutlass.Int32(bf) + if b < cutlass.Int32(0): + b = cutlass.Int32(0) + if b > cutlass.Int32(NB - 1): + b = cutlass.Int32(NB - 1) + atomicAdd(smem_hist.iterator + b, cutlass.Int32(1)) + ig = ig + cutlass.Int32(num_threads) + cute.arch.barrier() + + if warp_id == cutlass.Int32(0): + top = cutlass.Int32(NB - 1) - lane * cutlass.Int32(SEG) + seg_frag = cute.make_fragment((SEG,), cutlass.Int32) + part = cutlass.Int32(0) + for j in cutlass.range_constexpr(SEG): + v8 = smem_hist[top - cutlass.Int32(j)] + seg_frag[j] = v8 + part = part + v8 + tp = part + for off_i in cutlass.range_constexpr(5): + off_v = cutlass.const_expr(1 << off_i) + other = cute.arch.shuffle_sync_up(tp, off_v, mask_and_clamp=0) + if lane >= cutlass.Int32(off_v): + tp = tp + other + excl = tp - part + total = cute.arch.shuffle_sync(tp, cutlass.Int32(self.WARP_SIZE - 1)) + run = cutlass.Int32(0) + for j in cutlass.range_constexpr(SEG): + run = run + seg_frag[j] + cum_at = excl + run + cum_before = cum_at - seg_frag[j] + for m in cutlass.range_constexpr(M): + if cum_at >= cutlass.Int32(self.qneeds[m]) and cum_before < cutlass.Int32( + self.qneeds[m] + ): + s_mt_thr[m] = v_lo + cutlass.Float32(top - cutlass.Int32(j)) * width + if lane == 0: + for m in cutlass.range_constexpr(M): + if total < cutlass.Int32(self.qneeds[m]): + s_mt_thr[m] = v_lo + cute.arch.barrier() + # ------------------------------------------------------------------ # block_count_ge — GE-count of input vs threshold (shared by P2/P3). # Per-thread strided accumulate → smem_ptcnt[tid] (for P3 prefix sum) @@ -960,6 +1362,163 @@ def block_count_ge( return cutlass.Int32(0) + # ------------------------------------------------------------------ + # block_count_ge_multi — GE-count of the input row against M + # thresholds in ONE vectorized scan, reusing block_count_ge's memory + # 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 + # 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 + # cluster (cluster_size>1, do_cluster_sync) with a release cluster_arrive + # mirroring block_count_ge; at cs==1 (or short-row degrade) the local + # totals are the answer. smem_ptcnt_multi holds slice-local per-thread + # columns (the accepted rung's column seeds Phase 3 per CTA). + # ------------------------------------------------------------------ + @cute.jit + def block_count_ge_multi( + self, + input_row, + slice_start, + slice_end, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_cluster_partial_m, + do_cluster_sync, + tidx, + warp_id, + lane, + smem_ptcnt=None, # vseed: last column's per-thread counts land here + ): + M = cutlass.const_expr(self.M_thr) + num_threads = cutlass.const_expr(self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + cluster_size = cutlass.const_expr(self.cluster_size) + vec_w = cutlass.const_expr(self.vec_bits // self.dtype.width) + elem_bytes = cutlass.const_expr(self.dtype.width // 8) + vec_align = cutlass.const_expr(self.vec_align_bytes) + copy_atom = self._make_load_copy_atom() + step_elem = cutlass.const_expr(num_threads * vec_w) + + thr_frag = cute.make_fragment((M,), cutlass.Float32) + cnt_frag = cute.make_fragment((M,), cutlass.Int32) + for m in cutlass.range_constexpr(M): + thr_frag[m] = s_mt_thr[m] + cnt_frag[m] = cutlass.Int32(0) + + row_addr = input_row.iterator.toint() + slice_len = slice_end - slice_start + n_aligned = slice_start + (slice_len // cutlass.Int32(vec_w)) * cutlass.Int32(vec_w) + i = slice_start + tidx * cutlass.Int32(vec_w) + step = cutlass.Int32(step_elem) + + if self.enable_unroll_4: + rng_frag = cute.make_fragment((vec_w,), self.dtype) + big_iters = cutlass.Int32(0) + if slice_end > i + cutlass.Int32(vec_w - 1): + big_iters = (slice_end - i - cutlass.Int32(vec_w)) // cutlass.Int32( + step_elem + ) + cutlass.Int32(1) + for k in cutlass.range(big_iters, unroll=self.mt_unroll): + i_local = i + k * cutlass.Int32(step_elem) + src_ptr_k = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(i_local) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + src_k = cute.make_tensor(src_ptr_k, cute.make_layout((vec_w,))) + cute.copy(copy_atom, src_k, rng_frag) + for j in cutlass.range_constexpr(vec_w): + if cutlass.const_expr(self.dtype == cutlass.Float32): + vj = rng_frag[j] + else: + vj = cutlass.Float32(rng_frag[j]) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) + i = i + big_iters * cutlass.Int32(step_elem) + + tail_frag = cute.make_fragment((vec_w,), self.dtype) + while i + cutlass.Int32(vec_w - 1) < slice_end: + src_ptr = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(i) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + src = cute.make_tensor(src_ptr, cute.make_layout((vec_w,))) + cute.copy(copy_atom, src, tail_frag) + for j in cutlass.range_constexpr(vec_w): + if cutlass.const_expr(self.dtype == cutlass.Float32): + vj = tail_frag[j] + else: + vj = cutlass.Float32(tail_frag[j]) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) + i = i + step + + it = n_aligned + tidx + while it < slice_end: + v = self._load_fp32(input_row, it) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(v >= thr_frag[m]) + it = it + cutlass.Int32(num_threads) + + for m in cutlass.range_constexpr(M): + if cutlass.const_expr(self.r0_vseed and m == self.M_qf): + smem_ptcnt[tidx] = cnt_frag[m] + else: + smem_ptcnt_multi[m * num_threads + tidx] = cnt_frag[m] + + for m in cutlass.range_constexpr(M): + wc = self.warp_reduce_sum_i32(cnt_frag[m]) + if lane == 0: + smem_wcnt_multi[m * num_warps + warp_id] = wc + cute.arch.barrier() + # Block-reduce the M warp counts to this CTA's slice totals. Stage + # into DSMEM scratch at cs>1 (for the cluster merge below), else + # write straight to s_mt_cnt. + if warp_id == cutlass.Int32(0): + for m in cutlass.range_constexpr(M): + v = cutlass.Int32(0) + if lane < cutlass.Int32(num_warps): + v = smem_wcnt_multi[m * num_warps + lane] + total = self.warp_reduce_sum_i32(v) + if lane == cutlass.Int32(0): + if cutlass.const_expr(cluster_size > 1): + s_cluster_partial_m[m] = total + else: + s_mt_cnt[m] = total + cute.arch.barrier() + if cutlass.const_expr(cluster_size > 1): + if do_cluster_sync: + # Release arrive (NOT relaxed): pairs with the peer + # cluster_wait acquire so the staged M totals are visible + # before any CTA reads them over DSMEM. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if tidx == cutlass.Int32(0): + local_ptr = s_cluster_partial_m.iterator + for m in cutlass.range_constexpr(M): + total = cutlass.Int32(0) + for peer in cutlass.range_constexpr(cluster_size): + peer_addr = mapa_shared_cluster( + local_ptr + cutlass.Int32(m), cutlass.Int32(peer) + ) + total = total + ld_shared_cluster_i32(peer_addr) + s_mt_cnt[m] = total + cute.arch.barrier() + else: + # short-row degrade: this CTA's local totals are the answer. + if tidx == cutlass.Int32(0): + for m in cutlass.range_constexpr(M): + s_mt_cnt[m] = s_cluster_partial_m[m] + cute.arch.barrier() + # ------------------------------------------------------------------ # Phase 2: Secant-interpolation threshold search # Refines threshold to bring cand_count into [kK, kCC] using secant @@ -1887,6 +2446,753 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane sel_out = cute.arch.shuffle_sync(sel_loc, src) return thr_out, sel_out + # ------------------------------------------------------------------ + # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter). + # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py. + # ------------------------------------------------------------------ + @cute.jit + def phase4_rank_scatter( + self, + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count, + tidx, + warp_id, + lane, + ): + kK = cutlass.const_expr(self.top_k) + kBins = cutlass.const_expr(self.kNumBins) + num_threads = cutlass.const_expr(self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + bins_per_warp = cutlass.const_expr(kBins // self.num_warps) + + if cand_count == cutlass.Int32(kK): + i4 = tidx + while i4 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i4] = self.dtype(smem_keys[i4]) + output_indices_row[i4] = smem_vals[i4] + i4 = i4 + cutlass.Int32(num_threads) + elif cand_count > cutlass.Int32(kK): + # ---- block min/max over candidates ---- + local_cmin = cutlass.Float32(self.FLT_MAX) + local_cmax = cutlass.Float32(self.NEG_FLT_MAX) + i5 = tidx + while i5 < cand_count: + v = smem_keys[i5] + local_cmin = _fmin_f32_inline(local_cmin, v) + local_cmax = cute.arch.fmax(local_cmax, v) + i5 = i5 + cutlass.Int32(num_threads) + cmin = self.warp_reduce_min_f32(local_cmin) + cmax = self.warp_reduce_max_f32(local_cmax) + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = float_as_uint32(cmin) + smem_hist[warp_id] = float_as_uint32(cmax) + cute.arch.barrier() + bmin_r = cutlass.Float32(self.FLT_MAX) + bmax_r = cutlass.Float32(self.NEG_FLT_MAX) + for w in cutlass.range_constexpr(self.num_warps): + vmin = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value()) + ) + vmax = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[w].ir_value()) + ) + bmin_r = _fmin_f32_inline(bmin_r, vmin) + bmax_r = cute.arch.fmax(bmax_r, vmax) + if bmax_r <= bmin_r: + bmax_r = bmin_r + cutlass.Float32(1e-6) + cute.arch.barrier() + # ---- zero + build histogram ---- + i6 = tidx + while i6 < cutlass.Int32(kBins): + smem_hist[i6] = cutlass.Int32(0) + i6 = i6 + cutlass.Int32(num_threads) + cute.arch.barrier() + range1 = bmax_r - bmin_r + inv1 = (cutlass.Float32(kBins - 1) + cutlass.Float32(0.99)) / range1 + i7 = tidx + while i7 < cand_count: + vk = smem_keys[i7] + bin_i = cutlass.Int32((vk - bmin_r) * inv1) + if bin_i < cutlass.Int32(0): + bin_i = cutlass.Int32(0) + if bin_i > cutlass.Int32(kBins - 1): + bin_i = cutlass.Int32(kBins - 1) + atomicAdd(smem_hist.iterator + bin_i, cutlass.Int32(1)) + i7 = i7 + cutlass.Int32(num_threads) + cute.arch.barrier() + # ---- 3-step high→low bin search → straddling bin b* + rank_above ---- + warp_bin_sum = cutlass.Int32(0) + for jb in cutlass.range_constexpr(bins_per_warp): + bidx_s = ( + cutlass.Int32(kBins - 1) + - warp_id * cutlass.Int32(bins_per_warp) + - cutlass.Int32(jb) + ) + warp_bin_sum = warp_bin_sum + smem_hist[bidx_s] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = warp_bin_sum + cute.arch.barrier() + if tidx == cutlass.Int32(0): + cum = cutlass.Int32(0) + tw = cutlass.Int32(num_warps - 1) + found = cutlass.Int32(0) + for w2 in cutlass.range_constexpr(self.num_warps): + cum = cum + smem_wcnt[w2] + if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0): + tw = cutlass.Int32(w2) + found = cutlass.Int32(1) + cum2 = cutlass.Int32(0) + for w3 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w3) < tw: + cum2 = cum2 + smem_wcnt[w3] + s_iscalars[2] = cum2 # prefix-count before target warp + s_iscalars[3] = tw + cute.arch.barrier() + target_warp = s_iscalars[3] + if warp_id == target_warp and lane == cutlass.Int32(0): + base_cum = s_iscalars[2] + b_star = cutlass.Int32(kBins - 1) + rank_above = base_cum + set_d = cutlass.Int32(0) + for jb2 in cutlass.range_constexpr(bins_per_warp): + bidx2 = ( + cutlass.Int32(kBins - 1) + - target_warp * cutlass.Int32(bins_per_warp) + - cutlass.Int32(jb2) + ) + ra_before = base_cum + base_cum = base_cum + smem_hist[bidx2] + if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0): + b_star = bidx2 + rank_above = ra_before # count in bins strictly above b* + set_d = cutlass.Int32(1) + s_iscalars[2] = rank_above + s_iscalars[3] = b_star + s_iscalars[4] = cutlass.Int32(0) # cnt_above + s_iscalars[1] = cutlass.Int32(0) # cnt_straddle + cute.arch.barrier() + b_star = s_iscalars[3] + rank_above = s_iscalars[2] + + # ---- EXACT: one fine-histogram recursion on the straddling bin b* ---- + if cutlass.const_expr(self.enable_p4_rank_scatter_exact): + # FIXED small fine-bin count (independent of kNumBins) — cuts the + # re-zero + 3-step cost (esp. K=2048 where kNumBins=2048); 256 + # sub-bins over bin b* gives kNumBins×256 effective resolution, + # enough to resolve the straddling bin to ≤1 distinct value. + fbins = cutlass.const_expr(256) + fbpw = cutlass.const_expr(256 // self.num_warps) + # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1) + f_lo = bmin_r + cutlass.Float32(b_star) / inv1 + finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1 + # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands + iz = tidx + while iz < cutlass.Int32(fbins): + smem_hist[iz] = cutlass.Int32(0) + iz = iz + cutlass.Int32(num_threads) + cute.arch.barrier() + ifb = tidx + while ifb < cand_count: + vf = smem_keys[ifb] + cb = cutlass.Int32((vf - bmin_r) * inv1) + if cb < cutlass.Int32(0): + cb = cutlass.Int32(0) + if cb > cutlass.Int32(kBins - 1): + cb = cutlass.Int32(kBins - 1) + if cb == b_star: + sb = cutlass.Int32((vf - f_lo) * finv) + if sb < cutlass.Int32(0): + sb = cutlass.Int32(0) + if sb > cutlass.Int32(fbins - 1): + sb = cutlass.Int32(fbins - 1) + atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1)) + ifb = ifb + cutlass.Int32(num_threads) + cute.arch.barrier() + # fine 3-step search seeded at rank_above (over fbins bins) + fws = cutlass.Int32(0) + for jbf in cutlass.range_constexpr(fbpw): + bif = ( + cutlass.Int32(fbins - 1) + - warp_id * cutlass.Int32(fbpw) + - cutlass.Int32(jbf) + ) + fws = fws + smem_hist[bif] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = fws + cute.arch.barrier() + if tidx == cutlass.Int32(0): + cumf = rank_above + twf = cutlass.Int32(num_warps - 1) + fnd = cutlass.Int32(0) + for w2 in cutlass.range_constexpr(self.num_warps): + cumf = cumf + smem_wcnt[w2] + if cumf >= cutlass.Int32(kK) and fnd == cutlass.Int32(0): + twf = cutlass.Int32(w2) + fnd = cutlass.Int32(1) + pre = rank_above + for w3 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w3) < twf: + pre = pre + smem_wcnt[w3] + # Stage prefix/target-warp metadata in spare s_iscalars + # slots, NOT smem_hist[0]/[1]: the last fine warp's reverse + # scan below walks fine bins down to 0/1, so reusing those + # histogram bins as scratch would corrupt sb_star/ra_fine + # when twf2 == num_warps-1. Slots [4]/[1] are dead here + # (re-zeroed at the cnt_above/cnt_strad reset below). + s_iscalars[4] = pre # prefix into target fine warp + s_iscalars[1] = twf # target fine warp + cute.arch.barrier() + pre_f = s_iscalars[4] + twf2 = s_iscalars[1] + if warp_id == twf2 and lane == cutlass.Int32(0): + base_f = pre_f + sb_star = cutlass.Int32(fbins - 1) + ra_fine = base_f + sd = cutlass.Int32(0) + for jb3 in cutlass.range_constexpr(fbpw): + sbi = ( + cutlass.Int32(fbins - 1) + - twf2 * cutlass.Int32(fbpw) + - cutlass.Int32(jb3) + ) + ra_b = base_f + base_f = base_f + smem_hist[sbi] + if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0): + sb_star = sbi + ra_fine = ra_b + sd = cutlass.Int32(1) + smem_hist[2] = sb_star + smem_hist[3] = ra_fine + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # cnt_above + s_iscalars[0] = cutlass.Int32(0) # cnt_mid (b*, sub>sb*) + s_iscalars[1] = cutlass.Int32(0) # cnt_strad (b*, sub==sb*) + cute.arch.barrier() + sb_star = smem_hist[2] + rank_above_fine = smem_hist[3] + isc = tidx + while isc < cand_count: + v = smem_keys[isc] + bin_i = cutlass.Int32((v - bmin_r) * inv1) + if bin_i < cutlass.Int32(0): + bin_i = cutlass.Int32(0) + if bin_i > cutlass.Int32(kBins - 1): + bin_i = cutlass.Int32(kBins - 1) + if bin_i > b_star: + pos = atomicAdd(s_iscalars.iterator + cutlass.Int32(4), cutlass.Int32(1)) + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(v) + output_indices_row[pos] = smem_vals[isc] + elif bin_i == b_star: + sb = cutlass.Int32((v - f_lo) * finv) + if sb < cutlass.Int32(0): + sb = cutlass.Int32(0) + if sb > cutlass.Int32(fbins - 1): + sb = cutlass.Int32(fbins - 1) + if sb > sb_star: + o = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1)) + pos = rank_above + o + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(v) + output_indices_row[pos] = smem_vals[isc] + elif sb == sb_star: + o = atomicAdd(s_iscalars.iterator + cutlass.Int32(1), cutlass.Int32(1)) + pos = rank_above_fine + o + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(v) + output_indices_row[pos] = smem_vals[isc] + isc = isc + cutlass.Int32(num_threads) + cute.arch.barrier() + cnt_strad = s_iscalars[1] + filled = rank_above_fine + cnt_strad + if filled > cutlass.Int32(kK): + filled = cutlass.Int32(kK) + ipad = filled + tidx + while ipad < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[ipad] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[ipad] = cutlass.Int32(-1) + ipad = ipad + cutlass.Int32(num_threads) + + # ---- EXACT-TAIL repair (p4_exact_tail, fp32): the fine bin + # resolves values to range/(kBins*fbins); two candidates + # closer than that can straddle the kK boundary inside ONE + # fine bin, and the arrival-order fill above then keeps an + # arbitrary subset. Gated on the ONLY case where that is + # ambiguous — the tie set overfills the remaining slots — this + # re-ranks the (b*, sb*) tie set exactly via an MSB-first + # 8-bit-digit radix select over the order-preserving integer + # keys (4 levels = bit-exact for fp32) and rewrites the tie + # slot range [rank_above_fine, kK). Unambiguous rows (the + # overwhelming majority) pay two scalar compares; the counters + # and the fine histogram are reused, so SMEM does not grow. + # [p4tt] tiny-tie fast path: when the exact-tail gate fires + # with a small (b*, sb*) tie class (cnt_strad <= 128 — the + # real firing cells hold 2), ONE candidate pass collects the + # class and thread0 selects the top-need exactly, replacing + # the 4 unconditional radix passes. Larger classes take the + # UNMODIFIED radix select below (verbatim copy). + if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast): # [p4tt] + need0 = cutlass.Int32(kK) - rank_above_fine + if cnt_strad > need0 and need0 > cutlass.Int32(0): + if cnt_strad <= cutlass.Int32(128): + # [p4tt] SMEM: (value_bits, cand_idx) pairs at + # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255). + # The 256 digit bins are dead here (the fast path + # replaces the radix levels that used them); the + # sb_star/ra staging in slots 2/3 was read by + # every thread before the pre-scatter barrier. + # Persistent radix scalars [256..258] untouched. + # Collect counter = s_iscalars[0] (dead after the + # scatter; same reuse as the radix rewrite pass). + if tidx == cutlass.Int32(0): + s_iscalars[0] = cutlass.Int32(0) + cute.arch.barrier() + itc = tidx + while itc < cand_count: + tv = smem_keys[itc] + tb = cutlass.Int32((tv - bmin_r) * inv1) + if tb < cutlass.Int32(0): + tb = cutlass.Int32(0) + if tb > cutlass.Int32(kBins - 1): + tb = cutlass.Int32(kBins - 1) + if tb == b_star: + ts = cutlass.Int32((tv - f_lo) * finv) + if ts < cutlass.Int32(0): + ts = cutlass.Int32(0) + if ts > cutlass.Int32(fbins - 1): + ts = cutlass.Int32(fbins - 1) + if ts == sb_star: + to = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1) + ) + if to < cutlass.Int32(128): + smem_hist[to + to] = float_as_int32(tv) + smem_hist[to + to + cutlass.Int32(1)] = smem_vals[itc] + itc = itc + cutlass.Int32(num_threads) + cute.arch.barrier() + # [p4tt] thread0 exact top-need0 select rewriting + # positions [rank_above_fine, kK). Consumed flag = + # the cand_idx slot set to -1 (indices are always + # >= 0), so a genuine -FLT_MAX value in the class + # remains selectable (no value sentinel). Ties + # (bit-equal values) pick arbitrarily: value-set + # exact. + if tidx == cutlass.Int32(0): + tj = cutlass.Int32(0) + while tj < need0: + tbv = cutlass.Float32(self.NEG_FLT_MAX) + tbi = cutlass.Int32(-1) + ti = cutlass.Int32(0) + while ti < cnt_strad: + tvi = smem_hist[ti + ti + cutlass.Int32(1)] + if tvi >= cutlass.Int32(0): + tvb = smem_hist[ti + ti] + tvv = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + tvb.ir_value(), + ) + ) + take = cutlass.Int32(0) + if tbi < cutlass.Int32(0): + take = cutlass.Int32(1) + elif tvv > tbv: + take = cutlass.Int32(1) + if take == cutlass.Int32(1): + tbv = tvv + tbi = ti + ti = ti + cutlass.Int32(1) + pos = rank_above_fine + tj + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(tbv) + output_indices_row[pos] = smem_hist[ + tbi + tbi + cutlass.Int32(1) + ] + smem_hist[tbi + tbi + cutlass.Int32(1)] = cutlass.Int32(-1) + 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() + 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() + else: + # ---- APPROX rank-and-scatter (single pass), arbitrary straddling order ---- + isc = tidx + while isc < cand_count: + v = smem_keys[isc] + bin_i = cutlass.Int32((v - bmin_r) * inv1) + if bin_i < cutlass.Int32(0): + bin_i = cutlass.Int32(0) + if bin_i > cutlass.Int32(kBins - 1): + bin_i = cutlass.Int32(kBins - 1) + if bin_i > b_star: + pos = atomicAdd(s_iscalars.iterator + cutlass.Int32(4), cutlass.Int32(1)) + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(v) + output_indices_row[pos] = smem_vals[isc] + elif bin_i == b_star: + off = atomicAdd(s_iscalars.iterator + cutlass.Int32(1), cutlass.Int32(1)) + pos = rank_above + off + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(v) + output_indices_row[pos] = smem_vals[isc] + isc = isc + cutlass.Int32(num_threads) + cute.arch.barrier() + cnt_strad = s_iscalars[1] + filled = rank_above + cnt_strad + if filled > cutlass.Int32(kK): + filled = cutlass.Int32(kK) + ipad = filled + tidx + while ipad < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[ipad] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[ipad] = cutlass.Int32(-1) + ipad = ipad + cutlass.Int32(num_threads) + else: + i10 = tidx + while i10 < cand_count: + if cutlass.const_expr(self.return_output_values): + 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) + # ------------------------------------------------------------------ # Phase 4: Histogram-based k-th selection + two-pass writeback # ------------------------------------------------------------------ @@ -2624,23 +3930,24 @@ def run_one_row( # slot k&1 — closes the straggler-read-vs-next-write DSMEM race), # slot 2 = tid0-private call counter. mapa.shared::cluster relies # on every CTA holding this block at the SAME SMEM offset, so it's - # allocated once here. Only needed at cs>1; skipped at cs=1 (all - # uses are gated by const_expr(cs>1) and None propagates - # harmlessly through the unused parameters). + # allocated once here. Only USED at cs>1 (uses are gated by + # const_expr(cs>1)), but ALLOCATED unconditionally: the LB hybrid + # kernel inlines a cs>1 and a cs=1 instance into one launch, and the + # DSL sizes the launch SMEM from the last-traced SmemAllocator only — + # the layouts must stay byte-identical across cluster_size (16B cost + # at cs=1). + s_cluster_partial = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((3,), order=(0,)), + byte_alignment=16, + ) if cutlass.const_expr(cluster_size > 1): - s_cluster_partial = smem.allocate_tensor( - element_type=cutlass.Int32, - layout=cute.make_ordered_layout((3,), order=(0,)), - byte_alignment=16, - ) # Zero the call counter before any block_count_ge call. tid0- # private (same thread reads/increments it), so program order # suffices — but parity must start at 0 on EVERY CTA of the # cluster for lockstep alignment. if tidx == cutlass.Int32(0): s_cluster_partial[2] = cutlass.Int32(0) - else: - s_cluster_partial = None # SMEM slice cache (optional). Sized in ``self.dtype`` so the same # vec_w-wide LDG→STS→LDS pipeline works for fp32/bf16/fp16. @@ -2655,6 +3962,74 @@ def run_one_row( else: smem_input = None + # op#26 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 + # / smem_input above). smem_ptcnt_multi caches M per-thread count + # columns; s_r0col carries the accepted rung index tid0 -> all. + if cutlass.const_expr(self.enable_r0): + M_r0 = cutlass.const_expr(self.M_thr) + # vseed (v3): the pmean column's per-thread counts reuse the + # existing single-column smem_ptcnt buffer, so the BIG multi + # buffer only holds the M_qf rung columns -> zero smem growth + # (the round-1 +2-4KB column pushed 16-bit mb3/T1024 configs over + # an occupancy cliff: K2048 fp16 BS1024 -26%). + M_r0_pt = cutlass.const_expr(self.M_qf) + s_mt_thr = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((M_r0,), order=(0,)), + byte_alignment=16, + ) + smem_ptcnt_multi = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((M_r0_pt * num_threads,), order=(0,)), + byte_alignment=128, + ) + smem_wcnt_multi = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((M_r0 * num_warps,), order=(0,)), + byte_alignment=64, + ) + s_mt_cnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((M_r0,), order=(0,)), + byte_alignment=16, + ) + s_r0col = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((1,), order=(0,)), + byte_alignment=16, + ) + # DSMEM scratch for the M-way cluster all-reduce of the R0 rung + # counts (mapa.shared::cluster needs the same offset on every + # CTA). Only USED at cs>1; allocated unconditionally so the + # cs=1 / cs>1 SMEM layouts stay byte-identical for the LB + # hybrid kernel (see s_cluster_partial above). + s_cluster_partial_m = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((M_r0,), order=(0,)), + byte_alignment=16, + ) + # p1b_cache: P1 stashes the K gathered preIdx values here so P1b + # skips a second GMEM random gather (dtype-gated: 16-bit only). + if cutlass.const_expr(self.p1b_cache): + smem_gath = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((self.top_k,), order=(0,)), + byte_alignment=128, + ) + else: + smem_gath = None + else: + s_mt_thr = None + smem_ptcnt_multi = None + smem_wcnt_multi = None + s_mt_cnt = None + s_r0col = None + s_cluster_partial_m = None + smem_gath = None + # ---- Per-row dispatch ---- # Three branches: # 1. Degenerate (N <= top_k): no GVR work, leader emits identity. @@ -2729,6 +4104,13 @@ def run_one_row( s_iscalars, s_cluster_partial, smem_input, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_r0col, + s_cluster_partial_m, + smem_gath, tidx, warp_id, lane, @@ -2763,6 +4145,13 @@ def run_one_row( s_iscalars, s_cluster_partial, smem_input, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_r0col, + s_cluster_partial_m, + smem_gath, tidx, warp_id, lane, @@ -2794,6 +4183,13 @@ def run_one_row( s_iscalars, s_cluster_partial, smem_input, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_r0col, + s_cluster_partial_m, + smem_gath, tidx, warp_id, lane, @@ -2828,6 +4224,13 @@ def _run_phases( s_iscalars, s_cluster_partial, smem_input, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_r0col, + s_cluster_partial_m, + smem_gath, tidx, warp_id, lane, @@ -2861,6 +4264,8 @@ def _run_phases( tidx, warp_id, lane, + smem_gath=smem_gath, # p1b_cache: stash gathered values (None-op OFF) + s_mt_thr=s_mt_thr, # r0_vseed: park pmean in the last rung column ) # Degenerate threshold init: val_hi <= -self.FLT_MAX or val_lo >= val_hi. @@ -2905,23 +4310,274 @@ def _run_phases( tidx, ) - # ---- Phase 2: secant threshold search ---- - self.phase2_secant_search( - input_row, - N, - slice_start, - slice_end, - smem_ptcnt, - smem_wcnt, - s_thr, - s_iscalars, - s_cluster_partial, - tidx, - warp_id, - lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, - ) + # ---- 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. + 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 + # the inline log-falsi R1 shot / fb_fix refine. At cs>1 each + # CTA scans its slice and block_count_ge_multi cluster-merges + # the rung counts (phase1b rungs are per-CTA identical since + # preIdx stats are full-row). + if cutlass.const_expr(self.p1b_cache): + # rungs from the SMEM gather-cache P1 stashed (no 2nd + # GMEM gather); 16-bit only. + self.phase1b_hspace_rungs_cached( + pre_idx_count, smem_gath, smem_hist, s_thr, s_mt_thr, tidx, warp_id, lane + ) + else: + self.phase1b_hspace_rungs( + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + self.block_count_ge_multi( + input_row, + slice_start, + slice_end, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_cluster_partial_m, + do_cluster_sync, + tidx, + warp_id, + lane, + smem_ptcnt=smem_ptcnt, + ) + cute.arch.barrier() + if tidx == 0: + # tightest admissible rung = SMALLEST count in [K, kC]. + # (Explicit argmin: with r0_vseed the pmean column is not + # sorted into the rung order; for sorted rungs this is + # equivalent to the old "last m in window" rule.) + best_m = cutlass.Int32(-1) + best_c = cutlass.Int32(2147483647) + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + cm = s_mt_cnt[m] + if ( + cm >= cutlass.Int32(self.top_k) + and cm <= cutlass.Int32(self.kC) + and cm < best_c + ): + best_m = cutlass.Int32(m) + best_c = cm + s_r0col[0] = best_m + if best_m >= cutlass.Int32(0): + s_thr[0] = s_mt_thr[best_m] + s_iscalars[0] = s_mt_cnt[best_m] + # done=1: the threshold is admitted, so Phase 3 must + # SKIP its retry-shrink and honor s_thr[0]. (block_count + # _ge / secant leave done via their own path; the R0 + # admission must set it explicitly or Phase 3 re-searches + # and the cluster collect diverges -> wrong output.) + s_iscalars[1] = cutlass.Int32(1) + # Snapshot this CTA's LOCAL slice count for the chosen + # rung into s_iscalars[5] — the per-CTA cand_count that + # Phase 3/4's cluster gather consumes (block_count_ge + # sets it too; the R0 admission must match). Without it + # the cluster collect under-counts -> wrong output. + if cutlass.const_expr(cluster_size > 1): + s_iscalars[5] = s_cluster_partial_m[best_m] + cute.arch.barrier() + bc = s_r0col[0] + if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf): + # accepted rung column: copy its cached per-thread counts + # into the secant hand-off buffer (zero rescan). The vseed + # column (bc == M_qf) is ALREADY in smem_ptcnt (v3 reuse). + smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx] + cute.arch.barrier() + # ---- R0 miss: SEEDED bounded log-falsi refine ---- + # At large N the M2D rungs straddle [K, kC]; the refine must + # find a threshold with count in [K, kC] between the measured + # 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 + # accept so Phase 3 skips its retry-shrink. + if bc < cutlass.Int32(0): + if cutlass.const_expr(self.fb_fix): + if tidx == cutlass.Int32(0): + M = cutlass.const_expr(self.M_thr) + blo = v_lo + bhi = v_hi + clo = cutlass.Int32(-1) + chi = cutlass.Int32(-1) + for m in cutlass.range_constexpr(M): + cm = s_mt_cnt[m] + tm = s_mt_thr[m] + if cm > cutlass.Int32(self.kC) and ( + clo < cutlass.Int32(0) or tm > blo + ): + blo = tm + clo = cm + if cm < cutlass.Int32(self.top_k) and ( + chi < cutlass.Int32(0) or tm < bhi + ): + bhi = tm + chi = cm + s_thr[1] = blo + s_thr[2] = bhi + s_iscalars[2] = clo # SEED known rung counts + s_iscalars[3] = chi + s_iscalars[1] = cutlass.Int32(0) # done=0 + cand = (blo + bhi) * cutlass.Float32(0.5) + if clo > cutlass.Int32(0) and chi >= cutlass.Int32(0): + chic = chi + if chic < cutlass.Int32(1): + chic = cutlass.Int32(1) + l_lo = cmath.log2(cutlass.Float32(clo), fastmath=True) + l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) + den = l_lo - l_hi + if den > cutlass.Float32(0.0): + t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den + cnd3 = bhi + t3 * (blo - bhi) + if cnd3 > blo and cnd3 < bhi: + cand = cnd3 + elif chi < cutlass.Int32(0): + cand = bhi + elif clo < cutlass.Int32(0): + cand = blo + s_thr[0] = cand + cute.arch.barrier() + rs = cutlass.Int32(0) + while rs < cutlass.Int32(8) and s_iscalars[1] == cutlass.Int32(0): + if rs > cutlass.Int32(0): + if tidx == cutlass.Int32(0): + lo3 = s_thr[1] + hi3 = s_thr[2] + clo3 = s_iscalars[2] + chi3 = s_iscalars[3] + cand = (lo3 + hi3) * cutlass.Float32(0.5) + if chi3 < cutlass.Int32(0): + cand = hi3 + elif clo3 < cutlass.Int32(0): + cand = lo3 + else: + chic = chi3 + if chic < cutlass.Int32(1): + chic = cutlass.Int32(1) + l_lo = cmath.log2(cutlass.Float32(clo3), fastmath=True) + l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) + den3 = l_lo - l_hi + if den3 > cutlass.Float32(0.0): + t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den3 + cnd3 = hi3 + t3 * (lo3 - hi3) + if cnd3 > lo3 and cnd3 < hi3: + cand = cnd3 + s_thr[0] = cand + cute.arch.barrier() + 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): + c3 = s_iscalars[0] + t3v = s_thr[0] + if c3 >= cutlass.Int32(self.top_k) and c3 <= cutlass.Int32(self.kC): + s_iscalars[1] = cutlass.Int32(1) # accept + elif c3 > cutlass.Int32(self.kC): + s_thr[1] = t3v + s_iscalars[2] = c3 + if t3v >= s_thr[2]: + rng3 = s_thr[2] - s_thr[1] + if rng3 < cutlass.Float32(1.0): + rng3 = cutlass.Float32(1.0) + s_thr[2] = s_thr[2] + rng3 * cutlass.Float32(8.0) + s_iscalars[3] = cutlass.Int32(-1) + else: + s_thr[2] = t3v + s_iscalars[3] = c3 + if t3v <= s_thr[1]: + rng3 = s_thr[2] - s_thr[1] + if rng3 < cutlass.Float32(1.0): + rng3 = cutlass.Float32(1.0) + s_thr[1] = s_thr[1] - rng3 * cutlass.Float32(8.0) + s_iscalars[2] = cutlass.Int32(-1) + 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() + else: + self.phase2_secant_search( + input_row, + N, + slice_start, + slice_end, + smem_ptcnt, + smem_wcnt, + s_thr, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + else: + self.phase2_secant_search( + input_row, + N, + slice_start, + slice_end, + smem_ptcnt, + smem_wcnt, + s_thr, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) # Cluster handoff #1 (end of Phase 2). Skipped when # do_cluster_sync is False (cs=1 or short-row degrade). @@ -2967,20 +4623,36 @@ def _run_phases( if cutlass.const_expr(cluster_size == 1): # cs=1: the single CTA per row IS the leader. cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) - self.phase4_histogram_snap( - smem_keys, - smem_vals, - smem_hist, - smem_wcnt, - s_thr, - s_iscalars, - output_values_row, - output_indices_row, - cand_count_p4, - tidx, - warp_id, - lane, - ) + if cutlass.const_expr(self.enable_p4_rank_scatter): + self.phase4_rank_scatter( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) + else: + self.phase4_histogram_snap( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) else: # cs>1: only the leader (CTA 0 in cluster) runs Phase 4. if is_leader: @@ -3028,20 +4700,36 @@ def _run_phases( # ---- Phase 4: histogram snap + writeback ---- cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) - self.phase4_histogram_snap( - smem_keys, - smem_vals, - smem_hist, - smem_wcnt, - s_thr, - s_iscalars, - output_values_row, - output_indices_row, - cand_count_p4, - tidx, - warp_id, - lane, - ) + if cutlass.const_expr(self.enable_p4_rank_scatter): + self.phase4_rank_scatter( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) + else: + self.phase4_histogram_snap( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) # Final cluster barrier: keep peer CTAs (and their SMEM) alive # until the leader's gather + Phase 4 finish. Skipped at @@ -3095,5 +4783,213 @@ def __call__( min_blocks_per_mp=self.min_blocks_per_mp, ) + # ------------------------------------------------------------------ # + # Host-side launch-shape policy + self-contained launcher # + # ------------------------------------------------------------------ # + # cluster_size / num_threads / min_blocks_per_mp / use_256bit_load are + # compile-time ctor knobs: a compiled kernel cannot change its own grid + # or cluster shape, so batch-size adaptation MUST happen at launch time + # by picking a different compiled variant. ``pick_config`` is that + # 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. + + _NUM_SMS: Optional[int] = None + _LAUNCH_CACHE: dict = {} + + @staticmethod + def _device_num_sms() -> int: + if GvrTopKKernel._NUM_SMS is None: + import torch # local: keep the module importable without torch + + GvrTopKKernel._NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return GvrTopKKernel._NUM_SMS + + @staticmethod + def pick_config( + torch_dtype, + num_rows: int, + num_candidates: int, + max_seq_len: Optional[int] = None, + num_sms: Optional[int] = None, + ) -> 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. + + ``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``. + """ + 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 + 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). + 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 + # 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 + vec_w = vec_bits // (32 if is_fp32 else 16) + n_vec_iters = max(1, n_per_cta // (num_threads * vec_w)) + if is_fp32: + 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: + 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( + 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, + ) + + @classmethod + def launch( + cls, + logits, + pre_idx, + seq_lens, + output_indices, + top_k: int, + next_n: int = 1, + compress_ratio: int = 1, + max_seq_len: Optional[int] = None, + num_sms: Optional[int] = None, + **kernel_overrides, + ) -> None: + """Compile-and-launch with ``pick_config`` shapes (indices-only path). + + Owns a class-level compiled-variant cache keyed by every ctor knob, + so repeated calls at any (BS, N, dtype) reuse the right variant. + ``kernel_overrides`` (e.g. ``enable_r0=False``, ``cluster_size=8``) + override the picked config and participate in the cache key. + Mirrors the custom op's compile contract: sym_int shapes, tvm-ffi + env stream (launches on the ambient torch stream), fixed + ``return_output_values=False`` / ``seqlen_sorted=False``. + """ + import torch # local: keep the module importable without torch + from cutlass.cute import runtime as _crt + + _cute_dt = { + torch.float32: cutlass.Float32, + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + } + num_rows, num_candidates = logits.shape + cfg = cls.pick_config( + logits.dtype, num_rows, num_candidates, max_seq_len=max_seq_len, num_sms=num_sms + ) + cfg.update(kernel_overrides) + if cfg["cluster_size"] > 1: + try: + from .single_pass_multi_cta_radix_topk_cluster import _query_max_cluster_size + + cfg["cluster_size"] = min(cfg["cluster_size"], _query_max_cluster_size()) + except ImportError: + pass # standalone snapshot: trust the [1, 16] ctor bound + if cfg.get("use_256bit_load") and logits.data_ptr() % 32 != 0: + cfg["use_256bit_load"] = False # 256-bit vec loads need 32B alignment + + key = (logits.dtype, top_k, next_n, compress_ratio) + tuple(sorted(cfg.items())) + compiled = cls._LAUNCH_CACHE.get(key) + if compiled is None: + kernel = cls( + dtype=_cute_dt[logits.dtype], + top_k=top_k, + next_n=next_n, + compress_ratio=compress_ratio, + return_output_values=False, + **cfg, + ) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + in_align = 32 if cfg["use_256bit_load"] else 16 + input_fake = _crt.make_fake_compact_tensor( + kernel.dtype, (n_rows, n_cols), stride_order=(1, 0), assumed_align=in_align + ) + pre_idx_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_batch, top_k), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_batch,), stride_order=(0,) + ) + out_indices_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, + input_fake, + pre_idx_fake, + seq_lens_fake, + None, + out_indices_fake, + None, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + cls._LAUNCH_CACHE[key] = compiled + compiled(logits, pre_idx, seq_lens, None, output_indices, None) + __all__ = ["GvrTopKKernel", "GvrParams"] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_load_balance.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_load_balance.py index 166ff80394e0..fd9339620b08 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_load_balance.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_load_balance.py @@ -247,9 +247,45 @@ def __init__( enable_warp_parallel_reduce=enable_warp_parallel_reduce, compress_ratio=compress_ratio, return_output_values=return_output_values, + # The two instances are inlined into ONE launch, and the CuTe DSL + # sizes the launch's dynamic SMEM from the last-traced + # SmemAllocator only (CuTeDSL.track_smem_allocator holds a single + # slot). Every smem-affecting knob must therefore resolve + # identically for both instances, or the larger branch reads/ + # writes past the reserved SMEM (CI IMA, 2026-07-21: fp32 K2048 + # long branch overflowed by the p1b_cache smem_gath the single + # instance's default had skipped). Pin the two knobs whose + # defaults diverge on cluster_size: + # - p1b_cache: cs>1 default True vs cs=1 fp32 default False. + # - kc_diet: cs=1 K512 default shrinks kC 4096→3072. + p1b_cache=True, + kc_diet=False, ) self._cluster_kernel = GvrTopKKernel(cluster_size=cluster_size, **common_kwargs) self._single_kernel = GvrTopKKernel(cluster_size=1, **common_kwargs) + # Drift guard: fail fast at construction if the derived smem-layout + # attributes ever diverge again (new knobs must be added here AND to + # the pinned kwargs above). + for attr in ( + "kC", + "kNumBins", + "M_thr", + "M_qf", + "p1b_cache", + "enable_smem_cache", + "smem_cache_elems", + "num_threads", + "p2_warp_redundant", + "p4_warp_redundant", + ): + a = getattr(self._cluster_kernel, attr) + b = getattr(self._single_kernel, attr) + assert a == b, ( + f"GvrTopKLBKernel: smem-layout attribute {attr!r} diverges " + f"between the cluster ({a}) and single ({b}) member kernels; " + f"the DSL sizes the launch from one SmemAllocator, so the " + f"layouts must be byte-identical." + ) # Prepare is decoupled from main: callers run it once per # decode step (seq_lens is layer-invariant). Use # GvrTopKLBPrepareKernel directly. 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 c29fd0bb232f..f8f7215eb6ee 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 @@ -15,10 +15,16 @@ from typing import Optional +import cutlass +import cutlass.cute as cute import pytest import torch +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.gvr_topk_decode import ( + GvrTopKKernel as _GvrTopKKernel, +) from tensorrt_llm._utils import get_sm_version skip_not_sm100 = pytest.mark.skipif( @@ -27,7 +33,7 @@ ) -def _make_inputs( +def _make_inputs_impl( num_rows: int, N: int, top_k: int, @@ -96,8 +102,7 @@ def _make_inputs( if preidx_hit_rate <= 0.0: # Worst-case: only slot 0 is meaningful, rest are junk arange. - for j in range(1, top_k): - pre_idx[:, j] = j + pre_idx[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") else: # Realistic: mix ``preidx_hit_rate`` real torch.topk indices with # random in-range fillers. Tests the Guess-phase short-circuit @@ -114,6 +119,66 @@ def _make_inputs( return logits, pre_idx, seq_lens +# Module-level input memoization. ``_make_inputs_impl`` is fully deterministic +# (seed-keyed RNG), and the parametrized sweeps below request the same +# (shape, dtype, hit-rate, ...) combination once per cluster_size / dispatch +# variant — regenerating logits + the reference topk dominated suite +# wall-clock, not the kernel under test. Cached tensors are returned WITHOUT +# cloning under a strict read-only convention: the op writes only +# ``out_indices`` (allocated fresh by every test), never its inputs. +_inputs_cache: dict = {} +# Reference top-K values memoized per cached-inputs identity (see +# ``_tie_aware_check``). Keyed on object ids, which is safe only because the +# keying tensors are pinned for the process lifetime by ``_inputs_cache``. +_ref_vals_cache: dict = {} + + +def _make_inputs( + num_rows: int, + N: int, + top_k: int, + dtype: torch.dtype, + next_n: int, + seed: int, + compress_ratio: int = 1, + preidx_hit_rate: float = 0.0, + varlen: bool = False, + seq_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Memoizing wrapper around ``_make_inputs_impl`` (same signature). + + Caller-provided ``seq_lens`` bypasses the cache (the tensor identity is + not part of a hashable key). + """ + if seq_lens is not None: + return _make_inputs_impl( + num_rows, + N, + top_k, + dtype, + next_n, + seed, + compress_ratio=compress_ratio, + preidx_hit_rate=preidx_hit_rate, + varlen=varlen, + seq_lens=seq_lens, + ) + key = (num_rows, N, top_k, dtype, next_n, seed, compress_ratio, preidx_hit_rate, varlen) + if key not in _inputs_cache: + _inputs_cache[key] = _make_inputs_impl( + num_rows, + N, + top_k, + dtype, + next_n, + seed, + compress_ratio=compress_ratio, + preidx_hit_rate=preidx_hit_rate, + varlen=varlen, + ) + return _inputs_cache[key] + + def _tie_aware_check( out_indices: torch.Tensor, logits: torch.Tensor, @@ -154,13 +219,23 @@ def _tie_aware_check( actual_kv_len = seq_lens_per_row - next_n + ofs + 1 N_eff = actual_kv_len // compress_ratio # [num_rows] - # Mask logits beyond per-row N_eff to -inf so torch.topk ignores tails. - 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")) - - # Reference per-row top-K, sorted descending. - ref_vals, _ = torch.topk(masked_logits, k=top_k, largest=True, sorted=True, dim=-1) + # 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]) @@ -227,7 +302,6 @@ def _tie_aware_check( @pytest.mark.parametrize("compress_ratio", [1, 4]) @pytest.mark.parametrize("preidx_hit_rate", [0.0, 0.5]) @pytest.mark.parametrize("cluster_size", [1, 4]) -@pytest.mark.parametrize("seqlen_sorted", [False, True]) def test_cute_dsl_gvr_topk_decode( dtype, top_k, @@ -238,7 +312,6 @@ def test_cute_dsl_gvr_topk_decode( compress_ratio, preidx_hit_rate, cluster_size, - seqlen_sorted, ): """Compare custom op output against torch.topk reference (tie-aware). @@ -250,14 +323,9 @@ def test_cute_dsl_gvr_topk_decode( ``varlen=False`` uses uniform seq_lens=N*cr across the batch; ``varlen=True`` draws per-row seq_lens uniformly in [N/2, N]*cr. - ``seqlen_sorted=True`` exercises the LJF host-side dispatch order: - we build ``order_row`` as a descending argsort over ``seq_lens`` and - pass it through the custom op. The - kernel must produce the same per-row top-K (rows are still written - back at their original positions, since the kernel uses - ``row_idx = order_row[req] * next_n + nn`` for both reads and - writes). The reference comparison is unchanged — it asserts that - each row's output is a valid top-K of that row's masked logits. + The LJF host-side dispatch order (``order_row``) is covered by the + dedicated ``test_cute_dsl_gvr_topk_decode_seqlen_sorted`` below on + representative cells instead of doubling this whole sweep. """ if N - next_n + 1 < top_k: pytest.skip(f"N_eff < top_k ({N - next_n + 1} < {top_k}) is a degenerate path") @@ -279,12 +347,66 @@ def test_cute_dsl_gvr_topk_decode( out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") - # LJF dispatch order — request-level descending argsort of seq_lens. - order_row = ( - torch.argsort(seq_lens, descending=True, stable=False).to(torch.int32) - if seqlen_sorted - else None + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + next_n=next_n, + compress_ratio=compress_ratio, + cluster_size=cluster_size, ) + torch.cuda.synchronize() + + _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + + +@skip_not_sm100 +@pytest.mark.parametrize( + "dtype,top_k,N,batch_size,varlen,next_n,compress_ratio,cluster_size", + [ + # Representative cells for the LJF host-side dispatch order + # (previously a full extra dimension on the sweep above): varlen + # batches so the argsort is a real permutation, both SMEM layout + # endpoints (bf16/K=512, fp32/K=2048), next_n=2 for the + # ``order_row[req] * next_n + nn`` row expansion, cr=4 for the + # order_row + compressed seq_lens interaction, cluster and + # single-CTA paths, and a batch_size=1 trivial-permutation smoke. + (torch.bfloat16, 512, 65536, 32, True, 1, 1, 4), + (torch.float32, 2048, 65536, 32, True, 2, 1, 1), + (torch.bfloat16, 1024, 4096, 32, True, 2, 4, 1), + (torch.float16, 1024, 65536, 1, False, 1, 1, 1), + ], +) +def test_cute_dsl_gvr_topk_decode_seqlen_sorted( + dtype, top_k, N, batch_size, varlen, next_n, compress_ratio, cluster_size +): + """LJF host-side dispatch order: ``order_row`` = descending argsort of + ``seq_lens`` passed through the custom op. + + The kernel must produce the same per-row top-K as the unsorted launch + (rows are still written back at their original positions, since the + kernel uses ``row_idx = order_row[req] * next_n + nn`` for both reads + and writes). The reference comparison is unchanged — it asserts that + each row's output is a valid top-K of that row's masked logits. + """ + num_rows = batch_size * next_n + logits, pre_idx, seq_lens = _make_inputs( + num_rows, + N, + top_k, + dtype, + next_n, + seed=42, + compress_ratio=compress_ratio, + preidx_hit_rate=0.5, + varlen=varlen, + ) + out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + + # LJF dispatch order — request-level descending argsort of seq_lens. + order_row = torch.argsort(seq_lens, descending=True, stable=False).to(torch.int32) torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits, @@ -588,9 +710,11 @@ def test_lb_main_branches(dtype, top_k, scenario, N, seq_lens_mode, batch_size, @pytest.mark.parametrize( "dtype,top_k", [ + # SMEM-layout endpoints only. LB dispatch (prepare partition + + # long/short branch selection) is dtype-insensitive; the full + # dtype x K production map stays covered by the main sweep above + # and by test_cute_dsl_gvr_topk_decode_r0_equivalence. (torch.bfloat16, 512), - (torch.bfloat16, 1024), - (torch.float16, 1024), (torch.float32, 2048), ], ) @@ -663,3 +787,351 @@ def test_lb_vs_reference( next_n, compress_ratio=compress_ratio, ) + + +# =========================================================================== +# R0 histogram-ladder admission equivalence tests. +# +# ``enable_r0=True`` (the GvrTopKKernel default) replaces the Phase-2 secant +# threshold search with a single-pass multi-threshold "rung ladder" admission +# seeded by a 256-bin histogram over the prev-topK gathered values. This must +# select the SAME top-K as the retained secant baseline (``enable_r0=False``). +# +# top-K is order-independent, so correctness is checked by INDEX SET (not +# position): for continuous fp32 logits (tie-free with probability 1) the R0 +# and base index sets must be identical; for bf16/fp16 boundary value-ties can +# make two equally-valid selections differ in index, so there the guarantee is +# value-set (multiset) equality against the tie-aware torch.topk reference. +# +# The custom op does not plumb ``enable_r0`` (activation / dispatch land in a +# follow-up PR), so these tests drive ``GvrTopKKernel`` directly. This is also +# the only remaining coverage of the secant fallback path, since every op-level +# test above now inherits the ``enable_r0=True`` default. +# =========================================================================== + +_R0_DT = { + torch.float32: cutlass.Float32, + torch.bfloat16: cutlass.BFloat16, + torch.float16: cutlass.Float16, +} +# Compiled-kernel cache keyed on (enable_r0, dtype, top_k, cluster_size, T, +# min_blocks_per_mp). Shapes (num_rows / N / batch) are symbolic, so one +# compile covers every N and batch_size within a bucket (mirrors the runner). +_r0_kernel_cache: dict = {} + + +def _compile_gvr_direct(kernel): + """Compile a ``GvrTopKKernel`` with symbolic shapes, mirroring the + production runner's fake-tensor construction (128-bit loads, no + ``order_row`` / ``output_values``).""" + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + in_f = _crt.make_fake_compact_tensor( + kernel.dtype, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + pi_f = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_batch, kernel.top_k), stride_order=(1, 0), assumed_align=16 + ) + sl_f = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + oi_f = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, kernel.top_k), stride_order=(1, 0), assumed_align=16 + ) + fs = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + # __call__(input, pre_idx, seq_lens, output_values, output_indices, order_row, stream) + return cute.compile( + kernel, in_f, pi_f, sl_f, None, oi_f, None, stream=fs, options="--enable-tvm-ffi" + ) + + +def _run_gvr_direct(logits, pre_idx, seq_lens, top_k, enable_r0, cluster_size): + """Drive ``GvrTopKKernel`` directly (bypassing the custom op, which does + not expose ``enable_r0``). Fixed at ``next_n=1``, ``compress_ratio=1``, + 128-bit loads. When ``enable_r0=True`` the ctor auto-derives the shipped + R0 config (r0_qfracs=M2D, cs-aware p1b_cache, K512 kC-diet, P4 + rank-scatter) — i.e. the exact default arm. Returns int32 + ``[num_rows, top_k]`` indices.""" + num_rows, N = logits.shape + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + num_threads = 1024 if (num_rows <= num_sms and N >= 65536) else 512 + min_blocks_per_mp = 1 if num_rows <= num_sms else 3 + key = (enable_r0, logits.dtype, top_k, cluster_size, num_threads, min_blocks_per_mp) + if key not in _r0_kernel_cache: + kernel = _GvrTopKKernel( + dtype=_R0_DT[logits.dtype], + top_k=top_k, + next_n=1, + num_threads=num_threads, + compress_ratio=1, + use_256bit_load=False, + min_blocks_per_mp=min_blocks_per_mp, + cluster_size=cluster_size, + return_output_values=False, + enable_r0=enable_r0, + ) + _r0_kernel_cache[key] = _compile_gvr_direct(kernel) + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + _r0_kernel_cache[key](logits, pre_idx, seq_lens, None, out, None) + torch.cuda.synchronize() + return out + + +def _assert_index_sets_equal_tie_aware(out_base, out_r0, logits): + """Assert two arms' top-K index sets match, modulo boundary value-ties. + + fp32 randn logits DO collide bit-exactly at these sample counts; when the + duplicated value sits on the top-K boundary, each arm may legitimately + keep a different member of the tie class. Indices in the symmetric + difference must all carry the row's boundary (minimum kept) value — + anything else is a genuine divergence. + """ + base_sorted, _ = out_base.sort(dim=-1) + r0_sorted, _ = out_r0.sort(dim=-1) + mismatch = (base_sorted != r0_sorted).any(dim=-1) + for bad in mismatch.nonzero().flatten().tolist(): + base_set = set(out_base[bad].tolist()) + r0_set = set(out_r0[bad].tolist()) + diff = sorted(base_set.symmetric_difference(r0_set)) + row_vals = logits[bad].float() + kth = row_vals[out_base[bad].long()].min() + diff_vals = row_vals[torch.tensor(diff, device=logits.device, dtype=torch.long)] + if not bool((diff_vals == kth).all().item()): + raise AssertionError( + f"row={bad}: R0 index set != secant-base index set beyond a " + f"boundary value-tie (kth={kth.item()}, " + f"diff={[(i, row_vals[i].item()) for i in diff]}, " + f"base={sorted(base_set)}, r0={sorted(r0_set)})" + ) + + +def _make_r0_pre_idx(logits, top_k, hint, seed): + """Build ``pre_idx`` in the kernel's native cr=1 convention: the kernel + reads ``logits[pre_idx + 1]``, so store ``true_index - 1``. + + ``hint='real'`` seeds a warm hint (near-topK indices) so R0's admission + ladder hits on the first pass; ``hint='rand'`` seeds a cold hint (random + in-range indices) that misses admission and forces the R0-miss inline + log-falsi (R1) + fb_fix fallback.""" + num_rows, N = logits.shape + g = torch.Generator(device="cuda").manual_seed(seed) + if hint == "real": + noised = logits.float() + 0.15 * torch.randn(num_rows, N, generator=g, device="cuda") + pre = noised.topk(top_k, dim=1).indices.int() + else: + pre = torch.randint(0, N, (num_rows, top_k), generator=g, device="cuda").int() + return (pre - 1).contiguous() + + +@skip_not_sm100 +@pytest.mark.parametrize( + "dtype,top_k", + [ + (torch.bfloat16, 512), + (torch.bfloat16, 1024), + (torch.float16, 1024), + (torch.float32, 2048), + ], +) +@pytest.mark.parametrize("N", [8192, 65536]) +@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): + """R0 admission (``enable_r0=True``, the new default) selects the same + top-K as the secant baseline (``enable_r0=False``), by index set. + + ``hint='real'`` exercises the R0 admission-hit fast path; ``hint='rand'`` + forces the R0-miss log-falsi (R1) + fb_fix fallback. ``cluster_size=4`` + confirms R0 gates to single-CTA and the ``None`` R0 buffers propagate + cleanly through the cluster path; ``cluster_size=8`` covers the runner's + tiny-grid large-N pick (BS<=4, N>=128K -> cs=8): 7-peer DSMEM + aggregation, large-N only (per-CTA slice too short below 64K). + """ + if N < top_k * 2: + pytest.skip(f"N ({N}) < 2*top_k ({2 * top_k}): GVR histogram bucket too coarse") + if cluster_size == 8 and N < 65536: + pytest.skip("cs=8 is a large-N production config (runner picks it only at N >= 131072)") + + num_rows = batch_size # next_n = 1 + torch.manual_seed(0) + torch.cuda.manual_seed(0) + logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(dtype).contiguous() + seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda") + pre_idx = _make_r0_pre_idx(logits, top_k, hint, seed=1) + + out_base = _run_gvr_direct( + logits, pre_idx, seq_lens, top_k, enable_r0=False, cluster_size=cluster_size + ) + out_r0 = _run_gvr_direct( + logits, pre_idx, seq_lens, top_k, enable_r0=True, cluster_size=cluster_size + ) + + # 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) + + # 2. Equivalence. fp32 logits are ALMOST tie-free, so R0 and base must + # return the identical index set (order-independent) — but randn + # quantized to fp32 does collide (bs=16, N=8192, seed 0: + # logits[3,2956] == logits[3,4949] bit-exactly, straddling the K=2048 + # boundary), and a boundary value-tie makes the arms' distinct index + # picks equally valid. Where the sets differ, require every differing + # index to carry the row's boundary (k-th) value; anything else is a + # real divergence. bf16/fp16 boundary ties are common, so equivalence + # there is the value-set equality already established in step 1 + # (both == torch.topk reference). + if dtype == torch.float32: + _assert_index_sets_equal_tie_aware(out_base, out_r0, logits) + + +@skip_not_sm100 +@pytest.mark.parametrize( + "dtype,top_k,N,batch_size,cluster_size", + [ + # Multi-wave single-CTA grids (batch_size > num_sms): exercises the + # occupancy regime where rows alone oversubscribe the device. + (torch.bfloat16, 512, 16384, 256, 1), + (torch.float32, 2048, 65536, 256, 1), + # Multi-wave cluster grid (batch_size * cs > num_sms): DSMEM handoff + # correctness across wave boundaries. + (torch.bfloat16, 512, 65536, 64, 4), + ], +) +@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 +): + """Big-batch R0-vs-secant equivalence: multi-wave grids only. + + The main equivalence grid tops out at batch_size=16 (single wave). + These cells lock R0 + cluster correctness when the grid spans several + waves — the throughput-bound regime of the BS-scaling study.""" + num_rows = batch_size + torch.manual_seed(0) + torch.cuda.manual_seed(0) + logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(dtype).contiguous() + seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda") + pre_idx = _make_r0_pre_idx(logits, top_k, hint, seed=1) + + out_base = _run_gvr_direct( + logits, pre_idx, seq_lens, top_k, enable_r0=False, cluster_size=cluster_size + ) + 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) + if dtype == torch.float32: + _assert_index_sets_equal_tie_aware(out_base, out_r0, logits) + + +@skip_not_sm100 +@pytest.mark.parametrize( + "top_k,N,cluster_size", + [(512, 16384, 1), (1024, 131072, 4), (2048, 131072, 4)], +) +@pytest.mark.parametrize("band", ["sub_resolution", "one_ulp"]) +def test_cute_dsl_gvr_topk_decode_p4_exact_tail_ties(top_k, N, cluster_size, band): + """fp32 near-tie adversarial exactness (``p4_exact_tail``). + + The P4 rank-scatter fine recursion resolves candidate values to + range/(kNumBins*256); distinct values spaced below that which straddle + the top-K boundary land in ONE fine bin and were previously kept in + arrival order (observed on real DSv4-Pro 512k-ISL captures as |miss|=1 + with dv ~ 3e-6). ``sub_resolution`` plants ~2.4k distinct values spaced + 5e-8 around the boundary; ``one_ulp`` plants a two-value bitwise plateau + (``nextafter`` pairs). Both stay within the kC candidate budget (tie + sets wider than kC are outside the kernel's contract). The default fp32 + kernel must return the exact top-K value set; natural random data never + triggers this, so the adversarial construction is the only regression + coverage.""" + torch.manual_seed(3) + torch.cuda.manual_seed(3) + logits = (torch.randn(1, N, device="cuda") * 2.0).float().contiguous() + boundary = torch.topk(logits[0], top_k).values[top_k - 1].item() + n_tie = 2400 if band == "sub_resolution" else 2000 + plant = torch.randperm(N)[:n_tie] + if band == "sub_resolution": + tie_vals = ( + boundary + (torch.arange(n_tie, dtype=torch.float32, device="cuda") - n_tie // 2) * 5e-8 + ) + else: + tie_vals = torch.full((n_tie,), boundary, device="cuda") + tie_vals[::2] = torch.nextafter(tie_vals[::2], torch.tensor(float("inf"), device="cuda")) + logits[0, plant] = tie_vals + seq_lens = torch.full((1,), N, dtype=torch.int32, device="cuda") + pre_idx = _make_r0_pre_idx(logits, top_k, "real", seed=4) + + out = _run_gvr_direct( + logits, pre_idx, seq_lens, top_k, enable_r0=True, cluster_size=cluster_size + ) + + # Value-multiset exactness (the boundary index set is not unique under + # bitwise plateaus, so indices are compared through their values). + sel = logits[0][out[0].long()].sort().values + ref = torch.topk(logits[0], top_k).values.sort().values + torch.testing.assert_close(sel, ref, rtol=0.0, atol=0.0) + + +@skip_not_sm100 +def test_cute_dsl_gvr_topk_decode_pick_config_policy(): + """``pick_config`` returns the runner-equivalent launch shapes. + + Locks the (BS, N) -> cluster_size map and the BS-aware occupancy knobs + (the 2026-07-15 big-BS triage: a config frozen at the BS=1 optimum is + geomean 2.27x slower than these picks at BS in {64, 256, 1024}).""" + sms = 148 # policy is expressed against a fixed SM count for determinism + pc = _GvrTopKKernel.pick_config + + # cluster_size policy: N<64K -> 1; tiny grid large-N -> 8; single-wave + # -> 4/2; multi-wave -> 1. + assert pc(torch.float32, 1, 32768, num_sms=sms)["cluster_size"] == 1 + assert pc(torch.float32, 2, 131072, num_sms=sms)["cluster_size"] == 8 + assert pc(torch.float32, 16, 65536, num_sms=sms)["cluster_size"] == 4 + assert pc(torch.float32, 64, 65536, num_sms=sms)["cluster_size"] == 2 + assert pc(torch.float32, 256, 65536, num_sms=sms)["cluster_size"] == 1 + + # Occupancy knobs at multi-wave BS: T=512 + mbpm>=2 (NOT the BS=1 + # frozen T=1024/mbpm=1 that loses 2.3-6x at big BS). + big = pc(torch.float32, 1024, 65536, num_sms=sms) + assert big["num_threads"] == 512 and big["min_blocks_per_mp"] == 2 + big16 = pc(torch.bfloat16, 1024, 65536, num_sms=sms) + assert big16["num_threads"] == 512 and big16["min_blocks_per_mp"] == 3 + + # Graph-capture contract: max_seq_len (peak N) overrides the capture N. + cap = pc(torch.bfloat16, 1, 8192, max_seq_len=131072, num_sms=sms) + assert cap["cluster_size"] == 8 # picked for the replay shape + + +@skip_not_sm100 +@pytest.mark.parametrize( + "dtype,top_k,N,batch_size", + [ + (torch.float32, 2048, 32768, 1), # cs=1 small-N + (torch.bfloat16, 512, 65536, 16), # cs=4 single-wave + (torch.float32, 1024, 131072, 2), # cs=8 tiny grid large-N + (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): + """``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.""" + num_rows = batch_size + torch.manual_seed(0) + torch.cuda.manual_seed(0) + logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(dtype).contiguous() + seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda") + pre_idx = _make_r0_pre_idx(logits, top_k, "real", seed=1) + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + + _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) + + # 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) + if dtype == torch.float32: + assert torch.equal(out.sort(dim=-1).values, out_sec.sort(dim=-1).values)