From c06c8cd46ae98ddc6097262bb5cc9d6d0865526a Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 26 Jul 2026 06:10:51 +0000 Subject: [PATCH 01/19] [None][perf] Port op43 BSX CuTe DSL top-K tiers (direct/reg/tp) onto GVR decode op Adds the op43 bsx CuTe DSL kernel family as a guarded fast path inside trtllm::cute_dsl_gvr_topk_decode (op signature unchanged): - gvr_topk_decode_bsx_tp.py: throughput GVR tier (gvr_topk_tp port) with ragged-N lane masking, pre_idx clamping, in-kernel degenerate emit and the mandatory cluster exit rendezvous. - gvr_topk_decode_bsx_reg.py: register-resident tier (gvr_topk_reg port); ragged-N handled by extending the OOR-lane -FLT_MAX idiom at the single register load; packed u64 candidate pushes and aligned cluster barriers preserved from the op43 convergence rounds. - gvr_topk_decode_bsx_direct.py: short-row (npad <= 12288) exact radix tier with ragged-N key substitution and degenerate emit. - gvr_topk_decode_bsx_dispatch.py: route table transcribed from op42 gvr_bsx.cu; env knobs renamed GVR_BSX_* -> TRTLLM_BSX_*; v1 guard (fp32, next_n=1, cr=4, no order_row/counters, K in {512,1024,2048}, npad <= 262144, npad % 64 == 0) plus a hardware cluster-size cap check that falls back to the in-tree kernel instead of degrading silently. - CI-sized exactness test covering every reg launch-table instance once, direct/tp tiers, ragged rows with poisoned tails, quantized ties, degenerate rows, pre_idx hardening, route-table asserts and dispatcher fallback; registered in l0_b300.yml. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 17 + .../blackwell/top_k/__init__.py | 9 + .../top_k/gvr_topk_decode_bsx_direct.py | 593 +++++++ .../top_k/gvr_topk_decode_bsx_dispatch.py | 256 +++ .../top_k/gvr_topk_decode_bsx_reg.py | 872 ++++++++++ .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 1478 +++++++++++++++++ .../test_lists/test-db/l0_b300.yml | 3 +- .../sparse/test_cute_dsl_bsx_topk_decode.py | 481 ++++++ 8 files changed, 3708 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py create mode 100644 tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index e53d58693a4b..32642a84cac4 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7271,6 +7271,10 @@ def warmup_cute_dsl_radix_topk_decode( # ------------------------------------------------------------------ # from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import \ GvrTopKKernel as _GvrTopKKernel + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_bsx_dispatch import \ + bsx_topk as _bsx_topk + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_bsx_dispatch import \ + is_bsx_supported as _is_bsx_supported class CuteDSLGvrTopKDecodeRunner: """Runner for the GVR Top-K cuTe DSL kernel (Blackwell SM100). @@ -7569,6 +7573,19 @@ def forward( ``counters`` without ``order_row`` is rejected. """ + # BSX tier fast path (op43 port): fp32 / next_n=1 / cr=4 / + # npad <= 262144 decode rows route to the direct/reg/tp CuTe DSL + # tiers; everything else (half-prec, LB, sort-indirect, V3.2, + # oversize npad, hw cluster cap) falls through to the in-tree + # kernel below. Host-only guard — no device sync. The op + # signature and output contract are unchanged (unordered int32 + # indices, -1 pad only for degenerate rows). + if _is_bsx_supported(logits, pre_idx, seq_lens, output_indices, + top_k, next_n, compress_ratio, order_row, + counters): + _bsx_topk(logits, pre_idx, seq_lens, output_indices, top_k) + return + cute_dtype = _TORCH_TO_CUTLASS_DTYPE[logits.dtype] num_rows = logits.shape[0] # seq_lens is request-level, logits is row-level (next_n diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index f9ee2ddde4e2..042e8e70fe2f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -17,6 +17,10 @@ from .filtered_top_k_decode_varlen import FilteredTopKKernelVarlenDecode from .filtered_top_k_varlen_util import FilteredTopKKernelVarlen from .gvr_topk_decode import GvrParams, GvrTopKKernel +from .gvr_topk_decode_bsx_direct import DirectTopKKernel +from .gvr_topk_decode_bsx_dispatch import bsx_topk, is_bsx_supported +from .gvr_topk_decode_bsx_reg import GvrRegKernel +from .gvr_topk_decode_bsx_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel __all__ = [ @@ -25,4 +29,9 @@ "FilteredTopKKernelVarlenDecode", "GvrParams", "GvrTopKKernel", + "GvrTpKernel", + "GvrRegKernel", + "DirectTopKKernel", + "bsx_topk", + "is_bsx_supported", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py new file mode 100644 index 000000000000..959355fe9dbb --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py @@ -0,0 +1,593 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BSX direct (short-row) top-K tier — CuTe DSL, Blackwell SM100. + +Port of the op43 ``ct_direct.py`` CuTe DSL translation of the CUDA +``direct_topk_kernel`` (op42 gvr_bsx.cu), adapted for the production +``trtllm::cute_dsl_gvr_topk_decode`` contract (ragged N via a device +``seq_lens`` tensor + per-row degenerate identity emit; see the module +docstring of ``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory). + +Exact top-K indices for short padded rows (npad <= DKCMAX = 12288), one CTA +per row, TB = 1024 threads. No threshold solve: the whole row is collected +into SMEM as packed (f2u order key << 32 | index) u64 candidates with an +in-flight 2048-bin radix histogram over the top 11 key bits, then an +11/11/10-bit radix select with whole-bin early exit and boundary-bin +compaction. Tie-aware emit (strict-greater mandatory slots + tie tickets +filling to K) makes the output value multiset equal torch.topk exactly. + +Ragged N: elements at index >= N_eff contribute key = f2u(-FLT_MAX) (the +minimum key) instead of their stale value; with the degenerate rows +(N_eff <= K) peeled off in-kernel, a masked element can never enter the +top-K, so out-of-range indices are never emitted. + +Public entry: ``direct_topk(logits, seq_lens, out, K)`` with torch tensors +(logits [BS, npad] fp32 contiguous, seq_lens [BS] int32, out [BS, K] int32). +Compiled variants are cached per (K, TB); BS and npad are dynamic (sym_int) +so one variant serves every shape. +""" + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import llvm +from cutlass.utils.smem_allocator import SmemAllocator + +FLT_MAX = 3.4028234663852886e38 +DKCMAX = 12288 # direct-path candidate capacity (mirrors gvr_bsx.cu) + + +@cute.jit +def _f2u(v): + """Order-preserving fp32 -> uint32 radix key (CUDA f2u). + + u ^ (sign ? 0xFFFFFFFF : 0x80000000): unsigned order of the result equals + fp32 order (NaN-free inputs). All shifts on the returned Uint32 are + logical, all compares unsigned. + """ + u = cutlass.Uint32(llvm.bitcast(cutlass.Uint32.mlir_type, v.ir_value())) + mask = (cutlass.Uint32(0) - (u >> cutlass.Uint32(31))) | cutlass.Uint32(0x80000000) + return u ^ mask + + +class DirectTopKKernel: + """One-CTA-per-row exact radix top-K for npad <= DKCMAX. + + Ctor knobs (compile-time): ``top_k`` in {512, 1024, 2048}, ``num_threads`` + (TB, default 1024), ``next_n`` / ``compress_ratio`` for the per-row + N_eff arithmetic (v1 dispatcher guard pins next_n=1, cr=4). SMEM layout + mirrors gvr_bsx.cu's DSmem with the same packed (key << 32 | idx) + u64 layout (op43 S-E round 6: split key/idx arrays cost 2x smem + transactions in collect + emits vs CUDA's single STS.64 / LDS.64 per + candidate): + + cand[DKCMAX] u64 | side[SIDECAP] u64 (boundary-bin compaction) | + hist[2048] i32 | iwred[NWARP] i32 | sel[2] i32 | cnt[3] i32 + """ + + WARP_SIZE = 32 + + def __init__( + self, top_k: int, num_threads: int = 1024, next_n: int = 1, compress_ratio: int = 4 + ): + if top_k not in (512, 1024, 2048): + raise ValueError(f"unsupported top_k {top_k}") + if num_threads % 32 != 0: + raise ValueError("num_threads must be a multiple of 32") + self.top_k = top_k + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + # DCAP_RUNGS * TB / 2 in CUDA (side[] unions the dead ptcnt[]); the + # direct path never allocates ptcnt so we keep only side[]. + self.sidecap = 16 * num_threads // 2 + + # ------------------------------------------------------------------ + # Per-row valid length (ragged N) — mirrors the in-tree run_one_row + # arithmetic exactly (see gvr_topk_decode_bsx_tp._row_n_eff). + # ------------------------------------------------------------------ + @cute.jit + def _row_n_eff(self, seq_lens: cute.Tensor, row): + NN = cutlass.const_expr(self.next_n) + seq_len = seq_lens[row // cutlass.Int32(NN)] + actual_kv_len = seq_len - cutlass.Int32(NN) + (row % cutlass.Int32(NN)) + cutlass.Int32(1) + if cutlass.const_expr(self.compress_ratio == 1): + n_eff = actual_kv_len + else: + n_eff = actual_kv_len // cutlass.Int32(self.compress_ratio) + return n_eff + + # ------------------------------------------------------------------ + # Warp suffix-inclusive scan (higher lane = higher bins), the + # __shfl_down_sync ladder of the CUDA bin_select. + # ------------------------------------------------------------------ + @cute.jit + def _warp_suffix_scan(self, val, lane): + x = val + for d in cutlass.range_constexpr(5): + off = cutlass.const_expr(1 << d) + src = lane + cutlass.Int32(off) + if src > cutlass.Int32(31): + src = cutlass.Int32(31) + up = cute.arch.shuffle_sync(x, src) + if lane + cutlass.Int32(off) < cutlass.Int32(32): + x = x + up + return x + + # ------------------------------------------------------------------ + # bin_select: find bin b such that above(b) < want <= above(b) + # + hist[b] where above(b) = sum_{j>b} hist[j]. Writes smem_sel = + # {b, above(b)}. Trailing barrier. + # ------------------------------------------------------------------ + @cute.jit + def _bin_select( + self, + nbins: cutlass.Constexpr, + want, + smem_hist, + smem_iwred, + smem_sel, + tidx, + lane, + warp_id, + ): + per = cutlass.const_expr(nbins // self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + h = cute.make_fragment((per,), cutlass.Int32) + s_sum = cutlass.Int32(0) + for q in cutlass.range_constexpr(per): + h[q] = smem_hist[tidx * cutlass.Int32(per) + cutlass.Int32(q)] + s_sum = s_sum + h[q] + # warp suffix-inclusive over lanes (higher lane = higher bins) + x = self._warp_suffix_scan(s_sum, lane) + if lane == cutlass.Int32(0): + smem_iwred[warp_id] = x # warp total + cute.arch.barrier() + if tidx < cutlass.Int32(32): + wt = cutlass.Int32(0) + if tidx < cutlass.Int32(num_warps): + wt = smem_iwred[tidx] + wx = self._warp_suffix_scan(wt, tidx) + if tidx < cutlass.Int32(num_warps): + smem_iwred[tidx] = wx - wt # strictly-above-warp sum + cute.arch.barrier() + a_cnt = smem_iwred[warp_id] + (x - s_sum) # bins strictly above chunk + if a_cnt < want and want <= a_cnt + s_sum: + run = a_cnt + for qr in cutlass.range_constexpr(per): + q = cutlass.const_expr(per - 1 - qr) + if run < want and want <= run + h[q]: + smem_sel[0] = tidx * cutlass.Int32(per) + cutlass.Int32(q) + smem_sel[1] = run + run = run + h[q] + cute.arch.barrier() + + # ------------------------------------------------------------------ + # Warp-aggregated slot allocation: ballot + leader atomicAdd + popc. + # Whole-warp uniform trip counts are guaranteed by Cpad % 32 == 0 and + # TB % 32 == 0 (same argument as the CUDA original). + # ------------------------------------------------------------------ + @cute.jit + def _emit_prefix_ge(self, cand_cnt, shift, prefix, smem_cand, smem_cnt, out_row, tidx, lane): + """Emit indices of all candidates with (key >> shift) >= prefix; + the caller guarantees exactly K of them.""" + num_threads = cutlass.const_expr(self.num_threads) + if tidx == cutlass.Int32(0): + smem_cnt[0] = cutlass.Int32(0) + cute.arch.barrier() + cpad = (cand_cnt + cutlass.Int32(31)) & cutlass.Int32(-32) + i = tidx + while i < cpad: + valid = i < cand_cnt + kv = cutlass.Uint64(0) + if valid: + kv = smem_cand[i] + key = cutlass.Uint32(kv >> cutlass.Uint64(32)) + e = valid and ((key >> shift) >= prefix) + bal = cute.arch.vote_ballot_sync(e) + if bal != cutlass.Uint32(0): + tz = cutlass.Int32( + cute.arch.popc((bal & (cutlass.Uint32(0) - bal)) - cutlass.Uint32(1)) + ) + base = cutlass.Int32(0) + if lane == tz: + base = cute.arch.atomic_add( + smem_cnt.iterator, cutlass.Int32(cute.arch.popc(bal)), scope="cta" + ) + base = cute.arch.shuffle_sync(base, tz) + if e: + rk = cutlass.Int32( + cute.arch.popc(bal & cutlass.Uint32(cute.arch.lanemask_lt())) + ) + out_row[base + rk] = cutlass.Int32( + cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF)) + ) + i = i + cutlass.Int32(num_threads) + + @cute.jit + def _emit_final(self, cand_cnt, kth, m, nt, smem_cand, smem_cnt, out_row, tidx, lane): + """keys > kth are mandatory (slots [0, m)); keys == kth fill the + remaining nt tie tickets (slots [m, m+nt)).""" + num_threads = cutlass.const_expr(self.num_threads) + if tidx == cutlass.Int32(0): + smem_cnt[0] = cutlass.Int32(0) + smem_cnt[1] = cutlass.Int32(0) + cute.arch.barrier() + cpad = (cand_cnt + cutlass.Int32(31)) & cutlass.Int32(-32) + i = tidx + while i < cpad: + valid = i < cand_cnt + kv = cutlass.Uint64(0) + if valid: + kv = smem_cand[i] + key = cutlass.Uint32(kv >> cutlass.Uint64(32)) + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + man = valid and (key > kth) + tie = valid and (key == kth) + bm = cute.arch.vote_ballot_sync(man) + bt = cute.arch.vote_ballot_sync(tie) + if bm != cutlass.Uint32(0): + tzm = cutlass.Int32( + cute.arch.popc((bm & (cutlass.Uint32(0) - bm)) - cutlass.Uint32(1)) + ) + base_m = cutlass.Int32(0) + if lane == tzm: + base_m = cute.arch.atomic_add( + smem_cnt.iterator, cutlass.Int32(cute.arch.popc(bm)), scope="cta" + ) + base_m = cute.arch.shuffle_sync(base_m, tzm) + if man: + rk = cutlass.Int32(cute.arch.popc(bm & cutlass.Uint32(cute.arch.lanemask_lt()))) + out_row[base_m + rk] = idx + if bt != cutlass.Uint32(0): + tzt = cutlass.Int32( + cute.arch.popc((bt & (cutlass.Uint32(0) - bt)) - cutlass.Uint32(1)) + ) + base_t = cutlass.Int32(0) + if lane == tzt: + base_t = cute.arch.atomic_add( + smem_cnt.iterator + 1, cutlass.Int32(cute.arch.popc(bt)), scope="cta" + ) + base_t = cute.arch.shuffle_sync(base_t, tzt) + if tie: + p = base_t + cutlass.Int32( + cute.arch.popc(bt & cutlass.Uint32(cute.arch.lanemask_lt())) + ) + if p < nt: + out_row[m + p] = idx + i = i + cutlass.Int32(num_threads) + + # ------------------------------------------------------------------ + # Kernel + # ------------------------------------------------------------------ + @cute.kernel + def direct_topk_kernel( + self, + logits: cute.Tensor, # [BS, npad] fp32; tail beyond N_eff may be garbage + seq_lens: cute.Tensor, # [BS] int32 (uncompressed-token space) + out: cute.Tensor, # [BS, K] int32 + ): + num_threads = cutlass.const_expr(self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + kK = cutlass.const_expr(self.top_k) + sidecap = cutlass.const_expr(self.sidecap) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + lane = tidx & cutlass.Int32(31) + warp_id = tidx // cutlass.Int32(32) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[bidx, None] + out_row = out[bidx, None] + + # Ragged N: per-row valid length from seq_lens. + n_eff = self._row_n_eff(seq_lens, bidx) + + # ---- SMEM (allocated unconditionally, before the dynamic branch, + # so the DSL sizes the launch SMEM identically on every path) ---- + smem = SmemAllocator() + s_cand = smem.allocate_tensor( + cutlass.Uint64, cute.make_layout((DKCMAX,)), byte_alignment=128 + ) + s_side = smem.allocate_tensor( + cutlass.Uint64, cute.make_layout((sidecap,)), byte_alignment=128 + ) + s_hist = smem.allocate_tensor(cutlass.Int32, cute.make_layout((2048,)), byte_alignment=128) + s_iwred = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((num_warps,)), byte_alignment=128 + ) + s_sel = smem.allocate_tensor(cutlass.Int32, cute.make_layout((2,)), byte_alignment=16) + s_cnt = smem.allocate_tensor(cutlass.Int32, cute.make_layout((3,)), byte_alignment=16) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad ---- + # (mirrors the in-tree kernel; CuTe DSL has no runtime return, so the + # main body lives in the else branch). + if n_eff <= cutlass.Int32(kK): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(num_threads) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(kK): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(num_threads) + else: + # ---- Collect: vectorized load + f2u keys + in-flight 2048-bin hist + for z in cutlass.range_constexpr(2048 // num_threads): + s_hist[tidx + cutlass.Int32(z * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128, invariant=True + ) + row_addr = logits_row.iterator.toint() + vcnt = npad >> cutlass.Int32(2) # npad % 4 == 0 (host-asserted) + # op43 S-E round 6: npad <= DKCMAX bounds the per-thread trip count + # at U = DKCMAX/4/TB (= 3 for TB 1024), so the whole grid-stride + # loop is flattened into one predicated register batch: issue ALL + # (<=3) float4 loads back-to-back, then consume. A dynamic + # `cutlass.range(iters, unroll=4)` loop never reaches its unrolled + # body (trip <= 3) and leaves ONE load in flight per iteration — + # 1.8x the long-scoreboard stall of the CUDA arm and a stable + # ~1.03 cold-kernel gap at npad 8256. Keep the flat batch. + n_batch = cutlass.const_expr((DKCMAX // 4 + num_threads - 1) // num_threads) + frags = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(n_batch) + ] # Python-unrolled register batch + for u in cutlass.range_constexpr(n_batch): + i_vec = tidx + cutlass.Int32(u * num_threads) + if i_vec < vcnt: + src_ptr = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(i_vec) * cutlass.Int64(16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + src = cute.make_tensor(src_ptr, cute.make_layout((4,))) + cute.copy(copy_atom, src, frags[u]) + for u in cutlass.range_constexpr(n_batch): + i_vec = tidx + cutlass.Int32(u * num_threads) + if i_vec < vcnt: + gi = i_vec * cutlass.Int32(4) + for q in cutlass.range_constexpr(4): + # Ragged N: stale values beyond N_eff contribute the + # minimum key f2u(-FLT_MAX); with n_eff > K here they + # can never enter the top-K, so their (out-of-range) + # indices are never emitted. + v = frags[u][q] + if gi + cutlass.Int32(q) >= n_eff: + v = cutlass.Float32(-FLT_MAX) + key = _f2u(v) + s_cand[gi + cutlass.Int32(q)] = ( + cutlass.Uint64(key) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(gi + cutlass.Int32(q))) + cute.arch.atomic_add( + s_hist.iterator + cutlass.Int32(key >> cutlass.Uint32(21)), + cutlass.Int32(1), + scope="cta", + ) + cute.arch.barrier() + + # ---- radix_select_emit (11/11/10-bit, whole-bin early exit) ---- + if npad == cutlass.Int32(kK): + # C == K: every candidate is admitted; identity emit. (With + # seq_lens present this is unreachable — n_eff <= npad == K + # lands in the degenerate branch — but kept for structural + # parity with the op43/CUDA source.) + io = tidx + while io < npad: + out_row[io] = cutlass.Int32( + cutlass.Uint32(s_cand[io] & cutlass.Uint64(0xFFFFFFFF)) + ) + io = io + cutlass.Int32(num_threads) + else: + want = cutlass.Int32(kK) + m = cutlass.Int32(0) + # level 0: top 11 bits (hist prebuilt during collect) + self._bin_select(2048, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b0 = s_sel[0] + a0 = s_sel[1] + h0 = s_hist[b0] + if want == a0 + h0: + # k-th boundary == bin edge: whole bin admitted + self._emit_prefix_ge( + npad, + cutlass.Uint32(21), + cutlass.Uint32(b0), + s_cand, + s_cnt, + out_row, + tidx, + lane, + ) + else: + m = m + a0 + want = want - a0 + docompact = h0 <= cutlass.Int32(sidecap) + cute.arch.barrier() # all reads of hist/sel done before reuse + if tidx == cutlass.Int32(0): + s_cnt[2] = cutlass.Int32(0) + for z1 in cutlass.range_constexpr(2048 // num_threads): + s_hist[tidx + cutlass.Int32(z1 * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + # level 1 sweep: mid 11 bits of boundary-bin members; compact + ub0 = cutlass.Uint32(b0) + i1 = tidx + while i1 < npad: + kv1 = s_cand[i1] + key1 = cutlass.Uint32(kv1 >> cutlass.Uint64(32)) + if (key1 >> cutlass.Uint32(21)) == ub0: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32( + (key1 >> cutlass.Uint32(10)) & cutlass.Uint32(0x7FF) + ), + cutlass.Int32(1), + scope="cta", + ) + if docompact: + p1 = cute.arch.atomic_add( + s_cnt.iterator + 2, cutlass.Int32(1), scope="cta" + ) + s_side[p1] = kv1 + i1 = i1 + cutlass.Int32(num_threads) + cute.arch.barrier() + self._bin_select(2048, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b1 = s_sel[0] + a1 = s_sel[1] + h1 = s_hist[b1] + p01 = (ub0 << cutlass.Uint32(11)) | cutlass.Uint32(b1) + if want == a1 + h1: + self._emit_prefix_ge( + npad, + cutlass.Uint32(10), + p01, + s_cand, + s_cnt, + out_row, + tidx, + lane, + ) + else: + m = m + a1 + want = want - a1 + cute.arch.barrier() + for z2 in cutlass.range_constexpr(1024 // num_threads): + s_hist[tidx + cutlass.Int32(z2 * num_threads)] = cutlass.Int32(0) + cute.arch.barrier() + # level 2 sweep: low 10 bits + ub1 = cutlass.Uint32(b1) + if docompact: + i2 = tidx + while i2 < h0: + key2 = cutlass.Uint32(s_side[i2] >> cutlass.Uint64(32)) + if ((key2 >> cutlass.Uint32(10)) & cutlass.Uint32(0x7FF)) == ub1: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32(key2 & cutlass.Uint32(0x3FF)), + cutlass.Int32(1), + scope="cta", + ) + i2 = i2 + cutlass.Int32(num_threads) + else: + i3 = tidx + while i3 < npad: + key3 = cutlass.Uint32(s_cand[i3] >> cutlass.Uint64(32)) + if (key3 >> cutlass.Uint32(10)) == p01: + cute.arch.atomic_add( + s_hist.iterator + + cutlass.Int32(key3 & cutlass.Uint32(0x3FF)), + cutlass.Int32(1), + scope="cta", + ) + i3 = i3 + cutlass.Int32(num_threads) + cute.arch.barrier() + self._bin_select(1024, want, s_hist, s_iwred, s_sel, tidx, lane, warp_id) + b2 = s_sel[0] + a2 = s_sel[1] + m = m + a2 + want = want - a2 + kth = (p01 << cutlass.Uint32(10)) | cutlass.Uint32(b2) + self._emit_final(npad, kth, m, want, s_cand, s_cnt, out_row, tidx, lane) + + # ------------------------------------------------------------------ + # Host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__(self, logits: cute.Tensor, seq_lens: cute.Tensor, out: cute.Tensor, stream): + num_rows = logits.shape[0] + self.direct_topk_kernel(logits, seq_lens, out).launch( + grid=(num_rows, 1, 1), + block=(self.num_threads, 1, 1), + stream=stream, + min_blocks_per_mp=1, # __launch_bounds__(TB, 1) + ) + + +# --------------------------------------------------------------------------- +# torch-facing entry with a compiled-variant cache keyed (K, TB). +# --------------------------------------------------------------------------- +_COMPILE_CACHE: dict = {} + + +def _get_compiled(top_k: int, num_threads: int = 1024): + key = (top_k, num_threads) + compiled = _COMPILE_CACHE.get(key) + if compiled is None: + from cutlass.cute import runtime as _crt + + kernel = DirectTopKKernel(top_k=top_k, num_threads=num_threads) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, top_k), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kernel, + logits_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +_CHECKED_SIGS: set = set() + + +def _check_contract(logits, seq_lens, out, K): + assert logits.dtype == torch.float32 and logits.is_contiguous() + assert out.dtype == torch.int32 and out.is_contiguous() + assert seq_lens.dtype == torch.int32 + bs, npad = logits.shape + assert npad <= DKCMAX, f"direct path requires npad <= {DKCMAX}, got {npad}" + assert npad % 4 == 0, f"npad must be a multiple of 4, got {npad}" + assert out.shape[0] == bs and out.shape[1] == K + assert logits.data_ptr() % 16 == 0 and out.data_ptr() % 16 == 0 + + +def direct_topk(logits: torch.Tensor, seq_lens: torch.Tensor, out: torch.Tensor, K: int) -> None: + """Exact top-K indices of ``logits`` rows into ``out`` (direct path). + + logits: [BS, npad] fp32 contiguous; the tail beyond each row's N_eff + may be stale garbage (masked in-kernel). + seq_lens: [BS] int32, request-level, uncompressed-token space. + out: [BS, K] int32 contiguous. + K: 512 / 1024 / 2048. + + Contract checks run once per (shape, K) signature to keep the hot + launch path at bare tvm-ffi cost. + """ + sig = (logits.shape, out.shape, K) + if sig not in _CHECKED_SIGS: + _check_contract(logits, seq_lens, out, K) + _CHECKED_SIGS.add(sig) + compiled = _COMPILE_CACHE.get((K, 1024)) + if compiled is None: + compiled = _get_compiled(K) + compiled(logits, seq_lens, out) + + +__all__ = ["DirectTopKKernel", "direct_topk", "DKCMAX"] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py new file mode 100644 index 000000000000..8eb3a8bfe963 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BSX top-K tier dispatcher — routes ``cute_dsl_gvr_topk_decode`` calls to +the op43 CuTe DSL tiers (direct / reg / tp). + +Port of the op43 ``ct_bsx.py`` unified launcher, itself an exact +transcription of the CUDA dispatch in op42 ``gvr_bsx.cu`` +(``gvr_topk_launch_batched`` + ``launch_dense`` + ``launch_tp<512>``). +The gvr streaming tier (cs=16 fallback for npad > 262144) is intentionally +NOT ported: it is unreachable inside the deployment envelope +(npad <= 262144), which :func:`is_bsx_supported` enforces. + +v1 dispatcher guard (anything else falls back to the in-tree +``GvrTopKKernel`` path in ``CuteDSLGvrTopKDecodeRunner.forward``): + dtype == fp32, next_n == 1, compress_ratio == 4, order_row is None, + counters is None, K in {512, 1024, 2048}, npad <= 262144, npad % 64 == 0, + contiguous 16B-aligned tensors, and the routed tier's cluster size within + the queried hardware max (never silently degrade the cluster shape). +Per-row degeneracy (N_eff <= K) and ragged N are handled INSIDE the tiers, +so the guard needs no device sync. + +Env knobs (identical semantics to the CUDA arm's ``GVR_BSX_*`` static locals, +renamed for the production tree; cached at first use like the CUDA static +locals; :func:`_reset_env_cache` re-reads for tests): + TRTLLM_BSX_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); + else the bs threshold at which the tp tier takes over. + TRTLLM_BSX_DENSE_BS same, for the dense (tb=1024) reg tiers. +""" + +import os + +import torch + +from .gvr_topk_decode_bsx_direct import DKCMAX, direct_topk +from .gvr_topk_decode_bsx_reg import reg_topk +from .gvr_topk_decode_bsx_tp import tp_cluster_size, tp_topk +from .single_pass_multi_cta_radix_topk_cluster import _query_max_cluster_size + +_BIG = 1 << 30 + +_ENV = {} # cached thresholds, mirrors the CUDA static-local caching + + +def _env_threshold(name): + t = _ENV.get(name) + if t is None: + e = os.environ.get(name) + t = int(e) if e not in (None, "") else -1 + if t == 0: + t = _BIG + _ENV[name] = t + return t + + +def _reset_env_cache(): + _ENV.clear() + _DISPATCH_CACHE.clear() # routes depend on the env thresholds + + +def _thresholds(npad): + tpb = _env_threshold("TRTLLM_BSX_TP_BS") + if tpb < 0: + tpb = 256 if npad <= 20480 else (128 if npad < 32768 else 16) + dnb = _env_threshold("TRTLLM_BSX_DENSE_BS") + if dnb < 0: + dnb = 8 if npad >= 163840 else (64 if npad < 32768 else _BIG) + return tpb, dnb + + +# tier name -> (kind, params). reg params = (cs, tb, maxv, ar). +def route(bs: int, npad: int, K: int) -> str: + """Tier the CUDA gvr_topk_launch_batched would take. Returns + 'tp' | 'direct' | 'reg(cs=C,tb=T,maxv=M,ar=A)' | 'gvr(cs=16,tb=512)' + ('gvr' is unreachable through :func:`bsx_topk` — see module docstring).""" + tpb, dnb = _thresholds(npad) + if bs >= tpb: + return "tp" + if npad > DKCMAX and bs >= dnb: + # launch_dense table + if npad <= 20480: + return "reg(cs=1,tb=1024,maxv=5,ar=8)" + if npad < 32768: + return "reg(cs=1,tb=1024,maxv=8,ar=8)" + if npad <= 65536: + return "reg(cs=2,tb=1024,maxv=8,ar=8)" + if npad <= 131072: + return "reg(cs=4,tb=1024,maxv=8,ar=8)" + if npad <= 262144: + return "reg(cs=8,tb=1024,maxv=8,ar=8)" + return "gvr(cs=16,tb=512)" + # latency ladder + if npad <= DKCMAX: + return "direct" + if npad < 16384: + return "reg(cs=1,tb=512,maxv=8,ar=8)" + if npad < 32768: + return "reg(cs=4,tb=512,maxv=4,ar=8)" + if npad <= 49152: + return "reg(cs=8,tb=512,maxv=3,ar=8)" + if npad <= 65536: + return "reg(cs=8,tb=512,maxv=4,ar=8)" + if npad <= 131072: + return "reg(cs=8,tb=512,maxv=8,ar=8)" + if npad <= 163840: + if K >= 2048: + return "reg(cs=16,tb=512,maxv=5,ar=6)" + return "reg(cs=16,tb=512,maxv=5,ar=8)" + if npad <= 262144: + if K == 2048: + return "reg(cs=16,tb=512,maxv=8,ar=8)" + return "reg(cs=16,tb=512,maxv=8,ar=6)" + return "gvr(cs=16,tb=512)" + + +def _parse_reg(tier): + body = tier[tier.index("(") + 1 : -1] + d = dict(kv.split("=") for kv in body.split(",")) + return int(d["cs"]), int(d["tb"]), int(d["maxv"]), int(d["ar"]) + + +def route_cluster_size(bs: int, npad: int, K: int) -> int: + """Cluster size the routed tier would launch with (host-only helper for + the hardware cluster-cap guard).""" + tier = route(bs, npad, K) + if tier == "tp": + return tp_cluster_size(bs, npad) + if tier == "direct": + return 1 + if tier.startswith("reg"): + return _parse_reg(tier)[0] + return 16 # gvr + + +# Per-call route()+_parse_reg() string work costs ~2.5-3us host-submit wall +# on <20us kernels (op43 S-E). The routing decision is pure in (bs, npad, K) +# [env thresholds are cached at first use, like the CUDA static locals], so +# bind it once per key to a closure. +_DISPATCH_CACHE = {} # (bs, npad, K) -> callable(logits, pre, seq_lens, out) + + +def _bind(bs, npad, K): + tier = route(bs, npad, K) + if tier == "tp": + + def fn(lg, pre, sl, out): + tp_topk(lg, pre, sl, out, K) + elif tier == "direct": + + def fn(lg, pre, sl, out): + direct_topk(lg, sl, out, K) + elif tier.startswith("gvr"): + raise ValueError( + f"bsx gvr tier is not ported (npad beyond the deployment " + f"envelope); is_bsx_supported must gate this out (bs={bs}, " + f"npad={npad}, K={K})" + ) + else: + cs, tb, maxv, ar = _parse_reg(tier) + + def fn(lg, pre, sl, out): + reg_topk(lg, pre, sl, out, K, cs, tb, maxv, ar) + + return fn + + +def is_bsx_supported( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, + next_n: int, + compress_ratio: int, + order_row, + counters, +) -> bool: + """Host-only v1 guard for the bsx tiers (no device sync; see module + docstring). Returns False -> caller uses the in-tree kernel.""" + if logits.dtype != torch.float32: + return False + if next_n != 1 or compress_ratio != 4: + return False + if order_row is not None or counters is not None: + return False + if top_k not in (512, 1024, 2048): + return False + bs, npad = logits.shape + if npad > 262144 or npad % 64 != 0: + return False + if not ( + logits.is_contiguous() + and pre_idx.is_contiguous() + and output_indices.is_contiguous() + and seq_lens.is_contiguous() + ): + return False + if pre_idx.shape != (bs, top_k) or output_indices.shape != (bs, top_k): + return False + if seq_lens.shape != (bs,) or seq_lens.dtype != torch.int32: + return False + if pre_idx.dtype != torch.int32 or output_indices.dtype != torch.int32: + return False + if ( + logits.data_ptr() % 16 != 0 + or pre_idx.data_ptr() % 16 != 0 + or output_indices.data_ptr() % 16 != 0 + ): + return False + # Cluster cap: fall back to the in-tree kernel (dispatcher-level) rather + # than silently degrading the tier's cluster shape. + if route_cluster_size(bs, npad, top_k) > _query_max_cluster_size(): + return False + return True + + +def bsx_topk( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, +) -> None: + """Unified bsx tier dispatch, replicating gvr_topk_launch_batched. + + logits [BS, npad] fp32 (npad multiple of 64; per-row tail beyond N_eff + may be garbage — masked in-kernel), pre_idx [BS, K] int32, seq_lens + [BS] int32 (request-level, uncompressed-token space), output_indices + [BS, K] int32. Caller must have passed :func:`is_bsx_supported`. + """ + bs, npad = logits.shape + key = (bs, npad, top_k) + fn = _DISPATCH_CACHE.get(key) + if fn is None: + fn = _DISPATCH_CACHE[key] = _bind(bs, npad, top_k) + fn(logits, pre_idx, seq_lens, output_indices) + + +__all__ = [ + "bsx_topk", + "is_bsx_supported", + "route", + "route_cluster_size", + "_reset_env_cache", +] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py new file mode 100644 index 000000000000..e6dd0d37c1db --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py @@ -0,0 +1,872 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BSX register-resident (reg) GVR Top-K tier — CuTe DSL, Blackwell SM100. + +Port of the op43 ``ct_reg.py`` CuTe DSL translation of the CUDA +``gvr_topk_reg`` register-resident GVR top-K kernel +(op42 iter12-F1b), adapted for the production +``trtllm::cute_dsl_gvr_topk_decode`` contract (see the module docstring of +``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory: ragged-N +masking, pre_idx clamping, per-row degenerate identity emit). + +The row is loaded ONCE from GMEM into per-thread registers (MAXV float4s per +thread; out-of-slice lanes AND lanes beyond the row's N_eff are filled +-FLT_MAX — the ragged-N adaptation extends the existing OOR-lane idiom, so a +single mask at load time covers every downstream count / max-below / plateau +emit / collect over the register array). All subsequent passes are pure ALU — +global traffic pinned at 4*npad bytes/row. + +Phase skeleton (shared with the tp tier — helpers REUSED from +``gvr_topk_decode_bsx_tp``): + P1 : hint gather + two-stage 64-bin histogram -> rung ladder (tp.phase1) + P2 : full AR-rung register count + cluster exchange, secant refine / + max-below plateau descent (count_reg here) + P3 : rank-scatter collect (NO atomics): cached per-thread counts -> warp + inclusive scan + block scan + cluster-peer prefix -> deterministic + scatter position; keys f2u'd at push; P4's round-0 radix histogram is + built IN FLIGHT with LOCAL smem atomics, non-rank0 CTAs merge their + histograms into CTA0 with <=256 remote atomics. + P4 : CTA0-solo radix select starting directly at bin-select for round 0, + then tie-aware ticketed emit. + +``__launch_bounds__(TB, 1)`` -> ``.launch(min_blocks_per_mp=1)``. +CS=16 relies on CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, which the +CuTe DSL sets unconditionally on every kernel — the dispatcher additionally +verifies the queried hardware max cluster size and falls back to the in-tree +kernel when the tier's CS exceeds it (never degrades silently). + +CRITICAL: CuTe DSL cluster kernels need EXPLICIT trailing cluster +rendezvous where remote (DSMEM) ops can still be in flight — nvcc inserts an +implicit cluster barrier before ret in DSMEM kernels, the DSL does not +(missing it => timing-dependent CUDA 719 faults). + +op43 S-E convergence notes preserved in this port (do not "simplify away"): +1. ``cutlass.utils.distributed.atomicAdd`` lowers to sys-scoped + ``atom.relaxed.SYS.shared`` which blocks ptxas' warp aggregation; the + CTA-scope relaxed ``_atomic_add_cta`` restores the fast path. +2. ``cute.arch.cluster_arrive/wait`` emit the NON-aligned barrier forms; + ``cg::cluster.sync()`` parity needs the ALIGNED pair + (``_cluster_sync_aligned``). +3. Candidates are a single packed (key<<32|idx) u64 array — one 8B DSMEM + push per candidate (halves remote-store transactions at CS>1). +""" + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import nvvm +from cutlass.cute import runtime as _crt +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from .gvr_topk_decode_bsx_tp import ( + FLT_MAX, + INF, + MAXPASS, + RUNGS, + GvrTpKernel, + _atom_shared_cluster_add_i32, + _exp2i, + _f2u_bits, + _f32_bits_u32, + _ld_shared_cluster_f32, + _ld_shared_cluster_i32, + _mapa_shared_cluster, + _shfl_down_add, + _shfl_up_add, + _st_shared_cluster_u64, +) + + +@dsl_user_op +def _atomic_add_cta(dst_ptr, val, *, loc=None, ip=None): + """CTA-scope relaxed atomicAdd on shared memory; returns the old value. + Drop-in for cutlass.utils.distributed.atomicAdd (which is sys-scoped and + blocks ptxas warp-aggregation — see module docstring note 1).""" + return cute.arch.atomic_add( + ptr=dst_ptr.llvm_ptr, val=val, sem="relaxed", scope="cta", loc=loc, ip=ip + ) + + +@dsl_user_op +def _cluster_arrive_aligned(*, loc=None, ip=None): + nvvm.cluster_arrive(aligned=True, loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_wait_aligned(*, loc=None, ip=None): + nvvm.cluster_wait(aligned=True, loc=loc, ip=ip) + + +@cute.jit +def _cluster_sync_aligned(): + _cluster_arrive_aligned() + _cluster_wait_aligned() + + +class GvrRegKernel(GvrTpKernel): + """CuTe DSL port of gvr_topk_reg (fp32, B200/B300). + + Inherits phase1 (already N_eff-aware), the DSMEM/scan idioms and + ``_row_n_eff`` from the tp port (identical CUDA source); adds + register-resident count/max-below/collect and the rank-scatter + + in-flight-histogram P3/P4. + """ + + def __init__( + self, + top_k: int, + kC: int, + cluster_size: int = 1, + ar: int = RUNGS, + maxv: int = 8, + num_threads: int = 512, + next_n: int = 1, + compress_ratio: int = 4, + ): + assert num_threads % 32 == 0 + assert ar in (6, 8) + assert cluster_size in (1, 2, 4, 8, 16) + assert 1 <= maxv <= 8 + self.top_k = top_k + self.kC = kC + self.cluster_size = cluster_size + self.ar = ar + self.uf = 4 # unused by the reg path; kept for inherited helpers + self.maxv = maxv + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + + # ------------------------------------------------------------------ + # count_reg: R-rung count over the register array. + # Per-thread counts -> s_ptcnt[r*TB + tid] (P3 rank-scatter offsets). + # Dummy (-FLT_MAX) lanes — both out-of-slice AND ragged-N-masked — + # only pass a rung == -FLT_MAX, which can never be `chosen` (its count + # > kC by construction since npad > kC). + # ------------------------------------------------------------------ + @cute.jit + def count_reg(self, R: cutlass.Constexpr, a, tidx, s_rungs, s_ptcnt): + TB = cutlass.const_expr(self.num_threads) + MAXV = cutlass.const_expr(self.maxv) + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + for u in cutlass.range_constexpr(MAXV): + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # exchange_counts override: identical to the tp tier's, but the cluster + # sync is the ALIGNED barrier pair (cg::cluster.sync() parity — module + # docstring note 2). The tp tier keeps the non-aligned form it was + # validated with. + # ------------------------------------------------------------------ + @cute.jit + def exchange_counts( + self, R: cutlass.Constexpr, par, tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + cute.arch.barrier() # ptcnt final + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if wid < cutlass.Int32(R): + s = cutlass.Int32(0) + for k in cutlass.range_constexpr(TB // 32): + s = s + s_ptcnt[wid * TB + lane + cutlass.Int32(32 * k)] + s = cute.arch.warp_redux_sync(s, "add") + if cutlass.const_expr(CS == 1): + if lane == cutlass.Int32(0): + s_rcnt[wid] = s + s_rpre[wid] = cutlass.Int32(0) + else: + if lane == cutlass.Int32(0): + s_ipartial[par * cutlass.Int32(RUNGS) + wid] = s + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + else: + _cluster_sync_aligned() + if tidx < cutlass.Int32(R): + tot = cutlass.Int32(0) + pre = cutlass.Int32(0) + local_ptr = s_ipartial.iterator + (par * cutlass.Int32(RUNGS) + tidx) + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + v = _ld_shared_cluster_i32(a) + tot = tot + v + if cutlass.Int32(rr) < rank: + pre = pre + v + s_rcnt[tidx] = tot + s_rpre[tidx] = pre + cute.arch.barrier() + + # ------------------------------------------------------------------ + # max_below_reg: largest value strictly below t_hi_bound across the + # whole row (cluster-reduced). No index gate needed: dummy lanes hold + # -FLT_MAX which never raises the max. + # ------------------------------------------------------------------ + @cute.jit + def max_below_reg(self, a, t_hi_bound, par, tidx, s_fwred, s_fpartial): + CS = cutlass.const_expr(self.cluster_size) + MAXV = cutlass.const_expr(self.maxv) + NWARP = cutlass.const_expr(self.num_warps) + m = cutlass.Float32(-FLT_MAX) + for u in cutlass.range_constexpr(MAXV): + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + m = cute.arch.warp_redux_sync(m, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = m + cute.arch.barrier() + if tidx == cutlass.Int32(0): + mm = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + mm = cute.arch.fmax(mm, s_fwred[w]) + s_fpartial[par] = mm + res = cutlass.Float32(-FLT_MAX) + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + res = s_fpartial[par] + else: + _cluster_sync_aligned() + local_ptr = s_fpartial.iterator + par + for rr in cutlass.range_constexpr(CS): + ad = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + res = cute.arch.fmax(res, _ld_shared_cluster_f32(ad)) + return res + + # ------------------------------------------------------------------ + # kernel + # ------------------------------------------------------------------ + @cute.kernel + def gvr_reg_kernel( + self, logits: cute.Tensor, pre_idx: cute.Tensor, seq_lens: cute.Tensor, out_idx: cute.Tensor + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + AR = cutlass.const_expr(self.ar) + MAXV = cutlass.const_expr(self.maxv) + K = cutlass.const_expr(self.top_k) + kC = cutlass.const_expr(self.kC) + NWARP = cutlass.const_expr(self.num_warps) + + if cutlass.const_expr(CS > 1): + row = bidx // cutlass.Int32(CS) + rank = cute.arch.block_idx_in_cluster() + else: + row = bidx + rank = cutlass.Int32(0) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[row, None] + pre_idx_row = pre_idx[row, None] + out_row = out_idx[row, None] + row_addr = logits_row.iterator.toint() + + # Ragged N: per-row valid length from seq_lens (inherited helper). + n_eff = self._row_n_eff(seq_lens, row) + + # ---- shared memory (order must be identical across CTAs for mapa) ---- + smem = SmemAllocator() + # Single packed (key<<32 | idx) u64 candidate array — CUDA's + # `unsigned long long cand[kC]`; one 8B DSMEM push per candidate in + # the P3 scatter (halves remote-store transactions under 8-cluster + # contention at BS>=8 — op43 S-E round 2). + s_cand = smem.allocate_tensor( + element_type=cutlass.Uint64, + layout=cute.make_ordered_layout((kC,), order=(0,)), + byte_alignment=128, + ) + s_ptcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS * self.num_threads,), order=(0,)), + byte_alignment=128, + ) + s_hist = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((256,), order=(0,)), + byte_alignment=128, + ) + s_rungs = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rpre = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_ipartial = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((2 * RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_fpartial = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + s_fwred = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2 * self.num_warps,), order=(0,)), + byte_alignment=64, + ) + s_hminmax = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + # iscalars: [0]=sel_bin [1]=sel_above [2]=sel_count [3]=cnt_m [4]=cnt_t + s_isc = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((8,), order=(0,)), + byte_alignment=32, + ) + # P3 rank-scatter block-scan scratch (reg path only) + s_iwred = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((self.num_warps,), order=(0,)), + byte_alignment=64, + ) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad. ---- + # Cluster-uniform branch (all CTAs own the same row); no DSMEM op is + # issued on this path, so no exit rendezvous is required either. + if n_eff <= cutlass.Int32(K): + if rank == cutlass.Int32(0): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(TB) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(K): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(TB) + else: + # ---- slice + one-time register row load (launcher: vpc <= MAXV*TB) + V4 = npad >> cutlass.Int32(2) + vpc = (V4 + cutlass.Int32(CS) - cutlass.Int32(1)) // cutlass.Int32(CS) + v0 = rank * vpc + v1 = v0 + vpc + if v1 > V4: + v1 = V4 + + copy_atom = self._copy_atom() + a = cute.make_fragment((MAXV * 4,), cutlass.Float32) + frag4 = cute.make_fragment((4,), cutlass.Float32) + for u in cutlass.range_constexpr(MAXV): + i = v0 + tidx + cutlass.Int32(u * TB) + if i < v1: + self._ld_float4(copy_atom, row_addr, i, frag4) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + # Ragged N: extend the OOR-lane idiom to the valid + # length — elements at global index >= N_eff become + # -FLT_MAX dummy lanes, exactly like out-of-slice + # lanes. One mask here covers ALL downstream register + # passes (count/max-below/plateau emit/P3 collect). + val = frag4[q] + if gi + cutlass.Int32(q) >= n_eff: + val = cutlass.Float32(-FLT_MAX) + a[4 * u + q] = val + else: + for q in cutlass.range_constexpr(4): + a[4 * u + q] = cutlass.Float32(-FLT_MAX) + + xch = cutlass.Int32(0) + thr = cutlass.Float32(0.0) + chosen = cutlass.Int32(-1) + C = cutlass.Int32(0) + cbase = cutlass.Int32(0) + m_gt = cutlass.Int32(-1) + + self.phase1(logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs) + # Pre-zero the P4 round-0 radix histogram now that P1 is done with + # hist. The first exchange's cluster barrier (release) publishes it + # before any CTA reaches the P3 collect that increments it in flight. + if tidx < cutlass.Int32(256): + s_hist[tidx] = cutlass.Int32(0) + self.count_reg(AR, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + + span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) + + # ---- P2: secant refine driver (redundant on every thread) ---- + t_lo = cutlass.Float32(-FLT_MAX) + t_hi = cutlass.Float32(INF) + c_hi = cutlass.Int32(0) + Rcur = cutlass.Int32(AR) + passno = cutlass.Int32(0) + running = cutlass.Int32(1) + descend_break = cutlass.Int32(0) + while running != cutlass.Int32(0): + # first rung index j with rcnt[j] >= K (rcnt ascending in j) + j = Rcur + for r_ in range(AR - 1, -1, -1): + if cutlass.Int32(r_) < Rcur: + if s_rcnt[r_] >= cutlass.Int32(K): + j = cutlass.Int32(r_) + jj = j + if jj > Rcur - cutlass.Int32(1): + jj = Rcur - cutlass.Int32(1) + cj = s_rcnt[jj] + rj = s_rungs[jj] + bj = s_rpre[jj] + found = cutlass.Int32(0) + if j < Rcur: + if cj <= cutlass.Int32(kC): + found = cutlass.Int32(1) + if found != cutlass.Int32(0): + chosen = j + thr = rj + C = cj + cbase = bj + running = cutlass.Int32(0) + else: + if j < Rcur: + if rj >= t_lo: + t_lo = rj + if j > cutlass.Int32(0): + jm = j - cutlass.Int32(1) + if s_rungs[jm] <= t_hi: + t_hi = s_rungs[jm] + c_hi = s_rcnt[jm] + descend = cutlass.Int32(0) + if passno >= cutlass.Int32(MAXPASS): + descend = cutlass.Int32(1) + e3 = passno * cutlass.Int32(3) + if e3 > cutlass.Int32(24): + e3 = cutlass.Int32(24) + step = span0 * _exp2i(e3) + dt = cutlass.Float32(0.0) + mode = cutlass.Int32(2) # 0=up-ladder, 1=down-ladder, 2=secant + if t_hi == cutlass.Float32(INF): + mode = cutlass.Int32(0) + nr0 = t_lo + step * cutlass.Float32(float(1 << (AR - 1))) + if nr0 == cutlass.Float32(INF): + descend = cutlass.Int32(1) + else: + if t_lo == cutlass.Float32(-FLT_MAX): + mode = cutlass.Int32(1) + else: + dt = (t_hi - t_lo) * cutlass.Float32(1.0 / float(AR + 1)) + nr_last = t_hi - dt * cutlass.Float32(float(AR)) + nr_first = t_hi - dt + ok = cutlass.Int32(0) + if nr_last > t_lo: + if nr_first < t_hi: + ok = cutlass.Int32(1) + if ok == cutlass.Int32(0): + descend = cutlass.Int32(1) + if descend != cutlass.Int32(0): + running = cutlass.Int32(0) + descend_break = cutlass.Int32(1) + else: + cute.arch.barrier() + if tidx == cutlass.Int32(0): + for r_ in cutlass.range_constexpr(AR): + nrv = cutlass.Float32(0.0) + if mode == cutlass.Int32(0): + nrv = t_lo + step * cutlass.Float32(float(1 << (AR - 1 - r_))) + else: + if mode == cutlass.Int32(1): + nrv = t_hi - step * cutlass.Float32(float(1 << r_)) + else: + nrv = t_hi - dt * cutlass.Float32(float(r_ + 1)) + s_rungs[r_] = nrv + cute.arch.barrier() + self.count_reg(AR, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + passno = passno + cutlass.Int32(1) + Rcur = cutlass.Int32(AR) + + # ---- plateau descent (exact max-below stepping) ---- + if descend_break != cutlass.Int32(0): + pl = cutlass.Int32(1) + while pl != cutlass.Int32(0): + vstar = self.max_below_reg( + a, t_hi, xch & cutlass.Int32(1), tidx, s_fwred, s_fpartial + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_rungs[0] = vstar + cute.arch.barrier() + self.count_reg(1, a, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + c = s_rcnt[0] + okc = cutlass.Int32(0) + if c >= cutlass.Int32(K): + if c <= cutlass.Int32(kC): + okc = cutlass.Int32(1) + if okc != cutlass.Int32(0): + chosen = cutlass.Int32(0) + thr = vstar + C = c + cbase = s_rpre[0] + pl = cutlass.Int32(0) + else: + if c < cutlass.Int32(K): + t_hi = vstar + c_hi = c + else: + thr = vstar + m_gt = c_hi + pl = cutlass.Int32(0) + + if m_gt >= cutlass.Int32(0): + # ---- plateau direct emit from registers ---- + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) # cnt_m + s_isc[4] = cutlass.Int32(0) # cnt_t + if cutlass.const_expr(CS > 1): + _cluster_sync_aligned() + a_m = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(0)) + a_t = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(4), cutlass.Int32(0)) + else: + cute.arch.barrier() + nt = cutlass.Int32(K) - m_gt + for u in cutlass.range_constexpr(MAXV): + i = v0 + tidx + cutlass.Int32(u * TB) + if i < v1: + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v > thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_m, cutlass.Int32(1)) + else: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1) + ) + out_row[p] = gi + cutlass.Int32(q) + else: + if v == thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_t, cutlass.Int32(1)) + else: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m_gt + p] = gi + cutlass.Int32(q) + else: + # ---- P3: rank-scatter collect from registers (NO atomics) ---- + myc = s_ptcnt[chosen * cutlass.Int32(TB) + tidx] + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + incl = myc + for o in [1, 2, 4, 8, 16]: + incl = _shfl_up_add(incl, lane, o) + if lane == cutlass.Int32(31): + s_iwred[wid] = incl + cute.arch.barrier() + if wid == cutlass.Int32(0): + v_ = cutlass.Int32(0) + if lane < cutlass.Int32(NWARP): + v_ = s_iwred[lane] + iv = v_ + for o in [x for x in [1, 2, 4, 8, 16] if x < self.num_warps]: + iv = _shfl_up_add(iv, lane, o) + if lane < cutlass.Int32(NWARP): + s_iwred[lane] = iv - v_ + cute.arch.barrier() + pos = cbase + s_iwred[wid] + (incl - myc) + + # scatter: f2u key at push + IN-FLIGHT local P4 round-0 + # histogram. thr > -FLT_MAX whenever we get here (chosen count + # <= kC < npad), so dummy lanes never pass. + if cutlass.const_expr(CS > 1): + a_cand = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + for u in cutlass.range_constexpr(MAXV): + gi = (v0 + tidx + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = a[4 * u + q] + if v >= thr: + ux = _f2u_bits(_f32_bits_u32(v)) + _atomic_add_cta( + s_hist.iterator + cutlass.Int32(ux >> cutlass.Uint32(24)), + cutlass.Int32(1), + ) + kv = (cutlass.Uint64(ux) << cutlass.Uint64(32)) | cutlass.Uint64( + cutlass.Uint32(gi + cutlass.Int32(q)) + ) + if cutlass.const_expr(CS > 1): + _st_shared_cluster_u64(a_cand + pos * cutlass.Int32(8), kv) + else: + s_cand[pos] = kv + pos = pos + cutlass.Int32(1) + + # non-rank0: merge local round-0 histogram into CTA0 + # (<=256 remote atomics) + if cutlass.const_expr(CS > 1): + if rank != cutlass.Int32(0): + cute.arch.barrier() # local hist final + a_hist = _mapa_shared_cluster(s_hist.iterator, cutlass.Int32(0)) + b = cutlass.Int32(tidx) + while b < cutlass.Int32(256): + hv = s_hist[b] + if hv != cutlass.Int32(0): + _atom_shared_cluster_add_i32(a_hist + b * cutlass.Int32(4), hv) + b = b + cutlass.Int32(TB) + _cluster_sync_aligned() + else: + cute.arch.barrier() + + # ---- P4 (CTA0 solo when CS>1) ---- + if rank == cutlass.Int32(0): + if C == cutlass.Int32(K): + ie = cutlass.Int32(tidx) + while ie < C: + out_row[ie] = cutlass.Int32( + cutlass.Uint32(s_cand[ie] & cutlass.Uint64(0xFFFFFFFF)) + ) + ie = ie + cutlass.Int32(TB) + else: + # 4x8-bit radix select; keys ALREADY f2u'd; round-0 + # histogram already built in flight -> round 0 starts + # at bin-select. + pref = cutlass.Uint32(0) + want = cutlass.Int32(K) + m = cutlass.Int32(0) + final_shift = cutlass.Uint32(0) + active = cutlass.Int32(1) + for r_ in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * r_) + if active != cutlass.Int32(0): + if cutlass.const_expr(r_ > 0): + ih = cutlass.Int32(tidx) + while ih < C: + u = cutlass.Uint32(s_cand[ih] >> cutlass.Uint64(32)) + if (u >> cutlass.Uint32(shift + 8)) == pref: + bb = cutlass.Int32( + (u >> cutlass.Uint32(shift)) & cutlass.Uint32(0xFF) + ) + _atomic_add_cta(s_hist.iterator + bb, cutlass.Int32(1)) + ih = ih + cutlass.Int32(TB) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + b8 = tidx * cutlass.Int32(8) + h = cute.make_fragment((8,), cutlass.Int32) + Ssum = cutlass.Int32(0) + for q in cutlass.range_constexpr(8): + h[q] = s_hist[b8 + cutlass.Int32(q)] + Ssum = Ssum + h[q] + s_hist[b8 + cutlass.Int32(q)] = cutlass.Int32(0) + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_down_add(x, tidx, o) + A = x - Ssum + if A < want: + if want <= A + Ssum: + run = A + for q in range(7, -1, -1): + if run < want: + if want <= run + h[q]: + s_isc[0] = b8 + cutlass.Int32(q) + s_isc[1] = run + s_isc[2] = h[q] + run = run + h[q] + cute.arch.barrier() + sel = s_isc[0] + above = s_isc[1] + m = m + above + want = want - above + pref = (pref << cutlass.Uint32(8)) | cutlass.Uint32(sel) + if s_isc[2] == want: + final_shift = cutlass.Uint32(shift) + active = cutlass.Int32(0) + kth = pref + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) + s_isc[4] = cutlass.Int32(0) + cute.arch.barrier() + nt = cutlass.Int32(K) - m + ie = cutlass.Int32(tidx) + while ie < C: + kv = s_cand[ie] + u = cutlass.Uint32(kv >> cutlass.Uint64(32)) >> final_shift + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + if u > kth: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1) + ) + out_row[p] = idx + else: + if u == kth: + p = _atomic_add_cta( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m + p] = idx + ie = ie + cutlass.Int32(TB) + + # Exit rendezvous — PLATEAU PATH ONLY (op43 S-E wave-overlap fix). + # On the P3/P4 path the post-merge _cluster_sync_aligned() above is + # the LAST remote-op rendezvous (P4 is CTA0-local smem + global + # stores only), so non-rank0 CTAs exit right after it — mirroring + # CUDA's `if (rank != 0) return;` and freeing 15/16 SMs one + # P4-duration earlier (~+5us/wave at BS>=8 otherwise). The + # plateau-emit path still has remote atomics in flight up to its + # end (timing-dependent CUDA 719 territory without this barrier — + # the DSL emits NO implicit pre-ret cluster barrier) and therefore + # KEEPS the trailing rendezvous. m_gt is cluster-uniform + # (redundantly computed driver), so this is a uniform branch + # around the cluster barrier. Degenerate rows (outer branch) + # issue no remote ops, so skipping it there is safe too. + if cutlass.const_expr(CS > 1): + if m_gt >= cutlass.Int32(0): + _cluster_sync_aligned() + + # ------------------------------------------------------------------ + # host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + seq_lens: cute.Tensor, + out_idx: cute.Tensor, + stream, + ): + num_rows = logits.shape[0] + CS = cutlass.const_expr(self.cluster_size) + self.gvr_reg_kernel(logits, pre_idx, seq_lens, out_idx).launch( + grid=(num_rows * CS, 1, 1), + block=(self.num_threads, 1, 1), + cluster=(CS, 1, 1) if cutlass.const_expr(CS > 1) else None, + stream=stream, + min_blocks_per_mp=1, + ) + + +# --------------------------------------------------------------------------- +# compile cache + public entry (explicit variant params; dispatcher selects) +# --------------------------------------------------------------------------- +_LAUNCH_CACHE = {} + + +def _get_compiled(K: int, CS: int, TB: int, MAXV: int, AR: int): + key = (K, CS, TB, MAXV, AR) + compiled = _LAUNCH_CACHE.get(key) + if compiled is None: + kC = 8192 if K >= 2048 else 6144 + kern = GvrRegKernel(top_k=K, kC=kC, cluster_size=CS, ar=AR, maxv=MAXV, num_threads=TB) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + pre_idx_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kern, + logits_fake, + pre_idx_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _LAUNCH_CACHE[key] = compiled + return compiled + + +_FAST = {} # (K, cs, tb, maxv, ar, npad) -> compiled (contract checked once) + + +def reg_topk( + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + cs: int, + tb: int, + maxv: int, + ar: int, +) -> None: + """CuTe DSL gvr_topk_reg. logits [BS, npad] fp32 (npad + mult of 64; tail beyond each row's N_eff may be garbage — masked at the + register load), pre_idx [BS, K] int32 hint, seq_lens [BS] int32 + (uncompressed-token space), out [BS, K] int32. Variant params are + explicit — the dispatcher selects; compile cache keyed + (K, cs, tb, maxv, ar).""" + npad = logits.shape[1] + key = (K, cs, tb, maxv, ar, npad) + fn = _FAST.get(key) + if fn is None: + assert npad % 64 == 0, f"npad {npad} not a multiple of 64" + vpc = (npad // 4 + cs - 1) // cs + assert vpc <= maxv * tb, ( + f"slice {vpc} float4s exceeds MAXV*TB={maxv * tb} (npad={npad}, cs={cs})" + ) + kc = 8192 if K >= 2048 else 6144 + assert npad > kc, f"npad {npad} <= kC {kc}: reg path has no trivial branch" + fn = _get_compiled(K, cs, tb, maxv, ar) + _FAST[key] = fn + fn(logits, pre_idx, seq_lens, out) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py new file mode 100644 index 000000000000..681c67793e0c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -0,0 +1,1478 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BSX throughput (tp) GVR Top-K tier — CuTe DSL, Blackwell SM100. + +Port of the op43 ``ct_tp.py`` CuTe DSL translation of the CUDA +``gvr_topk_tp`` throughput GVR top-K kernel (op42 iter12-F1b), +adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract: + +* RAGGED N: production logits tails beyond the per-row valid length are + stale garbage (NOT -FLT_MAX pad, unlike the op-bench harness). A device + ``seq_lens`` tensor is threaded into every tier and every global row read + is predicated at the LANE level: element index >= N_eff substitutes + -FLT_MAX AFTER the (always in-bounds, buffer width = logits.shape[1]) + float4 load. N_eff mirrors the in-tree ``run_one_row`` arithmetic: + ``(seq_lens[req] - next_n + row % next_n + 1) // compress_ratio``. +* pre_idx hardening: the P1 hint gather clamps out-of-range hint indices + into [0, N_eff-1] (production cold-start is all zeros; arbitrary garbage + must neither fault nor corrupt the rung ladder). +* Degenerate rows (N_eff <= K): per-row in-kernel identity emit + [0..N_eff-1] plus -1 padding to K, matching the in-tree kernel. + +Faithful phase-by-phase port otherwise: + P1 : hint gather + minmax + two-stage 64-bin histogram -> CCDF rung ladder + P2a : uniform every-32nd-float4 sampled multi-rung count -> 3-stage pivot + P2b : ONE fused streaming pass: exact counts at {pivot, hmin} + optimistic + collect of packed (key<<32|idx) u64 candidates >= pivot into CTA0 + smem (capped kC) + P2c : multi-rung secant refine / max-below plateau descent when the pivot + count misses [K, kC] + P3 : candidate reuse (thr == tpush, no overflow) or one re-stream collect + P4 : CTA0-solo 4x8-bit radix select + tie-aware ticketed emit + (plus the trivial npad <= kC path and the plateau direct-emit path) + +Templates -> compile-time ctor knobs; compile cache keyed (K, CS, AR, UF, TB). +``__launch_bounds__(TB, 2)`` -> ``.launch(min_blocks_per_mp=2)``. +grid dim3(CS, BS) -> 1-D grid BS*CS with cluster=(CS,1,1); row = bidx // CS. + +DSMEM (CS > 1) via inline PTX: mapa.shared::cluster + ld/st.shared::cluster ++ atom.relaxed.cluster.shared::cluster.add.u32. Writer-side cluster syncs use +FULL cluster_arrive (release) — never relaxed (a relaxed arrive has no +release semantics, so peer CTAs could observe stale DSMEM: see +``cluster_arrive_relaxed`` DSMEM-race history in the GVR campaign notes). +""" + +import cutlass +import cutlass.cute as cute +import cutlass.cute.math as cmath +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute import runtime as _crt +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils.distributed import atomicAdd +from cutlass.utils.smem_allocator import SmemAllocator + +RUNGS = 8 +MAXPASS = 8 +SS = 32 # P2a sample stride (float4s) +FLT_MAX = 3.4028234663852886e38 +INF = float("inf") + + +# --------------------------------------------------------------------------- +# DSMEM primitives (inline PTX) +# --------------------------------------------------------------------------- +@dsl_user_op +def _mapa_shared_cluster(smem_ptr, peer_rank, *, loc=None, ip=None): + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [smem_ptr_i32, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_i32(mapped_addr, *, loc=None, ip=None): + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.u32 $0, [$1];", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_f32(mapped_addr, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_shared_cluster_i32(mapped_addr, val, *, loc=None, ip=None): + llvm.inline_asm( + res=None, + operands_=[mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + asm_string="st.shared::cluster.u32 [$0], $1;", + constraints="r,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_shared_cluster_u64(mapped_addr, val, *, loc=None, ip=None): + """One 8B DSMEM candidate push (CUDA: dst[pos] = (u64)key<<32 | idx). + + op43 S-E round 4: a single packed u64 store halves remote-store + transactions vs two 4B pushes under CS>1 cluster contention (matches + the CUDA arm's ``unsigned long long cand[kC]``). Do NOT split back + into key/idx 4B stores — the split form measurably regresses the + K-scaled CS>1 band (pro/v32 BS16-64).""" + llvm.inline_asm( + res=None, + operands_=[mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + asm_string="st.shared::cluster.u64 [$0], $1;", + constraints="r,l", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_shared_cluster_add_i32(mapped_addr, val, *, loc=None, ip=None): + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [mapped_addr.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.relaxed.cluster.shared::cluster.add.u32 $0, [$1], $2;", + "=r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _f32_bits_u32(float_val): + """Raw fp32 bits as Uint32 (bit-cast).""" + return cutlass.Uint32(llvm.bitcast(cutlass.Uint32.mlir_type, float_val.ir_value())) + + +def _i32_bits_f32(int_val): + return cutlass.Float32(llvm.bitcast(cutlass.Float32.mlir_type, int_val.ir_value())) + + +@dsl_user_op +def _fmin_f32(a, b, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _f2u_bits(u): + """Order-preserving fp32-bits -> Uint32 key: u ^ (sign ? 0xFFFFFFFF : 0x80000000).""" + neg = cutlass.Uint32(0) - (u >> cutlass.Uint32(31)) # 0 or 0xFFFFFFFF + return u ^ (neg | cutlass.Uint32(0x80000000)) + + +@cute.jit +def _exp2i(e): + """2.0**e for dynamic int e in [0, 127] via exponent-bit construction.""" + bits = (e + cutlass.Int32(127)) << cutlass.Int32(23) + return _i32_bits_f32(bits) + + +@cute.jit +def _shfl_up_add(val, lane, offset: cutlass.Constexpr): + """Inclusive-scan step: val += shfl_up(val, offset) gated by lane >= offset.""" + src = lane - cutlass.Int32(offset) + if src < 0: + src = cutlass.Int32(0) + other = cute.arch.shuffle_sync(val, src) + if lane >= cutlass.Int32(offset): + val = val + other + return val + + +@cute.jit +def _shfl_down_add(val, lane, offset: cutlass.Constexpr): + """Suffix-scan step: val += shfl_down(val, offset) gated by lane + offset < 32.""" + src = lane + cutlass.Int32(offset) + if src > 31: + src = cutlass.Int32(31) + other = cute.arch.shuffle_sync(val, src) + if lane + cutlass.Int32(offset) < cutlass.Int32(32): + val = val + other + return val + + +@cute.jit +def _mask_tail(v, gidx, n_eff): + """Ragged-N lane predicate: value at global element index ``gidx`` is + replaced by -FLT_MAX when ``gidx >= n_eff``. + + The float4 LOAD itself is always in-bounds (the logits buffer is + allocated to npad = logits.shape[1]); only the stale VALUE beyond the + row's valid length must be masked so garbage can never enter counts, + max-below, candidate pushes, or emits.""" + r = v + if gidx >= n_eff: + r = cutlass.Float32(-FLT_MAX) + return r + + +class GvrTpKernel: + """CuTe DSL port of gvr_topk_tp (fp32, B200/B300). + + Ctor knobs mirror the CUDA template params plus the production + ``next_n`` / ``compress_ratio`` contract (v1 dispatcher guard pins + next_n=1, cr=4; the N_eff arithmetic is kept general to mirror the + in-tree ``run_one_row`` formula exactly). + """ + + WARP_SIZE = 32 + + def __init__( + self, + top_k: int, + kC: int, + cluster_size: int = 1, + ar: int = RUNGS, + uf: int = 4, + num_threads: int = 512, + next_n: int = 1, + compress_ratio: int = 4, + ): + assert num_threads % 32 == 0 + assert ar in (6, 8) + assert cluster_size in (1, 2, 4, 8) + self.top_k = top_k + self.kC = kC + self.cluster_size = cluster_size + self.ar = ar + self.uf = uf + self.num_threads = num_threads + self.num_warps = num_threads // 32 + self.next_n = next_n + self.compress_ratio = compress_ratio + + # ------------------------------------------------------------------ + # float4 (128-bit) strided streaming load helper pieces + # ------------------------------------------------------------------ + def _copy_atom(self): + return cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128 + ) + + @cute.jit + def _ld_float4(self, copy_atom, row_addr, v_idx, frag): + """Load float4 #v_idx of the row into frag[0..3].""" + p = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(v_idx) * cutlass.Int64(16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + src = cute.make_tensor(p, cute.make_layout((4,))) + cute.copy(copy_atom, src, frag) + + # ------------------------------------------------------------------ + # Per-row valid length (ragged N). Mirrors the in-tree run_one_row + # arithmetic EXACTLY (gvr_topk_decode.py): seq_lens is request-level + # and in uncompressed-token space; logits live in compressed space + # when cr > 1. + # ------------------------------------------------------------------ + @cute.jit + def _row_n_eff(self, seq_lens: cute.Tensor, row): + NN = cutlass.const_expr(self.next_n) + seq_len = seq_lens[row // cutlass.Int32(NN)] + actual_kv_len = seq_len - cutlass.Int32(NN) + (row % cutlass.Int32(NN)) + cutlass.Int32(1) + if cutlass.const_expr(self.compress_ratio == 1): + n_eff = actual_kv_len + else: + n_eff = actual_kv_len // cutlass.Int32(self.compress_ratio) + return n_eff + + # ------------------------------------------------------------------ + # count_pass: R-rung exact count over [v0, v1) float4s, + # per-thread counts -> s_ptcnt[r*TB + tid]. + # ------------------------------------------------------------------ + @cute.jit + def count_pass( + self, + R: cutlass.Constexpr, + U: cutlass.Constexpr, + row_addr, + v0, + v1, + n_eff, + tidx, + s_rungs, + s_ptcnt, + ): + # op43 S-E: explicit U-batched loads (CUDA `float4 a[U]` idiom) — a + # `cutlass.range(unroll=U)` loop leaves ONE load in flight per iter + # and costs ~11% kernel time in the DRAM-bound regimes. Keep the + # Python-unrolled register batch. + TB = cutlass.const_expr(self.num_threads) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + frags = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(U) + ] # Python-unrolled register batch + i = v0 + tidx + while i + cutlass.Int32((U - 1) * TB) < v1: + for u in cutlass.range_constexpr(U): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(U): + gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(U * TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # sample_count: uniform every-SS-th float4 over the slice. + # ------------------------------------------------------------------ + @cute.jit + def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt): + TB = cutlass.const_expr(self.num_threads) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + frag = cute.make_fragment((4,), cutlass.Float32) + j = cutlass.Int32(tidx) + while v0 + j * cutlass.Int32(SS) < v1: + self._ld_float4(copy_atom, row_addr, v0 + j * cutlass.Int32(SS), frag) + gi = (v0 + j * cutlass.Int32(SS)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frag[q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + j = j + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # exchange_counts: warp-fused CTA reduce -> cluster sum + prefix. + # rcnt[r] = cluster-wide count, rpre[r] = exclusive prefix over lower ranks. + # ------------------------------------------------------------------ + @cute.jit + def exchange_counts( + self, R: cutlass.Constexpr, par, tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + cute.arch.barrier() # ptcnt final + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if wid < cutlass.Int32(R): + s = cutlass.Int32(0) + for k in cutlass.range_constexpr(TB // 32): + s = s + s_ptcnt[wid * TB + lane + cutlass.Int32(32 * k)] + s = cute.arch.warp_redux_sync(s, "add") + if cutlass.const_expr(CS == 1): + if lane == cutlass.Int32(0): + s_rcnt[wid] = s + s_rpre[wid] = cutlass.Int32(0) + else: + if lane == cutlass.Int32(0): + s_ipartial[par * cutlass.Int32(RUNGS) + wid] = s + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + else: + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if tidx < cutlass.Int32(R): + tot = cutlass.Int32(0) + pre = cutlass.Int32(0) + local_ptr = s_ipartial.iterator + (par * cutlass.Int32(RUNGS) + tidx) + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + v = _ld_shared_cluster_i32(a) + tot = tot + v + if cutlass.Int32(rr) < rank: + pre = pre + v + s_rcnt[tidx] = tot + s_rpre[tidx] = pre + cute.arch.barrier() + + # ------------------------------------------------------------------ + # max_below_pass: largest value strictly below t_hi_bound (cluster-reduced). + # ------------------------------------------------------------------ + @cute.jit + def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred, s_fpartial): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + NWARP = cutlass.const_expr(self.num_warps) + copy_atom = self._copy_atom() + m = cutlass.Float32(-FLT_MAX) + # op43 S-E: explicit U=4 batched loads (CUDA `float4 a[4]` idiom). + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + i = v0 + tidx + while i + cutlass.Int32(3 * TB) < v1: + for u in cutlass.range_constexpr(4): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(4): + gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(4 * TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(TB) + m = cute.arch.warp_redux_sync(m, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = m + cute.arch.barrier() + if tidx == cutlass.Int32(0): + mm = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + mm = cute.arch.fmax(mm, s_fwred[w]) + s_fpartial[par] = mm + res = cutlass.Float32(-FLT_MAX) + if cutlass.const_expr(CS == 1): + cute.arch.barrier() + res = s_fpartial[par] + else: + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + local_ptr = s_fpartial.iterator + par + for rr in cutlass.range_constexpr(CS): + a = _mapa_shared_cluster(local_ptr, cutlass.Int32(rr)) + res = cute.arch.fmax(res, _ld_shared_cluster_f32(a)) + return res + + # ------------------------------------------------------------------ + # phase1: hint gather + stats + rung ladder from hint-value CCDF. + # ------------------------------------------------------------------ + @cute.jit + def phase1(self, logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs): + TB = cutlass.const_expr(self.num_threads) + AR = cutlass.const_expr(self.ar) + NWARP = cutlass.const_expr(self.num_warps) + K = cutlass.const_expr(self.top_k) + if tidx < cutlass.Int32(64): + s_hist[tidx] = cutlass.Int32(0) + hv = cute.make_fragment((4,), cutlass.Float32) + hok = cute.make_fragment((4,), cutlass.Int32) + mn = cutlass.Float32(FLT_MAX) + mx = cutlass.Float32(-FLT_MAX) + for jj in cutlass.range_constexpr(4): + j = tidx + cutlass.Int32(jj * TB) + hok[jj] = cutlass.Int32(0) + hv[jj] = cutlass.Float32(0.0) + if j < cutlass.Int32(K): + # pre_idx hardening: clamp hints into [0, N_eff-1] (not + # [0, npad-1] as in the op-bench port). Production + # cold-start is all zeros and arbitrary garbage must not + # crash or corrupt; clamping (rather than skipping, which + # the in-tree phase1_preidx_stats does) also keeps the CCDF + # ladder seeded with real in-range values so an all-OOR + # hint cannot collapse the ladder onto FLT_MAX and force a + # pathological one-value-per-pass plateau descent. + hidx = pre_idx_row[j] + if hidx < cutlass.Int32(0): + hidx = cutlass.Int32(0) + if hidx > n_eff - cutlass.Int32(1): + hidx = n_eff - cutlass.Int32(1) + v = logits_row[hidx] + hv[jj] = v + hok[jj] = cutlass.Int32(1) + mn = _fmin_f32(mn, v) + mx = cute.arch.fmax(mx, v) + mn = cute.arch.warp_redux_sync(mn, "fmin") + mx = cute.arch.warp_redux_sync(mx, "fmax") + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + if lane == cutlass.Int32(0): + s_fwred[wid] = mn + s_fwred[cutlass.Int32(NWARP) + wid] = mx + cute.arch.barrier() + if tidx == cutlass.Int32(0): + a = cutlass.Float32(FLT_MAX) + b = cutlass.Float32(-FLT_MAX) + for w in cutlass.range_constexpr(NWARP): + a = _fmin_f32(a, s_fwred[w]) + b = cute.arch.fmax(b, s_fwred[NWARP + w]) + s_hminmax[0] = a + s_hminmax[1] = b + cute.arch.barrier() + hmin = s_hminmax[0] + hmax = s_hminmax[1] + if hmax - hmin > cutlass.Float32(0.0): + # Stage 1: coarse 64-bin hist over [hmin, hmax] -> 97% trim point. + scale = cutlass.Float32(64.0) / (hmax - hmin) + for jj in cutlass.range_constexpr(4): + if hok[jj] != cutlass.Int32(0): + b1 = cutlass.Int32((hv[jj] - hmin) * scale) + if b1 < cutlass.Int32(0): + b1 = cutlass.Int32(0) + if b1 > cutlass.Int32(63): + b1 = cutlass.Int32(63) + atomicAdd(s_hist.iterator + b1, cutlass.Int32(1)) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + h0 = s_hist[cutlass.Int32(62) - cutlass.Int32(2) * tidx] + h1 = s_hist[cutlass.Int32(63) - cutlass.Int32(2) * tidx] + Ssum = h0 + h1 + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_up_add(x, tidx, o) + A = x - Ssum + binw = (hmax - hmin) * cutlass.Float32(1.0 / 64.0) + qtrim = cutlass.const_expr((K * 97) // 100) + if A < cutlass.Int32(qtrim): + if x >= cutlass.Int32(qtrim): + b2 = cutlass.Int32(62) - cutlass.Int32(2) * tidx + if A + h1 >= cutlass.Int32(qtrim): + b2 = cutlass.Int32(63) - cutlass.Int32(2) * tidx + s_hminmax[0] = hmin + binw * cutlass.Float32(b2) + if tidx == cutlass.Int32(31): + if x < cutlass.Int32(qtrim): + s_hminmax[0] = hmin + if tidx == cutlass.Int32(0): + s_rungs[AR - 1] = hmin + # barrier fold: re-zero both hist halves for stage 2 + s_hist[tidx] = cutlass.Int32(0) + s_hist[tidx + cutlass.Int32(32)] = cutlass.Int32(0) + cute.arch.barrier() + tlow = s_hminmax[0] + if hmax - tlow > cutlass.Float32(0.0): + # Stage 2: fine 64-bin hist over [tlow, hmax] -> rung quantiles. + scale2 = cutlass.Float32(64.0) / (hmax - tlow) + for jj in cutlass.range_constexpr(4): + if hok[jj] != cutlass.Int32(0): + if hv[jj] >= tlow: + b3 = cutlass.Int32((hv[jj] - tlow) * scale2) + if b3 < cutlass.Int32(0): + b3 = cutlass.Int32(0) + if b3 > cutlass.Int32(63): + b3 = cutlass.Int32(63) + atomicAdd(s_hist.iterator + b3, cutlass.Int32(1)) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + h0 = s_hist[cutlass.Int32(62) - cutlass.Int32(2) * tidx] + h1 = s_hist[cutlass.Int32(63) - cutlass.Int32(2) * tidx] + Ssum = h0 + h1 + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_up_add(x, tidx, o) + A = x - Ssum + binw2 = (hmax - tlow) * cutlass.Float32(1.0 / 64.0) + if cutlass.const_expr(AR == 6): + qt = ((K * 15) // 100, (K * 40) // 100, (K * 70) // 100, (K * 92) // 100) + else: + qt = ( + (K * 10) // 100, + (K * 25) // 100, + (K * 45) // 100, + (K * 65) // 100, + (K * 82) // 100, + (K * 94) // 100, + ) + tot = cute.arch.shuffle_sync(x, cutlass.Int32(31)) + for r in cutlass.range_constexpr(AR - 2): + qtr = cutlass.const_expr(qt[r]) + if A < cutlass.Int32(qtr): + if x >= cutlass.Int32(qtr): + b4 = cutlass.Int32(62) - cutlass.Int32(2) * tidx + if A + h1 >= cutlass.Int32(qtr): + b4 = cutlass.Int32(63) - cutlass.Int32(2) * tidx + s_rungs[r + 1] = tlow + binw2 * cutlass.Float32(b4) + if tidx == cutlass.Int32(31): + if tot < cutlass.Int32(qtr): + s_rungs[r + 1] = tlow + if tidx == cutlass.Int32(0): + s_rungs[0] = hmax + (hmax - tlow) + cute.arch.barrier() + else: + # degenerate trim: all rungs above the floor = hmax + if tidx < cutlass.Int32(AR - 1): + s_rungs[tidx] = hmax + cute.arch.barrier() + else: + # degenerate: all hint values equal + if tidx < cutlass.Int32(AR): + s_rungs[tidx] = hmin + cute.arch.barrier() + + # ------------------------------------------------------------------ + # fused_count_collect: exact counts at {rungs[0]=pivot, + # rungs[1]=hmin} + push (key,idx) >= tpush into CTA0 cand, capped kcap. + # Keys are RAW fp32 bits (f2u happens in P4 round 0), matching CUDA. + # Candidates are a single packed (key<<32 | idx) u64 array — CUDA's + # `unsigned long long cand[kC]` — one 8B push per candidate. + # ------------------------------------------------------------------ + @cute.jit + def _push_cand( + self, + a_cnt, + a_st, + s_cand, + s_isc, + kcap: cutlass.Constexpr, + capped: cutlass.Constexpr, + v, + gidx, + ): + """One candidate push: kv = (raw fp32 bits << 32) | index, single 8B + store (st.shared::cluster.u64 remote at CS>1, local u64 store at CS1). + a_cnt/a_st are the pre-mapa'd CTA0 addresses of cnt_c / cand[0].""" + CS = cutlass.const_expr(self.cluster_size) + kv = (cutlass.Uint64(_f32_bits_u32(v)) << cutlass.Uint64(32)) | cutlass.Uint64( + cutlass.Uint32(gidx) + ) + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_cnt, cutlass.Int32(1)) + if cutlass.const_expr(capped): + if p < cutlass.Int32(kcap): + _st_shared_cluster_u64(a_st + p * cutlass.Int32(8), kv) + else: + _st_shared_cluster_u64(a_st + p * cutlass.Int32(8), kv) + else: + p = atomicAdd(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(1)) + if cutlass.const_expr(capped): + if p < cutlass.Int32(kcap): + s_cand[p] = kv + else: + s_cand[p] = kv + + @cute.jit + def fused_count_collect( + self, + U: cutlass.Constexpr, + row_addr, + v0, + v1, + n_eff, + tpush, + tidx, + s_rungs, + s_ptcnt, + s_cand, + s_isc, + ): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + kcap = cutlass.const_expr(self.kC) + R = cutlass.const_expr(2) + copy_atom = self._copy_atom() + tr = cute.make_fragment((R,), cutlass.Float32) + cnt = cute.make_fragment((R,), cutlass.Int32) + for r in cutlass.range_constexpr(R): + tr[r] = s_rungs[r] + cnt[r] = cutlass.Int32(0) + a_cnt = cutlass.Int32(0) + a_st = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) + a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + # Explicit U-batched loads (CUDA `float4 a[U]` idiom), main + + # vec-tail while-loops (same op43 S-E fix as count_pass). + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(U)] + i = v0 + tidx + while i + cutlass.Int32((U - 1) * TB) < v1: + for u in cutlass.range_constexpr(U): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(U): + gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, True, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(U * TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, True, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(TB) + for r in cutlass.range_constexpr(R): + s_ptcnt[r * TB + tidx] = cnt[r] + + # ------------------------------------------------------------------ + # collect_at: plain streaming collect at thr (uncapped; caller + # guarantees count(thr) <= kC). + # ------------------------------------------------------------------ + @cute.jit + def collect_at(self, row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc): + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + kcap = cutlass.const_expr(self.kC) + copy_atom = self._copy_atom() + a_cnt = cutlass.Int32(0) + a_st = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) + a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) + # op43 S-E round 3: explicit U=4 batched loads. nvcc auto-unrolls the + # CUDA collect_at loop 4x with 4 LDG.E.128 issued back-to-back; a + # 1-deep loop was the dominant stall site (36% of all warp-stall + # samples) on cells whose reuse check fails at big npad x big BS. + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + i = v0 + tidx + while i + cutlass.Int32(3 * TB) < v1: + for u in cutlass.range_constexpr(4): + self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) + for u in cutlass.range_constexpr(4): + gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + if v >= thr: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, False, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(4 * TB) + while i < v1: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + gi = i << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frags[0][q], gi + cutlass.Int32(q), n_eff) + if v >= thr: + self._push_cand( + a_cnt, a_st, s_cand, s_isc, kcap, False, v, gi + cutlass.Int32(q) + ) + i = i + cutlass.Int32(TB) + + # ------------------------------------------------------------------ + # kernel + # ------------------------------------------------------------------ + @cute.kernel + def gvr_tp_kernel( + self, logits: cute.Tensor, pre_idx: cute.Tensor, seq_lens: cute.Tensor, out_idx: cute.Tensor + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + TB = cutlass.const_expr(self.num_threads) + CS = cutlass.const_expr(self.cluster_size) + AR = cutlass.const_expr(self.ar) + UF = cutlass.const_expr(self.uf) + K = cutlass.const_expr(self.top_k) + kC = cutlass.const_expr(self.kC) + + if cutlass.const_expr(CS > 1): + row = bidx // cutlass.Int32(CS) + rank = cute.arch.block_idx_in_cluster() + else: + row = bidx + rank = cutlass.Int32(0) + + npad = cutlass.Int32(logits.shape[1]) + logits_row = logits[row, None] + pre_idx_row = pre_idx[row, None] + out_row = out_idx[row, None] + row_addr = logits_row.iterator.toint() + + # Ragged N: per-row valid length from seq_lens (see _row_n_eff). + n_eff = self._row_n_eff(seq_lens, row) + + # ---- shared memory (order must be identical across CTAs for mapa) ---- + smem = SmemAllocator() + # Packed (key<<32 | idx) u64 candidates = CUDA's + # `unsigned long long cand[kC]` (one 8B push per candidate). + s_cand = smem.allocate_tensor( + element_type=cutlass.Uint64, + layout=cute.make_ordered_layout((kC,), order=(0,)), + byte_alignment=128, + ) + s_ptcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS * self.num_threads,), order=(0,)), + byte_alignment=128, + ) + s_hist = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((256,), order=(0,)), + byte_alignment=128, + ) + s_rungs = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rcnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_rpre = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_ipartial = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((2 * RUNGS,), order=(0,)), + byte_alignment=32, + ) + s_fpartial = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + s_fwred = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2 * self.num_warps,), order=(0,)), + byte_alignment=64, + ) + s_hminmax = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout((2,), order=(0,)), + byte_alignment=16, + ) + # iscalars: [0]=sel_bin [1]=sel_above [2]=sel_count [3]=cnt_m [4]=cnt_t [5]=cnt_c + s_isc = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((8,), order=(0,)), + byte_alignment=32, + ) + + # slice in 64-float units (npad multiple of 64); v0/v1 are float4 indices + units = npad >> cutlass.Int32(6) + u0 = (units * rank) // cutlass.Int32(CS) + u1 = (units * (rank + cutlass.Int32(1))) // cutlass.Int32(CS) + v0 = u0 << cutlass.Int32(4) + v1 = u1 << cutlass.Int32(4) + + xch = cutlass.Int32(0) + thr = cutlass.Float32(0.0) + tpush = cutlass.Float32(0.0) + C = cutlass.Int32(0) + m_gt = cutlass.Int32(-1) + span0 = cutlass.Float32(1e-3) + + # ---- Degenerate rows (N_eff <= K): identity emit + -1 pad. ---- + # Mirrors the in-tree kernel's degenerate branch. n_eff is uniform + # across the cluster (all CTAs own the same row), so this is a + # cluster-uniform branch; non-leader CTAs fall through to the exit + # rendezvous (CuTe DSL has no runtime return). + if n_eff <= cutlass.Int32(K): + if rank == cutlass.Int32(0): + jd = cutlass.Int32(tidx) + while jd < n_eff: + out_row[jd] = jd + jd = jd + cutlass.Int32(TB) + jp = n_eff + cutlass.Int32(tidx) + if jp < cutlass.Int32(0): + jp = cutlass.Int32(tidx) # n_eff < 0 (defensive): pad all + while jp < cutlass.Int32(K): + out_row[jp] = cutlass.Int32(-1) + jp = jp + cutlass.Int32(TB) + else: + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[5] = cutlass.Int32(0) # cnt_c + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + if npad <= cutlass.Int32(kC): + # trivial: whole row fits the candidate buffer. Masked tail + # lanes push -FLT_MAX candidates; they can never displace a + # real value because n_eff > K here (kth value > -FLT_MAX). + if tidx < cutlass.Int32(RUNGS): + s_rungs[tidx] = cutlass.Float32(-FLT_MAX) + cute.arch.barrier() + self.count_pass(1, 8, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, cutlass.Int32(0), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + thr = cutlass.Float32(-FLT_MAX) + C = s_rcnt[0] + self.collect_at(row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc) + cute.arch.barrier() + else: + self.phase1( + logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs + ) + hmin_floor = s_rungs[AR - 1] + span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) + # P2a: sampled ladder count -> pivot pick + self.sample_count(AR, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + # 3-stage pivot pick (iter7 band -> iter8b 2-sigma -> + # iter12-F1b gamble) + lo = cutlass.const_expr((3 * K) // 2) + hi = cutlass.const_expr((6 * kC) // 10) + tgt_py = min(max(3 * K, (3 * K) // 2), (6 * kC) // 10) + tgt = cutlass.Int32(tgt_py) + best = cutlass.Int32(AR - 1) + bestd = cutlass.Int32(0x7FFFFFFF) + for j in cutlass.range_constexpr(AR): + est = s_rcnt[j] * cutlass.Int32(SS) + if est >= cutlass.Int32(lo): + if est <= cutlass.Int32(hi): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < bestd: + bestd = dd + best = cutlass.Int32(j) + if bestd == cutlass.Int32(0x7FFFFFFF): + for j in cutlass.range_constexpr(AR): + est = s_rcnt[j] * cutlass.Int32(SS) + if est > cutlass.Int32(0): + g = cutlass.Float32(2.0) * cmath.sqrt( + cutlass.Float32(cutlass.Int32(SS) * est) + ) + fest = cutlass.Float32(est) + if fest - g >= cutlass.Float32(float(K)): + if fest + g <= cutlass.Float32(float(kC)): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < bestd: + bestd = dd + best = cutlass.Int32(j) + if bestd == cutlass.Int32(0x7FFFFFFF): + if s_rcnt[AR - 1] * cutlass.Int32(SS) > cutlass.Int32(kC): + for j in cutlass.range_constexpr(AR): + est = s_rcnt[j] * cutlass.Int32(SS) + if est > cutlass.Int32(0): + dd = est - tgt + if dd < cutlass.Int32(0): + dd = tgt - est + if dd < bestd: + bestd = dd + best = cutlass.Int32(j) + tp_ = s_rungs[best] + s_rungs[0] = tp_ + s_rungs[1] = hmin_floor + cute.arch.barrier() + tpush = s_rungs[0] + self.fused_count_collect( + UF, row_addr, v0, v1, n_eff, tpush, tidx, s_rungs, s_ptcnt, s_cand, s_isc + ) + self.exchange_counts( + 2, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial + ) + xch = xch + cutlass.Int32(1) + + # ---- P2c: secant refine driver (redundant on every thread) ---- + t_lo = cutlass.Float32(-FLT_MAX) + t_hi = cutlass.Float32(INF) + c_hi = cutlass.Int32(0) + Rcur = cutlass.Int32(2) + passno = cutlass.Int32(0) + running = cutlass.Int32(1) + descend_break = cutlass.Int32(0) + while running != cutlass.Int32(0): + # first rung index j with rcnt[j] >= K (rcnt ascending in j) + j = Rcur + for r_ in range(RUNGS - 1, -1, -1): + if cutlass.Int32(r_) < Rcur: + if s_rcnt[r_] >= cutlass.Int32(K): + j = cutlass.Int32(r_) + jj = j + if jj > Rcur - cutlass.Int32(1): + jj = Rcur - cutlass.Int32(1) + cj = s_rcnt[jj] + rj = s_rungs[jj] + found = cutlass.Int32(0) + if j < Rcur: + if cj <= cutlass.Int32(kC): + found = cutlass.Int32(1) + if found != cutlass.Int32(0): + thr = rj + C = cj + running = cutlass.Int32(0) + else: + if j < Rcur: + if rj >= t_lo: + t_lo = rj + if j > cutlass.Int32(0): + jm = j - cutlass.Int32(1) + if s_rungs[jm] <= t_hi: + t_hi = s_rungs[jm] + c_hi = s_rcnt[jm] + descend = cutlass.Int32(0) + if passno >= cutlass.Int32(MAXPASS): + descend = cutlass.Int32(1) + # ladder params (uniform; recomputed per thread) + e3 = passno * cutlass.Int32(3) + if e3 > cutlass.Int32(24): + e3 = cutlass.Int32(24) + step = span0 * _exp2i(e3) + dt = cutlass.Float32(0.0) + mode = cutlass.Int32(2) # 0=up-ladder,1=down-ladder,2=secant + if t_hi == cutlass.Float32(INF): + mode = cutlass.Int32(0) + nr0 = t_lo + step * cutlass.Float32(float(1 << (AR - 1))) + if nr0 == cutlass.Float32(INF): + descend = cutlass.Int32(1) + else: + if t_lo == cutlass.Float32(-FLT_MAX): + mode = cutlass.Int32(1) + else: + dt = (t_hi - t_lo) * cutlass.Float32(1.0 / float(AR + 1)) + nr_last = t_hi - dt * cutlass.Float32(float(AR)) + nr_first = t_hi - dt + ok = cutlass.Int32(0) + if nr_last > t_lo: + if nr_first < t_hi: + ok = cutlass.Int32(1) + if ok == cutlass.Int32(0): + descend = cutlass.Int32(1) + if descend != cutlass.Int32(0): + running = cutlass.Int32(0) + descend_break = cutlass.Int32(1) + else: + cute.arch.barrier() + if tidx == cutlass.Int32(0): + for r_ in cutlass.range_constexpr(AR): + nrv = cutlass.Float32(0.0) + if mode == cutlass.Int32(0): + nrv = t_lo + step * cutlass.Float32( + float(1 << (AR - 1 - r_)) + ) + else: + if mode == cutlass.Int32(1): + nrv = t_hi - step * cutlass.Float32(float(1 << r_)) + else: + nrv = t_hi - dt * cutlass.Float32(float(r_ + 1)) + s_rungs[r_] = nrv + cute.arch.barrier() + self.count_pass(AR, UF, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + AR, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + passno = passno + cutlass.Int32(1) + Rcur = cutlass.Int32(AR) + + # ---- plateau descent (exact max-below stepping) ---- + if descend_break != cutlass.Int32(0): + pl = cutlass.Int32(1) + while pl != cutlass.Int32(0): + vstar = self.max_below_pass( + row_addr, + v0, + v1, + n_eff, + t_hi, + xch & cutlass.Int32(1), + tidx, + s_fwred, + s_fpartial, + ) + xch = xch + cutlass.Int32(1) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_rungs[0] = vstar + cute.arch.barrier() + self.count_pass(1, 8, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) + self.exchange_counts( + RUNGS, + xch & cutlass.Int32(1), + tidx, + rank, + s_ptcnt, + s_rcnt, + s_rpre, + s_ipartial, + ) + xch = xch + cutlass.Int32(1) + c = s_rcnt[0] + okc = cutlass.Int32(0) + if c >= cutlass.Int32(K): + if c <= cutlass.Int32(kC): + okc = cutlass.Int32(1) + if okc != cutlass.Int32(0): + thr = vstar + C = c + pl = cutlass.Int32(0) + else: + if c < cutlass.Int32(K): + t_hi = vstar + c_hi = c + else: + thr = vstar + m_gt = c_hi + pl = cutlass.Int32(0) + + # ---- candidate reuse check / re-stream collect ---- + if m_gt < cutlass.Int32(0): + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + if cutlass.const_expr(CS > 1): + a_cnt0 = _mapa_shared_cluster( + s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0) + ) + dcnt = _ld_shared_cluster_i32(a_cnt0) + else: + dcnt = s_isc[5] + reuse = cutlass.Int32(0) + if thr == tpush: + if dcnt == C: + reuse = cutlass.Int32(1) + if reuse == cutlass.Int32(0): + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[5] = cutlass.Int32(0) + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + self.collect_at(row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc) + cute.arch.barrier() + + # ---- plateau direct emit (all CTAs stream their slice) ---- + if m_gt >= cutlass.Int32(0): + if rank == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) # cnt_m + s_isc[4] = cutlass.Int32(0) # cnt_t + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + a_m = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(0)) + a_t = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(4), cutlass.Int32(0)) + else: + cute.arch.barrier() + nt = cutlass.Int32(K) - m_gt + copy_atom = self._copy_atom() + frag = cute.make_fragment((4,), cutlass.Float32) + ii = v0 + tidx + while ii < v1: + self._ld_float4(copy_atom, row_addr, ii, frag) + gi = ii << cutlass.Int32(2) + for q in cutlass.range_constexpr(4): + v = _mask_tail(frag[q], gi + cutlass.Int32(q), n_eff) + if v > thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_m, cutlass.Int32(1)) + else: + p = atomicAdd(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1)) + out_row[p] = gi + cutlass.Int32(q) + else: + if v == thr: + if cutlass.const_expr(CS > 1): + p = _atom_shared_cluster_add_i32(a_t, cutlass.Int32(1)) + else: + p = atomicAdd( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m_gt + p] = gi + cutlass.Int32(q) + ii = ii + cutlass.Int32(TB) + else: + # ---- P4 (CTA0 solo when CS>1) ---- + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if rank == cutlass.Int32(0): + if C == cutlass.Int32(K): + ie = cutlass.Int32(tidx) + while ie < C: + out_row[ie] = cutlass.Int32( + cutlass.Uint32(s_cand[ie] & cutlass.Uint64(0xFFFFFFFF)) + ) + ie = ie + cutlass.Int32(TB) + else: + # 4x8-bit radix select over cand keys (f2u in round 0) + if tidx < cutlass.Int32(256): + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() + pref = cutlass.Uint32(0) + want = cutlass.Int32(K) + m = cutlass.Int32(0) + final_shift = cutlass.Uint32(0) + active = cutlass.Int32(1) + for r_ in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * r_) + if active != cutlass.Int32(0): + ih = cutlass.Int32(tidx) + while ih < C: + if cutlass.const_expr(r_ == 0): + kv = s_cand[ih] + raw = cutlass.Uint32(kv >> cutlass.Uint64(32)) + u = _f2u_bits(raw) + s_cand[ih] = (cutlass.Uint64(u) << cutlass.Uint64(32)) | ( + kv & cutlass.Uint64(0xFFFFFFFF) + ) + atomicAdd( + s_hist.iterator + + cutlass.Int32(u >> cutlass.Uint32(24)), + cutlass.Int32(1), + ) + else: + u = cutlass.Uint32(s_cand[ih] >> cutlass.Uint64(32)) + if (u >> cutlass.Uint32(shift + 8)) == pref: + b = cutlass.Int32( + (u >> cutlass.Uint32(shift)) & cutlass.Uint32(0xFF) + ) + atomicAdd(s_hist.iterator + b, cutlass.Int32(1)) + ih = ih + cutlass.Int32(TB) + cute.arch.barrier() + if tidx < cutlass.Int32(32): + b8 = tidx * cutlass.Int32(8) + h = cute.make_fragment((8,), cutlass.Int32) + Ssum = cutlass.Int32(0) + for q in cutlass.range_constexpr(8): + h[q] = s_hist[b8 + cutlass.Int32(q)] + Ssum = Ssum + h[q] + s_hist[b8 + cutlass.Int32(q)] = cutlass.Int32(0) + x = Ssum + for o in [1, 2, 4, 8, 16]: + x = _shfl_down_add(x, tidx, o) + A = x - Ssum + if A < want: + if want <= A + Ssum: + run = A + for q in range(7, -1, -1): + if run < want: + if want <= run + h[q]: + s_isc[0] = b8 + cutlass.Int32(q) + s_isc[1] = run + s_isc[2] = h[q] + run = run + h[q] + cute.arch.barrier() + sel = s_isc[0] + above = s_isc[1] + m = m + above + want = want - above + pref = (pref << cutlass.Uint32(8)) | cutlass.Uint32(sel) + if s_isc[2] == want: + final_shift = cutlass.Uint32(shift) + active = cutlass.Int32(0) + kth = pref + if tidx == cutlass.Int32(0): + s_isc[3] = cutlass.Int32(0) + s_isc[4] = cutlass.Int32(0) + cute.arch.barrier() + nt = cutlass.Int32(K) - m + ie = cutlass.Int32(tidx) + while ie < C: + kv = s_cand[ie] + u = cutlass.Uint32(kv >> cutlass.Uint64(32)) >> final_shift + idx = cutlass.Int32(cutlass.Uint32(kv & cutlass.Uint64(0xFFFFFFFF))) + if u > kth: + p = atomicAdd(s_isc.iterator + cutlass.Int32(3), cutlass.Int32(1)) + out_row[p] = idx + else: + if u == kth: + p = atomicAdd( + s_isc.iterator + cutlass.Int32(4), cutlass.Int32(1) + ) + if p < nt: + out_row[m + p] = idx + ie = ie + cutlass.Int32(TB) + + # Exit rendezvous: DSMEM (mapa'd peer smem) must stay valid until every + # CTA of the cluster is done issuing remote ops (plateau emit's remote + # cnt_m/cnt_t atomics can land after CTA0 finishes its own slice). nvcc + # inserts an implicit cluster barrier before ret for DSMEM kernels; CuTe + # DSL does not — without this, plateau rows fault with a + # timing-dependent cudaErrorLaunchFailure (CUDA 719, hidden under + # compute-sanitizer synccheck). DO NOT remove or make conditional. + if cutlass.const_expr(CS > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # ------------------------------------------------------------------ + # host launcher + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + seq_lens: cute.Tensor, + out_idx: cute.Tensor, + stream, + ): + num_rows = logits.shape[0] + CS = cutlass.const_expr(self.cluster_size) + self.gvr_tp_kernel(logits, pre_idx, seq_lens, out_idx).launch( + grid=(num_rows * CS, 1, 1), + block=(self.num_threads, 1, 1), + cluster=(CS, 1, 1) if cutlass.const_expr(CS > 1) else None, + stream=stream, + min_blocks_per_mp=2, + ) + + +# --------------------------------------------------------------------------- +# compile cache + public entry (mirrors launch_tp<512> selection exactly) +# --------------------------------------------------------------------------- +_LAUNCH_CACHE = {} + + +def _p2floor(x: int) -> int: + p = 1 + while p * 2 <= x: + p *= 2 + return p + + +def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int): + key = (K, CS, AR, UF, TB) + compiled = _LAUNCH_CACHE.get(key) + if compiled is None: + kC = 8192 if K >= 2048 else 6144 + kern = GvrTpKernel(top_k=K, kC=kC, cluster_size=CS, ar=AR, uf=UF, num_threads=TB) + n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 + ) + pre_idx_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kern, + logits_fake, + pre_idx_fake, + seq_lens_fake, + out_fake, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + _LAUNCH_CACHE[key] = compiled + return compiled + + +_FAST = {} # (K, bs, npad) -> compiled variant (hot-path single dict hit) + + +def tp_cluster_size(bs: int, npad: int) -> int: + """CS selection of launch_tp<512>: co-residency + slice floor.""" + cs = 1 + if bs < 128: + cs = _p2floor(296 // bs) if bs <= 296 else 1 + capn = _p2floor(npad // 8192 if npad // 8192 > 0 else 1) + if cs > capn: + cs = capn + if cs > 8: + cs = 8 + return cs + + +def _dispatch(K: int, bs: int, npad: int): + """launch_tp<512> selection: CS by co-residency + slice floor, UF by depth.""" + assert npad % 64 == 0 + TB = 512 + AR = RUNGS + cs = tp_cluster_size(bs, npad) + if cs > 1: + uf = 4 + else: + uf = 8 if npad >= 16384 else 4 + return _get_compiled(K, cs, AR, uf, TB) + + +def tp_topk( + logits: torch.Tensor, pre_idx: torch.Tensor, seq_lens: torch.Tensor, out: torch.Tensor, K: int +) -> None: + """CuTe DSL gvr_topk_tp. logits [BS, npad] fp32 (npad mult of 64; + tail beyond each row's N_eff may be garbage — masked in-kernel), + pre_idx [BS, K] int32 hint, seq_lens [BS] int32 (uncompressed-token + space), out [BS, K] int32. Launch-time CS/UF selection replicates + launch_tp<512>.""" + sh = logits.shape + key = (K, sh[0], sh[1]) + fn = _FAST.get(key) + if fn is None: + fn = _dispatch(K, sh[0], sh[1]) + _FAST[key] = fn + fn(logits, pre_idx, seq_lens, out) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 95938a31cd98..b9db233d6d19 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -16,10 +16,11 @@ l0_b300: backend: pytorch tests: # ------------- PyTorch tests --------------- - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/thop/serial - unittest/_torch/executor # 250s diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py new file mode 100644 index 000000000000..4848b82ca0f7 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""BSX (op43 direct/reg/tp CuTe DSL tiers) top-K decode tests. + +CI-sized exactness grid for the guarded fp32/next_n=1/cr=4 fast path inside +``trtllm::cute_dsl_gvr_topk_decode``: every reg launch-table instance once, +the direct and tp tiers at a few npad each, ragged (varlen) rows with +POISONED tails (stale-garbage simulation: +1e30 beyond N_eff, which would +dominate the top-K if the ragged-N masking were broken), quantized-tie +inputs, degenerate rows, pre_idx hardening, host-only route-table asserts, +and a dispatcher-fallback check (bf16 / next_n=2 route to the in-tree +kernel). +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_bsx_dispatch as bsx_dispatch, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.single_pass_multi_cta_radix_topk_cluster import ( # noqa: E501 + _query_max_cluster_size, +) +from tensorrt_llm._utils import get_sm_version + +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason=f"CuTe DSL BSX Top-K only supports SM 100/103, got SM {get_sm_version()}", +) + +CR = 4 # v1 dispatcher guard: compress_ratio == 4 (DSv4) +NEXT_N = 1 # v1 dispatcher guard: next_n == 1 + + +# --------------------------------------------------------------------------- +# Shared helpers. ``_tie_aware_check`` is a local copy of the one in +# test_cute_dsl_gvr_topk_decode.py (same directory): the sibling test module +# is not reliably importable across the repo's pytest invocation styles +# (rootdir-dependent package resolution), so the checker is duplicated here +# verbatim-in-spirit with the same semantics. Keep the two in sync. +# --------------------------------------------------------------------------- +def _tie_aware_check( + out_indices: torch.Tensor, + logits: torch.Tensor, + seq_lens: torch.Tensor, + top_k: int, + next_n: int = NEXT_N, + compress_ratio: int = CR, +) -> None: + """Vectorized multi-row tie-aware correctness check (strict sort+allclose). + + Per row r the scan range is ``logits[r, :N_eff(r)]`` where + ``N_eff = (seq_lens[r // next_n] - next_n + r % next_n + 1) // cr`` + (the kernels' exact formula). Checks: in-range indices, no duplicates, + no selected value below the K-th reference value, and sorted-value + multiset equality against torch.topk. + """ + num_rows, top_k_out = out_indices.shape + assert top_k_out == top_k + device = logits.device + logits_f32 = logits.to(torch.float32) + N = logits.shape[1] + + row_idx = torch.arange(num_rows, device=device) + group_idx = row_idx // next_n + ofs = row_idx % next_n + seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] + actual_kv_len = seq_lens_per_row - next_n + ofs + 1 + N_eff = actual_kv_len // compress_ratio # [num_rows] + + col_idx = torch.arange(N, device=device) + in_range_mask = col_idx[None, :] < N_eff[:, None] + 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) + + out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) + if bool(out_of_range.any().item()): + bad_row = int(out_of_range.any(dim=1).int().argmax().item()) + raise AssertionError( + f"row={bad_row}: out-of-range index " + f"(N_eff={int(N_eff[bad_row].item())}, " + f"indices={out_indices[bad_row].cpu().tolist()})" + ) + + sorted_idx, _ = out_indices.sort(dim=-1) + has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) + if bool(has_dup.any().item()): + bad_row = int(has_dup.int().argmax().item()) + raise AssertionError( + f"row={bad_row}: duplicate indices: {out_indices[bad_row].cpu().tolist()}" + ) + + sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) + kth_vals = ref_vals[:, -1:] + n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) + if bool((n_below_per_row > 0).any().item()): + bad_row = int(n_below_per_row.argmax().item()) + raise AssertionError( + f"row={bad_row}: {int(n_below_per_row[bad_row].item())} selected " + f"values < Kth-rank value ({float(kth_vals[bad_row, 0].item()):.6f})" + ) + + sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) + if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): + per_row_max = (sel_sorted - ref_vals).abs().max(dim=-1).values + bad_row = int(per_row_max.argmax().item()) + raise AssertionError( + f"row={bad_row}: sorted-value mismatch — max diff " + f"{float(per_row_max[bad_row].item()):.4e}" + ) + + +def _make_bsx_inputs( + bs: int, + npad: int, + top_k: int, + seed: int, + kind: str = "randn", + varlen: bool = True, + preidx: str = "mixed", +): + """Build (logits fp32, pre_idx int32, seq_lens int32) for the bsx path. + + ``kind='ties'`` quantizes the logits to 0.25 steps so massive value + plateaus straddle the K-th boundary (exercises the tie-ticket emits and + the max-below plateau-descent path). Ragged rows: the tail beyond each + row's N_eff is POISONED with +1e30 — stale garbage that would win the + top-K if the lane masking were broken (production tails are stale + values, not -FLT_MAX pad). + + ``preidx``: 'mixed' = argmax slot 0 + ~50% real topk hints; + 'zeros' = all-zero cold start; 'oor' = out-of-range garbage + (negative / >= npad) that the kernels must clamp harmlessly. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + logits = torch.randn(bs, npad, dtype=torch.float32, device="cuda") * 2.0 + if kind == "ties": + logits = (logits * 4.0).round() * 0.25 + + if varlen: + lo = (top_k + 1) * CR # keep every row non-degenerate (N_eff > K) + seq_lens = torch.randint(lo, npad * CR + 1, (bs,), dtype=torch.int32, device="cuda") + else: + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + + n_eff = (seq_lens.long() // CR).clamp(max=npad) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + valid_logits = torch.where(tail, torch.full_like(logits, float("-inf")), logits) + argmax_idx = valid_logits.argmax(dim=-1).int() + if preidx == "zeros": + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + elif preidx == "oor": + pre_idx = torch.randint(-npad, 4 * npad, (bs, top_k), dtype=torch.int32, device="cuda") + else: + ref_topk = valid_logits.topk(top_k, dim=-1).indices.int() + keep = torch.rand(ref_topk.shape, device="cuda") < 0.5 + junk = torch.arange(top_k, dtype=torch.int32, device="cuda").expand(bs, -1) + # junk arange is in-range for every row: N_eff > top_k here. + pre_idx = torch.where(keep, ref_topk, junk).contiguous() + pre_idx[:, 0] = argmax_idx + return logits, pre_idx, seq_lens + + +def _run_op(logits, pre_idx, seq_lens, top_k): + out = torch.empty(logits.shape[0], top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out, + top_k=top_k, + next_n=NEXT_N, + compress_ratio=CR, + ) + torch.cuda.synchronize() + return out + + +def _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected_tier): + """Host-only: the dispatcher must accept this call and route it to + ``expected_tier``.""" + out = torch.empty(logits.shape[0], top_k, dtype=torch.int32, device="cuda") + assert bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ), "expected the bsx fast path to accept this call" + bs, npad = logits.shape + tier = bsx_dispatch.route(bs, npad, top_k) + assert tier == expected_tier, f"route({bs}, {npad}, {top_k}) = {tier} != {expected_tier}" + + +def _skip_if_cluster_capped(bs, npad, top_k): + cs = bsx_dispatch.route_cluster_size(bs, npad, top_k) + torch.zeros(1, device="cuda") # the driver-API query needs a live context + hw_max = _query_max_cluster_size() + if cs > hw_max: + pytest.skip(f"tier cluster size {cs} exceeds device max {hw_max}") + + +# --------------------------------------------------------------------------- +# Host-only route-table asserts (mirror of the op42 gvr_bsx.cu dispatch). +# --------------------------------------------------------------------------- +def test_bsx_route_table(): + r = bsx_dispatch.route + # latency ladder + assert r(1, 4096, 512) == "direct" + assert r(1, 12288, 2048) == "direct" + assert r(1, 14336, 512) == "reg(cs=1,tb=512,maxv=8,ar=8)" + assert r(1, 24576, 512) == "reg(cs=4,tb=512,maxv=4,ar=8)" + assert r(1, 49152, 512) == "reg(cs=8,tb=512,maxv=3,ar=8)" + assert r(1, 65536, 1024) == "reg(cs=8,tb=512,maxv=4,ar=8)" + assert r(8, 131072, 512) == "reg(cs=8,tb=512,maxv=8,ar=8)" + assert r(1, 163840, 2048) == "reg(cs=16,tb=512,maxv=5,ar=6)" + assert r(4, 163840, 512) == "reg(cs=16,tb=512,maxv=5,ar=8)" + assert r(1, 262144, 2048) == "reg(cs=16,tb=512,maxv=8,ar=8)" + assert r(4, 262144, 1024) == "reg(cs=16,tb=512,maxv=8,ar=6)" + # dense table + assert r(64, 20480, 512) == "reg(cs=1,tb=1024,maxv=5,ar=8)" + assert r(64, 28672, 512) == "reg(cs=1,tb=1024,maxv=8,ar=8)" + assert r(8, 262144, 2048) == "reg(cs=8,tb=1024,maxv=8,ar=8)" + # tp takeover + assert r(256, 20480, 512) == "tp" + assert r(128, 28672, 512) == "tp" + assert r(16, 65536, 1024) == "tp" + assert r(128, 262144, 2048) == "tp" + # gvr tier only beyond the deployment envelope (guarded out upstream) + assert r(1, 262208, 512).startswith("gvr") + + +def test_bsx_route_env_knobs(monkeypatch): + """TRTLLM_BSX_TP_BS / TRTLLM_BSX_DENSE_BS keep the GVR_BSX_* semantics: + unset/-1 -> baked bands, 0 -> disable, else explicit bs threshold.""" + r = bsx_dispatch.route + try: + monkeypatch.setenv("TRTLLM_BSX_DENSE_BS", "8") + bsx_dispatch._reset_env_cache() + assert r(8, 65536, 512) == "reg(cs=2,tb=1024,maxv=8,ar=8)" + assert r(8, 131072, 1024) == "reg(cs=4,tb=1024,maxv=8,ar=8)" + + monkeypatch.setenv("TRTLLM_BSX_TP_BS", "0") # 0 -> disable tp + bsx_dispatch._reset_env_cache() + assert r(1024, 65536, 512) != "tp" + + monkeypatch.setenv("TRTLLM_BSX_TP_BS", "4") + bsx_dispatch._reset_env_cache() + assert r(4, 65536, 512) == "tp" + finally: + monkeypatch.delenv("TRTLLM_BSX_TP_BS", raising=False) + monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) + bsx_dispatch._reset_env_cache() + + +# --------------------------------------------------------------------------- +# reg tier: every launch-table instance once (cs, tb, maxv, ar), ragged rows +# with poisoned tails. The two dense cs=2/cs=4 instances are reachable only +# through the TRTLLM_BSX_DENSE_BS knob (default bands route around them), +# which doubles as the env-knob dispatch test on a live launch. +# --------------------------------------------------------------------------- +_REG_INSTANCES = [ + # (npad, bs, K, dense_env, expected tier) + (14336, 1, 512, None, "reg(cs=1,tb=512,maxv=8,ar=8)"), + (24576, 1, 512, None, "reg(cs=4,tb=512,maxv=4,ar=8)"), + (49152, 1, 512, None, "reg(cs=8,tb=512,maxv=3,ar=8)"), + (65536, 1, 1024, None, "reg(cs=8,tb=512,maxv=4,ar=8)"), + (131072, 8, 512, None, "reg(cs=8,tb=512,maxv=8,ar=8)"), + (163840, 1, 2048, None, "reg(cs=16,tb=512,maxv=5,ar=6)"), + (163840, 4, 512, None, "reg(cs=16,tb=512,maxv=5,ar=8)"), + (262144, 1, 2048, None, "reg(cs=16,tb=512,maxv=8,ar=8)"), + (262144, 4, 1024, None, "reg(cs=16,tb=512,maxv=8,ar=6)"), + (20480, 64, 512, None, "reg(cs=1,tb=1024,maxv=5,ar=8)"), + (28672, 64, 512, None, "reg(cs=1,tb=1024,maxv=8,ar=8)"), + (65536, 8, 512, "8", "reg(cs=2,tb=1024,maxv=8,ar=8)"), + (131072, 8, 1024, "8", "reg(cs=4,tb=1024,maxv=8,ar=8)"), + (262144, 8, 2048, None, "reg(cs=8,tb=1024,maxv=8,ar=8)"), +] + + +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,dense_env,expected", + _REG_INSTANCES, + ids=[t[4] + f"_n{t[0]}_bs{t[1]}_k{t[2]}" for t in _REG_INSTANCES], +) +@pytest.mark.parametrize("kind", ["randn", "ties"]) +def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, kind, monkeypatch): + _skip_if_cluster_capped(bs, npad, top_k) + try: + if dense_env is not None: + monkeypatch.setenv("TRTLLM_BSX_DENSE_BS", dense_env) + bsx_dispatch._reset_env_cache() + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=npad + bs + top_k, kind=kind, varlen=True + ) + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected) + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + finally: + if dense_env is not None: + monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) + bsx_dispatch._reset_env_cache() + + +# --------------------------------------------------------------------------- +# direct tier (npad <= DKCMAX): BS 1/8/64, randn + quantized ties, ragged + +# uniform seq_lens. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,kind,varlen", + [ + (4096, 1, 512, "randn", True), + (8256, 8, 512, "ties", True), + (12288, 64, 2048, "randn", True), + (4096, 8, 512, "randn", False), # uniform N_eff == npad (no tail) + ], +) +def test_bsx_direct(npad, bs, top_k, kind, varlen): + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=npad * 3 + bs, kind=kind, varlen=varlen + ) + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "direct") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +# --------------------------------------------------------------------------- +# tp tier (bs >= tp threshold): trivial npad<=kC path, cs>1 cluster path and +# cs=1 big-npad path. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,kind", + [ + (4096, 256, 512, "randn"), # cs=1, trivial npad <= kC path + (65536, 16, 1024, "ties"), # cs=8 cluster path + tie plateaus + (262144, 128, 2048, "randn"), # cs=1, uf=8 streaming path + ], +) +def test_bsx_tp(npad, bs, top_k, kind): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=npad + 7 * bs, kind=kind, varlen=True + ) + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +# --------------------------------------------------------------------------- +# Degenerate rows (N_eff <= K): in-kernel identity emit [0..N_eff-1] + -1 +# pad, mixed with normal rows in the same batch, on all three tiers. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize( + "npad,bs,top_k,expected_kind", + [ + (4096, 8, 2048, "direct"), + (65536, 8, 512, "reg"), + (65536, 16, 1024, "tp"), + ], +) +def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=11, kind="randn", varlen=True + ) + # Rows 0..3: degenerate (N_eff = 0 / 1 / K-1 / K); the rest keep their + # non-degenerate varlen lengths. Poison the tails accordingly. + for i, sl in enumerate([NEXT_N, CR + NEXT_N, (top_k - 1) * CR, top_k * CR]): + seq_lens[i] = sl + n_eff = (seq_lens.long() // CR).clamp(max=npad) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + tier = bsx_dispatch.route(bs, npad, top_k) + assert tier.startswith(expected_kind) or tier == expected_kind + out = _run_op(logits, pre_idx, seq_lens, top_k) + + for i in range(bs): + ne = int(n_eff[i].item()) + if ne <= top_k: + expect = torch.full((top_k,), -1, dtype=torch.int32, device="cuda") + expect[:ne] = torch.arange(ne, dtype=torch.int32, device="cuda") + assert torch.equal(out[i], expect), ( + f"row {i} (N_eff={ne}): degenerate identity emit mismatch: " + f"{out[i, : max(ne + 2, 8)].cpu().tolist()}..." + ) + else: + _tie_aware_check(out[i : i + 1], logits[i : i + 1], seq_lens[i : i + 1], top_k) + + +# --------------------------------------------------------------------------- +# pre_idx hardening: all-zero (production cold start) and out-of-range +# garbage hints must neither fault nor corrupt the output. +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize("preidx", ["zeros", "oor"]) +@pytest.mark.parametrize( + "npad,bs,top_k", + [ + (65536, 4, 512), # reg tier + (65536, 16, 1024), # tp tier + ], +) +def test_bsx_preidx_hardening(npad, bs, top_k, preidx): + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=13, kind="randn", varlen=True, preidx=preidx + ) + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +# --------------------------------------------------------------------------- +# Dispatcher fallback: unsupported inputs must route to the in-tree kernel +# and still produce its results (op contract unchanged). +# --------------------------------------------------------------------------- +@skip_not_sm100 +def test_bsx_dispatcher_fallback_bf16(): + bs, npad, top_k = 4, 65536, 512 + torch.manual_seed(21) + torch.cuda.manual_seed(21) + logits = (torch.randn(bs, npad, device="cuda") * 2.0).to(torch.bfloat16) + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre_idx[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + pre_idx[:, 0] = logits.float().argmax(dim=-1).int() + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + assert not bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ), "bf16 must NOT take the bsx fast path" + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=NEXT_N, compress_ratio=CR + ) + torch.cuda.synchronize() + _tie_aware_check(out, logits, seq_lens, top_k) + + +@skip_not_sm100 +def test_bsx_dispatcher_fallback_next_n2(): + """next_n=2 (cr=1, V3.2-style MTP rows) must fall back to the in-tree + kernel and produce its results.""" + bs, npad, top_k, next_n, cr = 4, 65536, 2048, 2, 1 + num_rows = bs * next_n + torch.manual_seed(23) + torch.cuda.manual_seed(23) + logits = (torch.randn(num_rows, npad, device="cuda") * 2.0).contiguous() + seq_lens = torch.full((bs,), npad, dtype=torch.int32, device="cuda") + # cr=1 kernel convention: it reads logits[pre_idx + (row % next_n) + 1]. + eff = npad - next_n # safe hint range for every row + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre_idx[:, 0] = logits[::next_n, :eff].argmax(dim=-1).int() - 1 + pre_idx[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + + assert not bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None + ), "next_n=2 must NOT take the bsx fast path" + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr + ) + torch.cuda.synchronize() + _tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) From 666efffe061185f292bfa9532c4f74c949f651ac Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:35:53 +0000 Subject: [PATCH 02/19] [None][perf] bsx tp: hint-ladder admission fast path (R0 parity on high-hit-rate rows) Port the in-tree R0 histogram-ladder ADMISSION concept (op#26, PR #16457) into the bsx throughput (tp) tier's pivot selection, closing the pr1 full-grid losses against the in-tree GVR kernel on real-capture rows at BS >= 16 while keeping the fused one-pass structure and every existing exactness invariant. Mechanism (measured on pr1 real-capture cells): the old 3-stage pivot pick targets ~3K sampled candidates inside a narrow [1.5K, 0.6kC] band. On real rows this either (a) picks a FAT rung (2-4x more P3 pushes and P4 candidates than needed - the v32/pro BS>=16 band), or (b) trusts a clustering-inflated sampled estimate and undershoots K, which fails the fused-pass reuse check and re-streams / secant-loops the full row (the flash_512k 1.6-1.8x losses; spatially clustered rows inflate a float4-sampled estimate up to 2.5x over the true count). Changes (gvr_topk_decode_bsx_tp.py only): * P2a stage-0 ADMISSION pick (R0 parity): accept the TIGHTEST ladder rung whose sampled-count confidence interval sits inside the [K, kC] acceptance window - the same "smallest exact count in [K, kC]" rule as the in-tree R0 admission, applied to the pre-pass estimates. The legacy 3-stage pick is unchanged as the fallback when no rung qualifies (cold-start / degenerate ladders take exactly the old path). * Clustering-aware sigma: sample_count now also tracks per-rung float4 OCCUPANCY, packed into the same per-thread accumulator (occ << 16 | cnt), giving the compound-Poisson sigma cnt/sqrt(occ) (equals the classic Poisson sigma on IID rows). Packing keeps registers, SMEM and the exchange at their pre-change sizes - an unpacked occ implementation measured a 14-25% whole-kernel regression (register spill in the streaming loops; A/B/bisect on v32_64k_L20). No field overflow: bsx guards npad <= 262144 => cluster-total cnt <= 8192. * Lower margin 2-sigma; K2048 uses 1.5-sigma (its [K, 4K] window is too narrow for 2-sigma to fire tight) backed by a RESCUE rung: the fused pass's second count column is now the next-fatter ladder rung instead of hmin (identical cost), so a pivot undershoot is caught with ONE collect re-stream instead of the multi-pass secant loop. Exactness machinery (accept window, P4 + tie tickets, plateau descent, ragged-N masking, degenerate emit, exit rendezvous) is untouched; admission never drops a top-K element because acceptance still requires the EXACT fused count in [K, kC]. Perf (nsys cold-L2 paired, same-rep A/B vs in-tree GVR kernel, real pr1 capture cells, B200; ratio = bsx/in-tree, lower is better): target cells before after flash_512k_L34 BS16 1.57 1.13 flash_512k_L34 BS64 1.72 1.11 flash_512k_L34 BS128 1.83 1.04 flash_512k_L34 BS256 1.54 0.73 flash_512k_L34 BS512 1.64 0.74 flash_512k_L34 BS1024 1.58 0.74 v32_32k_L50 BS16 1.40 1.21 v32_32k_L50 BS64 1.54 1.30 v32_32k_L50 BS256-1024 1.29-1.32 1.10-1.11 pro_128k_L54 BS16 1.27 1.19 pro_128k_L54 BS64 1.35 1.26 pro_128k_L54 BS128-1024 1.21-1.32 1.12-1.25 no-regression cells before after flash_512k_L36 BS1-128 0.55-0.94 0.56-0.92 flash_512k_L36 BS256-1024 0.67-0.68 0.66-0.67 v32_64k_L20 BS1-8 (reg) 0.63-0.65 0.63 v32_64k_L20 BS16-128 1.09-1.19 1.00-1.10 v32_64k_L20 BS256-1024 1.00-1.05 0.93-0.97 v32_32k_L04 BS16-1024 0.88-1.06 0.90-1.10 pro_256k_L30 BS16-1024 0.89-1.07 0.91-1.12 flash_16k_L26 BS1-1024 0.61-0.86 0.62-0.86 (reg/direct tiers untouched; the three +3-5% readings - pro_256k_L30 BS16/64, v32_32k_L04 BS64 - are inside the +-4.5% session noise floor measured on the fully untouched direct tier, e.g. flash_16k_L26 BS1 0.607 -> 0.634 with zero code change) The remaining v32_32k/pro_128k BS16-128 gap (1.19-1.30) is NOT the admission/pass-count mechanism: with admission the reuse check fires (one fused pass, admitted set slimmed 6678->2917 on v32_32k_L50) and a CS in {1,2,4} launch-shape sweep moves <= 7%. On those L2-resident shapes the in-tree kernel's whole-row-per-CTA structure is simply faster than the tp cluster split; closing it needs a tier-structure change, out of scope for this admission port. Exactness: bsx suite 56/56 (incl. new admission cases: hit-rate extremes on all three production shapes, tie plateau AT the admission threshold, count>kC overflow fallback, mixed admit/fallback ragged batch), in-tree gvr suite 671 passed / 144 skipped, 11/11 targeted screen (clustered/ties/zeros/oor/uniform/trivial), all pr1 A/B runs value-set-exact. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 123 +++++++++++++++--- .../sparse/test_cute_dsl_bsx_topk_decode.py | 109 +++++++++++++++- 2 files changed, 211 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index 681c67793e0c..5b7882e61c0d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -33,10 +33,16 @@ Faithful phase-by-phase port otherwise: P1 : hint gather + minmax + two-stage 64-bin histogram -> CCDF rung ladder - P2a : uniform every-32nd-float4 sampled multi-rung count -> 3-stage pivot - P2b : ONE fused streaming pass: exact counts at {pivot, hmin} + optimistic - collect of packed (key<<32|idx) u64 candidates >= pivot into CTA0 - smem (capped kC) + P2a : uniform every-32nd-float4 sampled multi-rung count (+ per-rung + float4 occupancy for a clustering-aware sigma) -> 4-stage pivot: + stage 0 is the hint-ladder ADMISSION (in-tree R0 op#26 parity) — + tightest rung whose sampled CI sits inside [K, kC] — followed by + the legacy 3-stage pick when no rung qualifies + P2b : ONE fused streaming pass: exact counts at {pivot, rescue rung} + + optimistic collect of packed (key<<32|idx) u64 candidates + >= pivot into CTA0 smem (capped kC). Pivot count in [K, kC] => + the candidates are reused as-is (1-pass admission); pivot under K + but rescue in-window => ONE collect re-stream P2c : multi-rung secant refine / max-below plateau descent when the pivot count misses [K, kC] P3 : candidate reuse (thr == tpush, no overflow) or one re-stream collect @@ -384,6 +390,17 @@ def count_pass( # ------------------------------------------------------------------ # sample_count: uniform every-SS-th float4 over the slice. + # Each per-thread accumulator PACKS the per-rung hit count (low 16 + # bits) with the per-rung float4 OCCUPANCY (high 16 bits) — the number + # of sampled float4s with at least one hit. On spatially-clustered + # real rows the 4 values of one float4 are strongly correlated, so the + # effective independent sample count is the occupancy, not the hit + # count: the admission stage sizes its confidence interval with the + # compound-Poisson sigma cnt/sqrt(occ) (equal to the classic + # sqrt(cnt) Poisson sigma on IID data where occ == cnt). Packing + # keeps registers, SMEM and the exchange at their pre-admission + # sizes; no field overflow: the bsx envelope pins npad <= 262144, so + # cluster-total cnt <= npad/32 = 8192 < 2^16 and occ <= npad/128. # ------------------------------------------------------------------ @cute.jit def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt): @@ -399,10 +416,18 @@ def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_ru while v0 + j * cutlass.Int32(SS) < v1: self._ld_float4(copy_atom, row_addr, v0 + j * cutlass.Int32(SS), frag) gi = (v0 + j * cutlass.Int32(SS)) << cutlass.Int32(2) - for q in cutlass.range_constexpr(4): - v = _mask_tail(frag[q], gi + cutlass.Int32(q), n_eff) - for r in cutlass.range_constexpr(R): - cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + v0m = _mask_tail(frag[0], gi, n_eff) + v1m = _mask_tail(frag[1], gi + cutlass.Int32(1), n_eff) + v2m = _mask_tail(frag[2], gi + cutlass.Int32(2), n_eff) + v3m = _mask_tail(frag[3], gi + cutlass.Int32(3), n_eff) + for r in cutlass.range_constexpr(R): + c4 = ( + cutlass.Int32(v0m >= tr[r]) + + cutlass.Int32(v1m >= tr[r]) + + cutlass.Int32(v2m >= tr[r]) + + cutlass.Int32(v3m >= tr[r]) + ) + cnt[r] = cnt[r] + c4 + (cutlass.Int32(c4 > cutlass.Int32(0)) << cutlass.Int32(16)) j = j + cutlass.Int32(TB) for r in cutlass.range_constexpr(R): s_ptcnt[r * TB + tidx] = cnt[r] @@ -410,6 +435,8 @@ def sample_count(self, R: cutlass.Constexpr, row_addr, v0, v1, n_eff, tidx, s_ru # ------------------------------------------------------------------ # exchange_counts: warp-fused CTA reduce -> cluster sum + prefix. # rcnt[r] = cluster-wide count, rpre[r] = exclusive prefix over lower ranks. + # P2a rows carry packed (occ << 16 | cnt) accumulators; the integer sum + # distributes over the packed fields (no overflow: see sample_count). # ------------------------------------------------------------------ @cute.jit def exchange_counts( @@ -662,7 +689,12 @@ def phase1(self, logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminma # ------------------------------------------------------------------ # fused_count_collect: exact counts at {rungs[0]=pivot, - # rungs[1]=hmin} + push (key,idx) >= tpush into CTA0 cand, capped kcap. + # rungs[1]=rescue} + push (key,idx) >= tpush into CTA0 cand, capped + # kcap. The rescue rung (next fatter ladder rung) replaces HEAD's hmin + # column at identical register cost: on a pivot undershoot the driver + # accepts the rescue with ONE collect re-stream instead of the + # multi-pass secant loop (hmin was a bracket-only data point; the + # secant fallback re-derives its bracket from the rescue instead). # Keys are RAW fp32 bits (f2u happens in P4 round 0), matching CUDA. # Candidates are a single packed (key<<32 | idx) u64 array — CUDA's # `unsigned long long cand[kC]` — one 8B push per candidate. @@ -962,7 +994,9 @@ def gvr_tp_kernel( ) hmin_floor = s_rungs[AR - 1] span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) - # P2a: sampled ladder count -> pivot pick + # P2a: sampled ladder count -> pivot pick. The exchanged rows + # carry packed (occ << 16 | cnt) values (admission sigma + # input); every consumer below unpacks with & 0xFFFF / >> 16. self.sample_count(AR, row_addr, v0, v1, n_eff, tidx, s_rungs, s_ptcnt) self.exchange_counts( AR, xch & cutlass.Int32(1), tidx, rank, s_ptcnt, s_rcnt, s_rpre, s_ipartial @@ -970,16 +1004,60 @@ def gvr_tp_kernel( xch = xch + cutlass.Int32(1) cute.arch.barrier() if tidx == cutlass.Int32(0): - # 3-stage pivot pick (iter7 band -> iter8b 2-sigma -> - # iter12-F1b gamble) + # 4-stage pivot pick (R0-admission tightest-in-window -> + # iter7 band -> iter8b 2-sigma -> iter12-F1b gamble) lo = cutlass.const_expr((3 * K) // 2) hi = cutlass.const_expr((6 * kC) // 10) tgt_py = min(max(3 * K, (3 * K) // 2), (6 * kC) // 10) tgt = cutlass.Int32(tgt_py) best = cutlass.Int32(AR - 1) bestd = cutlass.Int32(0x7FFFFFFF) + # Stage 0 — hint-ladder ADMISSION (in-tree R0 parity, + # op#26 lineage): accept the TIGHTEST ladder rung + # (highest threshold = smallest admitted-candidate set) + # whose sampled-count confidence interval lies inside + # the [K, kC] acceptance window, mirroring the R0 rule + # "smallest exact count in [K, kC]". The old stage-1 + # band ([1.5K, 0.6kC], target ~3K) systematically picks + # a FAT rung on real high-hit-rate rows (2-4x more P3 + # pushes + P4 candidates than needed — the v32/pro + # BS>=16 losses vs the in-tree R0 kernel) or, when + # spatially-clustered real data inflates a sampled + # estimate into the band while the true count is < K, + # an undershooting one (whole-row secant + re-stream: + # the flash_512k 1.6-1.8x losses). + # Sigma is clustering-aware: est = SS*cnt, sigma = + # SS*cnt/sqrt(occ) (compound Poisson over occupied + # float4 clumps; equals the classic sqrt(SS*est) when + # occ == cnt, i.e. IID rows). Lower margin is 2 sigma; + # K2048 uses 1.5 (its [K, 4K] window is too narrow for + # 2-sigma to ever fire tight, and the rescue rung + # bounds a miss at one extra collect pass). Rungs are + # descending in j, so the first passing j is the + # tightest; later stages never override (their dd >= 0 + # cannot beat bestd == 0). + mlo = cutlass.const_expr(1.5 if K >= 2048 else 2.0) + for j in cutlass.range_constexpr(AR): + cnt_j = s_rcnt[j] & cutlass.Int32(0xFFFF) + if cnt_j > cutlass.Int32(0): + if bestd == cutlass.Int32(0x7FFFFFFF): + occ_j = s_rcnt[j] >> cutlass.Int32(16) + if occ_j < cutlass.Int32(1): + occ_j = cutlass.Int32(1) + sig = ( + cutlass.Float32(float(SS)) + * cutlass.Float32(cnt_j) + / cmath.sqrt(cutlass.Float32(occ_j)) + ) + fest = cutlass.Float32(cnt_j * cutlass.Int32(SS)) + if fest - cutlass.Float32(mlo) * sig >= cutlass.Float32(float(K)): + if fest + cutlass.Float32(2.0) * sig <= cutlass.Float32( + float(kC) + ): + bestd = cutlass.Int32(0) + best = cutlass.Int32(j) for j in cutlass.range_constexpr(AR): - est = s_rcnt[j] * cutlass.Int32(SS) + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) if est >= cutlass.Int32(lo): if est <= cutlass.Int32(hi): dd = est - tgt @@ -990,7 +1068,7 @@ def gvr_tp_kernel( best = cutlass.Int32(j) if bestd == cutlass.Int32(0x7FFFFFFF): for j in cutlass.range_constexpr(AR): - est = s_rcnt[j] * cutlass.Int32(SS) + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) if est > cutlass.Int32(0): g = cutlass.Float32(2.0) * cmath.sqrt( cutlass.Float32(cutlass.Int32(SS) * est) @@ -1005,9 +1083,11 @@ def gvr_tp_kernel( bestd = dd best = cutlass.Int32(j) if bestd == cutlass.Int32(0x7FFFFFFF): - if s_rcnt[AR - 1] * cutlass.Int32(SS) > cutlass.Int32(kC): + if (s_rcnt[AR - 1] & cutlass.Int32(0xFFFF)) * cutlass.Int32( + SS + ) > cutlass.Int32(kC): for j in cutlass.range_constexpr(AR): - est = s_rcnt[j] * cutlass.Int32(SS) + est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) if est > cutlass.Int32(0): dd = est - tgt if dd < cutlass.Int32(0): @@ -1016,8 +1096,17 @@ def gvr_tp_kernel( bestd = dd best = cutlass.Int32(j) tp_ = s_rungs[best] + # Rescue rung: the next FATTER ladder rung below the + # pivot. If the pivot's exact count lands under K + # (sampling error on a clustered row), the driver can + # accept the rescue with ONE collect re-stream instead + # of the multi-pass secant loop. Read before the ladder + # slots are overwritten. + resc_ = hmin_floor + if best < cutlass.Int32(AR - 1): + resc_ = s_rungs[best + cutlass.Int32(1)] s_rungs[0] = tp_ - s_rungs[1] = hmin_floor + s_rungs[1] = resc_ cute.arch.barrier() tpush = s_rungs[0] self.fused_count_collect( diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index 4848b82ca0f7..9c328aedd6ea 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -131,6 +131,7 @@ def _make_bsx_inputs( kind: str = "randn", varlen: bool = True, preidx: str = "mixed", + hit_rate: float = 0.5, ): """Build (logits fp32, pre_idx int32, seq_lens int32) for the bsx path. @@ -141,9 +142,9 @@ def _make_bsx_inputs( top-K if the lane masking were broken (production tails are stale values, not -FLT_MAX pad). - ``preidx``: 'mixed' = argmax slot 0 + ~50% real topk hints; - 'zeros' = all-zero cold start; 'oor' = out-of-range garbage - (negative / >= npad) that the kernels must clamp harmlessly. + ``preidx``: 'mixed' = argmax slot 0 + ``hit_rate`` real topk hints + (default ~50%); 'zeros' = all-zero cold start; 'oor' = out-of-range + garbage (negative / >= npad) that the kernels must clamp harmlessly. """ torch.manual_seed(seed) torch.cuda.manual_seed(seed) @@ -170,7 +171,7 @@ def _make_bsx_inputs( pre_idx = torch.randint(-npad, 4 * npad, (bs, top_k), dtype=torch.int32, device="cuda") else: ref_topk = valid_logits.topk(top_k, dim=-1).indices.int() - keep = torch.rand(ref_topk.shape, device="cuda") < 0.5 + keep = torch.rand(ref_topk.shape, device="cuda") < hit_rate junk = torch.arange(top_k, dtype=torch.int32, device="cuda").expand(bs, -1) # junk arange is in-range for every row: N_eff > top_k here. pre_idx = torch.where(keep, ref_topk, junk).contiguous() @@ -428,6 +429,106 @@ def test_bsx_preidx_hardening(npad, bs, top_k, preidx): _tie_aware_check(out, logits, seq_lens, top_k) +# --------------------------------------------------------------------------- +# tp-tier hint-ladder ADMISSION fast path (R0 parity): the P2a stage-0 pick +# pushes candidates at the tightest ladder rung whose sampled CI lies in +# [K, kC], so high-hit-rate rows finish in one fused streaming pass. These +# cases pin the admission decision surface: hit-rate extremes, tie plateaus +# AT the admission threshold, forced count > kC overflow fallback, and a +# batch mixing admitted and fallback rows (per-row escape independence). +# --------------------------------------------------------------------------- +@skip_not_sm100 +@pytest.mark.parametrize("hit_rate", [0.85, 0.15]) +@pytest.mark.parametrize( + "npad,bs,top_k", + [ + (65536, 16, 1024), # cs=8 cluster tp (pro-like) + (32832, 16, 2048), # cs=4 cluster tp (v32 32k production shape) + (131136, 128, 512), # cs=1 streaming tp (flash 512k production shape) + ], +) +def test_bsx_tp_admission_hitrate(npad, bs, top_k, hit_rate): + """High-hr rows must admit (1-pass) and low-hr rows must stay exact via + the pivot/secant fallback; both paths must produce exact top-K.""" + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=npad + bs + int(hit_rate * 100), varlen=True, hit_rate=hit_rate + ) + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +@skip_not_sm100 +def test_bsx_tp_admission_tie_plateau(): + """Tie plateau AT the admission threshold: K/2 distinct high values + + a 3K-wide exact-tie plateau straddling the K-th rank. The hint set is + the true top-K, so the ladder rungs land ON the plateau value and the + admitted candidate set is tie-degenerate; the P4 tie-ticket emit must + still return an exact value multiset.""" + bs, npad, top_k = 16, 65536, 1024 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + torch.manual_seed(31) + torch.cuda.manual_seed(31) + logits = -torch.rand(bs, npad, dtype=torch.float32, device="cuda") - 1.0 + for r in range(bs): + perm = torch.randperm(npad, device="cuda") + hi = perm[: top_k // 2] + plateau = perm[top_k // 2 : top_k // 2 + 3 * top_k] + logits[r, hi] = 5.0 + torch.arange(top_k // 2, device="cuda").float() * 0.01 + logits[r, plateau] = 1.0 # exact ties straddling rank K + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +@skip_not_sm100 +@pytest.mark.parametrize("top_k,n_plateau", [(1024, 8000), (2048, 12000)]) +def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau): + """count(>= any admissible threshold) > kC: K-1 distinct high values, + then an n_plateau-wide exact-tie plateau (n_plateau > kC = 6144/8192). + Every threshold at or below the plateau overflows the candidate buffer + (the uncapped push counter detects it) and every threshold above it + undershoots (K-1 < K), so admission can never accept and the kernel + must fall through to the max-below plateau-descent path and its direct + emit. Exact answer: the K-1 highs plus exactly one plateau member.""" + bs, npad = 16, 65536 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + torch.manual_seed(37) + torch.cuda.manual_seed(37) + logits = -torch.rand(bs, npad, dtype=torch.float32, device="cuda") - 1.0 + for r in range(bs): + perm = torch.randperm(npad, device="cuda") + hi = perm[: top_k - 1] + plateau = perm[top_k - 1 : top_k - 1 + n_plateau] + logits[r, hi] = 5.0 + torch.arange(top_k - 1, device="cuda").float() * 0.01 + logits[r, plateau] = 1.0 + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + +@skip_not_sm100 +def test_bsx_tp_admission_mixed_batch(): + """Ragged batch mixing rows that admit (true-top-K hints) with rows + that must fall back (all-zero cold-start hints): per-row admission + escape is independent, so every row must stay exact regardless of + which path its cluster takes.""" + bs, npad, top_k = 16, 65536, 1024 # cs=8 tp + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=41, varlen=True, hit_rate=0.9 + ) + pre_idx[1::2] = 0 # odd rows: cold-start (degenerate ladder -> fallback) + _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) + + # --------------------------------------------------------------------------- # Dispatcher fallback: unsupported inputs must route to the in-tree kernel # and still produce its results (op contract unchanged). From c8fc03ab5ac54b6112b76e3c73c7877acf255897 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:01:48 +0000 Subject: [PATCH 03/19] [None][perf] bsx tp admission: tighten acceptance bound to pivot-band target Fix round for the hint-ladder admission fast path (previous commit): the full-grid pr2-vs-pr1 A/B (9515 cases) showed the admission closed the deep tail (flash_512k_L34 1.83 -> 1.05) but regressed 396 cases by >5pp and flipped 47 former wins, netting the overall gm to a wash. Verified root cause (in-kernel decision probe + host-side P2a emulation on the regression rows) - TWO distinct failure modes, both in the PICK: 1. Fat admit: the [K, kC] acceptance window let the tightest CI-passing rung admit a candidate set 2.5x the legacy ~3K pivot-band target (pro_64k_L06: admitted 4134 vs legacy 1671; pro_32k_L28 4047 vs 1554; flash_128k_L10 1970 vs 788), inflating P3 pushes + P4 scan. 2. Undershoot admit: the K2048 1.5-sigma lower margin accepted rungs whose true count lands under K on clustered rows (v32_32k_L23 rung3 est 2720, sigma 374, TRUE count 1959 < 2048) - the rescue re-stream then costs more than the fat-but-valid legacy pick it displaced (22.6 -> 28.2us at BS64). The improved twin v32_32k_L50 has nearly identical rung3 stats (est 2784, sigma 382) with true count 2917: the pivot margin (+1.80 vs +1.92 sigma) is the only discriminant. Changes (pick logic only; fused pass, rescue and all exactness machinery are untouched from the admission commit): * Admission upper bound tightened from kC to the legacy pivot-band hi (0.6*kC), so admission only fires when the rung is genuinely tight. * K2048 lower margin raised 1.5 -> 1.85 sigma: rejects the L23-class undershoots (+1.80 sigma) while keeping the genuine tight admits (L50 +1.92 sigma). K512/K1024 stay at 2.0 sigma. * Stage 0b: the legacy band pick (min |est-tgt| in [1.5K, 0.6kC], the exact pr1 stage-1 rule) now OVERRIDES the admitted rung when it is strictly leaner and safe by its own clustering-aware 1.5-sigma lower CI; with no admission it is taken as-is (pr1 parity). * The pick's margin tests are now sqrt/div-free (squared comparisons): the pick is a thread-0 serial section between two CTA barriers, and the previous sqrt+fdiv chain measured 2-5% whole-kernel on L2-resident accept-path rows. Perf (nsys cold-L2 paired vs in-tree GVR kernel, real capture cells, B200; ratio = bsx/in-tree, lower is better; bar = pr1 ratio + 3% for the regression cells, absolute for the retention cells): regression cells pr1 pr2 fixed bar v32_32k_L23 BS16 1.202 1.521 1.198 <=1.238 PASS v32_32k_L23 BS64 1.334 1.662 1.325-1.352 <=1.374 PASS v32_32k_L23 BS1024 1.131 1.413 1.123 <=1.165 PASS pro_64k_L06 BS256 1.062 1.352 1.061 <=1.094 PASS pro_64k_L06 BS1024 1.099 1.414 1.121 <=1.132 PASS pro_32k_L28 BS1024 1.059 1.295 1.084 <=1.091 PASS flash_128k_L10 BS64 1.100 1.279 1.094 <=1.133 PASS pro_32k_L02 BS64 0.652 0.659 0.652 <=0.672 PASS retention cells pr1 pr2 fixed bar flash_512k_L34 BS128 1.831 1.052 1.020 <=1.08 PASS flash_512k_L34 BS256 1.537 0.730 0.725 <=0.80 PASS v32_32k_L50 BS64 1.537 1.290 1.267 <=1.35 PASS v32_32k_L50 BS16/1024 1.401/1.316 1.199/1.130 1.177/1.112 no-regression spots pr1 pr2 fixed v32_32k_L04 BS64 1.057 1.086 1.032-1.068 (3 reps) flash_16k_L26 BS1 (direct) 0.607 0.622 0.600 v32_64k_L20 BS8 (reg) 0.633 0.630 0.624 flash_512k_L36 BS64/256 0.944/0.674 0.898/0.656 0.883/0.655 pro_128k_L54 BS64/256 1.351/1.209 1.245/1.122 1.234/1.093 Known give-back: v32_64k_L20 BS16-128 (rung margin +1.53 sigma, below the new 1.85 K2048 cut) reverts from the pr2 admission pick to the legacy band pick; BS256-1024 keep the tight rung via stage 0b. Falsified along the way (kept out): stashing the [rescue, pivot) band in the candidate-buffer tail during the fused pass (per-element atomic, per-warp stripes, and warp-aggregated ballot variants) - the fused streaming loop cannot absorb ANY extra per-element code within the +-3% bar (measured +4-10% whole-kernel on accept-path rows for all three variants). Exactness: bsx suite 56/56 (admission cases included; acceptance-rule comment updated), in-tree gvr suite K2048 subset green, all A/B runs value-set-exact (exact=True on every measured cell x BS x arm). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 105 ++++++++++++++---- .../sparse/test_cute_dsl_bsx_topk_decode.py | 4 +- 2 files changed, 84 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index 5b7882e61c0d..ed32b5e858e8 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -36,8 +36,11 @@ P2a : uniform every-32nd-float4 sampled multi-rung count (+ per-rung float4 occupancy for a clustering-aware sigma) -> 4-stage pivot: stage 0 is the hint-ladder ADMISSION (in-tree R0 op#26 parity) — - tightest rung whose sampled CI sits inside [K, kC] — followed by - the legacy 3-stage pick when no rung qualifies + tightest rung whose sampled CI sits inside [K, 0.6*kC] (the + legacy pivot-band hi) — stage 0b prefers the legacy band pick + when it is strictly leaner and safe by its clustering-aware + 1.5-sigma lower CI, and the legacy 3-stage pick handles rows + where no rung qualifies P2b : ONE fused streaming pass: exact counts at {pivot, rescue rung} + optimistic collect of packed (key<<32|idx) u64 candidates >= pivot into CTA0 smem (capped kC). Pivot count in [K, kC] => @@ -1029,14 +1032,40 @@ def gvr_tp_kernel( # Sigma is clustering-aware: est = SS*cnt, sigma = # SS*cnt/sqrt(occ) (compound Poisson over occupied # float4 clumps; equals the classic sqrt(SS*est) when - # occ == cnt, i.e. IID rows). Lower margin is 2 sigma; - # K2048 uses 1.5 (its [K, 4K] window is too narrow for - # 2-sigma to ever fire tight, and the rescue rung - # bounds a miss at one extra collect pass). Rungs are - # descending in j, so the first passing j is the - # tightest; later stages never override (their dd >= 0 - # cannot beat bestd == 0). - mlo = cutlass.const_expr(1.5 if K >= 2048 else 2.0) + # occ == cnt, i.e. IID rows). + # Fix round (pr2 full-grid regressions), all changes on + # the PICK only — the fused pass / rescue / exactness + # machinery is byte-identical to the admission commit: + # * upper acceptance bound tightened from kC to the + # legacy pivot-band hi (0.6*kC): an in-window-but- + # fat admitted rung inflates P3 pushes + P4 + # candidates 2-4x over the ~3K legacy target; + # * K2048 lower margin raised 1.5 -> 1.85 sigma: at + # 1.5 the admitted rung's true count lands under K + # on clustered rows (v32_32k_L23 rung est 2720, + # sigma 374, true 1959 < K) and the rescue + # re-stream costs more than the fat legacy pick it + # displaced; 1.85 keeps the genuine tight admits + # (v32_32k_L50: margin +1.92 sigma, true 2917); + # * stage 0b below: when the legacy band pick is + # strictly LEANER than the admitted rung and safe + # by its own clustering-aware 1.5-sigma lower CI, + # prefer it (mlo=2.0 rows whose next-leaner rung + # sits at 1.8-2.0 sigma otherwise admit a 2.5x + # fatter set: pro_64k_L06 4134 vs legacy 1671). + # Rungs are descending in j, so the first passing j is + # the tightest. + # All margin tests below are sqrt/div-free: with + # sigma = est/sqrt(occ), "est - m*sigma >= K" is + # "(est-K)^2 * occ >= m^2 * est^2" (est >= K), and + # "est + 2*sigma <= U" is "4*est^2 <= (U-est)^2 * occ" + # (est <= U). The pick is a THREAD-0 serial section + # between two CTA barriers; the sqrt+fdiv chain of the + # first admission cut measured ~2-5% whole-kernel on + # L2-resident accept-path rows. + mlo2 = cutlass.const_expr(1.85 * 1.85 if K >= 2048 else 4.0) + ubnd = cutlass.const_expr(float((6 * kC) // 10)) + adm_est = cutlass.Int32(0x7FFFFFFF) for j in cutlass.range_constexpr(AR): cnt_j = s_rcnt[j] & cutlass.Int32(0xFFFF) if cnt_j > cutlass.Int32(0): @@ -1044,18 +1073,26 @@ def gvr_tp_kernel( occ_j = s_rcnt[j] >> cutlass.Int32(16) if occ_j < cutlass.Int32(1): occ_j = cutlass.Int32(1) - sig = ( - cutlass.Float32(float(SS)) - * cutlass.Float32(cnt_j) - / cmath.sqrt(cutlass.Float32(occ_j)) - ) + focc = cutlass.Float32(occ_j) fest = cutlass.Float32(cnt_j * cutlass.Int32(SS)) - if fest - cutlass.Float32(mlo) * sig >= cutlass.Float32(float(K)): - if fest + cutlass.Float32(2.0) * sig <= cutlass.Float32( - float(kC) - ): - bestd = cutlass.Int32(0) - best = cutlass.Int32(j) + fe2 = fest * fest + a_lo = fest - cutlass.Float32(float(K)) + bhi = cutlass.Float32(ubnd) - fest + if a_lo >= cutlass.Float32(0.0): + if a_lo * a_lo * focc >= cutlass.Float32(mlo2) * fe2: + if bhi >= cutlass.Float32(0.0): + if cutlass.Float32(4.0) * fe2 <= bhi * bhi * focc: + bestd = cutlass.Int32(0) + best = cutlass.Int32(j) + adm_est = cnt_j * cutlass.Int32(SS) + # Stage 0b — legacy band pick (min |est - tgt| in + # [1.5K, 0.6kC], EXACTLY the pr1 stage-1 rule), then + # combine: with no admission it is taken as-is (pr1 + # parity); with an admitted rung it OVERRIDES only + # when strictly leaner AND safe by its own + # clustering-aware 1.5-sigma lower CI. + jb = cutlass.Int32(AR - 1) + jbd = cutlass.Int32(0x7FFFFFFF) for j in cutlass.range_constexpr(AR): est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) if est >= cutlass.Int32(lo): @@ -1063,9 +1100,29 @@ def gvr_tp_kernel( dd = est - tgt if dd < cutlass.Int32(0): dd = tgt - est - if dd < bestd: - bestd = dd - best = cutlass.Int32(j) + if dd < jbd: + jbd = dd + jb = cutlass.Int32(j) + if jbd != cutlass.Int32(0x7FFFFFFF): + if bestd == cutlass.Int32(0x7FFFFFFF): + bestd = jbd + best = jb + else: + cnt_b = s_rcnt[jb] & cutlass.Int32(0xFFFF) + est_b = cnt_b * cutlass.Int32(SS) + if est_b < adm_est: + occ_b = s_rcnt[jb] >> cutlass.Int32(16) + if occ_b < cutlass.Int32(1): + occ_b = cutlass.Int32(1) + fb = cutlass.Float32(est_b) + ab = fb - cutlass.Float32(float(K)) + if ab >= cutlass.Float32(0.0): + if ( + ab * ab * cutlass.Float32(occ_b) + >= cutlass.Float32(2.25) * fb * fb + ): + bestd = jbd + best = jb if bestd == cutlass.Int32(0x7FFFFFFF): for j in cutlass.range_constexpr(AR): est = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index 9c328aedd6ea..ae129d9deb6a 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -432,7 +432,9 @@ def test_bsx_preidx_hardening(npad, bs, top_k, preidx): # --------------------------------------------------------------------------- # tp-tier hint-ladder ADMISSION fast path (R0 parity): the P2a stage-0 pick # pushes candidates at the tightest ladder rung whose sampled CI lies in -# [K, kC], so high-hit-rate rows finish in one fused streaming pass. These +# [K, 0.6*kC] (the legacy pivot-band hi; a leaner SAFE legacy band pick +# overrides it — stage 0b), so high-hit-rate rows finish in one fused +# streaming pass. These # cases pin the admission decision surface: hit-rate extremes, tie plateaus # AT the admission threshold, forced count > kC overflow fallback, and a # batch mixing admitted and fallback rows (per-row escape independence). From fb98f75d886accf9073ca13051fe8146fe5befb5 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 03:37:47 +0000 Subject: [PATCH 04/19] [None][feat] bsx tiers: MTP (next_n>1) and compress_ratio=1 support; guard narrowed Extend the op43 bsx CuTe DSL tier family (direct/reg/tp) to the full in-tree run_one_row contract: next_n >= 1 (MTP speculative-decode rows) and compress_ratio in {1, 4}, and drop the next_n==1 / cr==4 conditions from the dispatcher guard. Semantics mirror the in-tree kernel line-for-line: - request-level hint sharing: pre_idx row = row // next_n (pre_idx and seq_lens are [num_rows // next_n, ...], validated by the guard); - cr==1 temporal hint shift (row % next_n) + 1, with shifted-OOR hints falling into the tiers' existing clamp hardening; - per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr (already general in _row_n_eff). next_n / cr are ctor constexpr: the next_n==1 / cr==4 hot path traces identically to the v1 port (const_expr branches; no offset value is even computed on cr>1 builds). All compile / dispatch / runner caches key on (next_n, cr). Validation (umbriel-b200-074, torch nv26.05): - bsx suite extended with the MTP axis (next_n in {2,3,4} x cr in {1,4} x three tiers x {random-garbage, noised-hint, cold-start-zeros, ties}), each case checked against torch.topk host N_eff/offset simulation AND a differential in-tree arm (order_row-forced, per-row value-multiset equality): 140/140 passed. - full in-tree GVR suite (fp32 MTP/cr=1 cells now auto-route to bsx): 671 passed, 144 skipped. - next_n=1 perf anchors (ab_pr smoke, paired same-GPU back-to-back vs stashed baseline): bsxd/gvrpr ratio deltas within +/-3% on flash_16k_L26 BS1, v32_64k_L20 BS8, flash_512k_L34 BS128/256, v32_32k_L23 BS64, pro_256k_L30 BS256; all exact. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 17 +- .../top_k/gvr_topk_decode_bsx_direct.py | 28 +- .../top_k/gvr_topk_decode_bsx_dispatch.py | 54 ++-- .../top_k/gvr_topk_decode_bsx_reg.py | 61 +++- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 102 +++++-- .../sparse/test_cute_dsl_bsx_topk_decode.py | 267 ++++++++++++++++-- 6 files changed, 433 insertions(+), 96 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 32642a84cac4..2c3090b437d5 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7573,17 +7573,18 @@ def forward( ``counters`` without ``order_row`` is rejected. """ - # BSX tier fast path (op43 port): fp32 / next_n=1 / cr=4 / - # npad <= 262144 decode rows route to the direct/reg/tp CuTe DSL - # tiers; everything else (half-prec, LB, sort-indirect, V3.2, - # oversize npad, hw cluster cap) falls through to the in-tree - # kernel below. Host-only guard — no device sync. The op - # signature and output contract are unchanged (unordered int32 - # indices, -1 pad only for degenerate rows). + # BSX tier fast path (op43 port): fp32 / next_n >= 1 (MTP) / + # cr in {1, 4} / npad <= 262144 decode rows route to the + # direct/reg/tp CuTe DSL tiers; everything else (half-prec, LB, + # sort-indirect, oversize npad, hw cluster cap) falls through + # to the in-tree kernel below. Host-only guard — no device + # sync. The op signature and output contract are unchanged + # (unordered int32 indices, -1 pad only for degenerate rows). if _is_bsx_supported(logits, pre_idx, seq_lens, output_indices, top_k, next_n, compress_ratio, order_row, counters): - _bsx_topk(logits, pre_idx, seq_lens, output_indices, top_k) + _bsx_topk(logits, pre_idx, seq_lens, output_indices, top_k, + next_n, compress_ratio) return cute_dtype = _TORCH_TO_CUTLASS_DTYPE[logits.dtype] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py index 959355fe9dbb..6fb8558cafaf 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py @@ -67,7 +67,8 @@ class DirectTopKKernel: Ctor knobs (compile-time): ``top_k`` in {512, 1024, 2048}, ``num_threads`` (TB, default 1024), ``next_n`` / ``compress_ratio`` for the per-row - N_eff arithmetic (v1 dispatcher guard pins next_n=1, cr=4). SMEM layout + N_eff arithmetic (constexpr; the direct tier reads no hints, so MTP + support is the N_eff formula alone). SMEM layout mirrors gvr_bsx.cu's DSmem with the same packed (key << 32 | idx) u64 layout (op43 S-E round 6: split key/idx arrays cost 2x smem transactions in collect + emits vs CUDA's single STS.64 / LDS.64 per @@ -526,13 +527,15 @@ def __call__(self, logits: cute.Tensor, seq_lens: cute.Tensor, out: cute.Tensor, _COMPILE_CACHE: dict = {} -def _get_compiled(top_k: int, num_threads: int = 1024): - key = (top_k, num_threads) +def _get_compiled(top_k: int, num_threads: int = 1024, next_n: int = 1, cr: int = 4): + key = (top_k, num_threads, next_n, cr) compiled = _COMPILE_CACHE.get(key) if compiled is None: from cutlass.cute import runtime as _crt - kernel = DirectTopKKernel(top_k=top_k, num_threads=num_threads) + kernel = DirectTopKKernel( + top_k=top_k, num_threads=num_threads, next_n=next_n, compress_ratio=cr + ) n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 @@ -568,25 +571,32 @@ def _check_contract(logits, seq_lens, out, K): assert logits.data_ptr() % 16 == 0 and out.data_ptr() % 16 == 0 -def direct_topk(logits: torch.Tensor, seq_lens: torch.Tensor, out: torch.Tensor, K: int) -> None: +def direct_topk( + logits: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + next_n: int = 1, + cr: int = 4, +) -> None: """Exact top-K indices of ``logits`` rows into ``out`` (direct path). logits: [BS, npad] fp32 contiguous; the tail beyond each row's N_eff may be stale garbage (masked in-kernel). - seq_lens: [BS] int32, request-level, uncompressed-token space. + seq_lens: [BS // next_n] int32, request-level, uncompressed-token space. out: [BS, K] int32 contiguous. K: 512 / 1024 / 2048. Contract checks run once per (shape, K) signature to keep the hot launch path at bare tvm-ffi cost. """ - sig = (logits.shape, out.shape, K) + sig = (logits.shape, out.shape, K, next_n, cr) if sig not in _CHECKED_SIGS: _check_contract(logits, seq_lens, out, K) _CHECKED_SIGS.add(sig) - compiled = _COMPILE_CACHE.get((K, 1024)) + compiled = _COMPILE_CACHE.get((K, 1024, next_n, cr)) if compiled is None: - compiled = _get_compiled(K) + compiled = _get_compiled(K, next_n=next_n, cr=cr) compiled(logits, seq_lens, out) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index 8eb3a8bfe963..655881712de4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -22,14 +22,17 @@ NOT ported: it is unreachable inside the deployment envelope (npad <= 262144), which :func:`is_bsx_supported` enforces. -v1 dispatcher guard (anything else falls back to the in-tree +Dispatcher guard (anything else falls back to the in-tree ``GvrTopKKernel`` path in ``CuteDSLGvrTopKDecodeRunner.forward``): - dtype == fp32, next_n == 1, compress_ratio == 4, order_row is None, - counters is None, K in {512, 1024, 2048}, npad <= 262144, npad % 64 == 0, - contiguous 16B-aligned tensors, and the routed tier's cluster size within - the queried hardware max (never silently degrade the cluster shape). -Per-row degeneracy (N_eff <= K) and ragged N are handled INSIDE the tiers, -so the guard needs no device sync. + dtype == fp32, next_n >= 1 (MTP; num_rows divisible by next_n), + compress_ratio in {1, 4}, order_row is None, counters is None, + K in {512, 1024, 2048}, npad <= 262144, npad % 64 == 0, contiguous + 16B-aligned tensors, and the routed tier's cluster size within the + queried hardware max (never silently degrade the cluster shape). +pre_idx / seq_lens are request-level ([num_rows // next_n, K] / +[num_rows // next_n]) — the in-tree contract. Per-row degeneracy +(N_eff <= K) and ragged N are handled INSIDE the tiers, so the guard +needs no device sync. Env knobs (identical semantics to the CUDA arm's ``GVR_BSX_*`` static locals, renamed for the production tree; cached at first use like the CUDA static @@ -147,19 +150,19 @@ def route_cluster_size(bs: int, npad: int, K: int) -> int: # on <20us kernels (op43 S-E). The routing decision is pure in (bs, npad, K) # [env thresholds are cached at first use, like the CUDA static locals], so # bind it once per key to a closure. -_DISPATCH_CACHE = {} # (bs, npad, K) -> callable(logits, pre, seq_lens, out) +_DISPATCH_CACHE = {} # (bs, npad, K, next_n, cr) -> callable(logits, pre, seq_lens, out) -def _bind(bs, npad, K): +def _bind(bs, npad, K, next_n, cr): tier = route(bs, npad, K) if tier == "tp": def fn(lg, pre, sl, out): - tp_topk(lg, pre, sl, out, K) + tp_topk(lg, pre, sl, out, K, next_n, cr) elif tier == "direct": def fn(lg, pre, sl, out): - direct_topk(lg, sl, out, K) + direct_topk(lg, sl, out, K, next_n, cr) elif tier.startswith("gvr"): raise ValueError( f"bsx gvr tier is not ported (npad beyond the deployment " @@ -170,7 +173,7 @@ def fn(lg, pre, sl, out): cs, tb, maxv, ar = _parse_reg(tier) def fn(lg, pre, sl, out): - reg_topk(lg, pre, sl, out, K, cs, tb, maxv, ar) + reg_topk(lg, pre, sl, out, K, cs, tb, maxv, ar, next_n, cr) return fn @@ -186,17 +189,20 @@ def is_bsx_supported( order_row, counters, ) -> bool: - """Host-only v1 guard for the bsx tiers (no device sync; see module + """Host-only guard for the bsx tiers (no device sync; see module docstring). Returns False -> caller uses the in-tree kernel.""" if logits.dtype != torch.float32: return False - if next_n != 1 or compress_ratio != 4: + if next_n < 1 or compress_ratio not in (1, 4): return False if order_row is not None or counters is not None: return False if top_k not in (512, 1024, 2048): return False bs, npad = logits.shape + if bs % next_n != 0: + return False + n_req = bs // next_n if npad > 262144 or npad % 64 != 0: return False if not ( @@ -206,9 +212,9 @@ def is_bsx_supported( and seq_lens.is_contiguous() ): return False - if pre_idx.shape != (bs, top_k) or output_indices.shape != (bs, top_k): + if pre_idx.shape != (n_req, top_k) or output_indices.shape != (bs, top_k): return False - if seq_lens.shape != (bs,) or seq_lens.dtype != torch.int32: + if seq_lens.shape != (n_req,) or seq_lens.dtype != torch.int32: return False if pre_idx.dtype != torch.int32 or output_indices.dtype != torch.int32: return False @@ -231,19 +237,23 @@ def bsx_topk( seq_lens: torch.Tensor, output_indices: torch.Tensor, top_k: int, + next_n: int = 1, + compress_ratio: int = 4, ) -> None: """Unified bsx tier dispatch, replicating gvr_topk_launch_batched. - logits [BS, npad] fp32 (npad multiple of 64; per-row tail beyond N_eff - may be garbage — masked in-kernel), pre_idx [BS, K] int32, seq_lens - [BS] int32 (request-level, uncompressed-token space), output_indices - [BS, K] int32. Caller must have passed :func:`is_bsx_supported`. + logits [BS, npad] fp32 (BS = num_requests * next_n; npad multiple of + 64; per-row tail beyond N_eff may be garbage — masked in-kernel), + pre_idx [BS // next_n, K] int32 (request-level), seq_lens + [BS // next_n] int32 (request-level, uncompressed-token space), + output_indices [BS, K] int32. Caller must have passed + :func:`is_bsx_supported`. """ bs, npad = logits.shape - key = (bs, npad, top_k) + key = (bs, npad, top_k, next_n, compress_ratio) fn = _DISPATCH_CACHE.get(key) if fn is None: - fn = _DISPATCH_CACHE[key] = _bind(bs, npad, top_k) + fn = _DISPATCH_CACHE[key] = _bind(bs, npad, top_k, next_n, compress_ratio) fn(logits, pre_idx, seq_lens, output_indices) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py index e6dd0d37c1db..25d15aed6a35 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py @@ -285,7 +285,12 @@ def gvr_reg_kernel( npad = cutlass.Int32(logits.shape[1]) logits_row = logits[row, None] - pre_idx_row = pre_idx[row, None] + # Hint sharing: pre_idx is request-level ([num_rows // next_n, K]) — + # see the tp tier for the in-tree run_one_row parity notes. + if cutlass.const_expr(self.next_n == 1): + pre_idx_row = pre_idx[row, None] + else: + pre_idx_row = pre_idx[row // cutlass.Int32(self.next_n), None] out_row = out_idx[row, None] row_addr = logits_row.iterator.toint() @@ -414,7 +419,25 @@ def gvr_reg_kernel( cbase = cutlass.Int32(0) m_gt = cutlass.Int32(-1) - self.phase1(logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs) + # cr==1 hint temporal shift (in-tree pre_idx_offset parity); the + # cr>1 branch keeps the exact pre-MTP trace — see the tp tier. + if cutlass.const_expr(self.compress_ratio == 1): + hint_off = (row % cutlass.Int32(self.next_n)) + cutlass.Int32(1) + self.phase1( + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=hint_off, + ) + else: + self.phase1( + logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs + ) # Pre-zero the P4 round-0 radix histogram now that P1 is done with # hist. The first exchange's cluster barrier (release) publishes it # before any CTA reaches the P3 collect that increments it in flight. @@ -805,18 +828,30 @@ def __call__( _LAUNCH_CACHE = {} -def _get_compiled(K: int, CS: int, TB: int, MAXV: int, AR: int): - key = (K, CS, TB, MAXV, AR) +def _get_compiled(K: int, CS: int, TB: int, MAXV: int, AR: int, next_n: int = 1, cr: int = 4): + key = (K, CS, TB, MAXV, AR, next_n, cr) compiled = _LAUNCH_CACHE.get(key) if compiled is None: kC = 8192 if K >= 2048 else 6144 - kern = GvrRegKernel(top_k=K, kC=kC, cluster_size=CS, ar=AR, maxv=MAXV, num_threads=TB) + kern = GvrRegKernel( + top_k=K, + kC=kC, + cluster_size=CS, + ar=AR, + maxv=MAXV, + num_threads=TB, + next_n=next_n, + compress_ratio=cr, + ) n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + # pre_idx is request-level (num_rows // next_n rows); keep the shared + # n_rows sym at next_n == 1 (identical compiled artifact to v1). + n_pre = n_rows if next_n == 1 else cute.sym_int() logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 ) pre_idx_fake = _crt.make_fake_compact_tensor( - cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + cutlass.Int32, (n_pre, K), stride_order=(1, 0), assumed_align=16 ) seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) out_fake = _crt.make_fake_compact_tensor( @@ -849,15 +884,17 @@ def reg_topk( tb: int, maxv: int, ar: int, + next_n: int = 1, + cr: int = 4, ) -> None: """CuTe DSL gvr_topk_reg. logits [BS, npad] fp32 (npad mult of 64; tail beyond each row's N_eff may be garbage — masked at the - register load), pre_idx [BS, K] int32 hint, seq_lens [BS] int32 - (uncompressed-token space), out [BS, K] int32. Variant params are - explicit — the dispatcher selects; compile cache keyed - (K, cs, tb, maxv, ar).""" + register load), pre_idx [BS // next_n, K] int32 hint (request-level), + seq_lens [BS // next_n] int32 (uncompressed-token space), out [BS, K] + int32. Variant params are explicit — the dispatcher selects; compile + cache keyed (K, cs, tb, maxv, ar, next_n, cr).""" npad = logits.shape[1] - key = (K, cs, tb, maxv, ar, npad) + key = (K, cs, tb, maxv, ar, npad, next_n, cr) fn = _FAST.get(key) if fn is None: assert npad % 64 == 0, f"npad {npad} not a multiple of 64" @@ -867,6 +904,6 @@ def reg_topk( ) kc = 8192 if K >= 2048 else 6144 assert npad > kc, f"npad {npad} <= kC {kc}: reg path has no trivial branch" - fn = _get_compiled(K, cs, tb, maxv, ar) + fn = _get_compiled(K, cs, tb, maxv, ar, next_n, cr) _FAST[key] = fn fn(logits, pre_idx, seq_lens, out) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index ed32b5e858e8..bce236be4e57 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -270,9 +270,11 @@ class GvrTpKernel: """CuTe DSL port of gvr_topk_tp (fp32, B200/B300). Ctor knobs mirror the CUDA template params plus the production - ``next_n`` / ``compress_ratio`` contract (v1 dispatcher guard pins - next_n=1, cr=4; the N_eff arithmetic is kept general to mirror the - in-tree ``run_one_row`` formula exactly). + ``next_n`` / ``compress_ratio`` contract (compile-time constexpr: + N_eff arithmetic, request-level hint-row sharing and the cr==1 hint + temporal shift all mirror the in-tree ``run_one_row`` / + ``phase1_preidx_stats`` formulas exactly; the next_n==1 / cr==4 hot + path traces identically to the v1 port). """ WARP_SIZE = 32 @@ -540,7 +542,18 @@ def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred # phase1: hint gather + stats + rung ladder from hint-value CCDF. # ------------------------------------------------------------------ @cute.jit - def phase1(self, logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs): + def phase1( + self, + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=None, # cr==1 temporal shift ((row % next_n) + 1); None => 0 + ): TB = cutlass.const_expr(self.num_threads) AR = cutlass.const_expr(self.ar) NWARP = cutlass.const_expr(self.num_warps) @@ -565,6 +578,12 @@ def phase1(self, logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminma # hint cannot collapse the ladder onto FLT_MAX and force a # pathological one-value-per-pass plateau descent. hidx = pre_idx_row[j] + # cr==1 temporal shift, mirroring the in-tree run_one_row + # pre_idx_offset ((row % next_n) + 1); shifted-out-of-range + # hints fall into the same clamp below. const_expr'd out on + # cr>1 builds (hint_off is None => identical trace). + if cutlass.const_expr(hint_off is not None): + hidx = hidx + hint_off if hidx < cutlass.Int32(0): hidx = cutlass.Int32(0) if hidx > n_eff - cutlass.Int32(1): @@ -868,7 +887,13 @@ def gvr_tp_kernel( npad = cutlass.Int32(logits.shape[1]) logits_row = logits[row, None] - pre_idx_row = pre_idx[row, None] + # Hint sharing: pre_idx is request-level ([num_rows // next_n, K]); + # the next_n MTP rows of one request share the same hint row — + # mirrors the in-tree run_one_row's pre_idx_row_idx = row // next_n. + if cutlass.const_expr(self.next_n == 1): + pre_idx_row = pre_idx[row, None] + else: + pre_idx_row = pre_idx[row // cutlass.Int32(self.next_n), None] out_row = out_idx[row, None] row_addr = logits_row.iterator.toint() @@ -992,9 +1017,27 @@ def gvr_tp_kernel( self.collect_at(row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc) cute.arch.barrier() else: - self.phase1( - logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs - ) + # cr==1 hint temporal shift (in-tree pre_idx_offset parity: + # (row % next_n) + 1 maps prev-step indices into this step's + # KV space); cr>1 keeps the exact pre-MTP trace (no offset + # value is even computed). + if cutlass.const_expr(self.compress_ratio == 1): + hint_off = (row % cutlass.Int32(self.next_n)) + cutlass.Int32(1) + self.phase1( + logits_row, + pre_idx_row, + n_eff, + tidx, + s_hist, + s_fwred, + s_hminmax, + s_rungs, + hint_off=hint_off, + ) + else: + self.phase1( + logits_row, pre_idx_row, n_eff, tidx, s_hist, s_fwred, s_hminmax, s_rungs + ) hmin_floor = s_rungs[AR - 1] span0 = cute.arch.fmax(s_hminmax[1] - s_hminmax[0], cutlass.Float32(1e-3)) # P2a: sampled ladder count -> pivot pick. The exchanged rows @@ -1547,18 +1590,31 @@ def _p2floor(x: int) -> int: return p -def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int): - key = (K, CS, AR, UF, TB) +def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, cr: int = 4): + key = (K, CS, AR, UF, TB, next_n, cr) compiled = _LAUNCH_CACHE.get(key) if compiled is None: kC = 8192 if K >= 2048 else 6144 - kern = GvrTpKernel(top_k=K, kC=kC, cluster_size=CS, ar=AR, uf=UF, num_threads=TB) + kern = GvrTpKernel( + top_k=K, + kC=kC, + cluster_size=CS, + ar=AR, + uf=UF, + num_threads=TB, + next_n=next_n, + compress_ratio=cr, + ) n_rows, n_cols, n_batch = cute.sym_int(), cute.sym_int(), cute.sym_int() + # pre_idx is request-level: num_rows // next_n rows. At next_n == 1 + # keep the shared n_rows sym (identical compiled artifact to the v1 + # port); at next_n > 1 the row counts differ, so use a distinct sym. + n_pre = n_rows if next_n == 1 else cute.sym_int() logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (n_rows, n_cols), stride_order=(1, 0), assumed_align=16 ) pre_idx_fake = _crt.make_fake_compact_tensor( - cutlass.Int32, (n_rows, K), stride_order=(1, 0), assumed_align=16 + cutlass.Int32, (n_pre, K), stride_order=(1, 0), assumed_align=16 ) seq_lens_fake = _crt.make_fake_compact_tensor(cutlass.Int32, (n_batch,), stride_order=(0,)) out_fake = _crt.make_fake_compact_tensor( @@ -1594,7 +1650,7 @@ def tp_cluster_size(bs: int, npad: int) -> int: return cs -def _dispatch(K: int, bs: int, npad: int): +def _dispatch(K: int, bs: int, npad: int, next_n: int = 1, cr: int = 4): """launch_tp<512> selection: CS by co-residency + slice floor, UF by depth.""" assert npad % 64 == 0 TB = 512 @@ -1604,21 +1660,27 @@ def _dispatch(K: int, bs: int, npad: int): uf = 4 else: uf = 8 if npad >= 16384 else 4 - return _get_compiled(K, cs, AR, uf, TB) + return _get_compiled(K, cs, AR, uf, TB, next_n, cr) def tp_topk( - logits: torch.Tensor, pre_idx: torch.Tensor, seq_lens: torch.Tensor, out: torch.Tensor, K: int + logits: torch.Tensor, + pre_idx: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + K: int, + next_n: int = 1, + cr: int = 4, ) -> None: """CuTe DSL gvr_topk_tp. logits [BS, npad] fp32 (npad mult of 64; tail beyond each row's N_eff may be garbage — masked in-kernel), - pre_idx [BS, K] int32 hint, seq_lens [BS] int32 (uncompressed-token - space), out [BS, K] int32. Launch-time CS/UF selection replicates - launch_tp<512>.""" + pre_idx [BS // next_n, K] int32 hint (request-level), seq_lens + [BS // next_n] int32 (uncompressed-token space), out [BS, K] int32. + Launch-time CS/UF selection replicates launch_tp<512>.""" sh = logits.shape - key = (K, sh[0], sh[1]) + key = (K, sh[0], sh[1], next_n, cr) fn = _FAST.get(key) if fn is None: - fn = _dispatch(K, sh[0], sh[1]) + fn = _dispatch(K, sh[0], sh[1], next_n, cr) _FAST[key] = fn fn(logits, pre_idx, seq_lens, out) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index ae129d9deb6a..efbb0cef4061 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -14,14 +14,18 @@ # limitations under the License. """BSX (op43 direct/reg/tp CuTe DSL tiers) top-K decode tests. -CI-sized exactness grid for the guarded fp32/next_n=1/cr=4 fast path inside -``trtllm::cute_dsl_gvr_topk_decode``: every reg launch-table instance once, -the direct and tp tiers at a few npad each, ragged (varlen) rows with -POISONED tails (stale-garbage simulation: +1e30 beyond N_eff, which would -dominate the top-K if the ragged-N masking were broken), quantized-tie -inputs, degenerate rows, pre_idx hardening, host-only route-table asserts, -and a dispatcher-fallback check (bf16 / next_n=2 route to the in-tree -kernel). +CI-sized exactness grid for the guarded fp32 fast path inside +``trtllm::cute_dsl_gvr_topk_decode`` (next_n >= 1, cr in {1, 4}): every reg +launch-table instance once, the direct and tp tiers at a few npad each, +ragged (varlen) rows with POISONED tails (stale-garbage simulation: +1e30 +beyond N_eff, which would dominate the top-K if the ragged-N masking were +broken), quantized-tie inputs, degenerate rows, pre_idx hardening, +host-only route-table asserts, a dispatcher-fallback check (bf16 routes to +the in-tree kernel), and the MTP axis: next_n in {2, 3, 4} x cr in {1, 4} +x all three tiers, each case checked against BOTH the torch.topk host +N_eff/offset simulation (``_tie_aware_check``) and a differential in-tree +arm (same inputs through the in-tree kernel via ``order_row``, per-row +value-multiset equality). """ import pytest @@ -41,8 +45,8 @@ reason=f"CuTe DSL BSX Top-K only supports SM 100/103, got SM {get_sm_version()}", ) -CR = 4 # v1 dispatcher guard: compress_ratio == 4 (DSv4) -NEXT_N = 1 # v1 dispatcher guard: next_n == 1 +CR = 4 # default axis of the legacy (pre-MTP) cases: compress_ratio == 4 (DSv4) +NEXT_N = 1 # default axis of the legacy (pre-MTP) cases: next_n == 1 # --------------------------------------------------------------------------- @@ -558,27 +562,240 @@ def test_bsx_dispatcher_fallback_bf16(): @skip_not_sm100 -def test_bsx_dispatcher_fallback_next_n2(): - """next_n=2 (cr=1, V3.2-style MTP rows) must fall back to the in-tree - kernel and produce its results.""" - bs, npad, top_k, next_n, cr = 4, 65536, 2048, 2, 1 - num_rows = bs * next_n - torch.manual_seed(23) - torch.cuda.manual_seed(23) - logits = (torch.randn(num_rows, npad, device="cuda") * 2.0).contiguous() - seq_lens = torch.full((bs,), npad, dtype=torch.int32, device="cuda") - # cr=1 kernel convention: it reads logits[pre_idx + (row % next_n) + 1]. - eff = npad - next_n # safe hint range for every row +def test_bsx_dispatcher_fallback_bad_shapes(): + """Host-only: MTP contract violations must fall back to the in-tree + kernel (which asserts on them) rather than mis-launch a bsx tier.""" + bs, npad, top_k = 4, 65536, 512 + logits = torch.randn(bs * 2, npad, dtype=torch.float32, device="cuda") + seq_lens = torch.full((bs,), npad * 4, dtype=torch.int32, device="cuda") pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") - pre_idx[:, 0] = logits[::next_n, :eff].argmax(dim=-1).int() - 1 - pre_idx[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs * 2, top_k, dtype=torch.int32, device="cuda") + ok = bsx_dispatch.is_bsx_supported + # next_n=2 with request-level pre_idx/seq_lens: accepted. + assert ok(logits, pre_idx, seq_lens, out, top_k, 2, 4, None, None) + # cr outside {1, 4}: rejected. + assert not ok(logits, pre_idx, seq_lens, out, top_k, 2, 2, None, None) + # num_rows not divisible by next_n: rejected. + assert not ok(logits, pre_idx, seq_lens, out, top_k, 3, 4, None, None) + # row-level (non-request-level) pre_idx under next_n=2: rejected. + pre_row = torch.zeros(bs * 2, top_k, dtype=torch.int32, device="cuda") + assert not ok(logits, pre_row, seq_lens, out, top_k, 2, 4, None, None) + # row-level seq_lens under next_n=2: rejected. + sl_row = torch.full((bs * 2,), npad * 4, dtype=torch.int32, device="cuda") + assert not ok(logits, pre_idx, sl_row, out, top_k, 2, 4, None, None) + + +# --------------------------------------------------------------------------- +# MTP (next_n > 1) + cr in {1, 4}: exactness on all three tiers, checked +# against BOTH the torch.topk host N_eff/offset simulation +# (``_tie_aware_check``) and a differential in-tree arm — the SAME inputs +# through the in-tree kernel (forced via ``order_row``, which the bsx guard +# rejects), per-row value-multiset equality. Ragged: per-request varlen +# seq_lens make N_eff differ across requests AND (via row % next_n) across +# the MTP rows of one request; tails are poisoned with +1e30. +# --------------------------------------------------------------------------- +def _n_eff_rows(seq_lens, num_rows, next_n, cr): + r = torch.arange(num_rows, device="cuda") + sl = seq_lens.to(device="cuda", dtype=torch.long)[r // next_n] + return (sl - next_n + (r % next_n) + 1) // cr + + +def _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed, kind, preidx, hit_rate=0.5, argmax_slot0=True +): + """(logits [bs_req*next_n, npad] fp32, pre_idx [bs_req, K] int32, + seq_lens [bs_req] int32) with poisoned per-row tails. ``preidx``: + 'noised' = per-request true-top-K hints (host-simulated cr==1 temporal + offset: hint = ref_idx - 1) mixed with junk at ``hit_rate``; 'random' = + out-of-range garbage the kernels must clamp; 'zeros' = cold start. + + ``argmax_slot0=True`` (default) enforces the op contract + ``pre_idx[..., 0] = per-group argmax`` (over the min-N_eff window, + mirroring the in-tree test suite). The in-tree kernel REQUIRES this + invariant for exactness — the differential arm is only valid with it. + The bsx tiers do not require it (clamp hardening); pass False for + bsx-only robustness cases.""" + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + num_rows = bs_req * next_n + logits = torch.randn(num_rows, npad, dtype=torch.float32, device="cuda") * 2.0 + if kind == "ties": + logits = (logits * 4.0).round() * 0.25 + lo = (top_k + 1) * cr + next_n # every row non-degenerate (N_eff > K) + seq_lens = torch.randint(lo, npad * cr + 1, (bs_req,), dtype=torch.int32, device="cuda") + n_eff = _n_eff_rows(seq_lens, num_rows, next_n, cr) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + if preidx == "zeros": + pre_idx = torch.zeros(bs_req, top_k, dtype=torch.int32, device="cuda") + elif preidx == "random": + pre_idx = torch.randint(-npad, 4 * npad, (bs_req, top_k), dtype=torch.int32, device="cuda") + else: # noised + valid = torch.where(tail, torch.full_like(logits, float("-inf")), logits) + ref = valid[::next_n].topk(top_k, dim=-1).indices.int() + # host-side temporal-offset simulation: the cr==1 kernels read + # logits[hint + (row % next_n) + 1], so a prev-step hint for a + # top value at index i is (i - 1) for the first MTP row. + off = 1 if cr == 1 else 0 + good = (ref - off).clamp(min=0) + keep = torch.rand(good.shape, device="cuda") < hit_rate + junk = torch.arange(top_k, dtype=torch.int32, device="cuda").expand(bs_req, -1) + pre_idx = torch.where(keep, good, junk).contiguous() + if argmax_slot0: + min_ne = int(n_eff.min().item()) + pre_idx[:, 0] = logits[::next_n, :min_ne].argmax(dim=-1).int() + return logits, pre_idx, seq_lens + + +def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): + """bsx arm + in-tree differential arm (order_row forces the in-tree + sort path; the bsx guard rejects order_row). Returns (out_bsx, out_ref).""" + num_rows = logits.shape[0] out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + assert bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None + ), "expected the bsx fast path to accept this MTP call" + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr + ) + order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() + out_ref = torch.empty_like(out) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_ref, + top_k=top_k, + next_n=next_n, + compress_ratio=cr, + order_row=order_row, + ) + torch.cuda.synchronize() + return out, out_ref - assert not bsx_dispatch.is_bsx_supported( + +_MTP_TIER_CELLS = [ + # (tier, npad, bs_req_base, top_k) — bs_req is scaled so num_rows = + # bs_req * next_n stays in the tier's route band for every next_n. + ("direct", 4096, 4, 512), + ("reg", 24576, 1, 2048), # reg(cs=4,tb=512,maxv=4,ar=8) at num_rows < 16 + ("tp", 65536, 8, 1024), # tp takeover at num_rows >= 16 (npad >= 32768) +] + +_MTP_KINDS = [ + ("randn", "random"), # random logits + OOR-garbage hints (clamp hardening) + ("randn", "noised"), # realistic noised hints (+ host offset simulation) + ("randn", "zeros"), # all-zero cold-start hints + ("ties", "noised"), # quantized tie plateaus +] + + +@skip_not_sm100 +@pytest.mark.parametrize("next_n", [2, 3, 4]) +@pytest.mark.parametrize("cr", [1, 4]) +@pytest.mark.parametrize( + "tier,npad,bs_req,top_k", _MTP_TIER_CELLS, ids=[c[0] for c in _MTP_TIER_CELLS] +) +@pytest.mark.parametrize("kind,preidx", _MTP_KINDS, ids=[f"{k}-{p}" for k, p in _MTP_KINDS]) +def test_bsx_mtp_exactness(next_n, cr, tier, npad, bs_req, top_k, kind, preidx): + num_rows = bs_req * next_n + if tier == "tp" and num_rows < 16: + bs_req = (16 + next_n - 1) // next_n + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + assert bsx_dispatch.route(num_rows, npad, top_k).startswith(tier) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed=npad + 13 * next_n + cr, kind=kind, preidx=preidx + ) + out, out_ref = _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr) + # Independent torch.topk + host N_eff/offset simulation reference. + _tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + # Differential oracle: per-row value multiset equal to the in-tree arm. + sel = torch.gather(logits, -1, out.long()).sort(-1, descending=True).values + sel_ref = torch.gather(logits, -1, out_ref.long()).sort(-1, descending=True).values + assert torch.equal(sel, sel_ref), ( + f"bsx vs in-tree value-multiset mismatch (next_n={next_n}, cr={cr}, tier={tier})" + ) + + +@skip_not_sm100 +@pytest.mark.parametrize("next_n", [2, 4]) +@pytest.mark.parametrize("cr", [1, 4]) +@pytest.mark.parametrize("preidx", ["zeros", "random"]) +def test_bsx_mtp_preidx_hardening(next_n, cr, preidx): + """bsx-only MTP robustness: hint sets that VIOLATE the op's argmax- + slot-0 contract (pure zeros / out-of-range garbage) must neither fault + nor break exactness on the bsx tiers. No differential arm here: the + in-tree kernel requires the argmax invariant for exactness, the bsx + tiers deliberately do not (clamp hardening).""" + bs_req, npad, top_k = 2, 65536, 1024 # reg tier at num_rows 4/8 + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, + next_n, + cr, + npad, + top_k, + seed=19 + next_n + cr, + kind="randn", + preidx=preidx, + argmax_slot0=False, + ) + out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + assert bsx_dispatch.is_bsx_supported( logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None - ), "next_n=2 must NOT take the bsx fast path" + ) torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr ) torch.cuda.synchronize() _tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + + +@skip_not_sm100 +@pytest.mark.parametrize("next_n", [2, 3]) +@pytest.mark.parametrize("cr", [1, 4]) +def test_bsx_mtp_degenerate_rows(next_n, cr): + """Degenerate MTP requests (N_eff <= K for some or all of the next_n + rows): identity emit [0..N_eff-1] + -1 pad must hold per ROW, with the + per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr.""" + bs_req, npad, top_k = 8, 65536, 1024 + num_rows = bs_req * next_n + _skip_if_cluster_capped(num_rows, npad, top_k) + logits, pre_idx, seq_lens = _make_mtp_inputs( + bs_req, next_n, cr, npad, top_k, seed=17, kind="randn", preidx="noised" + ) + # Requests 0..3: degenerate / boundary — N_eff spans 0, 1, K-1, K, K+1 + # across their MTP rows. + for i, sl in enumerate([next_n, cr + next_n, (top_k - 1) * cr + next_n - 1, top_k * cr]): + seq_lens[i] = sl + n_eff = _n_eff_rows(seq_lens, num_rows, next_n, cr) + col = torch.arange(npad, device="cuda") + tail = col[None, :] >= n_eff.clamp(min=0)[:, None] + logits = torch.where(tail, torch.full_like(logits, 1e30), logits) + + out, out_ref = _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr) + for r in range(num_rows): + ne = int(n_eff[r].item()) + if ne <= top_k: + expect = torch.full((top_k,), -1, dtype=torch.int32, device="cuda") + if ne > 0: + expect[:ne] = torch.arange(ne, dtype=torch.int32, device="cuda") + assert torch.equal(out[r], expect), ( + f"row {r} (N_eff={ne}): degenerate identity emit mismatch" + ) + else: + # Per-row torch reference: collapse the row's MTP arithmetic + # into an equivalent next_n=1 seq_lens so the shared checker + # scans exactly N_eff(r) columns. + sl_row = seq_lens[r // next_n : r // next_n + 1] - next_n + (r % next_n) + 1 + _tie_aware_check( + out[r : r + 1], logits[r : r + 1], sl_row, top_k, next_n=1, compress_ratio=cr + ) + nd = (n_eff > top_k).nonzero().flatten() + if nd.numel(): + sel = torch.gather(logits[nd], -1, out[nd].long()).sort(-1, descending=True).values + sel_ref = torch.gather(logits[nd], -1, out_ref[nd].long()).sort(-1, descending=True).values + assert torch.equal(sel, sel_ref) From 9535fd9cfe0954994ea11aaeeb04696549770958 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:26:42 +0000 Subject: [PATCH 05/19] [None][perf] bsx tp: lean-pivot admission (stage-0c CI override + per-K interpolated tpush) and mask-hoisted streaming loops Ships the op43-pr6 campaign candidate (measured head 1a88ede77d) as a clean patch on the pr4 head: - stage 0c LEAN-PIVOT OVERRIDE (pick-only): on stage-0 CI-admission miss with npad in [16384, 262144], a strictly leaner ladder rung whose clustering-aware 1.5-sigma CI lands in [K, kC] becomes the pivot; undershoot costs one rescue re-stream (pr4 economics preserved). - per-K interpolated lean tpush (log2-count interpolation between pick and next-tighter rung; target 1.5K, 2.0K for K=512 sampling noise; npad <= 98304 gate) with original-pick rescue. - mask-hoisted streaming loops in count/max-below/fused-count-collect/ collect (unmasked main loop + masked vec-tail) - uniform micro-gain, independently exonerated of the pr5 regression by the C5 fingerprint ablation. - candidate budget pinned at the flat pr4 kC (8192 for K>=2048, else 6144); wide pr4 ladder quantiles pinned. Full-grid evidence (865 real decode cells x 11 BS, same-rep cold-L2 paired vs the in-tree kernel): gm 1.4170 (pr4 head 1.4120), win 88.0%, <0.909 cases 517 -> from 636. Exactness: 140-test production suite + adversarial + synthetic gates all PASS; measured module and this file are AST-identical modulo dead scaffolding removal. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 281 ++++++++++++++++-- 1 file changed, 264 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index bce236be4e57..859d8aa85146 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -63,6 +63,8 @@ ``cluster_arrive_relaxed`` DSMEM-race history in the GVR campaign notes). """ +import math + import cutlass import cutlass.cute as cute import cutlass.cute.math as cmath @@ -76,6 +78,30 @@ RUNGS = 8 MAXPASS = 8 SS = 32 # P2a sample stride (float4s) + +# pr6 DATA-ADAPTIVE admission (op43-pr6 campaign, shipped head). +# This is the pr4 admission machinery +# BYTE-PARITY on every streaming pass, the fused R=2 count-collect, the +# reuse rule, the P4 select and the P2c driver (iter1's dedicated window +# rungs and iter2's R=4 window both taxed the whole kernel 4-15% +# globally — even trivial-branch cells — via register pressure/extra +# compares/P4 prefilter, and iter1 additionally carried an adversarial +# under-emit bug). The adaptivity is a PICK-ONLY delta: +# * stage 0c LEAN-PIVOT OVERRIDE: when the stage-0 CI admission FAILED +# (the pick in hand is a band/2-sigma/overshoot fallback: fat or +# undershoot-prone), npad <= 262144, and a strictly LEANER ladder +# rung's sampled count lands in [K, kC] by its clustering-aware +# 1.5-sigma CI on both sides, that rung becomes the pivot (the +# rescue is recomputed from it by the standard next-fatter rule). +# A fat band pick (est ~3-4x K, the P3-push/P4-radix fatness the +# op43-pr6 D2 probe identified as the whole residual-band gap) is +# replaced by a lean CI-backed pivot; an undershoot costs ONE rescue +# re-stream — pr4's own economics. +# * kC pinned to the pr4 flat budget 8192 (K>=2048) / 6144 — the pr6 +# iron rule (the K-scaled diet was falsified as pure harm). +# * ladder quantiles pinned WIDE (pr4): abl3 showed the re-placed +# spread under this admission is a net residual harm. +# * the C4 occupancy CS cut keeps the fix-head default (neutral). FLT_MAX = 3.4028234663852886e38 INF = float("inf") @@ -213,6 +239,24 @@ def _fmin_f32(a, b, *, loc=None, ip=None): ) +@dsl_user_op +def _lg2_f32(a, *, loc=None, ip=None): + """lg2.approx.f32 — thread-0 serial pick section only (lean interp).""" + return cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [a.ir_value(loc=loc, ip=ip)], + "lg2.approx.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + @cute.jit def _f2u_bits(u): """Order-preserving fp32-bits -> Uint32 key: u ^ (sign ? 0xFFFFFFFF : 0x80000000).""" @@ -293,6 +337,10 @@ def __init__( assert num_threads % 32 == 0 assert ar in (6, 8) assert cluster_size in (1, 2, 4, 8) + # pr6 D2a-lean2 interpolation target (per-K: K512 rows carry + # 5-20x sparser P2a samples, so 2.0K keeps interpolation-error + # margin; K>=1024 takes the measured-win 1.5K). + self.lean_tgt = (2 * top_k) if top_k < 1024 else (3 * top_k) // 2 self.top_k = top_k self.kC = kC self.cluster_size = cluster_size @@ -371,17 +419,34 @@ def count_pass( frags = [ cute.make_fragment((4,), cutlass.Float32) for _ in range(U) ] # Python-unrolled register batch + # Mask hoist: float4s below vmain = min(v1, n_eff >> 2) are fully + # valid (4i+3 < n_eff), so the main loop runs mask-free (the + # per-element gi compare + select was ~10% of whole-kernel + # instructions); only [vmain, v1) — the n_eff boundary float4 and + # the pad tail — takes the masked epilogue. Bit-exact: the mask + # only changes values at gi >= n_eff, all of which live in + # [vmain, v1). + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull i = v0 + tidx - while i + cutlass.Int32((U - 1) * TB) < v1: + while i + cutlass.Int32((U - 1) * TB) < vmain: for u in cutlass.range_constexpr(U): self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) for u in cutlass.range_constexpr(U): - gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) for q in cutlass.range_constexpr(4): - v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + v = frags[u][q] for r in cutlass.range_constexpr(R): cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) i = i + cutlass.Int32(U * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + i = i + cutlass.Int32(TB) while i < v1: self._ld_float4(copy_atom, row_addr, i, frags[0]) gi = i << cutlass.Int32(2) @@ -495,17 +560,27 @@ def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred m = cutlass.Float32(-FLT_MAX) # op43 S-E: explicit U=4 batched loads (CUDA `float4 a[4]` idiom). frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull i = v0 + tidx - while i + cutlass.Int32(3 * TB) < v1: + while i + cutlass.Int32(3 * TB) < vmain: for u in cutlass.range_constexpr(4): self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) for u in cutlass.range_constexpr(4): - gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) for q in cutlass.range_constexpr(4): - v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + v = frags[u][q] if v < t_hi_bound: m = cute.arch.fmax(m, v) i = i + cutlass.Int32(4 * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + if v < t_hi_bound: + m = cute.arch.fmax(m, v) + i = i + cutlass.Int32(TB) while i < v1: self._ld_float4(copy_atom, row_addr, i, frags[0]) gi = i << cutlass.Int32(2) @@ -672,8 +747,16 @@ def phase1( x = _shfl_up_add(x, tidx, o) A = x - Ssum binw2 = (hmax - tlow) * cutlass.Float32(1.0 / 64.0) + # pr6 D1: WIDE (pr4) quantile spread, pinned — the + # P2A path must stay pr4-parity (abl3: the re-placed + # spread under this admission is a net residual harm). if cutlass.const_expr(AR == 6): - qt = ((K * 15) // 100, (K * 40) // 100, (K * 70) // 100, (K * 92) // 100) + qt = ( + (K * 15) // 100, + (K * 40) // 100, + (K * 70) // 100, + (K * 92) // 100, + ) else: qt = ( (K * 10) // 100, @@ -786,23 +869,53 @@ def fused_count_collect( a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) # Explicit U-batched loads (CUDA `float4 a[U]` idiom), main + - # vec-tail while-loops (same op43 S-E fix as count_pass). + # vec-tail while-loops (same op43 S-E fix as count_pass). Mask + # hoist as in count_pass: [v0, vmain) mask-free (gi computed only + # inside the rare push branch), [vmain, v1) masked epilogue. frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(U)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull i = v0 + tidx - while i + cutlass.Int32((U - 1) * TB) < v1: + while i + cutlass.Int32((U - 1) * TB) < vmain: for u in cutlass.range_constexpr(U): self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) for u in cutlass.range_constexpr(U): - gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) for q in cutlass.range_constexpr(4): - v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + v = frags[u][q] for r in cutlass.range_constexpr(R): cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) if v >= tpush: self._push_cand( - a_cnt, a_st, s_cand, s_isc, kcap, True, v, gi + cutlass.Int32(q) + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + True, + v, + ((i + cutlass.Int32(u * TB)) << cutlass.Int32(2)) + cutlass.Int32(q), ) i = i + cutlass.Int32(U * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + for r in cutlass.range_constexpr(R): + cnt[r] = cnt[r] + cutlass.Int32(v >= tr[r]) + if v >= tpush: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + True, + v, + (i << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(TB) while i < v1: self._ld_float4(copy_atom, row_addr, i, frags[0]) gi = i << cutlass.Int32(2) @@ -838,19 +951,45 @@ def collect_at(self, row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc): # 1-deep loop was the dominant stall site (36% of all warp-stall # samples) on cells whose reuse check fails at big npad x big BS. frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] + vmain = v1 + vfull = n_eff >> cutlass.Int32(2) + if vmain > vfull: + vmain = vfull i = v0 + tidx - while i + cutlass.Int32(3 * TB) < v1: + while i + cutlass.Int32(3 * TB) < vmain: for u in cutlass.range_constexpr(4): self._ld_float4(copy_atom, row_addr, i + cutlass.Int32(u * TB), frags[u]) for u in cutlass.range_constexpr(4): - gi = (i + cutlass.Int32(u * TB)) << cutlass.Int32(2) for q in cutlass.range_constexpr(4): - v = _mask_tail(frags[u][q], gi + cutlass.Int32(q), n_eff) + v = frags[u][q] if v >= thr: self._push_cand( - a_cnt, a_st, s_cand, s_isc, kcap, False, v, gi + cutlass.Int32(q) + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + False, + v, + ((i + cutlass.Int32(u * TB)) << cutlass.Int32(2)) + cutlass.Int32(q), ) i = i + cutlass.Int32(4 * TB) + while i < vmain: + self._ld_float4(copy_atom, row_addr, i, frags[0]) + for q in cutlass.range_constexpr(4): + v = frags[0][q] + if v >= thr: + self._push_cand( + a_cnt, + a_st, + s_cand, + s_isc, + kcap, + False, + v, + (i << cutlass.Int32(2)) + cutlass.Int32(q), + ) + i = i + cutlass.Int32(TB) while i < v1: self._ld_float4(copy_atom, row_addr, i, frags[0]) gi = i << cutlass.Int32(2) @@ -1128,6 +1267,13 @@ def gvr_tp_kernel( bestd = cutlass.Int32(0) best = cutlass.Int32(j) adm_est = cnt_j * cutlass.Int32(SS) + # pr6 D1: record the stage-0 CI-admission outcome + # BEFORE stages 0b/1/2 mutate bestd; a stage-0 admit + # is the "confident P2A" signal — those rows never + # take the stage-0c override. + s0_ok = cutlass.Int32(0) + if bestd == cutlass.Int32(0): + s0_ok = cutlass.Int32(1) # Stage 0b — legacy band pick (min |est - tgt| in # [1.5K, 0.6kC], EXACTLY the pr1 stage-1 rule), then # combine: with no admission it is taken as-is (pr1 @@ -1195,6 +1341,64 @@ def gvr_tp_kernel( if dd < bestd: bestd = dd best = cutlass.Int32(j) + # ---- pr6 D1 stage 0c: lean-pivot override ---- + # When the stage-0 CI admission FAILED, npad <= + # 262144, and some ladder rung is (i) strictly LEANER + # than the pick in hand and (ii) lands in [K, kC] by + # its own clustering-aware 1.5-sigma CI on BOTH sides + # (a looser window than stage 0's [K, 0.6kC] at + # 1.85/2.0-sigma — rows passing THAT never get here), + # make IT the pivot. Rungs are descending in j: scan + # ascending and keep the FIRST (tightest) qualifier. + # Hint-unrepresentative rows (the 44f0a208a9 harm + # band) have no CI-qualifying rung and keep the pr4 + # pick. Same sqrt/div-free margin algebra as stage 0; + # everything downstream (fused R=2, rescue, reuse, + # P4, secant) is byte-parity pr4. + # npad floor 16384: below it the stride-32 sample sees + # <= ~128 float4s, the occ-aware CI fires on noise and + # misfire rescues cost 3-5% (v32_8k probe, pr6d1d); + # the absolute win-room there is small anyway. + gate_ok = cutlass.Int32(0) + if npad >= cutlass.Int32(16384): + if npad <= cutlass.Int32(262144): + gate_ok = cutlass.Int32(1) + if gate_ok != cutlass.Int32(0): + if s0_ok == cutlass.Int32(0): + est_pick = cutlass.Int32(0x7FFFFFFF) + if bestd != cutlass.Int32(0x7FFFFFFF): + est_pick = (s_rcnt[best] & cutlass.Int32(0xFFFF)) * cutlass.Int32( + SS + ) + jw = cutlass.Int32(AR) + for w_ in cutlass.range_constexpr(AR): + if jw == cutlass.Int32(AR): + cnt_w = s_rcnt[w_] & cutlass.Int32(0xFFFF) + if cnt_w > cutlass.Int32(0): + estw = cnt_w * cutlass.Int32(SS) + if estw < est_pick: + occ_w = s_rcnt[w_] >> cutlass.Int32(16) + if occ_w < cutlass.Int32(1): + occ_w = cutlass.Int32(1) + foccw = cutlass.Float32(occ_w) + festw = cutlass.Float32(estw) + few2 = festw * festw + a_low = festw - cutlass.Float32(float(K)) + bhiw = cutlass.Float32(float(kC)) - festw + if a_low >= cutlass.Float32(0.0): + if ( + a_low * a_low * foccw + >= cutlass.Float32(2.25) * few2 + ): + if bhiw >= cutlass.Float32(0.0): + if ( + cutlass.Float32(2.25) * few2 + <= bhiw * bhiw * foccw + ): + jw = cutlass.Int32(w_) + if jw < cutlass.Int32(AR): + best = jw + bestd = cutlass.Int32(0) tp_ = s_rungs[best] # Rescue rung: the next FATTER ladder rung below the # pivot. If the pivot's exact count lands under K @@ -1205,6 +1409,43 @@ def gvr_tp_kernel( resc_ = hmin_floor if best < cutlass.Int32(AR - 1): resc_ = s_rungs[best + cutlass.Int32(1)] + # ---- pr6 D2a lean pivot (ported from pr6d2 lean2, + # composed AFTER stage 0c): the pick in hand — band, + # 2-sigma, overshoot fallback, or a stage-0c lean rung + # that is still fat — is structurally FAT on the + # 32k-128k residual band (exact Cp 2.4-4.7x K; P3 + # pushes + P4 radix pay for it). When the picked + # rung's sampled est >= 1.8K, interpolate tpush + # between the pick and the next tighter rung with + # est < lean_tgt, targeting count ~= lean_tgt + # (log2-count interpolation: tails are ~exponential, + # the linear form undershoots). Rescue = the ORIGINAL + # pick, whose exact count the fused pass computes + # anyway: an undershoot costs one collect re-stream, + # cheap under the npad <= 98304 gate; the 256k-1024k + # guard band keeps the stock pick bit-for-bit. + if npad <= cutlass.Int32(98304): + estp = (s_rcnt[best] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if estp >= cutlass.Int32((9 * K) // 5): + jm = cutlass.Int32(-1) + estm = cutlass.Int32(1) + for j in cutlass.range_constexpr(AR): + if cutlass.Int32(j) < best: + cj = (s_rcnt[j] & cutlass.Int32(0xFFFF)) * cutlass.Int32(SS) + if cj < cutlass.Int32(self.lean_tgt): + jm = cutlass.Int32(j) + estm = cj + if jm >= cutlass.Int32(0): + if estm < cutlass.Int32(1): + estm = cutlass.Int32(1) + fp = _lg2_f32(cutlass.Float32(estp)) + fm = _lg2_f32(cutlass.Float32(estm)) + den = fp - fm + if den < cutlass.Float32(1e-6): + den = cutlass.Float32(1e-6) + frac = (fp - cutlass.Float32(math.log2(self.lean_tgt))) / den + resc_ = tp_ + tp_ = tp_ + (s_rungs[jm] - tp_) * frac s_rungs[0] = tp_ s_rungs[1] = resc_ cute.arch.barrier() @@ -1594,6 +1835,8 @@ def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, c key = (K, CS, AR, UF, TB, next_n, cr) compiled = _LAUNCH_CACHE.get(key) if compiled is None: + # pr6 iron rule: pr4 flat candidate budget, never the K-scaled + # diet (falsified as pure harm). kC = 8192 if K >= 2048 else 6144 kern = GvrTpKernel( top_k=K, @@ -1638,7 +1881,11 @@ def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, c def tp_cluster_size(bs: int, npad: int) -> int: - """CS selection of launch_tp<512>: co-residency + slice floor.""" + """CS selection of launch_tp<512>: co-residency + slice floor. + Occupancy cut for mid-N large-BS (measured neutral-to-positive).""" + # C4 occupancy cut kept at fix-head default (ablation: neutral). + if npad < 65536 and bs >= 64: + return 1 cs = 1 if bs < 128: cs = _p2floor(296 // bs) if bs <= 296 else 1 From 22b7af18dcb3b9ab77b14174e8c603b6dddbd56c Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:29:17 +0000 Subject: [PATCH 06/19] [None][perf] bsx dispatch: measured fallback band table to the in-tree kernel Routes the (npad, bs) buckets where the op43-pr6 full-grid verdict found at least one production layer >1.10x slower on every bsx tier than the in-tree kernel (L2-resident mid-N regime: the in-tree exact-count ladder admits a leaner candidate set and its row-slice cluster split keeps all CTAs busy through P4). With the table the operator is a strict Pareto improvement over the in-tree kernel: routed buckets run at parity, everything else keeps the bsx win (full-grid gm 1.40 vs the in-tree head, worst case capped at 1.10x). npad keys resolve by nearest power of two; TRTLLM_BSX_FALLBACK_BANDS=0 disables the table. Band data: 865 real decode cells x 11 BS same-rep cold-L2 nsys pairs, 2026-07-28. Tests: kernel-contract cases pin the table off via an autouse fixture (they exist to exercise the bsx tiers); the table itself is covered by test_bsx_fallback_band_table (bucket membership, nearest-pow2 resolution, neighbour non-routing, kill-switch), and CUDA-graph capture/replay by test_bsx_cuda_graph_capture_replay on both sides of the table (host-side dispatch bakes at capture; replay stays exact over in-place rewritten inputs). Suite: 143 passed; sibling test_cute_dsl_gvr_topk_decode suite: 671 passed / 144 skipped. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_bsx_dispatch.py | 34 ++++++++ .../sparse/test_cute_dsl_bsx_topk_decode.py | 86 +++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index 655881712de4..a40b232eef88 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -82,6 +82,38 @@ def _thresholds(npad): return tpb, dnb +# Measured fallback band table (op43-pr6 full-grid verdict, 2026-07-28): +# (npad, bs) buckets where the in-tree kernel is faster than every bsx +# tier by >1.10x on at least one production layer (865 real decode cells +# x 11 BS, same-rep cold-L2 nsys pairs; data +# op43_bsx_cutedsl/results/pr6_band_table.json, recalibration recipe in +# the campaign notes). These are the L2-resident mid-N shapes where the +# in-tree exact-count ladder admits a leaner candidate set and its +# row-slice cluster split keeps every CTA busy through P4. Routing them +# to the in-tree kernel caps the worst case at 1.10x while keeping the +# bsx win elsewhere (full-grid gm 1.40 vs the in-tree head). +# Keys are the nearest power-of-two of npad; values are inclusive bs +# ranges. TRTLLM_BSX_FALLBACK_BANDS=0 disables the table (bsx serves +# every guarded shape). +_FALLBACK_BANDS = { + 8192: (256, 1 << 30), + 16384: (256, 1 << 30), + 32768: (16, 1 << 30), + 65536: (16, 1 << 30), + 131072: (16, 255), + 262144: (16, 127), +} + + +def _in_fallback_band(bs: int, npad: int) -> bool: + if _env_threshold("TRTLLM_BSX_FALLBACK_BANDS") == _BIG: # "0" -> off + return False + up = 1 << max(npad - 1, 1).bit_length() # pow2 >= npad + np2 = up if 4 * npad >= 3 * up else up >> 1 # arithmetic-midpoint nearest + band = _FALLBACK_BANDS.get(np2) + return band is not None and band[0] <= bs <= band[1] + + # tier name -> (kind, params). reg params = (cs, tb, maxv, ar). def route(bs: int, npad: int, K: int) -> str: """Tier the CUDA gvr_topk_launch_batched would take. Returns @@ -205,6 +237,8 @@ def is_bsx_supported( n_req = bs // next_n if npad > 262144 or npad % 64 != 0: return False + if _in_fallback_band(bs, npad): + return False if not ( logits.is_contiguous() and pre_idx.is_contiguous() diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index efbb0cef4061..b4dffc9fba3c 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -49,6 +49,92 @@ NEXT_N = 1 # default axis of the legacy (pre-MTP) cases: next_n == 1 +@pytest.fixture(autouse=True) +def _bsx_bands_off(monkeypatch): + """Kernel-contract tests must reach the bsx tiers: disable the measured + fallback band table (it legitimately routes several tested (npad, bs) + shapes to the in-tree kernel in production). The table itself is covered + by ``test_bsx_fallback_band_table``.""" + monkeypatch.setenv("TRTLLM_BSX_FALLBACK_BANDS", "0") + bsx_dispatch._reset_env_cache() + yield + monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) + bsx_dispatch._reset_env_cache() + + +@skip_not_sm100 +@pytest.mark.parametrize("bands", ["on", "off"]) +def test_bsx_cuda_graph_capture_replay(monkeypatch, bands): + """The op must be CUDA-graph capturable and replay-consistent on both + sides of the fallback band table. Shape (npad=32768, bs=64) sits inside + a fallback bucket: bands=on captures the in-tree branch, bands=off the + bsx tp branch — the dispatch decision is host-side per shape, so it + bakes into the graph at capture time and must hold across replays, + including replays over in-place rewritten inputs.""" + if bands == "on": + # the autouse fixture pinned the table off; restore the default + monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) + bsx_dispatch._reset_env_cache() + bs, npad, top_k = 64, 32768, 2048 + logits, pre_idx, seq_lens = _make_bsx_inputs(bs, npad, top_k, seed=1234) + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + def call(): + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre_idx, seq_lens, out, top_k=top_k, next_n=NEXT_N, compress_ratio=CR + ) + + # warmup outside capture: JIT compile + any lazy init, on a side stream + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(2): + call() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + call() + g.replay() + torch.cuda.synchronize() + _tie_aware_check(out, logits, seq_lens, top_k) + + # replay over in-place rewritten inputs (fresh logits + fresh hints) + logits2, pre_idx2, seq_lens2 = _make_bsx_inputs(bs, npad, top_k, seed=4321) + logits.copy_(logits2) + pre_idx.copy_(pre_idx2) + seq_lens.copy_(seq_lens2) + g.replay() + torch.cuda.synchronize() + _tie_aware_check(out, logits, seq_lens, top_k) + + +def test_bsx_fallback_band_table(monkeypatch): + """The measured (npad, bs) fallback buckets route to the in-tree kernel + by default; the kill-switch restores bsx service; neighbours are not + over-routed. Bands: op43-pr6 full-grid verdict (2026-07-28).""" + monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) + bsx_dispatch._reset_env_cache() + inb = bsx_dispatch._in_fallback_band + # routed buckets (bucket floor <0.909 vs the in-tree head) + assert inb(256, 8192) and inb(1024, 8192) + assert inb(16, 32768) and inb(1024, 65536) + assert inb(128, 131072) and inb(64, 262144) + assert inb(48, 40960) # off-grid npad resolves to the nearest pow2 + # neighbours stay on bsx + assert not inb(128, 8192) # direct/tp win band + assert not inb(8, 32768) # latency reg band + assert not inb(256, 131072) and not inb(128, 262144) # large-N tp win + assert not inb(1024, 4096) # small-npad tp win + # kill-switch + monkeypatch.setenv("TRTLLM_BSX_FALLBACK_BANDS", "0") + bsx_dispatch._reset_env_cache() + assert not inb(64, 32768) + monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) + bsx_dispatch._reset_env_cache() + + # --------------------------------------------------------------------------- # Shared helpers. ``_tie_aware_check`` is a local copy of the one in # test_cute_dsl_gvr_topk_decode.py (same directory): the sibling test module From 78013ba542588b20f7cec23591f3f676fbe160b0 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:56:43 +0000 Subject: [PATCH 07/19] [None][perf] bsx dispatch: recalibrate 131072 fallback band lower bound to bs=8 The original band-table calibration measured the reg tier through a harness whose GvrTpKernel base class carried a faster experimental streaming phase1 (GvrRegKernel inherits phase1 and the streaming-load helpers from GvrTpKernel; the measurement arm rebound only the tp entry point, so reg-routed shapes silently rode the experimental base). On this branch's reg tier the 128K x BS8 shapes measure 0.82-0.91x vs the in-tree kernel on 5 production layers (865-cell x 11-BS full-grid, same-rep cold-L2 nsys pairs), below the 0.909 floor the table guarantees. Routing the (131072, 8..15) bucket to the in-tree kernel restores the floor; full-grid gm 1.397 -> 1.390. All other buckets re-verified on the shipped kernel (direct 1.0031 / reg 0.9998 vs its true pr4 baseline / tp 1.0008). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_bsx_dispatch.py | 8 +++++++- .../attention/sparse/test_cute_dsl_bsx_topk_decode.py | 4 +++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index a40b232eef88..ebe6cab99cd4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -95,12 +95,18 @@ def _thresholds(npad): # Keys are the nearest power-of-two of npad; values are inclusive bs # ranges. TRTLLM_BSX_FALLBACK_BANDS=0 disables the table (bsx serves # every guarded shape). +# 2026-07-29 recalibration: the 131072 band lower bound moved 16 -> 8. +# The original calibration run measured the reg tier with a faster +# experimental streaming phase1 inherited from a side branch; on this +# branch's reg tier the 128K x BS8 shapes dip to 0.82-0.91x vs the +# in-tree kernel (5 production layers), so they are routed as well +# (full-grid gm 1.397 -> 1.390). _FALLBACK_BANDS = { 8192: (256, 1 << 30), 16384: (256, 1 << 30), 32768: (16, 1 << 30), 65536: (16, 1 << 30), - 131072: (16, 255), + 131072: (8, 255), 262144: (16, 127), } diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index b4dffc9fba3c..1b15fb3804f7 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -113,7 +113,8 @@ def call(): def test_bsx_fallback_band_table(monkeypatch): """The measured (npad, bs) fallback buckets route to the in-tree kernel by default; the kill-switch restores bsx service; neighbours are not - over-routed. Bands: op43-pr6 full-grid verdict (2026-07-28).""" + over-routed. Bands: op43-pr6 full-grid verdict (2026-07-28), 131072 + lower bound recalibrated to 8 (2026-07-29).""" monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) bsx_dispatch._reset_env_cache() inb = bsx_dispatch._in_fallback_band @@ -121,6 +122,7 @@ def test_bsx_fallback_band_table(monkeypatch): assert inb(256, 8192) and inb(1024, 8192) assert inb(16, 32768) and inb(1024, 65536) assert inb(128, 131072) and inb(64, 262144) + assert inb(8, 131072) # 128K x BS8 recalibration (reg-tier floor) assert inb(48, 40960) # off-grid npad resolves to the nearest pow2 # neighbours stay on bsx assert not inb(128, 8192) # direct/tp win band From e0401dd4d5b6feec1e8a0e43680c8924c7d223f5 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 02:34:15 +0000 Subject: [PATCH 08/19] [None][test] bsx suite: cut CI wall-clock 46% by deduplicating JIT compile variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite's cost is DSL JIT compiles (15-34s per constexpr variant), not case count. Measured on B200: 570s -> 306s, zero code-path coverage lost: - MTP exactness: the 3-tier x next_n{2,3,4} x cr{1,4} cross (18 variant compiles of BOTH arms, 385s) becomes an explicit 8-combo list — full {2,3} x {1,4} cross on tp (the most complex MTP arithmetic), complementary (next_n, cr) diagonals on direct/reg. next_n=4 dropped: 2 covers even row-sharing, 3 covers odd division (and is the production MTP depth). All 4 input kinds kept (they reuse compiles). - reg launch table: all 14 route asserts kept (host-only, free); live launches reduced to the 6 instances that together cover every codegen axis value (cs {1,2,8,16} x tb {512,1024} x ar {6,8} x maxv {5,8} + the dense-knob-only path). - MTP preidx hardening / admission hit-rate: cells re-pinned onto already-compiled variants (same route/cs class, npad is a runtime parameter) so they add zero compiles. - bf16 dispatcher fallback: guard-only (the bf16 in-tree execution it falls back to is exhaustively covered by the sibling gvr suite; recompiling that variant here cost ~10s for no added coverage). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../sparse/test_cute_dsl_bsx_topk_decode.py | 112 +++++++++++------- 1 file changed, 70 insertions(+), 42 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index 1b15fb3804f7..5ea242d805d4 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -360,38 +360,44 @@ def test_bsx_route_env_knobs(monkeypatch): # --------------------------------------------------------------------------- -# reg tier: every launch-table instance once (cs, tb, maxv, ar), ragged rows -# with poisoned tails. The two dense cs=2/cs=4 instances are reachable only -# through the TRTLLM_BSX_DENSE_BS knob (default bands route around them), -# which doubles as the env-knob dispatch test on a live launch. +# reg tier: every launch-table instance's ROUTE is asserted (host-only, +# free); a LAUNCH (JIT compile ~5-7s each) runs only for the 6 instances +# that together cover every codegen axis value (cs {1,2,8,16} x tb +# {512,1024} x ar {6,8} x maxv {5,8} + the dense-knob-only path) — the +# dropped launches differ only in axis combinations, not in code paths. +# The cs=2 dense instance is reachable only through the TRTLLM_BSX_DENSE_BS +# knob (default bands route around it), doubling as the env-knob dispatch +# test on a live launch. # --------------------------------------------------------------------------- _REG_INSTANCES = [ - # (npad, bs, K, dense_env, expected tier) - (14336, 1, 512, None, "reg(cs=1,tb=512,maxv=8,ar=8)"), - (24576, 1, 512, None, "reg(cs=4,tb=512,maxv=4,ar=8)"), - (49152, 1, 512, None, "reg(cs=8,tb=512,maxv=3,ar=8)"), - (65536, 1, 1024, None, "reg(cs=8,tb=512,maxv=4,ar=8)"), - (131072, 8, 512, None, "reg(cs=8,tb=512,maxv=8,ar=8)"), - (163840, 1, 2048, None, "reg(cs=16,tb=512,maxv=5,ar=6)"), - (163840, 4, 512, None, "reg(cs=16,tb=512,maxv=5,ar=8)"), - (262144, 1, 2048, None, "reg(cs=16,tb=512,maxv=8,ar=8)"), - (262144, 4, 1024, None, "reg(cs=16,tb=512,maxv=8,ar=6)"), - (20480, 64, 512, None, "reg(cs=1,tb=1024,maxv=5,ar=8)"), - (28672, 64, 512, None, "reg(cs=1,tb=1024,maxv=8,ar=8)"), - (65536, 8, 512, "8", "reg(cs=2,tb=1024,maxv=8,ar=8)"), - (131072, 8, 1024, "8", "reg(cs=4,tb=1024,maxv=8,ar=8)"), - (262144, 8, 2048, None, "reg(cs=8,tb=1024,maxv=8,ar=8)"), + # (npad, bs, K, dense_env, expected tier, launch) + (14336, 1, 512, None, "reg(cs=1,tb=512,maxv=8,ar=8)", True), + (24576, 1, 512, None, "reg(cs=4,tb=512,maxv=4,ar=8)", False), + (49152, 1, 512, None, "reg(cs=8,tb=512,maxv=3,ar=8)", False), + (65536, 1, 1024, None, "reg(cs=8,tb=512,maxv=4,ar=8)", False), + (131072, 8, 512, None, "reg(cs=8,tb=512,maxv=8,ar=8)", True), + (163840, 1, 2048, None, "reg(cs=16,tb=512,maxv=5,ar=6)", True), + (163840, 4, 512, None, "reg(cs=16,tb=512,maxv=5,ar=8)", False), + (262144, 1, 2048, None, "reg(cs=16,tb=512,maxv=8,ar=8)", False), + (262144, 4, 1024, None, "reg(cs=16,tb=512,maxv=8,ar=6)", False), + (20480, 64, 512, None, "reg(cs=1,tb=1024,maxv=5,ar=8)", False), + (28672, 64, 512, None, "reg(cs=1,tb=1024,maxv=8,ar=8)", True), + (65536, 8, 512, "8", "reg(cs=2,tb=1024,maxv=8,ar=8)", True), + (131072, 8, 1024, "8", "reg(cs=4,tb=1024,maxv=8,ar=8)", False), + (262144, 8, 2048, None, "reg(cs=8,tb=1024,maxv=8,ar=8)", True), ] @skip_not_sm100 @pytest.mark.parametrize( - "npad,bs,top_k,dense_env,expected", + "npad,bs,top_k,dense_env,expected,launch", _REG_INSTANCES, ids=[t[4] + f"_n{t[0]}_bs{t[1]}_k{t[2]}" for t in _REG_INSTANCES], ) @pytest.mark.parametrize("kind", ["randn", "ties"]) -def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, kind, monkeypatch): +def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, launch, kind, monkeypatch): + if not launch and kind == "ties": + pytest.skip("route-assert-only instance (single kind suffices)") _skip_if_cluster_capped(bs, npad, top_k) try: if dense_env is not None: @@ -401,8 +407,9 @@ def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, kind, monkey bs, npad, top_k, seed=npad + bs + top_k, kind=kind, varlen=True ) _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected) - out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + if launch: + out = _run_op(logits, pre_idx, seq_lens, top_k) + _tie_aware_check(out, logits, seq_lens, top_k) finally: if dense_env is not None: monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) @@ -537,7 +544,10 @@ def test_bsx_preidx_hardening(npad, bs, top_k, preidx): "npad,bs,top_k", [ (65536, 16, 1024), # cs=8 cluster tp (pro-like) - (32832, 16, 2048), # cs=4 cluster tp (v32 32k production shape) + # cs=8 cluster tp, K=2048 (v32-like); npad pinned to 65536 so the + # JIT variant is shared with the overflow test (route/cs and the + # compiled kernel key on npad-derived cs, not on npad itself) + (65536, 16, 2048), (131136, 128, 512), # cs=1 streaming tp (flash 512k production shape) ], ) @@ -642,11 +652,10 @@ def test_bsx_dispatcher_fallback_bf16(): assert not bsx_dispatch.is_bsx_supported( logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None ), "bf16 must NOT take the bsx fast path" - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - logits, pre_idx, seq_lens, out, top_k=top_k, next_n=NEXT_N, compress_ratio=CR - ) - torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k) + # Guard-only on purpose (host, no JIT): the bf16 in-tree execution the + # op falls back to is exhaustively covered by + # test_cute_dsl_gvr_topk_decode.py; compiling that variant here again + # costs ~10s of CI for no added coverage. @skip_not_sm100 @@ -763,12 +772,31 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): return out, out_ref -_MTP_TIER_CELLS = [ - # (tier, npad, bs_req_base, top_k) — bs_req is scaled so num_rows = +_MTP_TIER_CELLS = { + # tier -> (npad, bs_req_base, top_k) — bs_req is scaled so num_rows = # bs_req * next_n stays in the tier's route band for every next_n. - ("direct", 4096, 4, 512), - ("reg", 24576, 1, 2048), # reg(cs=4,tb=512,maxv=4,ar=8) at num_rows < 16 - ("tp", 65536, 8, 1024), # tp takeover at num_rows >= 16 (npad >= 32768) + "direct": (4096, 4, 512), + "reg": (24576, 1, 2048), # reg(cs=4,tb=512,maxv=4,ar=8) at num_rows < 16 + "tp": (65536, 8, 1024), # tp takeover at num_rows >= 16 (npad >= 32768) +} + +# (tier, next_n, cr) — every (next_n, cr) constexpr pair is a separate JIT +# compile of BOTH arms (~15-34s each), so the full 3-tier x {2,3,4} x {1,4} +# cross is prohibitively slow in CI. Coverage kept: the tp tier (the most +# complex MTP arithmetic: cluster exchange + admission escape) runs the +# full {2,3} x {1,4} cross; direct/reg run complementary (next_n, cr) +# diagonals so each tier still sees odd/even next_n and both cr values. +# next_n=4 is dropped: 2 covers the even/row-sharing arithmetic, 3 covers +# odd division (and is the production MTP depth). +_MTP_COMBOS = [ + ("tp", 2, 1), + ("tp", 2, 4), + ("tp", 3, 1), + ("tp", 3, 4), + ("direct", 2, 4), + ("direct", 3, 1), + ("reg", 2, 1), + ("reg", 3, 4), ] _MTP_KINDS = [ @@ -780,13 +808,12 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): @skip_not_sm100 -@pytest.mark.parametrize("next_n", [2, 3, 4]) -@pytest.mark.parametrize("cr", [1, 4]) @pytest.mark.parametrize( - "tier,npad,bs_req,top_k", _MTP_TIER_CELLS, ids=[c[0] for c in _MTP_TIER_CELLS] + "tier,next_n,cr", _MTP_COMBOS, ids=[f"{t}-{n}-{c}" for t, n, c in _MTP_COMBOS] ) @pytest.mark.parametrize("kind,preidx", _MTP_KINDS, ids=[f"{k}-{p}" for k, p in _MTP_KINDS]) -def test_bsx_mtp_exactness(next_n, cr, tier, npad, bs_req, top_k, kind, preidx): +def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx): + npad, bs_req, top_k = _MTP_TIER_CELLS[tier] num_rows = bs_req * next_n if tier == "tp" and num_rows < 16: bs_req = (16 + next_n - 1) // next_n @@ -808,16 +835,17 @@ def test_bsx_mtp_exactness(next_n, cr, tier, npad, bs_req, top_k, kind, preidx): @skip_not_sm100 -@pytest.mark.parametrize("next_n", [2, 4]) -@pytest.mark.parametrize("cr", [1, 4]) +@pytest.mark.parametrize("next_n,cr", [(2, 1), (3, 4)]) @pytest.mark.parametrize("preidx", ["zeros", "random"]) def test_bsx_mtp_preidx_hardening(next_n, cr, preidx): """bsx-only MTP robustness: hint sets that VIOLATE the op's argmax- slot-0 contract (pure zeros / out-of-range garbage) must neither fault nor break exactness on the bsx tiers. No differential arm here: the in-tree kernel requires the argmax invariant for exactness, the bsx - tiers deliberately do not (clamp hardening).""" - bs_req, npad, top_k = 2, 65536, 1024 # reg tier at num_rows 4/8 + tiers deliberately do not (clamp hardening). Cell and (next_n, cr) + pairs match the reg combos of ``test_bsx_mtp_exactness`` so the JIT + variants are reused (zero extra compiles).""" + bs_req, npad, top_k = 1, 24576, 2048 # reg tier at num_rows 2/3 num_rows = bs_req * next_n _skip_if_cluster_capped(num_rows, npad, top_k) logits, pre_idx, seq_lens = _make_mtp_inputs( From 2a1950cf388f277a1227e7bbfca49b7f4afd72b5 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 08:28:45 +0000 Subject: [PATCH 09/19] [None][chore] bsx: scrub internal development codenames from comments Comment-only: replace internal iteration/campaign identifiers with neutral engineering descriptions (measured/ablation/baseline), keep every load-bearing note (convergence constraints, measured costs, falsified alternatives, recalibration recipe) intact. No code change. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_bsx_direct.py | 10 +-- .../top_k/gvr_topk_decode_bsx_dispatch.py | 16 ++--- .../top_k/gvr_topk_decode_bsx_reg.py | 10 +-- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 64 +++++++++---------- .../sparse/test_cute_dsl_bsx_topk_decode.py | 6 +- 5 files changed, 53 insertions(+), 53 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py index 6fb8558cafaf..098a110cc1fc 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py @@ -14,8 +14,8 @@ """BSX direct (short-row) top-K tier — CuTe DSL, Blackwell SM100. -Port of the op43 ``ct_direct.py`` CuTe DSL translation of the CUDA -``direct_topk_kernel`` (op42 gvr_bsx.cu), adapted for the production +CuTe DSL translation of the original CUDA ``direct_topk_kernel``, +adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract (ragged N via a device ``seq_lens`` tensor + per-row degenerate identity emit; see the module docstring of ``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory). @@ -70,7 +70,7 @@ class DirectTopKKernel: N_eff arithmetic (constexpr; the direct tier reads no hints, so MTP support is the N_eff formula alone). SMEM layout mirrors gvr_bsx.cu's DSmem with the same packed (key << 32 | idx) - u64 layout (op43 S-E round 6: split key/idx arrays cost 2x smem + u64 layout (measured: split key/idx arrays cost 2x smem transactions in collect + emits vs CUDA's single STS.64 / LDS.64 per candidate): @@ -339,7 +339,7 @@ def direct_topk_kernel( ) row_addr = logits_row.iterator.toint() vcnt = npad >> cutlass.Int32(2) # npad % 4 == 0 (host-asserted) - # op43 S-E round 6: npad <= DKCMAX bounds the per-thread trip count + # npad <= DKCMAX bounds the per-thread trip count # at U = DKCMAX/4/TB (= 3 for TB 1024), so the whole grid-stride # loop is flattened into one predicated register batch: issue ALL # (<=3) float4 loads back-to-back, then consume. A dynamic @@ -390,7 +390,7 @@ def direct_topk_kernel( # C == K: every candidate is admitted; identity emit. (With # seq_lens present this is unreachable — n_eff <= npad == K # lands in the degenerate branch — but kept for structural - # parity with the op43/CUDA source.) + # parity with the CUDA source.) io = tidx while io < npad: out_row[io] = cutlass.Int32( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index ebe6cab99cd4..b44f4dd9ddb1 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -13,10 +13,10 @@ # limitations under the License. """BSX top-K tier dispatcher — routes ``cute_dsl_gvr_topk_decode`` calls to -the op43 CuTe DSL tiers (direct / reg / tp). +the BSX CuTe DSL tiers (direct / reg / tp). -Port of the op43 ``ct_bsx.py`` unified launcher, itself an exact -transcription of the CUDA dispatch in op42 ``gvr_bsx.cu`` +Exact transcription of the original CUDA implementation's dispatch +(``gvr_topk_launch_batched``) (``gvr_topk_launch_batched`` + ``launch_dense`` + ``launch_tp<512>``). The gvr streaming tier (cs=16 fallback for npad > 262144) is intentionally NOT ported: it is unreachable inside the deployment envelope @@ -82,12 +82,12 @@ def _thresholds(npad): return tpb, dnb -# Measured fallback band table (op43-pr6 full-grid verdict, 2026-07-28): +# Measured fallback band table (full-grid calibration, 2026-07-28): # (npad, bs) buckets where the in-tree kernel is faster than every bsx # tier by >1.10x on at least one production layer (865 real decode cells -# x 11 BS, same-rep cold-L2 nsys pairs; data -# op43_bsx_cutedsl/results/pr6_band_table.json, recalibration recipe in -# the campaign notes). These are the L2-resident mid-N shapes where the +# x 11 BS, same-rep cold-L2 nsys pairs; to recalibrate, re-run that +# paired sweep and route every (npad, bs) bucket whose per-case floor +# drops below 0.909). These are the L2-resident mid-N shapes where the # in-tree exact-count ladder admits a leaner candidate set and its # row-slice cluster split keeps every CTA busy through P4. Routing them # to the in-tree kernel caps the worst case at 1.10x while keeping the @@ -185,7 +185,7 @@ def route_cluster_size(bs: int, npad: int, K: int) -> int: # Per-call route()+_parse_reg() string work costs ~2.5-3us host-submit wall -# on <20us kernels (op43 S-E). The routing decision is pure in (bs, npad, K) +# on <20us kernels (measured). The routing decision is pure in (bs, npad, K) # [env thresholds are cached at first use, like the CUDA static locals], so # bind it once per key to a closure. _DISPATCH_CACHE = {} # (bs, npad, K, next_n, cr) -> callable(logits, pre, seq_lens, out) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py index 25d15aed6a35..c5dc1637bce3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py @@ -14,9 +14,9 @@ """BSX register-resident (reg) GVR Top-K tier — CuTe DSL, Blackwell SM100. -Port of the op43 ``ct_reg.py`` CuTe DSL translation of the CUDA +CuTe DSL translation of the CUDA ``gvr_topk_reg`` register-resident GVR top-K kernel -(op42 iter12-F1b), adapted for the production +(tuned CUDA head), adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract (see the module docstring of ``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory: ragged-N masking, pre_idx clamping, per-row degenerate identity emit). @@ -52,7 +52,7 @@ implicit cluster barrier before ret in DSMEM kernels, the DSL does not (missing it => timing-dependent CUDA 719 faults). -op43 S-E convergence notes preserved in this port (do not "simplify away"): +Convergence notes preserved in this port (do not "simplify away"): 1. ``cutlass.utils.distributed.atomicAdd`` lowers to sys-scoped ``atom.relaxed.SYS.shared`` which blocks ptxas' warp aggregation; the CTA-scope relaxed ``_atomic_add_cta`` restores the fast path. @@ -302,7 +302,7 @@ def gvr_reg_kernel( # Single packed (key<<32 | idx) u64 candidate array — CUDA's # `unsigned long long cand[kC]`; one 8B DSMEM push per candidate in # the P3 scatter (halves remote-store transactions under 8-cluster - # contention at BS>=8 — op43 S-E round 2). + # contention at BS>=8 — measured). s_cand = smem.allocate_tensor( element_type=cutlass.Uint64, layout=cute.make_ordered_layout((kC,), order=(0,)), @@ -782,7 +782,7 @@ def gvr_reg_kernel( out_row[m + p] = idx ie = ie + cutlass.Int32(TB) - # Exit rendezvous — PLATEAU PATH ONLY (op43 S-E wave-overlap fix). + # Exit rendezvous — PLATEAU PATH ONLY (wave-overlap fix). # On the P3/P4 path the post-merge _cluster_sync_aligned() above is # the LAST remote-op rendezvous (P4 is CTA0-local smem + global # stores only), so non-rank0 CTAs exit right after it — mirroring diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index 859d8aa85146..a569d5acf652 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -14,12 +14,12 @@ """BSX throughput (tp) GVR Top-K tier — CuTe DSL, Blackwell SM100. -Port of the op43 ``ct_tp.py`` CuTe DSL translation of the CUDA -``gvr_topk_tp`` throughput GVR top-K kernel (op42 iter12-F1b), +CuTe DSL translation of the CUDA ``gvr_topk_tp`` +throughput GVR top-K kernel (tuned CUDA head), adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract: * RAGGED N: production logits tails beyond the per-row valid length are - stale garbage (NOT -FLT_MAX pad, unlike the op-bench harness). A device + stale garbage (NOT -FLT_MAX pad, unlike the standalone bench harness). A device ``seq_lens`` tensor is threaded into every tier and every global row read is predicated at the LANE level: element index >= N_eff substitutes -FLT_MAX AFTER the (always in-bounds, buffer width = logits.shape[1]) @@ -60,7 +60,7 @@ + atom.relaxed.cluster.shared::cluster.add.u32. Writer-side cluster syncs use FULL cluster_arrive (release) — never relaxed (a relaxed arrive has no release semantics, so peer CTAs could observe stale DSMEM: see -``cluster_arrive_relaxed`` DSMEM-race history in the GVR campaign notes). +``cluster_arrive_relaxed`` DSMEM race observed in development). """ import math @@ -79,13 +79,13 @@ MAXPASS = 8 SS = 32 # P2a sample stride (float4s) -# pr6 DATA-ADAPTIVE admission (op43-pr6 campaign, shipped head). -# This is the pr4 admission machinery +# DATA-ADAPTIVE admission (measured full-grid verdict). +# This is the baseline admission machinery # BYTE-PARITY on every streaming pass, the fused R=2 count-collect, the -# reuse rule, the P4 select and the P2c driver (iter1's dedicated window -# rungs and iter2's R=4 window both taxed the whole kernel 4-15% +# reuse rule, the P4 select and the P2c driver (a dedicated-window-rung variant +# rungs and a later R=4 window variant both taxed the whole kernel 4-15% # globally — even trivial-branch cells — via register pressure/extra -# compares/P4 prefilter, and iter1 additionally carried an adversarial +# compares/P4 prefilter, and the former additionally carried an adversarial # under-emit bug). The adaptivity is a PICK-ONLY delta: # * stage 0c LEAN-PIVOT OVERRIDE: when the stage-0 CI admission FAILED # (the pick in hand is a band/2-sigma/overshoot fallback: fat or @@ -94,12 +94,12 @@ # 1.5-sigma CI on both sides, that rung becomes the pivot (the # rescue is recomputed from it by the standard next-fatter rule). # A fat band pick (est ~3-4x K, the P3-push/P4-radix fatness the -# op43-pr6 D2 probe identified as the whole residual-band gap) is +# mechanism probe identified as the whole residual-band gap) is # replaced by a lean CI-backed pivot; an undershoot costs ONE rescue -# re-stream — pr4's own economics. -# * kC pinned to the pr4 flat budget 8192 (K>=2048) / 6144 — the pr6 -# iron rule (the K-scaled diet was falsified as pure harm). -# * ladder quantiles pinned WIDE (pr4): abl3 showed the re-placed +# re-stream — the baseline's own economics. +# * kC pinned to the flat budget 8192 (K>=2048) / 6144 +# (the K-scaled diet was falsified as pure harm). +# * ladder quantiles pinned WIDE (baseline): ablation showed the re-placed # spread under this admission is a net residual harm. # * the C4 occupancy CS cut keeps the fix-head default (neutral). FLT_MAX = 3.4028234663852886e38 @@ -179,7 +179,7 @@ def _st_shared_cluster_i32(mapped_addr, val, *, loc=None, ip=None): def _st_shared_cluster_u64(mapped_addr, val, *, loc=None, ip=None): """One 8B DSMEM candidate push (CUDA: dst[pos] = (u64)key<<32 | idx). - op43 S-E round 4: a single packed u64 store halves remote-store + Measured: a single packed u64 store halves remote-store transactions vs two 4B pushes under CS>1 cluster contention (matches the CUDA arm's ``unsigned long long cand[kC]``). Do NOT split back into key/idx 4B stores — the split form measurably regresses the @@ -337,7 +337,7 @@ def __init__( assert num_threads % 32 == 0 assert ar in (6, 8) assert cluster_size in (1, 2, 4, 8) - # pr6 D2a-lean2 interpolation target (per-K: K512 rows carry + # Lean-pivot interpolation target (per-K: K512 rows carry # 5-20x sparser P2a samples, so 2.0K keeps interpolation-error # margin; K>=1024 takes the measured-win 1.5K). self.lean_tgt = (2 * top_k) if top_k < 1024 else (3 * top_k) // 2 @@ -405,7 +405,7 @@ def count_pass( s_rungs, s_ptcnt, ): - # op43 S-E: explicit U-batched loads (CUDA `float4 a[U]` idiom) — a + # Explicit U-batched loads (CUDA `float4 a[U]` idiom) — a # `cutlass.range(unroll=U)` loop leaves ONE load in flight per iter # and costs ~11% kernel time in the DRAM-bound regimes. Keep the # Python-unrolled register batch. @@ -558,7 +558,7 @@ def max_below_pass(self, row_addr, v0, v1, n_eff, t_hi_bound, par, tidx, s_fwred NWARP = cutlass.const_expr(self.num_warps) copy_atom = self._copy_atom() m = cutlass.Float32(-FLT_MAX) - # op43 S-E: explicit U=4 batched loads (CUDA `float4 a[4]` idiom). + # Explicit U=4 batched loads (CUDA `float4 a[4]` idiom). frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(4)] vmain = v1 vfull = n_eff >> cutlass.Int32(2) @@ -645,7 +645,7 @@ def phase1( hv[jj] = cutlass.Float32(0.0) if j < cutlass.Int32(K): # pre_idx hardening: clamp hints into [0, N_eff-1] (not - # [0, npad-1] as in the op-bench port). Production + # [0, npad-1] as in the standalone bench port). Production # cold-start is all zeros and arbitrary garbage must not # crash or corrupt; clamping (rather than skipping, which # the in-tree phase1_preidx_stats does) also keeps the CCDF @@ -747,8 +747,8 @@ def phase1( x = _shfl_up_add(x, tidx, o) A = x - Ssum binw2 = (hmax - tlow) * cutlass.Float32(1.0 / 64.0) - # pr6 D1: WIDE (pr4) quantile spread, pinned — the - # P2A path must stay pr4-parity (abl3: the re-placed + # WIDE (baseline) quantile spread, pinned — the + # P2A path must stay baseline-parity (ablation: the re-placed # spread under this admission is a net residual harm). if cutlass.const_expr(AR == 6): qt = ( @@ -869,7 +869,7 @@ def fused_count_collect( a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) # Explicit U-batched loads (CUDA `float4 a[U]` idiom), main + - # vec-tail while-loops (same op43 S-E fix as count_pass). Mask + # vec-tail while-loops (same fix as count_pass). Mask # hoist as in count_pass: [v0, vmain) mask-free (gi computed only # inside the rare push branch), [vmain, v1) masked epilogue. frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(U)] @@ -946,7 +946,7 @@ def collect_at(self, row_addr, v0, v1, n_eff, thr, tidx, s_cand, s_isc): if cutlass.const_expr(CS > 1): a_cnt = _mapa_shared_cluster(s_isc.iterator + cutlass.Int32(5), cutlass.Int32(0)) a_st = _mapa_shared_cluster(s_cand.iterator, cutlass.Int32(0)) - # op43 S-E round 3: explicit U=4 batched loads. nvcc auto-unrolls the + # Explicit U=4 batched loads. nvcc auto-unrolls the # CUDA collect_at loop 4x with 4 LDG.E.128 issued back-to-back; a # 1-deep loop was the dominant stall site (36% of all warp-stall # samples) on cells whose reuse check fails at big npad x big BS. @@ -1190,7 +1190,7 @@ def gvr_tp_kernel( cute.arch.barrier() if tidx == cutlass.Int32(0): # 4-stage pivot pick (R0-admission tightest-in-window -> - # iter7 band -> iter8b 2-sigma -> iter12-F1b gamble) + # band target -> 2-sigma bound -> gamble fallback) lo = cutlass.const_expr((3 * K) // 2) hi = cutlass.const_expr((6 * kC) // 10) tgt_py = min(max(3 * K, (3 * K) // 2), (6 * kC) // 10) @@ -1267,7 +1267,7 @@ def gvr_tp_kernel( bestd = cutlass.Int32(0) best = cutlass.Int32(j) adm_est = cnt_j * cutlass.Int32(SS) - # pr6 D1: record the stage-0 CI-admission outcome + # Record the stage-0 CI-admission outcome # BEFORE stages 0b/1/2 mutate bestd; a stage-0 admit # is the "confident P2A" signal — those rows never # take the stage-0c override. @@ -1341,7 +1341,7 @@ def gvr_tp_kernel( if dd < bestd: bestd = dd best = cutlass.Int32(j) - # ---- pr6 D1 stage 0c: lean-pivot override ---- + # ---- stage 0c: lean-pivot override ---- # When the stage-0 CI admission FAILED, npad <= # 262144, and some ladder rung is (i) strictly LEANER # than the pick in hand and (ii) lands in [K, kC] by @@ -1350,14 +1350,14 @@ def gvr_tp_kernel( # 1.85/2.0-sigma — rows passing THAT never get here), # make IT the pivot. Rungs are descending in j: scan # ascending and keep the FIRST (tightest) qualifier. - # Hint-unrepresentative rows (the 44f0a208a9 harm - # band) have no CI-qualifying rung and keep the pr4 + # Hint-unrepresentative rows (the measured harm + # band) have no CI-qualifying rung and keep the baseline # pick. Same sqrt/div-free margin algebra as stage 0; # everything downstream (fused R=2, rescue, reuse, - # P4, secant) is byte-parity pr4. + # P4, secant) is byte-parity with the baseline. # npad floor 16384: below it the stride-32 sample sees # <= ~128 float4s, the occ-aware CI fires on noise and - # misfire rescues cost 3-5% (v32_8k probe, pr6d1d); + # misfire rescues cost 3-5% (measured, npad~8K probe); # the absolute win-room there is small anyway. gate_ok = cutlass.Int32(0) if npad >= cutlass.Int32(16384): @@ -1409,7 +1409,7 @@ def gvr_tp_kernel( resc_ = hmin_floor if best < cutlass.Int32(AR - 1): resc_ = s_rungs[best + cutlass.Int32(1)] - # ---- pr6 D2a lean pivot (ported from pr6d2 lean2, + # ---- lean pivot (per-K interpolated push threshold, # composed AFTER stage 0c): the pick in hand — band, # 2-sigma, overshoot fallback, or a stage-0c lean rung # that is still fat — is structurally FAT on the @@ -1835,7 +1835,7 @@ def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, c key = (K, CS, AR, UF, TB, next_n, cr) compiled = _LAUNCH_CACHE.get(key) if compiled is None: - # pr6 iron rule: pr4 flat candidate budget, never the K-scaled + # Flat candidate budget, never the K-scaled # diet (falsified as pure harm). kC = 8192 if K >= 2048 else 6144 kern = GvrTpKernel( diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index 5ea242d805d4..9fd52095082f 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""BSX (op43 direct/reg/tp CuTe DSL tiers) top-K decode tests. +"""BSX (direct/reg/tp CuTe DSL tiers) top-K decode tests. CI-sized exactness grid for the guarded fp32 fast path inside ``trtllm::cute_dsl_gvr_topk_decode`` (next_n >= 1, cr in {1, 4}): every reg @@ -113,7 +113,7 @@ def call(): def test_bsx_fallback_band_table(monkeypatch): """The measured (npad, bs) fallback buckets route to the in-tree kernel by default; the kill-switch restores bsx service; neighbours are not - over-routed. Bands: op43-pr6 full-grid verdict (2026-07-28), 131072 + over-routed. Bands: full-grid calibration (2026-07-28), 131072 lower bound recalibrated to 8 (2026-07-29).""" monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) bsx_dispatch._reset_env_cache() @@ -307,7 +307,7 @@ def _skip_if_cluster_capped(bs, npad, top_k): # --------------------------------------------------------------------------- -# Host-only route-table asserts (mirror of the op42 gvr_bsx.cu dispatch). +# Host-only route-table asserts (mirror of the original CUDA dispatch). # --------------------------------------------------------------------------- def test_bsx_route_table(): r = bsx_dispatch.route From 41268973eb47fd613a2bf9865320e4ec3569056b Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:01:09 +0000 Subject: [PATCH 10/19] [None][chore] GVR decode: resolve the review follow-ups from #16457 Four of the five follow-up items committed to reviewers on #16457 (the fifth, the dispatch guard itself, is this PR's dispatcher): 1. Comment pruning (requested by @lfr-0531): measurement-history / tuning-provenance commentary reduced to invariants and contracts across the kernel and custom-op files. 2. Launch-shape policy single source (requested by @limin2021): pick_config is split into pick_cluster_size + pick_tuning on the kernel class (the single source of truth); the production runner's _pick_tuning becomes a thin adapter and its cluster auto-pick delegates to pick_cluster_size. The intentional shell divergence is kept and documented: the runner ASSERTS on a 32B-misaligned logits pointer (contract violation) while GvrTopKKernel.launch silently downgrades to 128-bit loads. New test test_..._pick_policy_single_source sweeps dtype x BS x N x graph-capture and pins runner == kernel policy. 3. 16-bit exact-tail (requested by @mingyangHao): p4_exact_tail now defaults ON for fp16/bf16 as well - candidate keys are ALWAYS fp32 (16-bit inputs are upcast injectively at collect), so the tail radix re-rank on the full fp32 order key is exact for every dtype; the overclaiming 'fully resolved' docstring is corrected. New adversarial test: two distinct 16-bit values (1.0 vs 1.25) in one fine bin straddling the K boundary under a wide Phase-2 bracket, fp16 + bf16. 4. P4 exact-tail radix de-duplication (requested by @mingyangHao): the two verbatim copies (tiny-tie fast path's large-class fallback and the plain exact-tail path; token-identical, 1162 tokens) collapse into one @cute.jit helper _p4_exact_tail_radix_select. Verified: the p4_tail_fast=False variant compiles to BYTE-IDENTICAL PTX before and after (465,875 bytes, CUTE_DSL_KEEP=ptx). Remaining item (plateau undershoot terminal routing to an exact tie-aware fallback) follows as its own commit: the audit found the rank-scatter path currently has no cand_count < K branch at all, so the fix is wider than the review comment assumed and deserves isolated review. Gates: full sparse-attention suite 674 passed / 144 skipped (includes the two new tests); PTX identity proof above. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 111 +-- .../blackwell/top_k/gvr_topk_decode.py | 744 ++++++++---------- .../sparse/test_cute_dsl_gvr_topk_decode.py | 68 ++ 3 files changed, 441 insertions(+), 482 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 2c3090b437d5..316d5a70297e 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7296,77 +7296,37 @@ def _pick_tuning( max_seq_len: Optional[int], data_ptr: int, ) -> dict: - """Pick T / V / min_blocks_per_mp tuning knobs shared by - single-CTA / sort and LB compile paths. Returned keys match - ``_compile`` / ``_compile_lb`` param names for ``**tuning`` - spreading. + """Adapter over :meth:`GvrTopKKernel.pick_tuning` (the single + source of truth for the T / V / min_blocks_per_mp / + warp-reduce policy), shared by the single-CTA / sort and LB + compile paths. Returned keys match ``_compile`` / + ``_compile_lb`` param names for ``**tuning`` spreading. + + Intentional shell divergence from ``GvrTopKKernel.launch``: + a 32B-misaligned logits pointer is a CONTRACT VIOLATION here + (assert), while ``launch`` silently downgrades to 128-bit + loads (dev convenience for ad-hoc tensors). """ - enable_unroll_4 = True - enable_phase3_unroll = True - use_constant_hint = False - - # T=1024 needs 1 CTA/SM grid AND enough per-CTA vec work. - # Under graph capture, raise the half-prec bar so a small - # capture-N doesn't force T=1024 on small-N replays - # (~14-16% regression). - if max_seq_len is not None and torch_dtype != torch.float32: - n_thresh_t = 131072 - else: - n_thresh_t = 65536 - num_threads_per_block = (1024 if - (num_rows <= num_sms - and N_per_cta >= n_thresh_t) else 512) - # V=256-bit only helps fp32 at large N. Half-prec cvt - # doubles reg pressure (5-11% loss at K=512/1024). Caller - # must hand a contiguous (32B-aligned) tensor — torch.empty - # / row slices satisfy this; column / stride-padded layouts - # may not. - use_256bit_load = (torch_dtype == torch.float32 - and N_per_cta >= 16384) - if use_256bit_load: + cfg = _GvrTopKKernel.pick_tuning( + torch_dtype, + num_rows, + N_per_cta, + num_sms, + graph_capture=max_seq_len is not None, + ) + if cfg["use_256bit_load"]: assert data_ptr % 32 == 0, ( f"use_256bit_load=True requires 32B-aligned " f"logits.data_ptr(), got {data_ptr} % 32 = " f"{data_ptr % 32}.") - # Warp-parallel reduce only pays at 32-warp (T=1024). - enable_warp_parallel_reduce = num_threads_per_block == 1024 - - # min_blocks_per_mp: reg-vs-occupancy 3-tier. Half-prec - # prefers extra CTA/SM (cvt-ILP fits in 40 regs); fp32 - # wants mb=2 (4-LDG ILP needs ~70 regs). - vec_bits_host = 256 if use_256bit_load else 128 - vec_w_host = vec_bits_host // (32 if torch_dtype == torch.float32 - else 16) - n_vec_iters = max(1, - N_per_cta // (num_threads_per_block * vec_w_host)) - if torch_dtype == torch.float32: - if n_vec_iters < 4: - min_blocks_per_mp = 0 - elif num_rows <= num_sms: - min_blocks_per_mp = 1 - elif (num_sms * 2 < num_rows <= num_sms * 3 - and N_per_cta <= 32768): - # mb=3 packs all CTAs in 1 wave; at N>=64K kernel - # is bandwidth-bound and mb=2 wins instead. - min_blocks_per_mp = 3 - else: - min_blocks_per_mp = 2 - else: - if num_rows > num_sms: - min_blocks_per_mp = 3 - elif n_vec_iters < 4: - min_blocks_per_mp = 0 - else: - min_blocks_per_mp = 1 - return dict( - enable_unroll_4=enable_unroll_4, - enable_phase3_unroll=enable_phase3_unroll, - use_constant_hint=use_constant_hint, - num_threads_per_block=num_threads_per_block, - use_256bit_load=use_256bit_load, - enable_warp_parallel_reduce=enable_warp_parallel_reduce, - min_blocks_per_mp=min_blocks_per_mp, + enable_unroll_4=True, + enable_phase3_unroll=True, + use_constant_hint=False, + num_threads_per_block=cfg["num_threads"], + use_256bit_load=cfg["use_256bit_load"], + enable_warp_parallel_reduce=cfg["enable_warp_parallel_reduce"], + min_blocks_per_mp=cfg["min_blocks_per_mp"], ) @classmethod @@ -7573,7 +7533,7 @@ def forward( ``counters`` without ``order_row`` is rejected. """ - # BSX tier fast path (op43 port): fp32 / next_n >= 1 (MTP) / + # BSX tier fast path: fp32 / next_n >= 1 (MTP) / # cr in {1, 4} / npad <= 262144 decode rows route to the # direct/reg/tp CuTe DSL tiers; everything else (half-prec, LB, # sort-indirect, oversize npad, hw cluster cap) falls through @@ -7639,21 +7599,8 @@ def forward( f"prepare, or use the single-CTA path.") else: if cluster_size is None: - # B200 SXM5 synth-data tuning, 2026-06-10: - # N < 64K -> 1 (sync unrecouped) - # N >= 128K, BS <= 4 -> 8 (tiny grid) - # BS * cs <= num_sms -> cs (single-wave) - # else -> 1 (multi-wave loses) - if N_row < 65536: - cluster_size = 1 - elif num_rows <= 4 and N_row >= 131072: - cluster_size = 8 - elif num_rows * 4 <= num_sms: - cluster_size = 4 - elif num_rows * 2 <= num_sms: - cluster_size = 2 - else: - cluster_size = 1 + cluster_size = _GvrTopKKernel.pick_cluster_size( + num_rows, N_row, num_sms) if cluster_size > 1: hw_max_cluster = _query_max_cluster_size() if cluster_size > hw_max_cluster: @@ -7715,7 +7662,7 @@ def forward( # ``num_rows >= 2 * num_sms``. Physical meaning: wave-2 must fit a # full SM-row's worth of CTAs so the sort has long-vs-short rows to # swap. Below that threshold the win is noise / can regress a few - # percent (B200 N∈{8K,16K,32K} sweep 2026-06-23). + # percent (measured, N in {8K,16K,32K}). @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode", mutates_args=("output_indices", ), device_types="cuda") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py index 1fe0dfa5eddc..0f8547cc1269 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py @@ -390,7 +390,7 @@ def __init__( self.FLT_MAX = 3.4028235e38 self.NEG_FLT_MAX = -self.FLT_MAX - # --- op#26 R0 histogram-ladder admission (default ON) --- + # --- R0 histogram-ladder admission (default ON) --- # enable_r0: replace the Phase-2 secant search with a single-pass # multi-threshold "rung ladder" admission seeded by a 256-bin # histogram over the prev-topK gathered values (P1b). @@ -423,9 +423,9 @@ def __init__( self.enable_r0 = bool(enable_r0) self.mt_unroll = int(mt_unroll) self.fb_fix = bool(fb_fix) - # C7 dispatch (op#26 host policy folded into the ctor; all gated on + # C7 dispatch (host policy folded into the ctor; all gated on # enable_r0 so an OFF kernel is byte-identical to the base): - # - qfracs default = M2D (0.85, 0.35): dispatch_r0_op26 ships M2D for + # - qfracs default = M2D (0.85, 0.35): the shipped dispatch uses M2D for # every (dtype, K, N); the M=2 pass is ~free and the R1 falsi shot # covers the 3-7% bracket misses. uh4 (M=4) was silicon-falsified # (mc geomean 0.956 — admission != latency). @@ -435,7 +435,7 @@ def __init__( # single-CTA path does NOT reproduce in the cluster kernel # (latency-bound, different SMEM budget). nsys cs=4: K1024 # ~1.01x / K2048 ~1.02x / K512 wash, 0 losses, exact. Matches - # op26 dispatch_p1bc_mc (unconditional ON). + # the multi-CTA dispatch (unconditional ON). # * cs=1 (single-CTA): (dtype != fp32). The gather-cache wins # +0.8-2.8% on 16-bit (random half-prec gather is the cost) but # is flat/negative on fp32 (occupancy at kC=6144), so OFF there. @@ -445,14 +445,14 @@ def __init__( if r0_vseed is None: r0_vseed = enable_r0 if enable_r0 and r0_qfracs is None: - # Per-K default (2026-07-16 vseed full-envelope audit, 2772 + # Per-K default (full-envelope audit, 2772 # cells): with the virtual seed rung on, pmean covers q.35's # admission region for K512/K1024 (2 count columns = zero # column tax); K2048 keeps q.35 (kC/K = 2.5 makes a fat admit # costlier than a slim 2-pass miss). Without vseed, q.35 must # stay for all K (it is the only slim rung). - # K2048 low rung 0.85 -> 0.6 (2026-07-19 real-content rung - # recalibration + paired nsys cold-L2 A/B, B200): the shipped + # K2048 low rung 0.85 -> 0.6 (real-content rung recalibration, + # paired cold-L2 A/B): the shipped # 0.85 rung's admission straddles [K, kC] on real V3.2 decode # captures (bracket on 86% of steps -> one extra falsi pass); # 0.6 lands the first pass. Measured: real V3.2 geomean @@ -479,8 +479,8 @@ def __init__( kc_diet = cluster_size == 1 if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072: self.kC = 3072 - # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (2026-07-19 - # paired nsys cold-L2 A/B on B200, all cells exact). The P4 zero / + # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (paired + # cold-L2 A/B, all cells exact). The P4 zero / # atomic build / serial scan all shrink 4x; the deeper boundary-bin # recursion costs less than the saved passes at kC=6144 candidates. # Measured vs this head: real V3.2 decode captures geomean +6.1% @@ -500,7 +500,7 @@ def __init__( "r0_qfracs must be descending h (ascending threshold value)" ) self.M_thr = len(self.r0_qfracs) - # --- vseed (2026-07-16): fold P1's pmean (the secant init + # --- vseed: fold P1's pmean (the secant init # probe) into the M-ary R0 count pass as one extra "virtual rung". # Fixes the flash-1M fat-admission regression (the coarse q.85 rung # admits ~4400 candidates where pmean admits ~630 -> 7x P3/P4 cand @@ -524,14 +524,14 @@ def __init__( else 0.0 ) - # --- op#7 P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- + # --- P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- # Replaces phase4_histogram_snap's k-th-bin search + 2-pass writeback - # with a single rank-and-scatter pass (op#7 PR#15709), cutting Phase-4 + # with a single rank-and-scatter pass (PR#15709), cutting Phase-4 # barriers ~14 -> ~7. On a latency-bound kernel that is a whole-kernel # win (~1.078x, HW-invariant). enable_p4_rank_scatter_exact adds ONE # fine-histogram recursion on the straddling coarse bin so the result is # bit-exact vs torch.topk (adds a few barriers back but still < snap). - # Default ON with R0: nsys over the op22 4k-1M BS=1 best/worst envelope + # Default ON with R0: measured over the 4k-1M BS=1 best/worst envelope # gives geomean ~1.09x (K1024 1.12 / K2048 1.12 / K512 1.05) with NO # cell regressing >2%. Resolves to OFF when enable_r0 is False, so the # base kernel stays byte-identical to upstream. @@ -542,17 +542,20 @@ def __init__( self.enable_p4_rank_scatter = bool(enable_p4_rank_scatter) self.enable_p4_rank_scatter_exact = bool(enable_p4_rank_scatter_exact) # p4_exact_tail: ambiguity-gated exact tie-resolution for the fine - # straddling bin (fp32 inputs only; see phase4_rank_scatter). The - # fine recursion resolves values to range/(kNumBins*256); two fp32 - # values closer than that straddling the kK boundary inside one fine - # bin were previously picked in arrival order (observed as |miss|=1 - # with |dv| ~ 3e-6 on real Pro 512k-ISL captures). Default ON for - # fp32 rank-scatter-exact kernels; 16-bit inputs keep the arrival - # fill (their upconverted keys are already fully resolved by the - # two-level histogram, and 16-bit tie plateaus are bitwise-equal, - # where arrival order is value-exact). + # straddling bin (see phase4_rank_scatter). The fine recursion + # resolves values to range/(kNumBins*256) — WINDOW-RELATIVE, so ANY + # dtype (fp32 or upconverted 16-bit) can leave distinct values in + # one fine bin whenever the Phase-2 bracket is wide relative to the + # boundary-local ULP (e.g. fp16 1.0 vs 1.25 under a [0, 65504] + # bracket); values straddling the kK boundary inside one fine bin + # were previously picked in arrival order (observed as |miss|=1 + # with |dv| ~ 3e-6 on real captures). The tail radix re-ranks on + # the full fp32 order key — candidate keys are ALWAYS fp32 (16-bit + # inputs are upcast injectively at collect), so the repair is exact + # for every supported dtype. Default ON for rank-scatter-exact + # kernels of all dtypes. if p4_exact_tail is None: - p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32 + p4_exact_tail = self.enable_p4_rank_scatter_exact 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 @@ -565,7 +568,7 @@ def __init__( # text (byte-identical PTX modulo kernel name) for A/B. # Default gate = p4_exact_tail AND top_k >= 1024: the non-firing # codegen tax concentrates at K512 cs=1 mid-N (flash 64k/128k - # -6.6/-9.1%, cross-GPU reproducible, 2026-07-20 b200-035) while the + # -6.6/-9.1%, cross-GPU reproducible) while the # fire census (pro/512k bench + 9 per-layer fixture cells) contains # NO K512 cell — so K512 keeps the original byte-identical kernel. if p4_tail_fast is None: # [p4tt] @@ -1368,7 +1371,7 @@ def block_count_ge( # path (same vec_w / 4-way-unroll / tail loops) with M static register # counters. Caches all M per-thread count columns in smem_ptcnt_multi so # the accepted rung's column seeds Phase 3 with zero rescan. This is the - # R0 admission primitive (op#18 multithresh lineage); it is only invoked + # R0 admission primitive (multi-threshold lineage); it is only invoked # from the enable_r0 path added in a later commit, so the base kernel is # unaffected. Slice + cluster form: each CTA scans [slice_start, # slice_end) and the M per-CTA totals are DSMEM all-reduced across the @@ -2447,9 +2450,199 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane return thr_out, sel_out # ------------------------------------------------------------------ - # Phase 4 (alt): op#7 fused rank-and-scatter (enable_p4_rank_scatter). + # Phase 4 (alt): fused rank-and-scatter (enable_p4_rank_scatter). # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py. # ------------------------------------------------------------------ + @cute.jit + def _p4_exact_tail_radix_select( + self, + kK: cutlass.Constexpr, + kBins: cutlass.Constexpr, + num_threads: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ): + """MSB-first 4x8-bit exact radix select over the straddling + fine-bin tie set (``p4_exact_tail``) — single source shared by + the tiny-tie fast path's large-class fallback and the plain + exact-tail path (previously two verbatim copies; ``@cute.jit`` + helpers inline, so codegen is unchanged).""" + # Persistent scalars live above the 256 digit bins + # (kNumBins >= 512 always): [256] key prefix (chosen + # digits, remaining bits 0), [257] slots still to fill + # inside the current equal-prefix set, [258] ties + # strictly above the prefix (their slots precede it). + if tidx == cutlass.Int32(0): + smem_hist[256] = cutlass.Int32(0) + smem_hist[257] = need0 + smem_hist[258] = cutlass.Int32(0) + cute.arch.barrier() + for lvl in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * lvl) + iz2 = tidx + while iz2 < cutlass.Int32(256): + smem_hist[iz2] = cutlass.Int32(0) + iz2 = iz2 + cutlass.Int32(num_threads) + cute.arch.barrier() + uthr_cur = smem_hist[256] + it2 = tidx + while it2 < cand_count: + vt = smem_keys[it2] + bt = cutlass.Int32((vt - bmin_r) * inv1) + if bt < cutlass.Int32(0): + bt = cutlass.Int32(0) + if bt > cutlass.Int32(kBins - 1): + bt = cutlass.Int32(kBins - 1) + if bt == b_star: + st2 = cutlass.Int32((vt - f_lo) * finv) + if st2 < cutlass.Int32(0): + st2 = cutlass.Int32(0) + if st2 > cutlass.Int32(fbins - 1): + st2 = cutlass.Int32(fbins - 1) + if st2 == sb_star: + uk = f32_order_key(vt) + pmatch = cutlass.Int32(1) + if cutlass.const_expr(lvl > 0): + if (uk >> cutlass.Int32(shift + 8)) != ( + uthr_cur >> cutlass.Int32(shift + 8) + ): + pmatch = cutlass.Int32(0) + if pmatch == cutlass.Int32(1): + dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) + atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) + it2 = it2 + cutlass.Int32(num_threads) + cute.arch.barrier() + # Two-stage descending digit scan (mirrors the + # fine 3-step search): per-warp partial sums, + # thread0 picks the target warp, its lane0 walks + # the warp's digit range — 2*num_warps serial + # steps instead of 256. + fdw = cutlass.const_expr(256 // self.num_warps) + wsum2 = cutlass.Int32(0) + for jd in cutlass.range_constexpr(fdw): + dix = ( + cutlass.Int32(255) + - warp_id * cutlass.Int32(fdw) + - cutlass.Int32(jd) + ) + wsum2 = wsum2 + smem_hist[dix] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = wsum2 + cute.arch.barrier() + if tidx == cutlass.Int32(0): + needl = smem_hist[257] + cw = cutlass.Int32(0) + tw3 = cutlass.Int32(num_warps - 1) + f3 = cutlass.Int32(0) + for w4 in cutlass.range_constexpr(self.num_warps): + cw = cw + smem_wcnt[w4] + if cw >= needl and f3 == cutlass.Int32(0): + tw3 = cutlass.Int32(w4) + f3 = cutlass.Int32(1) + pre3 = cutlass.Int32(0) + for w5 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w5) < tw3: + pre3 = pre3 + smem_wcnt[w5] + s_iscalars[4] = pre3 # prefix above target warp + s_iscalars[0] = tw3 # target warp + cute.arch.barrier() + pre4 = s_iscalars[4] + tw4 = s_iscalars[0] + if warp_id == tw4 and lane == cutlass.Int32(0): + needl2 = smem_hist[257] + base4 = pre4 + dstar = cutlass.Int32(0) + above_d = pre4 + sd4 = cutlass.Int32(0) + for jd2 in cutlass.range_constexpr(fdw): + dix2 = ( + cutlass.Int32(255) + - tw4 * cutlass.Int32(fdw) + - cutlass.Int32(jd2) + ) + ra4 = base4 + base4 = base4 + smem_hist[dix2] + if base4 >= needl2 and sd4 == cutlass.Int32(0): + dstar = dix2 + above_d = ra4 + sd4 = cutlass.Int32(1) + smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) + smem_hist[257] = needl2 - above_d + smem_hist[258] = smem_hist[258] + above_d + cute.arch.barrier() + # Rewrite the tie slot range: ties with key > u_thr + # first (there are exactly cnt_ab of them), then the + # first need_eq bitwise-equal-to-u_thr ties in arrival + # order (value-exact by construction). Signed compare + # needs the top bit flipped (unsigned-monotonic key). + u_thr = smem_hist[256] + cnt_ab = smem_hist[258] + need_eq = smem_hist[257] + ks_thr = u_thr ^ cutlass.Int32(-2147483648) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # above-writer ctr + s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr + cute.arch.barrier() + ir2 = tidx + while ir2 < cand_count: + vr = smem_keys[ir2] + br = cutlass.Int32((vr - bmin_r) * inv1) + if br < cutlass.Int32(0): + br = cutlass.Int32(0) + if br > cutlass.Int32(kBins - 1): + br = cutlass.Int32(kBins - 1) + if br == b_star: + sr = cutlass.Int32((vr - f_lo) * finv) + if sr < cutlass.Int32(0): + sr = cutlass.Int32(0) + if sr > cutlass.Int32(fbins - 1): + sr = cutlass.Int32(fbins - 1) + if sr == sb_star: + uk2 = f32_order_key(vr) + ks2 = uk2 ^ cutlass.Int32(-2147483648) + if ks2 > ks_thr: + o2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(4), + cutlass.Int32(1), + ) + pos = rank_above_fine + o2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + elif ks2 == ks_thr: + q2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if q2 < need_eq: + pos = rank_above_fine + cnt_ab + q2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + ir2 = ir2 + cutlass.Int32(num_threads) + cute.arch.barrier() + @cute.jit def phase4_rank_scatter( self, @@ -2825,325 +3018,61 @@ def phase4_rank_scatter( tj = tj + cutlass.Int32(1) cute.arch.barrier() else: - # Persistent scalars live above the 256 digit bins - # (kNumBins >= 512 always): [256] key prefix (chosen - # digits, remaining bits 0), [257] slots still to fill - # inside the current equal-prefix set, [258] ties - # strictly above the prefix (their slots precede it). - if tidx == cutlass.Int32(0): - smem_hist[256] = cutlass.Int32(0) - smem_hist[257] = need0 - smem_hist[258] = cutlass.Int32(0) - cute.arch.barrier() - for lvl in cutlass.range_constexpr(4): - shift = cutlass.const_expr(24 - 8 * lvl) - iz2 = tidx - while iz2 < cutlass.Int32(256): - smem_hist[iz2] = cutlass.Int32(0) - iz2 = iz2 + cutlass.Int32(num_threads) - cute.arch.barrier() - uthr_cur = smem_hist[256] - it2 = tidx - while it2 < cand_count: - vt = smem_keys[it2] - bt = cutlass.Int32((vt - bmin_r) * inv1) - if bt < cutlass.Int32(0): - bt = cutlass.Int32(0) - if bt > cutlass.Int32(kBins - 1): - bt = cutlass.Int32(kBins - 1) - if bt == b_star: - st2 = cutlass.Int32((vt - f_lo) * finv) - if st2 < cutlass.Int32(0): - st2 = cutlass.Int32(0) - if st2 > cutlass.Int32(fbins - 1): - st2 = cutlass.Int32(fbins - 1) - if st2 == sb_star: - uk = f32_order_key(vt) - pmatch = cutlass.Int32(1) - if cutlass.const_expr(lvl > 0): - if (uk >> cutlass.Int32(shift + 8)) != ( - uthr_cur >> cutlass.Int32(shift + 8) - ): - pmatch = cutlass.Int32(0) - if pmatch == cutlass.Int32(1): - dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32( - 0xFF - ) - atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) - it2 = it2 + cutlass.Int32(num_threads) - cute.arch.barrier() - # Two-stage descending digit scan (mirrors the - # fine 3-step search): per-warp partial sums, - # thread0 picks the target warp, its lane0 walks - # the warp's digit range — 2*num_warps serial - # steps instead of 256. - fdw = cutlass.const_expr(256 // self.num_warps) - wsum2 = cutlass.Int32(0) - for jd in cutlass.range_constexpr(fdw): - dix = ( - cutlass.Int32(255) - - warp_id * cutlass.Int32(fdw) - - cutlass.Int32(jd) - ) - wsum2 = wsum2 + smem_hist[dix] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = wsum2 - cute.arch.barrier() - if tidx == cutlass.Int32(0): - needl = smem_hist[257] - cw = cutlass.Int32(0) - tw3 = cutlass.Int32(num_warps - 1) - f3 = cutlass.Int32(0) - for w4 in cutlass.range_constexpr(self.num_warps): - cw = cw + smem_wcnt[w4] - if cw >= needl and f3 == cutlass.Int32(0): - tw3 = cutlass.Int32(w4) - f3 = cutlass.Int32(1) - pre3 = cutlass.Int32(0) - for w5 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w5) < tw3: - pre3 = pre3 + smem_wcnt[w5] - s_iscalars[4] = pre3 # prefix above target warp - s_iscalars[0] = tw3 # target warp - cute.arch.barrier() - pre4 = s_iscalars[4] - tw4 = s_iscalars[0] - if warp_id == tw4 and lane == cutlass.Int32(0): - needl2 = smem_hist[257] - base4 = pre4 - dstar = cutlass.Int32(0) - above_d = pre4 - sd4 = cutlass.Int32(0) - for jd2 in cutlass.range_constexpr(fdw): - dix2 = ( - cutlass.Int32(255) - - tw4 * cutlass.Int32(fdw) - - cutlass.Int32(jd2) - ) - ra4 = base4 - base4 = base4 + smem_hist[dix2] - if base4 >= needl2 and sd4 == cutlass.Int32(0): - dstar = dix2 - above_d = ra4 - sd4 = cutlass.Int32(1) - smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) - smem_hist[257] = needl2 - above_d - smem_hist[258] = smem_hist[258] + above_d - cute.arch.barrier() - # Rewrite the tie slot range: ties with key > u_thr - # first (there are exactly cnt_ab of them), then the - # first need_eq bitwise-equal-to-u_thr ties in arrival - # order (value-exact by construction). Signed compare - # needs the top bit flipped (unsigned-monotonic key). - u_thr = smem_hist[256] - cnt_ab = smem_hist[258] - need_eq = smem_hist[257] - ks_thr = u_thr ^ cutlass.Int32(-2147483648) - if tidx == cutlass.Int32(0): - s_iscalars[4] = cutlass.Int32(0) # above-writer ctr - s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr - cute.arch.barrier() - ir2 = tidx - while ir2 < cand_count: - vr = smem_keys[ir2] - br = cutlass.Int32((vr - bmin_r) * inv1) - if br < cutlass.Int32(0): - br = cutlass.Int32(0) - if br > cutlass.Int32(kBins - 1): - br = cutlass.Int32(kBins - 1) - if br == b_star: - sr = cutlass.Int32((vr - f_lo) * finv) - if sr < cutlass.Int32(0): - sr = cutlass.Int32(0) - if sr > cutlass.Int32(fbins - 1): - sr = cutlass.Int32(fbins - 1) - if sr == sb_star: - uk2 = f32_order_key(vr) - ks2 = uk2 ^ cutlass.Int32(-2147483648) - if ks2 > ks_thr: - o2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(4), - cutlass.Int32(1), - ) - pos = rank_above_fine + o2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - elif ks2 == ks_thr: - q2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), - cutlass.Int32(1), - ) - if q2 < need_eq: - pos = rank_above_fine + cnt_ab + q2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr( - self.return_output_values - ): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - ir2 = ir2 + cutlass.Int32(num_threads) - cute.arch.barrier() + self._p4_exact_tail_radix_select( + kK, + kBins, + num_threads, + num_warps, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ) elif cutlass.const_expr(self.p4_exact_tail): # [p4tt] if->elif only need0 = cutlass.Int32(kK) - rank_above_fine if cnt_strad > need0 and need0 > cutlass.Int32(0): - # Persistent scalars live above the 256 digit bins - # (kNumBins >= 512 always): [256] key prefix (chosen - # digits, remaining bits 0), [257] slots still to fill - # inside the current equal-prefix set, [258] ties - # strictly above the prefix (their slots precede it). - if tidx == cutlass.Int32(0): - smem_hist[256] = cutlass.Int32(0) - smem_hist[257] = need0 - smem_hist[258] = cutlass.Int32(0) - cute.arch.barrier() - for lvl in cutlass.range_constexpr(4): - shift = cutlass.const_expr(24 - 8 * lvl) - iz2 = tidx - while iz2 < cutlass.Int32(256): - smem_hist[iz2] = cutlass.Int32(0) - iz2 = iz2 + cutlass.Int32(num_threads) - cute.arch.barrier() - uthr_cur = smem_hist[256] - it2 = tidx - while it2 < cand_count: - vt = smem_keys[it2] - bt = cutlass.Int32((vt - bmin_r) * inv1) - if bt < cutlass.Int32(0): - bt = cutlass.Int32(0) - if bt > cutlass.Int32(kBins - 1): - bt = cutlass.Int32(kBins - 1) - if bt == b_star: - st2 = cutlass.Int32((vt - f_lo) * finv) - if st2 < cutlass.Int32(0): - st2 = cutlass.Int32(0) - if st2 > cutlass.Int32(fbins - 1): - st2 = cutlass.Int32(fbins - 1) - if st2 == sb_star: - uk = f32_order_key(vt) - pmatch = cutlass.Int32(1) - if cutlass.const_expr(lvl > 0): - if (uk >> cutlass.Int32(shift + 8)) != ( - uthr_cur >> cutlass.Int32(shift + 8) - ): - pmatch = cutlass.Int32(0) - if pmatch == cutlass.Int32(1): - dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) - atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) - it2 = it2 + cutlass.Int32(num_threads) - cute.arch.barrier() - # Two-stage descending digit scan (mirrors the - # fine 3-step search): per-warp partial sums, - # thread0 picks the target warp, its lane0 walks - # the warp's digit range — 2*num_warps serial - # steps instead of 256. - fdw = cutlass.const_expr(256 // self.num_warps) - wsum2 = cutlass.Int32(0) - for jd in cutlass.range_constexpr(fdw): - dix = ( - cutlass.Int32(255) - - warp_id * cutlass.Int32(fdw) - - cutlass.Int32(jd) - ) - wsum2 = wsum2 + smem_hist[dix] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = wsum2 - cute.arch.barrier() - if tidx == cutlass.Int32(0): - needl = smem_hist[257] - cw = cutlass.Int32(0) - tw3 = cutlass.Int32(num_warps - 1) - f3 = cutlass.Int32(0) - for w4 in cutlass.range_constexpr(self.num_warps): - cw = cw + smem_wcnt[w4] - if cw >= needl and f3 == cutlass.Int32(0): - tw3 = cutlass.Int32(w4) - f3 = cutlass.Int32(1) - pre3 = cutlass.Int32(0) - for w5 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w5) < tw3: - pre3 = pre3 + smem_wcnt[w5] - s_iscalars[4] = pre3 # prefix above target warp - s_iscalars[0] = tw3 # target warp - cute.arch.barrier() - pre4 = s_iscalars[4] - tw4 = s_iscalars[0] - if warp_id == tw4 and lane == cutlass.Int32(0): - needl2 = smem_hist[257] - base4 = pre4 - dstar = cutlass.Int32(0) - above_d = pre4 - sd4 = cutlass.Int32(0) - for jd2 in cutlass.range_constexpr(fdw): - dix2 = ( - cutlass.Int32(255) - - tw4 * cutlass.Int32(fdw) - - cutlass.Int32(jd2) - ) - ra4 = base4 - base4 = base4 + smem_hist[dix2] - if base4 >= needl2 and sd4 == cutlass.Int32(0): - dstar = dix2 - above_d = ra4 - sd4 = cutlass.Int32(1) - smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) - smem_hist[257] = needl2 - above_d - smem_hist[258] = smem_hist[258] + above_d - cute.arch.barrier() - # Rewrite the tie slot range: ties with key > u_thr - # first (there are exactly cnt_ab of them), then the - # first need_eq bitwise-equal-to-u_thr ties in arrival - # order (value-exact by construction). Signed compare - # needs the top bit flipped (unsigned-monotonic key). - u_thr = smem_hist[256] - cnt_ab = smem_hist[258] - need_eq = smem_hist[257] - ks_thr = u_thr ^ cutlass.Int32(-2147483648) - if tidx == cutlass.Int32(0): - s_iscalars[4] = cutlass.Int32(0) # above-writer ctr - s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr - cute.arch.barrier() - ir2 = tidx - while ir2 < cand_count: - vr = smem_keys[ir2] - br = cutlass.Int32((vr - bmin_r) * inv1) - if br < cutlass.Int32(0): - br = cutlass.Int32(0) - if br > cutlass.Int32(kBins - 1): - br = cutlass.Int32(kBins - 1) - if br == b_star: - sr = cutlass.Int32((vr - f_lo) * finv) - if sr < cutlass.Int32(0): - sr = cutlass.Int32(0) - if sr > cutlass.Int32(fbins - 1): - sr = cutlass.Int32(fbins - 1) - if sr == sb_star: - uk2 = f32_order_key(vr) - ks2 = uk2 ^ cutlass.Int32(-2147483648) - if ks2 > ks_thr: - o2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(4), - cutlass.Int32(1), - ) - pos = rank_above_fine + o2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - elif ks2 == ks_thr: - q2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), - cutlass.Int32(1), - ) - if q2 < need_eq: - pos = rank_above_fine + cnt_ab + q2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - ir2 = ir2 + cutlass.Int32(num_threads) - cute.arch.barrier() + self._p4_exact_tail_radix_select( + kK, + kBins, + num_threads, + num_warps, + need0, + cand_count, + rank_above_fine, + b_star, + sb_star, + bmin_r, + f_lo, + finv, + fbins, + inv1, + tidx, + lane, + warp_id, + smem_hist, + smem_keys, + smem_vals, + smem_wcnt, + s_iscalars, + output_indices_row, + output_values_row, + ) else: # ---- APPROX rank-and-scatter (single pass), arbitrary straddling order ---- isc = tidx @@ -3962,7 +3891,7 @@ def run_one_row( else: smem_input = None - # op#26 R0 admission scratch (single-CTA fast path). Allocated only + # R0 admission scratch (single-CTA fast path). Allocated only # when enable_r0; None otherwise so the base SMEM layout is byte-for- # byte unchanged and these propagate harmlessly through _run_phases' # const_expr(enable_r0)-gated branch (same idiom as s_cluster_partial @@ -4312,7 +4241,7 @@ def _run_phases( # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast # path) or the secant threshold search ---- - # enable_r0 gates to cluster_size==1 for now: op#26's R0 scans the + # enable_r0 gates to cluster_size==1 for now: 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. @@ -4406,7 +4335,7 @@ def _run_phases( # rungs. SEED the loop with the rung bracket AND its known # counts (clo/chi) so it does log-count regula-falsi from # iter 0 with no re-measure and no separate R1 shot -> ~2-3 - # count passes (op#26 efficiency) instead of ~6. done=1 on + # count passes instead of ~6. done=1 on # accept so Phase 3 skips its retry-shrink. if bc < cutlass.Int32(0): if cutlass.const_expr(self.fb_fix): @@ -4793,9 +4722,11 @@ def __call__( # policy as a pure function colocated with the kernel (single source of # truth), and ``launch`` is a thin variant-cache wrapper so direct-drive # users (tests, benchmarks) get the same shapes production would pick. - # The production custom op keeps its own equivalent inline policy for - # now; unifying it onto ``pick_config`` is a call-site change deferred - # to the dispatch-guard follow-up PR. + # The production custom op delegates here (``pick_cluster_size`` / + # ``pick_tuning``) — one policy, two shells. Intentional shell + # divergence: on a 32B-misaligned logits pointer the production runner + # ASSERTS (contract violation), while ``launch`` silently downgrades to + # 128-bit loads (dev convenience for ad-hoc tensors). _NUM_SMS: Optional[int] = None _LAUNCH_CACHE: dict = {} @@ -4811,69 +4742,51 @@ def _device_num_sms() -> int: return GvrTopKKernel._NUM_SMS @staticmethod - def pick_config( + def pick_cluster_size(num_rows: int, n_row: int, num_sms: int) -> int: + """Cluster-size policy: N < 64K -> 1 (sync unrecouped); tiny grid + at large N -> 8; single-wave -> 4/2; multi-wave -> 1 (row + parallelism already saturates the SMs; per-row splitting is pure + overhead past one wave).""" + if n_row < 65536: + return 1 + if num_rows <= 4 and n_row >= 131072: + return 8 + if num_rows * 4 <= num_sms: + return 4 + if num_rows * 2 <= num_sms: + return 2 + return 1 + + @staticmethod + def pick_tuning( torch_dtype, num_rows: int, - num_candidates: int, - max_seq_len: Optional[int] = None, - num_sms: Optional[int] = None, + n_per_cta: int, + num_sms: int, + graph_capture: bool, ) -> dict: - """Pick the launch-shape ctor kwargs for ``(dtype, BS, N)``. - - Mirrors the production runner policy (cluster_size auto-pick + - ``_pick_tuning``) so any caller instantiating the kernel directly - gets the same shapes the custom op would use. Rationale (B200, - nsys cold-L2, 2026-07-15 big-BS triage): a config frozen at the - BS=1 optimum (cs = N>=65536 ? 4 : 1, T=1024, mbpm=1) is geomean - 2.27x slower (max 6.0x) than the op-bench anchor at BS in - {64, 256, 1024}, while this policy is 0.95x (parity/better). - Multi-CTA splitting only pays while the grid is a single wave - (num_rows * cluster_size <= num_sms); past that, row parallelism - already saturates the SMs and per-row splitting is pure overhead. + """T / V / min_blocks_per_mp / warp-reduce policy at a given + per-CTA row width (cluster split already applied). - ``max_seq_len``: pass the peak runtime N under CUDA-graph capture - so the variant is picked for the replay shape, not the capture - shape (same contract as the custom op's ``_pick_tuning``). - - Returns kwargs for ``GvrTopKKernel(...)``: ``cluster_size``, - ``num_threads``, ``use_256bit_load``, ``min_blocks_per_mp``, - ``enable_warp_parallel_reduce``. + ``graph_capture``: raise the half-prec T=1024 bar so a small + capture-N does not pin T=1024 onto small-N replays. + Returns ``num_threads``, ``use_256bit_load``, + ``min_blocks_per_mp``, ``enable_warp_parallel_reduce``. """ import torch # local: keep the module importable without torch - if num_sms is None: - num_sms = GvrTopKKernel._device_num_sms() - n_row = max_seq_len if max_seq_len is not None else num_candidates is_fp32 = torch_dtype == torch.float32 - - # cluster_size: B200 SXM5 synth-data tuning (matches the custom - # op's auto-pick): N < 64K -> 1 (sync unrecouped); tiny grid at - # large N -> 8; single-wave -> 4/2; multi-wave -> 1. - if n_row < 65536: - cluster_size = 1 - elif num_rows <= 4 and n_row >= 131072: - cluster_size = 8 - elif num_rows * 4 <= num_sms: - cluster_size = 4 - elif num_rows * 2 <= num_sms: - cluster_size = 2 - else: - cluster_size = 1 - - # Cluster CTAs split the row, so tuning targets per-CTA work. - n_per_cta = n_row // cluster_size - # T=1024 needs 1 CTA/SM grid AND enough per-CTA vec work. Under - # graph capture, raise the half-prec bar so a small capture-N - # doesn't force T=1024 on small-N replays. - n_thresh_t = 131072 if (max_seq_len is not None and not is_fp32) else 65536 - num_threads = 1024 if (num_rows <= num_sms and n_per_cta >= n_thresh_t) else 512 + # T=1024 needs a 1 CTA/SM grid AND enough per-CTA vec work. + n_thresh_t = 131072 if (graph_capture and not is_fp32) else 65536 + num_threads = 1024 if (num_rows <= num_sms + and n_per_cta >= n_thresh_t) else 512 # V=256-bit only helps fp32 at large N; half-prec cvt doubles reg - # pressure. Caller must hand a 32B-aligned contiguous tensor - # (``launch`` downgrades on misalignment). + # pressure. Requires a 32B-aligned contiguous tensor (see the + # shell-divergence note above). use_256bit_load = is_fp32 and n_per_cta >= 16384 enable_warp_parallel_reduce = num_threads == 1024 - # min_blocks_per_mp: reg-vs-occupancy 3-tier (fp32 wants ~70 regs + # min_blocks_per_mp: reg-vs-occupancy tiers (fp32 wants ~70 regs # for 4-LDG ILP -> mb<=2; half-prec fits 40 regs -> mb=3 packs # 3 CTA/SM when rows oversubscribe the device). vec_bits = 256 if use_256bit_load else 128 @@ -4897,13 +4810,44 @@ def pick_config( min_blocks_per_mp = 1 return dict( - cluster_size=cluster_size, num_threads=num_threads, use_256bit_load=use_256bit_load, min_blocks_per_mp=min_blocks_per_mp, enable_warp_parallel_reduce=enable_warp_parallel_reduce, ) + @staticmethod + def pick_config( + torch_dtype, + num_rows: int, + num_candidates: int, + max_seq_len: Optional[int] = None, + num_sms: Optional[int] = None, + ) -> dict: + """Launch-shape ctor kwargs for ``(dtype, BS, N)`` — the single + source of truth shared by the production runner + (``CuteDSLGvrTopKDecodeRunner``) and direct-drive users (tests, + benchmarks): composition of :meth:`pick_cluster_size` and + :meth:`pick_tuning`. + + ``max_seq_len``: pass the peak runtime N under CUDA-graph capture + so the variant is picked for the replay shape, not the capture + shape. + """ + if num_sms is None: + num_sms = GvrTopKKernel._device_num_sms() + n_row = max_seq_len if max_seq_len is not None else num_candidates + cluster_size = GvrTopKKernel.pick_cluster_size(num_rows, n_row, num_sms) + cfg = GvrTopKKernel.pick_tuning( + torch_dtype, + num_rows, + n_row // cluster_size, + num_sms, + graph_capture=max_seq_len is not None, + ) + cfg["cluster_size"] = cluster_size + return cfg + @classmethod def launch( cls, diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py index f8f7215eb6ee..fc3c6a767227 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 @@ -1135,3 +1135,71 @@ def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size) _tie_aware_check(out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1) if dtype == torch.float32: assert torch.equal(out.sort(dim=-1).values, out_sec.sort(dim=-1).values) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): + """16-bit exact-tail adversarial: two DISTINCT half-precision values in + ONE fine bin straddling the K boundary (window-relative binning cannot + separate them under a wide Phase-2 bracket — e.g. fp16 1.0 vs 1.25 + under a [0, 65504]-scale bracket). With p4_exact_tail now default-on + for all dtypes the tail radix must keep every 1.25 above every 1.0.""" + torch.manual_seed(11) + top_k, n = 1024, 32768 + bs = 2 + lo = torch.full((bs, n), -1.0, dtype=dtype, device="cuda") + # wide bracket anchors (force a wide Phase-2 window) + lo[:, :8] = torch.tensor([60000.0, 40000.0, 20000.0, 10000.0, + 5000.0, 2500.0, 1200.0, 600.0], + device="cuda").to(dtype) + # boundary: (top_k - 8 - 512) high-tie values 1.25 and 1024 low-tie 1.0 + n_hi, n_lo = top_k - 8 - 512, 1024 + perm = torch.randperm(n - 8, device="cuda") + 8 + hi_pos, lo_pos = perm[:n_hi], perm[n_hi:n_hi + n_lo] + for r in range(bs): + lo[r, hi_pos] = torch.tensor(1.25, dtype=dtype, device="cuda") + lo[r, lo_pos] = torch.tensor(1.0, dtype=dtype, device="cuda") + seq_lens = torch.full((bs,), n, dtype=torch.int32, device="cuda") + pre = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre[:, 0] = lo.float().argmax(dim=-1).int() + pre[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1) + torch.cuda.synchronize() + sel = torch.gather(lo.float(), -1, out.long()) + # every 1.25 must be selected before any 1.0 fills the remainder + assert int((sel == 1.25).sum()) == n_hi * bs, ( + f"exact-tail 16-bit: {(sel == 1.25).sum()} of {n_hi * bs} " + f"high-tie values selected") + ref = torch.topk(lo.float(), top_k, dim=-1).values.sort(-1).values + got = sel.sort(-1).values + assert torch.equal(ref, got) + + +def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): + """The production runner's tuning adapter must agree with the kernel's + pick_cluster_size/pick_tuning single source across a shape sweep + (guards the de-duplicated launch-shape policy against drift).""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLGvrTopKDecodeRunner as R, ) + num_sms = 148 + for torch_dtype in (torch.float32, torch.bfloat16, torch.float16): + for num_rows in (1, 4, 32, 148, 300, 512): + for n in (4096, 16384, 65536, 131072, 262144): + for msl in (None, 262144): + cs = _GvrTopKKernel.pick_cluster_size( + num_rows, msl if msl is not None else n, num_sms) + cfg = _GvrTopKKernel.pick_config( + torch_dtype, num_rows, n, max_seq_len=msl, + num_sms=num_sms) + assert cfg["cluster_size"] == cs + tuning = R._pick_tuning( + torch_dtype, num_rows, + (msl if msl is not None else n) // cs, num_sms, + msl, 0) + assert tuning["num_threads_per_block"] == cfg["num_threads"] + assert tuning["use_256bit_load"] == cfg["use_256bit_load"] + assert tuning["min_blocks_per_mp"] == cfg["min_blocks_per_mp"] + assert (tuning["enable_warp_parallel_reduce"] + == cfg["enable_warp_parallel_reduce"]) From b82f6e03790c80f944491fe7c466b41bb40053f0 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:03:01 +0000 Subject: [PATCH 11/19] [None][fix] apply ruff-format to the follow-up changes The Release-Check stage of the previous CI run failed on ruff-format: three hand-wrapped expressions introduced by the follow-up commit (two digit-index computations in the P4 exact-tail helper, one conditional in pick_tuning) and the appended tests did not match the formatter's output. Formatting only, no behavior change; ruff format --check is clean on every touched file and the affected test families re-run green (95 passed / 16 skipped: r0-equivalence, exact-tail, policy-parity). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode.py | 15 ++------ .../sparse/test_cute_dsl_gvr_topk_decode.py | 38 +++++++++++-------- 2 files changed, 26 insertions(+), 27 deletions(-) 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 0f8547cc1269..6ddb4a1a6aa0 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 @@ -2539,11 +2539,7 @@ def _p4_exact_tail_radix_select( 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) - ) + 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 @@ -2574,11 +2570,7 @@ def _p4_exact_tail_radix_select( 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) - ) + 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): @@ -4778,8 +4770,7 @@ def pick_tuning( is_fp32 = torch_dtype == torch.float32 # T=1024 needs a 1 CTA/SM grid AND enough per-CTA vec work. n_thresh_t = 131072 if (graph_capture and not is_fp32) else 65536 - num_threads = 1024 if (num_rows <= num_sms - and n_per_cta >= n_thresh_t) else 512 + 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. Requires a 32B-aligned contiguous tensor (see the # shell-divergence note above). 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 fc3c6a767227..0ed81531eba3 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 @@ -1150,13 +1150,13 @@ def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): bs = 2 lo = torch.full((bs, n), -1.0, dtype=dtype, device="cuda") # wide bracket anchors (force a wide Phase-2 window) - lo[:, :8] = torch.tensor([60000.0, 40000.0, 20000.0, 10000.0, - 5000.0, 2500.0, 1200.0, 600.0], - device="cuda").to(dtype) + lo[:, :8] = torch.tensor( + [60000.0, 40000.0, 20000.0, 10000.0, 5000.0, 2500.0, 1200.0, 600.0], device="cuda" + ).to(dtype) # boundary: (top_k - 8 - 512) high-tie values 1.25 and 1024 low-tie 1.0 n_hi, n_lo = top_k - 8 - 512, 1024 perm = torch.randperm(n - 8, device="cuda") + 8 - hi_pos, lo_pos = perm[:n_hi], perm[n_hi:n_hi + n_lo] + hi_pos, lo_pos = perm[:n_hi], perm[n_hi : n_hi + n_lo] for r in range(bs): lo[r, hi_pos] = torch.tensor(1.25, dtype=dtype, device="cuda") lo[r, lo_pos] = torch.tensor(1.0, dtype=dtype, device="cuda") @@ -1170,8 +1170,8 @@ def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): sel = torch.gather(lo.float(), -1, out.long()) # every 1.25 must be selected before any 1.0 fills the remainder assert int((sel == 1.25).sum()) == n_hi * bs, ( - f"exact-tail 16-bit: {(sel == 1.25).sum()} of {n_hi * bs} " - f"high-tie values selected") + f"exact-tail 16-bit: {(sel == 1.25).sum()} of {n_hi * bs} high-tie values selected" + ) ref = torch.topk(lo.float(), top_k, dim=-1).values.sort(-1).values got = sel.sort(-1).values assert torch.equal(ref, got) @@ -1182,24 +1182,32 @@ def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): pick_cluster_size/pick_tuning single source across a shape sweep (guards the de-duplicated launch-shape policy against drift).""" from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - CuteDSLGvrTopKDecodeRunner as R, ) + CuteDSLGvrTopKDecodeRunner as R, + ) + num_sms = 148 for torch_dtype in (torch.float32, torch.bfloat16, torch.float16): for num_rows in (1, 4, 32, 148, 300, 512): for n in (4096, 16384, 65536, 131072, 262144): for msl in (None, 262144): cs = _GvrTopKKernel.pick_cluster_size( - num_rows, msl if msl is not None else n, num_sms) + num_rows, msl if msl is not None else n, num_sms + ) cfg = _GvrTopKKernel.pick_config( - torch_dtype, num_rows, n, max_seq_len=msl, - num_sms=num_sms) + torch_dtype, num_rows, n, max_seq_len=msl, num_sms=num_sms + ) assert cfg["cluster_size"] == cs tuning = R._pick_tuning( - torch_dtype, num_rows, - (msl if msl is not None else n) // cs, num_sms, - msl, 0) + torch_dtype, + num_rows, + (msl if msl is not None else n) // cs, + num_sms, + msl, + 0, + ) assert tuning["num_threads_per_block"] == cfg["num_threads"] assert tuning["use_256bit_load"] == cfg["use_256bit_load"] assert tuning["min_blocks_per_mp"] == cfg["min_blocks_per_mp"] - assert (tuning["enable_warp_parallel_reduce"] - == cfg["enable_warp_parallel_reduce"]) + assert ( + tuning["enable_warp_parallel_reduce"] == cfg["enable_warp_parallel_reduce"] + ) From 4642d92643900e130e59283ffde972dea52abd49 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:06:35 +0000 Subject: [PATCH 12/19] [None][fix] GVR decode: exact tie-aware terminal for boundary plateaus Resolves the last #16457 review follow-up: a bitwise-equal plateau wider than the candidate buffer (kC) straddling the K boundary has no admissible threshold - every threshold either overflows the buffer or undershoots K - so Phase 2 previously gave up on the undershoot side and Phase 4 padded the tail with -1, dropping entries that belong in the answer. Fix: on that terminal the driver first collapses the bracket by pure bisection until [lo, hi] are ADJACENT floats. Every value in [lo, hi) is then bitwise-equal, i.e. a genuine tie class, so the row can be completed exactly: Phase 4 emits the cnt(>= hi) sure winners and a ticketed fill takes any (K - count)-subset of the tie class - a valid tie-aware completion by definition. A count that lands in [K, kC] mid-collapse converges normally. Both terminals (the base secant driver and the admission path's tie-plateau fail-soft) take this route; the guard requires a coherent undershoot-overflow bracket with both counts current, so the admission retry's widened brackets - whose adjacency carries no tie-class meaning - are excluded. Non-plateau undershoot keeps the documented -1-pad encoding. Note for future work in Phase 4: the terminal is captured into a dedicated SMEM slot BEFORE Phase 4 runs, because Phase 4 reuses s_iscalars[1] as radix scratch - reading the terminal back from it afterwards yields a mid-radix value. Adversarial test: a plateau wider than kC straddling K, fp32 + fp16 x {rank-scatter cs=1, cs=4, histogram-snap}, 6/6. Full sparse-attention suite unchanged at 674 passed / 144 skipped. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode.py | 255 +++++++++++++++--- .../sparse/test_cute_dsl_gvr_topk_decode.py | 49 ++++ 2 files changed, 266 insertions(+), 38 deletions(-) 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 6ddb4a1a6aa0..ea68cd735d6c 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 @@ -1736,11 +1736,22 @@ def phase2_secant_search( nv = vhi - rng * cutlass.Float32(0.05) if nv == vlo or nv == vhi: - # Bracket exhausted — try midpoint, else give up. + # Bracket exhausted — try midpoint, else terminal. nv = (vlo + vhi) * cutlass.Float32(0.5) if nv == vlo or nv == vhi: - s_thr[0] = vlo - s_iscalars[1] = cutlass.Int32(2) # done = 2 (give up) + # ADJACENT-FLOAT bracket: every value in + # [vlo, vhi) is bitwise-equal to vlo. Low side + # overflowing the candidate buffer AND high side + # undershooting K means the boundary sits inside + # a bitwise-equal plateau wider than kC — record + # the plateau terminal (done = 3) and keep the + # sure-winner threshold vhi. + if clo > cutlass.Int32(kCC) and chi < cutlass.Int32(kK): + s_thr[0] = vhi + s_iscalars[1] = cutlass.Int32(3) # done = 3 (plateau) + else: + s_thr[0] = vlo + s_iscalars[1] = cutlass.Int32(2) # done = 2 (give up) else: s_thr[0] = nv else: @@ -3107,12 +3118,13 @@ def phase4_rank_scatter( output_values_row[i10] = self.dtype(smem_keys[i10]) output_indices_row[i10] = smem_vals[i10] i10 = i10 + cutlass.Int32(num_threads) - i11 = cand_count + tidx - while i11 < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) - output_indices_row[i11] = cutlass.Int32(-1) - i11 = i11 + cutlass.Int32(num_threads) + if s_iscalars[6] == cutlass.Int32(0): # plateau fill completes done=3 + i11 = cand_count + tidx + while i11 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[i11] = cutlass.Int32(-1) + i11 = i11 + cutlass.Int32(num_threads) # ------------------------------------------------------------------ # Phase 4: Histogram-based k-th selection + two-pass writeback @@ -3618,12 +3630,13 @@ def phase4_histogram_snap( ) output_indices_row[i10] = self._smem_ld(cutlass.Int32, vals_base, i10) i10 = i10 + cutlass.Int32(num_threads) - i11 = cand_count + tidx - while i11 < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) - output_indices_row[i11] = cutlass.Int32(-1) - i11 = i11 + cutlass.Int32(num_threads) + if s_iscalars[6] == cutlass.Int32(0): # plateau fill completes done=3 + i11 = cand_count + tidx + while i11 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i11] = self.dtype(self.NEG_FLT_MAX) + output_indices_row[i11] = cutlass.Int32(-1) + i11 = i11 + cutlass.Int32(num_threads) # ------------------------------------------------------------------ # Main kernel — one CTA per row @@ -3841,9 +3854,13 @@ def run_one_row( # [4] out_count # [5] local cand_count (per-CTA snapshot before cluster all-reduce; # consumed by the kernel-level cluster handoff) + # [6] plateau terminal flag, captured from [1] BEFORE Phase 4 + # (Phase 4 REUSES [1] as radix scratch, so the terminal must + # never be re-read from it afterwards) + # [7] plateau fill ticket (done == 3 only) s_iscalars = smem.allocate_tensor( element_type=cutlass.Int32, - layout=cute.make_ordered_layout((6,), order=(0,)), + layout=cute.make_ordered_layout((8,), order=(0,)), byte_alignment=16, ) # Per-CTA DSMEM scratch for the cluster all-reduce of cand_count: @@ -4443,28 +4460,115 @@ def _run_phases( cute.arch.barrier() rs = rs + cutlass.Int32(1) if s_iscalars[1] != cutlass.Int32(1): - # tie-plateau fail-soft: land on the measured - # undershoot side (count <= kC => no overflow). - self.block_count_ge( - input_row, - slice_start, - slice_end, - s_thr[2], - smem_ptcnt, - smem_wcnt, - s_iscalars, - s_cluster_partial, - tidx, - warp_id, - lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, - ) - cute.arch.barrier() - if tidx == cutlass.Int32(0): - s_thr[0] = s_thr[2] - s_iscalars[1] = cutlass.Int32(1) - cute.arch.barrier() + # The retry budget could not land in [K, kC]. + # ONLY the coherent undershoot-overflow corner + # (count(>= lo) > kC AND 0 <= count(>= hi) < K, + # both counts CURRENT — the retry's bracket + # widening marks a side stale with -1 and thus + # fails this guard) collapses the bracket by + # pure bisection to ADJACENT floats, where the + # plateau terminal (done = 3, threshold = hi) + # is exact: Phase 4 emits the sure winners and + # the plateau fill completes the row from the + # tie class. A mid-collapse count landing in + # [K, kC] converges normally; anything else + # (incl. an exhausted collapse budget) falls + # through to the fail-soft terminal below. + it4 = cutlass.Int32(0) + if ( + s_iscalars[2] <= cutlass.Int32(self.kC) + or s_iscalars[3] < cutlass.Int32(0) + or s_iscalars[3] >= cutlass.Int32(self.top_k) + ): + it4 = cutlass.Int32(40) # guard: skip collapse + while it4 < cutlass.Int32(40) and s_iscalars[1] == cutlass.Int32(0): + if tidx == cutlass.Int32(0): + lo4 = s_thr[1] + hi4 = s_thr[2] + mid4 = (lo4 + hi4) * cutlass.Float32(0.5) + if mid4 == lo4 or mid4 == hi4: + s_thr[0] = hi4 + s_iscalars[1] = cutlass.Int32(3) + else: + s_thr[0] = mid4 + cute.arch.barrier() + if s_iscalars[1] == cutlass.Int32(0): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + c4 = s_iscalars[0] + t4 = s_thr[0] + if c4 >= cutlass.Int32(self.top_k) and c4 <= cutlass.Int32( + self.kC + ): + s_iscalars[1] = cutlass.Int32(1) + elif c4 > cutlass.Int32(self.kC): + s_thr[1] = t4 + s_iscalars[2] = c4 + else: + s_thr[2] = t4 + s_iscalars[3] = c4 + cute.arch.barrier() + it4 = it4 + cutlass.Int32(1) + if s_iscalars[1] == cutlass.Int32(3): + # recount at the terminal threshold so P3's + # cached per-thread counts describe the + # sure-winner set the fill completes. + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + elif s_iscalars[1] != cutlass.Int32(1): + # fail-soft (non-plateau): land on the + # measured undershoot side (count <= kC => + # no overflow; -1 pad stays the documented + # non-convergence encoding). + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[2], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_thr[0] = s_thr[2] + s_iscalars[1] = cutlass.Int32(1) + cute.arch.barrier() else: self.phase2_secant_search( input_row, @@ -4543,6 +4647,13 @@ def _run_phases( cand_count_p4 = cutlass.Int32(0) if cutlass.const_expr(cluster_size == 1): # cs=1: the single CTA per row IS the leader. + # Capture the P2 terminal BEFORE Phase 4: P4 reuses + # s_iscalars[1] as radix scratch. + if tidx == cutlass.Int32(0): + s_iscalars[6] = cutlass.Int32(0) + if s_iscalars[1] == cutlass.Int32(3): + s_iscalars[6] = cutlass.Int32(1) + cute.arch.barrier() cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) if cutlass.const_expr(self.enable_p4_rank_scatter): self.phase4_rank_scatter( @@ -4574,6 +4685,36 @@ def _run_phases( warp_id, lane, ) + # ---- plateau fill (done == 3): complete the row from the + # bitwise-equal plateau class. The terminal is only set on an + # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) + # is bitwise-equal; Phase 4 has already emitted the + # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of + # the tie class is a valid tie-aware completion. Ticket counter + # lives in the DEDICATED s_iscalars[7]. + if s_iscalars[6] == cutlass.Int32(1): + pv_lo = s_thr[1] + pv_hi = s_thr[0] + if tidx == cutlass.Int32(0): + s_iscalars[7] = min(s_iscalars[0], cutlass.Int32(self.kC)) + cute.arch.barrier() + ifp = tidx + while ifp < N: + vfp = cutlass.Float32(0.0) + if cutlass.const_expr(self.dtype == cutlass.Float32): + vfp = input_row[ifp] + else: + vfp = cutlass.Float32(input_row[ifp]) + if vfp >= pv_lo and vfp < pv_hi: + pfill = atomicAdd( + s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1) + ) + if pfill < cutlass.Int32(self.top_k): + if cutlass.const_expr(self.return_output_values): + output_values_row[pfill] = self.dtype(vfp) + output_indices_row[pfill] = ifp + ifp = ifp + cutlass.Int32(self.num_threads) + cute.arch.barrier() else: # cs>1: only the leader (CTA 0 in cluster) runs Phase 4. if is_leader: @@ -4620,6 +4761,13 @@ def _run_phases( # smem_keys/smem_vals (no peers to gather from). # ---- Phase 4: histogram snap + writeback ---- + # Capture the P2 terminal BEFORE Phase 4: P4 + # reuses s_iscalars[1] as radix scratch. + if tidx == cutlass.Int32(0): + s_iscalars[6] = cutlass.Int32(0) + if s_iscalars[1] == cutlass.Int32(3): + s_iscalars[6] = cutlass.Int32(1) + cute.arch.barrier() cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) if cutlass.const_expr(self.enable_p4_rank_scatter): self.phase4_rank_scatter( @@ -4652,6 +4800,37 @@ def _run_phases( lane, ) + # ---- plateau fill (done == 3): complete the row from the + # bitwise-equal plateau class. The terminal is only set on an + # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) + # is bitwise-equal; Phase 4 has already emitted the + # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of + # the tie class is a valid tie-aware completion. Ticket counter + # lives in the DEDICATED s_iscalars[7]. + if s_iscalars[6] == cutlass.Int32(1): + pv_lo = s_thr[1] + pv_hi = s_thr[0] + if tidx == cutlass.Int32(0): + s_iscalars[7] = min(s_iscalars[0], cutlass.Int32(self.kC)) + cute.arch.barrier() + ifp = tidx + while ifp < N: + vfp = cutlass.Float32(0.0) + if cutlass.const_expr(self.dtype == cutlass.Float32): + vfp = input_row[ifp] + else: + vfp = cutlass.Float32(input_row[ifp]) + if vfp >= pv_lo and vfp < pv_hi: + pfill = atomicAdd( + s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1) + ) + if pfill < cutlass.Int32(self.top_k): + if cutlass.const_expr(self.return_output_values): + output_values_row[pfill] = self.dtype(vfp) + output_indices_row[pfill] = ifp + ifp = ifp + cutlass.Int32(self.num_threads) + cute.arch.barrier() + # Final cluster barrier: keep peer CTAs (and their SMEM) alive # until the leader's gather + Phase 4 finish. Skipped at # do_cluster_sync=False (no peers; short-row degrade non-leaders 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 0ed81531eba3..6cbc28238222 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 @@ -1211,3 +1211,52 @@ def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): assert ( tuning["enable_warp_parallel_reduce"] == cfg["enable_warp_parallel_reduce"] ) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize( + "variant", + ["rank_scatter_cs1", "rank_scatter_cs4", "snap_cs1"], +) +def test_cute_dsl_gvr_topk_decode_plateau_terminal(dtype, variant): + """Adversarial plateau terminal (done == 3): a bitwise-equal plateau + WIDER than the candidate buffer (kC) straddling the K boundary. Any + threshold either overflows the buffer (>= plateau) or undershoots K + (> plateau), so the Phase-2 bracket collapses to adjacent floats. + The plateau terminal must emit the sure winners plus plateau members + (exact tie-aware) instead of the old -1 pad / unfilled tail.""" + torch.manual_seed(23) + top_k, n, bs = 1024, 32768, 2 + n_hi, n_plateau = 512, 8500 # kC = 6144 for K=1024; 8500 > kC + lo = torch.full((bs, n), -1.0, dtype=dtype, device="cuda") + for r in range(bs): + perm = torch.randperm(n, device="cuda") + hi = perm[:n_hi] + pl = perm[n_hi : n_hi + n_plateau] + lo[r, hi] = (5.0 + torch.arange(n_hi, device="cuda").float() * 0.01).to(dtype) + lo[r, pl] = torch.tensor(1.0, dtype=dtype, device="cuda") + seq_lens = torch.full((bs,), n, dtype=torch.int32, device="cuda") + pre = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + pre[:, 0] = lo.float().argmax(dim=-1).int() + pre[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + overrides = {} + if variant == "rank_scatter_cs4": + overrides["cluster_size"] = 4 + elif variant == "snap_cs1": + overrides["enable_p4_rank_scatter"] = False + _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1, **overrides) + torch.cuda.synchronize() + assert int((out < 0).sum()) == 0, ( + f"{int((out < 0).sum())} unfilled/-1 slots in the plateau terminal" + ) + # in-range, unique + assert out.max() < n and out.min() >= 0 + for r in range(bs): + assert out[r].unique().numel() == top_k + sel = torch.gather(lo.float(), -1, out.long()) + assert int((sel >= 5.0).sum()) == n_hi * bs, "missing sure winners" + assert int((sel == 1.0).sum()) == (top_k - n_hi) * bs, "remaining slots must be plateau members" + ref = torch.topk(lo.float(), top_k, dim=-1).values.sort(-1).values + assert torch.equal(sel.sort(-1).values, ref) From cd6c10600a6fd946e4b6c476d015cf0d331dca74 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:36:30 +0000 Subject: [PATCH 13/19] [None][perf] bsx dispatch: scope the bs=8 fallback band to true 128K shapes A full-grid re-measure showed the earlier bs=8 extension of the 131072 band was too coarse. That bucket (nearest-power-of-two of npad) mixes two unrelated shape families: npad=131136, which holds the 5 production layers whose reg-tier ratio dips to 0.816-0.909 and therefore forces the routing, and npad~163776, which only ROUNDS into the bucket and runs 1.36-1.60x AHEAD of the in-tree kernel on the reg tier. Routing the whole bucket gave up 58 winning cells to protect 5 floor cells. _BAND_LOW_BS_NPAD_MAX caps the low-bs end of the 131072 entry at npad 147456; larger shapes keep their calibrated bs>=16 routing, so the floor guarantee is untouched. Verified by re-measuring all 58 affected cells x BS {8,16,32,64,128,256} (same-rep cold-L2 nsys): bs=8 back on the bsx path at geo 1.608 (min 1.385), bs=16..128 still ~1.0 (routed), bs=256 unchanged at 2.01. Full 9515-case grid with those cells substituted in: gm 1.3996, floor 0.9516, zero cases below 0.909. Band-table test extended with the new boundary. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_bsx_dispatch.py | 26 ++++++++++++++----- .../sparse/test_cute_dsl_bsx_topk_decode.py | 7 ++++- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index b44f4dd9ddb1..9d029fac9502 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -95,12 +95,12 @@ def _thresholds(npad): # Keys are the nearest power-of-two of npad; values are inclusive bs # ranges. TRTLLM_BSX_FALLBACK_BANDS=0 disables the table (bsx serves # every guarded shape). -# 2026-07-29 recalibration: the 131072 band lower bound moved 16 -> 8. -# The original calibration run measured the reg tier with a faster -# experimental streaming phase1 inherited from a side branch; on this -# branch's reg tier the 128K x BS8 shapes dip to 0.82-0.91x vs the -# in-tree kernel (5 production layers), so they are routed as well -# (full-grid gm 1.397 -> 1.390). +# 2026-07-29 recalibration: the 131072 band was extended down to bs=8 for +# TRUE 128K shapes only (see _BAND_LOW_BS_NPAD_MAX). At bs=8 the reg tier +# dips to 0.82-0.91x vs the in-tree kernel on 5 production layers at +# npad=131136, which the floor guarantee cannot admit; the same bucket at +# npad~163776 wins 1.36-1.60x, and those shapes only land in this bucket +# through the nearest-power-of-two rounding, so the extension excludes them. _FALLBACK_BANDS = { 8192: (256, 1 << 30), 16384: (256, 1 << 30), @@ -110,6 +110,12 @@ def _thresholds(npad): 262144: (16, 127), } +# The bs=8 end of the 131072 band applies only up to this npad: shapes above +# it (e.g. npad 163776) round INTO the 131072 bucket but behave like the next +# tier band, where the bsx reg tier is far ahead. They keep the calibrated +# bs>=16 routing. +_BAND_LOW_BS_NPAD_MAX = {131072: 147456} + def _in_fallback_band(bs: int, npad: int) -> bool: if _env_threshold("TRTLLM_BSX_FALLBACK_BANDS") == _BIG: # "0" -> off @@ -117,7 +123,13 @@ def _in_fallback_band(bs: int, npad: int) -> bool: up = 1 << max(npad - 1, 1).bit_length() # pow2 >= npad np2 = up if 4 * npad >= 3 * up else up >> 1 # arithmetic-midpoint nearest band = _FALLBACK_BANDS.get(np2) - return band is not None and band[0] <= bs <= band[1] + if band is None: + return False + lo, hi = band + npad_max = _BAND_LOW_BS_NPAD_MAX.get(np2) + if npad_max is not None and npad > npad_max and lo < 16: + lo = 16 + return lo <= bs <= hi # tier name -> (kind, params). reg params = (cs, tb, maxv, ar). diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index 9fd52095082f..cf24cef91696 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -122,7 +122,12 @@ def test_bsx_fallback_band_table(monkeypatch): assert inb(256, 8192) and inb(1024, 8192) assert inb(16, 32768) and inb(1024, 65536) assert inb(128, 131072) and inb(64, 262144) - assert inb(8, 131072) # 128K x BS8 recalibration (reg-tier floor) + assert inb(8, 131072) and inb(8, 131136) # 128K x BS8 (reg-tier floor) + # ... but shapes that only ROUND into the 131072 bucket keep the + # calibrated bs>=16 routing: at bs=8 the reg tier is 1.36-1.60x ahead + # there, so the low-bs extension must not capture them. + assert not inb(8, 163776) + assert inb(16, 163776) and inb(255, 163776) and not inb(256, 163776) assert inb(48, 40960) # off-grid npad resolves to the nearest pow2 # neighbours stay on bsx assert not inb(128, 8192) # direct/tp win band From e382f9814df5f525146019d36e18a10eea6a8f4e Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:50:02 +0000 Subject: [PATCH 14/19] [None][fix] gvr top-k: plateau terminal on the register-resident admission path Phase 2 has two secant drivers: the SMEM/leader driver and a register-resident redundant-warp driver used at cluster_size == 1. The plateau terminal (done == 3) added earlier only landed in the leader driver, so a bitwise-equal plateau wider than the candidate buffer that straddles the K boundary still fell through to the legacy give-up on the register-resident path and left -1 pads in the output. Mirror both terminals into the register-resident driver: - adjacent-float bracket inside the refine loop -> plateau terminal - budget-exhausted bisection collapse after the loop, then a recount at the terminal threshold so Phase 3 sees sure-winner counts Every thread replays the driver from identical registers, so the added branches stay warp-uniform and block_count_ge keeps its barrier cadence. Extend the adversarial plateau matrix from 3 to 5 variants, adding the classic secant admission (enable_r0=False) at cluster_size 1 and 4 - the route that exposed this gap. GVR top-k suite 684 passed / 0 failed, bsx top-k suite 91 passed / 0 failed. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode.py | 200 ++++++++++++++++-- .../sparse/test_cute_dsl_gvr_topk_decode.py | 15 +- 2 files changed, 196 insertions(+), 19 deletions(-) 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 ea68cd735d6c..103aacfe1066 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 @@ -398,13 +398,20 @@ def __init__( # workloads (25-cell seq-len scan) where R0 wins 24/25 vs the # secant baseline, geomean 1.33x (pro 128k 2.10x). Correctness is # value-set-exact vs torch.topk (186/186 across dtype/K/N/BS/cluster - # + tie plateaus). The secant path is retained verbatim and remains - # reachable via enable_r0=False; it is the exact fallback for the - # large-N / cold-hint (low preIdx hit-rate) regime where R0 can - # regress on the synthetic worst axis — a follow-up PR adds a - # data-driven dispatch guard to route between the two. All R0 fields - # are const-foldable, so an enable_r0=False kernel is byte-identical - # to the pre-R0 upstream base. + # + tie plateaus). + # THE SECANT PATH IS NOT DEAD CODE AND MUST NOT BE DELETED. It has + # two distinct roles: + # (a) EXACT FALLBACK, live at the default enable_r0=True: when the + # rung ladder admits no candidate (R0 miss) the row falls + # through to phase2_secant_search, so this code runs in + # production on every hint-unrepresentative row; + # (b) DIFFERENTIAL ORACLE, via enable_r0=False: it is the classic + # baseline the R0 admission is checked against + # (test_..._r0_equivalence), and the direct-drive entry used to + # bisect admission-vs-baseline regressions. + # All R0 fields are const-foldable, so an enable_r0=False kernel is + # byte-identical to the pre-R0 upstream base and the disabled branch + # costs nothing at runtime. # r0_qfracs: descending h-space quantile fractions defining the M # candidate rungs (ascending threshold values); None => no rungs. # r0_vseed: park P1's pmean (the secant init probe) as one extra @@ -1620,8 +1627,18 @@ def phase2_secant_search( if nv == vlo_r or nv == vhi_r: nv = (vlo_r + vhi_r) * cutlass.Float32(0.5) if nv == vlo_r or nv == vhi_r: - thr_r = vlo_r - done_r = cutlass.Int32(2) + # ADJACENT-FLOAT bracket, same terminal as the + # leader path: a low side over the candidate + # buffer plus a high side under K means the + # boundary sits inside a bitwise-equal plateau + # wider than kC. Keep the sure-winner threshold + # and let Phase 4's plateau fill finish the row. + if clo_r > cutlass.Int32(kCC) and chi_r < cutlass.Int32(kK): + thr_r = vhi_r + done_r = cutlass.Int32(3) + else: + thr_r = vlo_r + done_r = cutlass.Int32(2) if done_r == cutlass.Int32(0): thr_r = nv par_r = par_r ^ cutlass.Int32(1) @@ -1651,6 +1668,75 @@ def phase2_secant_search( vhi_r = thr_r chi_r = cnt_r it = it + cutlass.Int32(1) + # ---- Budget-exhausted plateau collapse (mirrors the leader + # path): the refine budget can run out while the bracket is + # still wide because a tie plateau wider than kC admits no + # threshold. On exactly that signature, bisect to adjacent + # floats so the plateau terminal is exact. Every thread + # replays this from identical registers, so the branch stays + # warp-uniform and block_count_ge keeps its barrier cadence. + if ( + done_r == cutlass.Int32(0) + and clo_r > cutlass.Int32(kCC) + and chi_r >= cutlass.Int32(0) + and chi_r < cutlass.Int32(kK) + ): + itc = cutlass.Int32(0) + while itc < cutlass.Int32(64) and done_r == cutlass.Int32(0): + mid_c = (vlo_r + vhi_r) * cutlass.Float32(0.5) + if mid_c == vlo_r or mid_c == vhi_r: + thr_r = vhi_r + done_r = cutlass.Int32(3) + else: + thr_r = mid_c + par_r = par_r ^ cutlass.Int32(1) + cnt_r = self.block_count_ge( + input_row, + slice_start, + slice_end, + thr_r, + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + cutlass.Boolean(False), # do_cluster_sync (cs==1) + smem_input=smem_input, + redundant=True, + wcnt_off=par_r * cutlass.Int32(nwp2), + ) + if cnt_r >= cutlass.Int32(kK) and cnt_r <= cutlass.Int32(kCC): + done_r = cutlass.Int32(1) + elif cnt_r > cutlass.Int32(kCC): + vlo_r = thr_r + clo_r = cnt_r + else: + vhi_r = thr_r + chi_r = cnt_r + itc = itc + cutlass.Int32(1) + if done_r == cutlass.Int32(3): + # recount at the terminal threshold so Phase 3 sees + # per-thread counts for the sure-winner set. + par_r = par_r ^ cutlass.Int32(1) + cnt_r = self.block_count_ge( + input_row, + slice_start, + slice_end, + thr_r, + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + cutlass.Boolean(False), # do_cluster_sync (cs==1) + smem_input=smem_input, + redundant=True, + wcnt_off=par_r * cutlass.Int32(nwp2), + ) if done_r == cutlass.Int32(0): if clo_r <= cutlass.Int32(kCC * 2): thr_r = vlo_r @@ -1791,6 +1877,84 @@ def phase2_secant_search( cute.arch.barrier() it = it + cutlass.Int32(1) + # ---- Budget-exhausted plateau collapse ---- + # The refine budget can run out while the bracket is still wide: the + # secant step keeps making progress (the bracket shrinks every + # iteration) but a tie plateau wider than kC admits no threshold, so + # the count never lands in [kK, kCC]. In exactly that signature - + # count(>= v_lo) > kCC AND count(>= v_hi) < kK, both counts current - + # collapse the bracket by pure bisection until the ends are ADJACENT + # floats; every value in [v_lo, v_hi) is then bitwise-equal, so the + # plateau terminal (done = 3) is exact and Phase 4 completes the row + # from that tie class. A count landing in [kK, kCC] mid-collapse + # converges normally. Anything else keeps the legacy give-up below. + if ( + s_iscalars[1] == cutlass.Int32(0) + and s_iscalars[2] > cutlass.Int32(kCC) + and s_iscalars[3] >= cutlass.Int32(0) + and s_iscalars[3] < cutlass.Int32(kK) + ): + itc = cutlass.Int32(0) + while itc < cutlass.Int32(64) and s_iscalars[1] == cutlass.Int32(0): + if tidx == 0: + vlo_c = s_thr[1] + vhi_c = s_thr[2] + mid_c = (vlo_c + vhi_c) * cutlass.Float32(0.5) + if mid_c == vlo_c or mid_c == vhi_c: + s_thr[0] = vhi_c + s_iscalars[1] = cutlass.Int32(3) # plateau terminal + else: + s_thr[0] = mid_c + cute.arch.barrier() + if s_iscalars[1] == cutlass.Int32(0): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + smem_input=smem_input, + do_cluster_sync=do_cluster_sync, + ) + if tidx == 0: + c_c = s_iscalars[0] + t_c = s_thr[0] + if c_c >= cutlass.Int32(kK) and c_c <= cutlass.Int32(kCC): + s_iscalars[1] = cutlass.Int32(1) + elif c_c > cutlass.Int32(kCC): + s_thr[1] = t_c + s_iscalars[2] = c_c + else: + s_thr[2] = t_c + s_iscalars[3] = c_c + cute.arch.barrier() + itc = itc + cutlass.Int32(1) + if s_iscalars[1] == cutlass.Int32(3): + # recount at the terminal threshold so Phase 3's cached + # per-thread counts describe the sure-winner set. + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + smem_input=smem_input, + do_cluster_sync=do_cluster_sync, + ) + cute.arch.barrier() + # ---- Post-loop fallback: if still not done, force threshold ---- if tidx == 0: if s_iscalars[1] == cutlass.Int32(0): @@ -4250,10 +4414,12 @@ def _run_phases( # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast # path) or the secant threshold search ---- - # enable_r0 gates to cluster_size==1 for now: R0 scans the - # full row in one CTA. The slice-parallel + cluster count-merge - # variant that lets R0 cover the cs>1 long-row branch lands in a - # later commit; until then cs>1 keeps the secant path. + # R0 covers every cluster size: at cs>1 each CTA scans its own + # slice and block_count_ge_multi cluster-merges the rung counts + # (the P1b rungs are per-CTA identical because the preIdx stats are + # full-row). The secant search below is the exact fallback taken + # when the ladder admits nothing, plus the enable_r0=False + # differential-oracle entry. if cutlass.const_expr(self.enable_r0): # P1b rung placement -> ONE M-ary R0 count pass -> accept the # tightest rung with count in [K, kC]. On a miss, fall back to @@ -4696,7 +4862,9 @@ def _run_phases( pv_lo = s_thr[1] pv_hi = s_thr[0] if tidx == cutlass.Int32(0): - s_iscalars[7] = min(s_iscalars[0], cutlass.Int32(self.kC)) + # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] + # is radix scratch by now (same hazard as the flag). + s_iscalars[7] = cand_count_p4 cute.arch.barrier() ifp = tidx while ifp < N: @@ -4811,7 +4979,9 @@ def _run_phases( pv_lo = s_thr[1] pv_hi = s_thr[0] if tidx == cutlass.Int32(0): - s_iscalars[7] = min(s_iscalars[0], cutlass.Int32(self.kC)) + # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] + # is radix scratch by now (same hazard as the flag). + s_iscalars[7] = cand_count_p4 cute.arch.barrier() ifp = tidx while ifp < N: 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 6cbc28238222..06b89b0382b8 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 @@ -1181,9 +1181,7 @@ def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): """The production runner's tuning adapter must agree with the kernel's pick_cluster_size/pick_tuning single source across a shape sweep (guards the de-duplicated launch-shape policy against drift).""" - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - CuteDSLGvrTopKDecodeRunner as R, - ) + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLGvrTopKDecodeRunner as R num_sms = 148 for torch_dtype in (torch.float32, torch.bfloat16, torch.float16): @@ -1217,7 +1215,11 @@ def test_cute_dsl_gvr_topk_decode_pick_policy_single_source(): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) @pytest.mark.parametrize( "variant", - ["rank_scatter_cs1", "rank_scatter_cs4", "snap_cs1"], + # base_r0off exercises the classic secant admission (the exact fallback the + # production path takes when the R0 ladder misses): its refine budget can + # run out while the bracket is still wide, which is a different route into + # the plateau terminal than the R0 ladder's. + ["rank_scatter_cs1", "rank_scatter_cs4", "snap_cs1", "base_r0off", "base_r0off_cs4"], ) def test_cute_dsl_gvr_topk_decode_plateau_terminal(dtype, variant): """Adversarial plateau terminal (done == 3): a bitwise-equal plateau @@ -1246,6 +1248,11 @@ def test_cute_dsl_gvr_topk_decode_plateau_terminal(dtype, variant): overrides["cluster_size"] = 4 elif variant == "snap_cs1": overrides["enable_p4_rank_scatter"] = False + elif variant == "base_r0off": + overrides["enable_r0"] = False + elif variant == "base_r0off_cs4": + overrides["enable_r0"] = False + overrides["cluster_size"] = 4 _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1, **overrides) torch.cuda.synchronize() assert int((out < 0).sum()) == 0, ( From d29db18566664d87e4f8051a7928016af354526c Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:11:20 +0000 Subject: [PATCH 15/19] [None][fix] bsx review round: order_row accept-and-ignore, device-derived wave budget, env soft-fail, shared tie-aware checker, 16-bit exact-tail default revert Address the five review comments on #16877: 1. order_row: the guard now ACCEPTS and ignores order_row (the LJF hint dsa.py computes for every batch with num_rows >= 2*num_sms). The bsx tiers launch per-row CTAs and never consume the permutation, so rejecting it silently turned bsx off for exactly the large-batch shapes the fallback-band table keeps in service; accepting it aligns deployment behavior with the benched natural-row-order mesh. New TRTLLM_BSX_DISABLE kill switch replaces order_row as the test-side (and operational) force-in-tree mechanism. New test asserts guard acceptance, per-row index-set equality with/without the permutation, and the num_rows >= 2*num_sms production shape. 2. tp_cluster_size: the 296 co-residency budget is now derived as 2 * multi_processor_count (cached; bit-identical on B200), with an explicit note that every other constant remains frozen B200 calibration (family measured HW-invariant in prior cross-arch A/Bs). 3. 16-bit p4_exact_tail default REVERTED to fp32-only (measured): 16-bit quantization puts value plateaus at the K boundary on virtually every input, so the ambiguity gate fires constantly - B200 envelope (bf16 K512/K1024 x 16k-262k x BS 1-256, same-process paired, 26 cells x 2 input kinds) costs gm 1.29-1.36x, worst 2.27x, while typical bf16 inputs are already value-exact without the tail (48/48 paired runs). The repair stays available via the knob; the adversarial 16-bit test now opts in explicitly. 4. env knobs: malformed values fail soft (warn once + baked default) instead of raising ValueError on the first decode step; whitespace tolerated. Covered by a new soft-fail test. 5. _tie_aware_check: single canonical implementation moved to conftest.py, injected as the tie_aware_check fixture (collected by pytest regardless of rootdir/package-resolution style); both test modules drop their copies. The gvr module keeps its pinned-inputs ref-vals memo via the checker's ref_vals_cache parameter. Both suites green locally on B200 (bsx 94 passed / 8 skipped, gvr 684 passed / 144 skipped). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-Authored-By: Claude Fable 5 --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 11 +- .../blackwell/top_k/gvr_topk_decode.py | 21 +- .../top_k/gvr_topk_decode_bsx_dispatch.py | 45 ++- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 18 +- .../_torch/attention/sparse/conftest.py | 136 +++++++++ .../sparse/test_cute_dsl_bsx_topk_decode.py | 269 ++++++++++-------- .../sparse/test_cute_dsl_gvr_topk_decode.py | 184 ++++-------- 7 files changed, 423 insertions(+), 261 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/conftest.py diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 316d5a70297e..7e55cc241b6b 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7536,10 +7536,13 @@ def forward( # BSX tier fast path: fp32 / next_n >= 1 (MTP) / # cr in {1, 4} / npad <= 262144 decode rows route to the # direct/reg/tp CuTe DSL tiers; everything else (half-prec, LB, - # sort-indirect, oversize npad, hw cluster cap) falls through - # to the in-tree kernel below. Host-only guard — no device - # sync. The op signature and output contract are unchanged - # (unordered int32 indices, -1 pad only for degenerate rows). + # oversize npad, hw cluster cap) falls through to the in-tree + # kernel below. ``order_row`` (the LJF hint dsa.py computes for + # num_rows >= 2 * num_sms) is accepted and ignored: the bsx + # tiers launch per-row CTAs and do not consume the permutation. + # Host-only guard — no device sync. The op signature and output + # contract are unchanged (unordered int32 indices, -1 pad only + # for degenerate rows). if _is_bsx_supported(logits, pre_idx, seq_lens, output_indices, top_k, next_n, compress_ratio, order_row, counters): 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 103aacfe1066..44275ab832c3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py @@ -556,13 +556,22 @@ def __init__( # boundary-local ULP (e.g. fp16 1.0 vs 1.25 under a [0, 65504] # bracket); values straddling the kK boundary inside one fine bin # were previously picked in arrival order (observed as |miss|=1 - # with |dv| ~ 3e-6 on real captures). The tail radix re-ranks on - # the full fp32 order key — candidate keys are ALWAYS fp32 (16-bit - # inputs are upcast injectively at collect), so the repair is exact - # for every supported dtype. Default ON for rank-scatter-exact - # kernels of all dtypes. + # with |dv| ~ 3e-6 on real fp32 captures). The tail radix re-ranks + # on the full fp32 order key — candidate keys are ALWAYS fp32 + # (16-bit inputs are upcast injectively at collect), so the repair + # is exact for every supported dtype WHEN ENABLED. Default ON for + # fp32 only: on fp32 the gate fires rarely and the fix is ~free, + # but 16-bit quantization puts value plateaus at the boundary on + # virtually every input, so the gate fires constantly — measured + # B200 envelope cost (bf16, K512/K1024 x 16k-262k x BS 1-256, + # same-process paired) is gm 1.29-1.36x, worst 2.27x, while typical + # bf16 inputs (randn and quantized-tie, 48 paired runs) are already + # value-exact without it: 16-bit misses need an adversarially wide + # Phase-2 bracket (distinct 16-bit values inside one fine bin). + # 16-bit callers that need the guarantee opt in via the knob (see + # test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit). if p4_exact_tail is None: - p4_exact_tail = self.enable_p4_rank_scatter_exact + 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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index 9d029fac9502..771e4a6b7d62 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -25,10 +25,18 @@ Dispatcher guard (anything else falls back to the in-tree ``GvrTopKKernel`` path in ``CuteDSLGvrTopKDecodeRunner.forward``): dtype == fp32, next_n >= 1 (MTP; num_rows divisible by next_n), - compress_ratio in {1, 4}, order_row is None, counters is None, + compress_ratio in {1, 4}, counters is None, K in {512, 1024, 2048}, npad <= 262144, npad % 64 == 0, contiguous 16B-aligned tensors, and the routed tier's cluster size within the queried hardware max (never silently degrade the cluster shape). +``order_row`` is accepted and IGNORED: it is the LJF scheduling hint the +caller (``dsa.py``) computes for the in-tree persistent kernel whenever +num_rows >= 2 * num_sms, and the bsx tiers launch per-row CTAs that the +hardware schedules directly — the permutation affects neither their +correctness nor their measured performance (the tier mesh was benched in +natural row order). Rejecting it would silently turn bsx off for every +large batch in production. ``counters`` (LB mode) still falls back: the +LB partition contract belongs to the in-tree kernel. pre_idx / seq_lens are request-level ([num_rows // next_n, K] / [num_rows // next_n]) — the in-tree contract. Per-row degeneracy (N_eff <= K) and ragged N are handled INSIDE the tiers, so the guard @@ -40,12 +48,19 @@ TRTLLM_BSX_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); else the bs threshold at which the tp tier takes over. TRTLLM_BSX_DENSE_BS same, for the dense (tb=1024) reg tiers. + TRTLLM_BSX_DISABLE any value other than unset/""/"0" -> the guard + rejects everything (kill switch: every call takes the + in-tree kernel path). +Malformed numeric values fail soft: a warning is logged once and the knob +falls back to unset (-1) instead of raising on the decode path. """ import os import torch +from tensorrt_llm.logger import logger + from .gvr_topk_decode_bsx_direct import DKCMAX, direct_topk from .gvr_topk_decode_bsx_reg import reg_topk from .gvr_topk_decode_bsx_tp import tp_cluster_size, tp_topk @@ -60,13 +75,33 @@ def _env_threshold(name): t = _ENV.get(name) if t is None: e = os.environ.get(name) - t = int(e) if e not in (None, "") else -1 + if e is None or e.strip() == "": + t = -1 + else: + try: + t = int(e.strip()) + except ValueError: + # Tuning override with a safe baked default: fail soft + # (treat as unset) instead of killing the decode path. + logger.warning( + f"{name}={e!r} is not an integer; ignoring the override " + f"and using the baked default." + ) + t = -1 if t == 0: t = _BIG _ENV[name] = t return t +def _env_flag(name): + t = _ENV.get(name) + if t is None: + t = os.environ.get(name, "").strip() not in ("", "0") + _ENV[name] = t + return t + + def _reset_env_cache(): _ENV.clear() _DISPATCH_CACHE.clear() # routes depend on the env thresholds @@ -241,11 +276,15 @@ def is_bsx_supported( ) -> bool: """Host-only guard for the bsx tiers (no device sync; see module docstring). Returns False -> caller uses the in-tree kernel.""" + if _env_flag("TRTLLM_BSX_DISABLE"): + return False if logits.dtype != torch.float32: return False if next_n < 1 or compress_ratio not in (1, 4): return False - if order_row is not None or counters is not None: + # order_row is accepted and ignored (in-tree scheduling hint; see the + # module docstring). counters (LB mode) keeps the in-tree path. + if counters is not None: return False if top_k not in (512, 1024, 2048): return False diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index a569d5acf652..f2d090462268 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -1831,6 +1831,21 @@ def _p2floor(x: int) -> int: return p +def _two_waves_rows() -> int: + """2 * SM count of the current device (cached): the co-residency row + budget behind the CS selection below (min_blocks_per_mp=2). Mirrors + ``_get_num_sms`` in ``cute_dsl_custom_ops.py`` (not importable from + here — the import runs the other way). NOTE: only this term scales + with the device; every other constant in ``tp_cluster_size`` and the + dispatch band tables is frozen B200 calibration (the kernel family + measured HW-invariant across B200/B300 in the op9/op17 cross-arch + A/Bs, so the bands are kept device-independent on purpose). + """ + if not hasattr(_two_waves_rows, "_value"): + _two_waves_rows._value = 2 * torch.cuda.get_device_properties().multi_processor_count + return _two_waves_rows._value + + def _get_compiled(K: int, CS: int, AR: int, UF: int, TB: int, next_n: int = 1, cr: int = 4): key = (K, CS, AR, UF, TB, next_n, cr) compiled = _LAUNCH_CACHE.get(key) @@ -1888,7 +1903,8 @@ def tp_cluster_size(bs: int, npad: int) -> int: return 1 cs = 1 if bs < 128: - cs = _p2floor(296 // bs) if bs <= 296 else 1 + two_waves = _two_waves_rows() # 296 on B200 (2 x 148 SMs) + cs = _p2floor(two_waves // bs) if bs <= two_waves else 1 capn = _p2floor(npad // 8192 if npad // 8192 > 0 else 1) if cs > capn: cs = capn diff --git a/tests/unittest/_torch/attention/sparse/conftest.py b/tests/unittest/_torch/attention/sparse/conftest.py new file mode 100644 index 000000000000..60e6b73f470b --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/conftest.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared checkers for the top-K decode test modules in this directory. + +The tie-aware checker lives here (single definition) and is handed to tests +via the ``tie_aware_check`` fixture: conftest is collected by pytest +regardless of rootdir/package-resolution style, so this sidesteps the +cross-module import problem that previously forced a keep-in-sync duplicate +of the checker in ``test_cute_dsl_bsx_topk_decode.py``. +""" + +import pytest +import torch + + +def _tie_aware_check_impl( + out_indices: torch.Tensor, + logits: torch.Tensor, + seq_lens: torch.Tensor, + top_k: int, + next_n: int, + compress_ratio: int = 1, + ref_vals_cache: dict = None, +) -> None: + """Vectorized multi-row tie-aware correctness check with strict sort+allclose. + + Per row r: scan range is ``logits[r, :N_eff(r)]`` where N_eff mirrors + the kernel's exact formula (see ``GvrTopKKernel.gvr_topk_kernel``): + + actual_kv_len = seq_lens[r // next_n] - next_n + (r % next_n) + 1 + N_eff = actual_kv_len // compress_ratio # cr=1 is identity + + Reference ``torch.topk`` is masked to this range so the reference and + kernel scan exactly the same columns under any (next_n, cr) combo + (including next_n>=3 + cr>=2 where the floor-division makes per-row + N_eff vary within a group). + + All checks (out-of-range, duplicates, n_below, sort+allclose) run as + batched GPU ops; only assertion-failure diagnostics fall back to host. + + ``ref_vals_cache``: optional id-keyed memo dict for the reference top-K + values. Pass it ONLY when (logits, seq_lens) are pinned alive for the + process lifetime by the caller (id-keying on transient tensors would + alias after GC); the caller owns that guarantee. + """ + num_rows, top_k_out = out_indices.shape + assert top_k_out == top_k + device = logits.device + logits_f32 = logits.to(torch.float32) + N = logits.shape[1] + + # Per-row N_eff mirroring the kernel formula. seq_lens is per-group + # (length num_rows // next_n); broadcast across the next_n rows of + # each group before computing actual_kv_len // cr. + row_idx = torch.arange(num_rows, device=device) + group_idx = row_idx // next_n + ofs = row_idx % next_n + seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] + actual_kv_len = seq_lens_per_row - next_n + ofs + 1 + N_eff = actual_kv_len // compress_ratio # [num_rows] + + # Reference per-row top-K, sorted descending, over logits masked beyond + # per-row N_eff. + ref_key = None + if ref_vals_cache is not None: + ref_key = (id(logits), id(seq_lens), top_k, next_n, compress_ratio) + if ref_key is not None and ref_key in ref_vals_cache: + ref_vals = ref_vals_cache[ref_key] + else: + col_idx = torch.arange(N, device=device) + in_range_mask = col_idx[None, :] < N_eff[:, None] # [num_rows, N] + masked_logits = torch.where(in_range_mask, logits_f32, float("-inf")) + ref_vals, _ = torch.topk(masked_logits, k=top_k, largest=True, sorted=True, dim=-1) + if ref_key is not None: + ref_vals_cache[ref_key] = ref_vals + + # ---- 1. Out-of-range / -1 placeholder check (single fused mask) ---- + out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) + if bool(out_of_range.any().item()): + bad_row = int(out_of_range.any(dim=1).int().argmax().item()) + bad_indices = out_indices[bad_row].cpu().tolist() + raise AssertionError( + f"row={bad_row}: kernel returned out-of-range index " + f"(N_eff={int(N_eff[bad_row].item())}, indices={bad_indices})" + ) + + # ---- 2. Duplicate-index check (sort each row, scan consecutive eq) ---- + sorted_idx, _ = out_indices.sort(dim=-1) + has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) + if bool(has_dup.any().item()): + bad_row = int(has_dup.int().argmax().item()) + raise AssertionError( + f"row={bad_row}: kernel returned duplicate indices: " + f"{out_indices[bad_row].cpu().tolist()}" + ) + + # ---- 3. Gather selected values (safe — already in range) ---- + sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) + + # ---- 4. n_below check vs per-row K-th value ---- + kth_vals = ref_vals[:, -1:] # [num_rows, 1] + n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) + if bool((n_below_per_row > 0).any().item()): + bad_row = int(n_below_per_row.argmax().item()) + n_below = int(n_below_per_row[bad_row].item()) + kth = float(kth_vals[bad_row, 0].item()) + raise AssertionError( + f"row={bad_row}: {n_below} selected values < Kth-rank value ({kth:.6f})" + ) + + # ---- 5. Strict: sorted-value multiset == torch.topk reference ---- + sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) + diff = (sel_sorted - ref_vals).abs() + if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): + per_row_max = diff.max(dim=-1).values + bad_row = int(per_row_max.argmax().item()) + max_diff = float(per_row_max[bad_row].item()) + raise AssertionError(f"row={bad_row}: sorted-value mismatch — max diff {max_diff:.4e}") + + +@pytest.fixture +def tie_aware_check(): + """The shared tie-aware top-K correctness checker (plain function).""" + return _tie_aware_check_impl diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index cf24cef91696..d80a2bc80dfa 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -23,11 +23,13 @@ host-only route-table asserts, a dispatcher-fallback check (bf16 routes to the in-tree kernel), and the MTP axis: next_n in {2, 3, 4} x cr in {1, 4} x all three tiers, each case checked against BOTH the torch.topk host -N_eff/offset simulation (``_tie_aware_check``) and a differential in-tree -arm (same inputs through the in-tree kernel via ``order_row``, per-row -value-multiset equality). +N_eff/offset simulation (the shared ``tie_aware_check`` conftest fixture) +and a differential in-tree arm (same inputs through the in-tree kernel via +the ``TRTLLM_BSX_DISABLE`` kill switch, per-row value-multiset equality). """ +import os + import pytest import torch @@ -64,7 +66,7 @@ def _bsx_bands_off(monkeypatch): @skip_not_sm100 @pytest.mark.parametrize("bands", ["on", "off"]) -def test_bsx_cuda_graph_capture_replay(monkeypatch, bands): +def test_bsx_cuda_graph_capture_replay(monkeypatch, bands, tie_aware_check): """The op must be CUDA-graph capturable and replay-consistent on both sides of the fallback band table. Shape (npad=32768, bs=64) sits inside a fallback bucket: bands=on captures the in-tree branch, bands=off the @@ -98,7 +100,7 @@ def call(): call() g.replay() torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # replay over in-place rewritten inputs (fresh logits + fresh hints) logits2, pre_idx2, seq_lens2 = _make_bsx_inputs(bs, npad, top_k, seed=4321) @@ -107,7 +109,7 @@ def call(): seq_lens.copy_(seq_lens2) g.replay() torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) def test_bsx_fallback_band_table(monkeypatch): @@ -142,84 +144,6 @@ def test_bsx_fallback_band_table(monkeypatch): bsx_dispatch._reset_env_cache() -# --------------------------------------------------------------------------- -# Shared helpers. ``_tie_aware_check`` is a local copy of the one in -# test_cute_dsl_gvr_topk_decode.py (same directory): the sibling test module -# is not reliably importable across the repo's pytest invocation styles -# (rootdir-dependent package resolution), so the checker is duplicated here -# verbatim-in-spirit with the same semantics. Keep the two in sync. -# --------------------------------------------------------------------------- -def _tie_aware_check( - out_indices: torch.Tensor, - logits: torch.Tensor, - seq_lens: torch.Tensor, - top_k: int, - next_n: int = NEXT_N, - compress_ratio: int = CR, -) -> None: - """Vectorized multi-row tie-aware correctness check (strict sort+allclose). - - Per row r the scan range is ``logits[r, :N_eff(r)]`` where - ``N_eff = (seq_lens[r // next_n] - next_n + r % next_n + 1) // cr`` - (the kernels' exact formula). Checks: in-range indices, no duplicates, - no selected value below the K-th reference value, and sorted-value - multiset equality against torch.topk. - """ - num_rows, top_k_out = out_indices.shape - assert top_k_out == top_k - device = logits.device - logits_f32 = logits.to(torch.float32) - N = logits.shape[1] - - row_idx = torch.arange(num_rows, device=device) - group_idx = row_idx // next_n - ofs = row_idx % next_n - seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] - actual_kv_len = seq_lens_per_row - next_n + ofs + 1 - N_eff = actual_kv_len // compress_ratio # [num_rows] - - col_idx = torch.arange(N, device=device) - in_range_mask = col_idx[None, :] < N_eff[:, None] - 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) - - out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) - if bool(out_of_range.any().item()): - bad_row = int(out_of_range.any(dim=1).int().argmax().item()) - raise AssertionError( - f"row={bad_row}: out-of-range index " - f"(N_eff={int(N_eff[bad_row].item())}, " - f"indices={out_indices[bad_row].cpu().tolist()})" - ) - - sorted_idx, _ = out_indices.sort(dim=-1) - has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) - if bool(has_dup.any().item()): - bad_row = int(has_dup.int().argmax().item()) - raise AssertionError( - f"row={bad_row}: duplicate indices: {out_indices[bad_row].cpu().tolist()}" - ) - - sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) - kth_vals = ref_vals[:, -1:] - n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) - if bool((n_below_per_row > 0).any().item()): - bad_row = int(n_below_per_row.argmax().item()) - raise AssertionError( - f"row={bad_row}: {int(n_below_per_row[bad_row].item())} selected " - f"values < Kth-rank value ({float(kth_vals[bad_row, 0].item()):.6f})" - ) - - sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) - if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): - per_row_max = (sel_sorted - ref_vals).abs().max(dim=-1).values - bad_row = int(per_row_max.argmax().item()) - raise AssertionError( - f"row={bad_row}: sorted-value mismatch — max diff " - f"{float(per_row_max[bad_row].item()):.4e}" - ) - - def _make_bsx_inputs( bs: int, npad: int, @@ -400,7 +324,9 @@ def test_bsx_route_env_knobs(monkeypatch): ids=[t[4] + f"_n{t[0]}_bs{t[1]}_k{t[2]}" for t in _REG_INSTANCES], ) @pytest.mark.parametrize("kind", ["randn", "ties"]) -def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, launch, kind, monkeypatch): +def test_bsx_reg_launch_table( + npad, bs, top_k, dense_env, expected, launch, kind, monkeypatch, tie_aware_check +): if not launch and kind == "ties": pytest.skip("route-assert-only instance (single kind suffices)") _skip_if_cluster_capped(bs, npad, top_k) @@ -414,7 +340,7 @@ def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, launch, kind _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected) if launch: out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) finally: if dense_env is not None: monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) @@ -435,13 +361,13 @@ def test_bsx_reg_launch_table(npad, bs, top_k, dense_env, expected, launch, kind (4096, 8, 512, "randn", False), # uniform N_eff == npad (no tail) ], ) -def test_bsx_direct(npad, bs, top_k, kind, varlen): +def test_bsx_direct(npad, bs, top_k, kind, varlen, tie_aware_check): logits, pre_idx, seq_lens = _make_bsx_inputs( bs, npad, top_k, seed=npad * 3 + bs, kind=kind, varlen=varlen ) _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "direct") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # --------------------------------------------------------------------------- @@ -457,14 +383,14 @@ def test_bsx_direct(npad, bs, top_k, kind, varlen): (262144, 128, 2048, "randn"), # cs=1, uf=8 streaming path ], ) -def test_bsx_tp(npad, bs, top_k, kind): +def test_bsx_tp(npad, bs, top_k, kind, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) logits, pre_idx, seq_lens = _make_bsx_inputs( bs, npad, top_k, seed=npad + 7 * bs, kind=kind, varlen=True ) _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # --------------------------------------------------------------------------- @@ -480,7 +406,7 @@ def test_bsx_tp(npad, bs, top_k, kind): (65536, 16, 1024, "tp"), ], ) -def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind): +def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) logits, pre_idx, seq_lens = _make_bsx_inputs( bs, npad, top_k, seed=11, kind="randn", varlen=True @@ -508,7 +434,14 @@ def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind): f"{out[i, : max(ne + 2, 8)].cpu().tolist()}..." ) else: - _tie_aware_check(out[i : i + 1], logits[i : i + 1], seq_lens[i : i + 1], top_k) + tie_aware_check( + out[i : i + 1], + logits[i : i + 1], + seq_lens[i : i + 1], + top_k, + next_n=NEXT_N, + compress_ratio=CR, + ) # --------------------------------------------------------------------------- @@ -524,13 +457,13 @@ def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind): (65536, 16, 1024), # tp tier ], ) -def test_bsx_preidx_hardening(npad, bs, top_k, preidx): +def test_bsx_preidx_hardening(npad, bs, top_k, preidx, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) logits, pre_idx, seq_lens = _make_bsx_inputs( bs, npad, top_k, seed=13, kind="randn", varlen=True, preidx=preidx ) out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # --------------------------------------------------------------------------- @@ -556,7 +489,7 @@ def test_bsx_preidx_hardening(npad, bs, top_k, preidx): (131136, 128, 512), # cs=1 streaming tp (flash 512k production shape) ], ) -def test_bsx_tp_admission_hitrate(npad, bs, top_k, hit_rate): +def test_bsx_tp_admission_hitrate(npad, bs, top_k, hit_rate, tie_aware_check): """High-hr rows must admit (1-pass) and low-hr rows must stay exact via the pivot/secant fallback; both paths must produce exact top-K.""" _skip_if_cluster_capped(bs, npad, top_k) @@ -565,11 +498,11 @@ def test_bsx_tp_admission_hitrate(npad, bs, top_k, hit_rate): ) _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 -def test_bsx_tp_admission_tie_plateau(): +def test_bsx_tp_admission_tie_plateau(tie_aware_check): """Tie plateau AT the admission threshold: K/2 distinct high values + a 3K-wide exact-tie plateau straddling the K-th rank. The hint set is the true top-K, so the ladder rungs land ON the plateau value and the @@ -590,12 +523,12 @@ def test_bsx_tp_admission_tie_plateau(): pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 @pytest.mark.parametrize("top_k,n_plateau", [(1024, 8000), (2048, 12000)]) -def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau): +def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau, tie_aware_check): """count(>= any admissible threshold) > kC: K-1 distinct high values, then an n_plateau-wide exact-tie plateau (n_plateau > kC = 6144/8192). Every threshold at or below the plateau overflows the candidate buffer @@ -618,11 +551,11 @@ def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau): pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 -def test_bsx_tp_admission_mixed_batch(): +def test_bsx_tp_admission_mixed_batch(tie_aware_check): """Ragged batch mixing rows that admit (true-top-K hints) with rows that must fall back (all-zero cold-start hints): per-row admission escape is independent, so every row must stay exact regardless of @@ -635,7 +568,7 @@ def test_bsx_tp_admission_mixed_batch(): pre_idx[1::2] = 0 # odd rows: cold-start (degenerate ladder -> fallback) _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) - _tie_aware_check(out, logits, seq_lens, top_k) + tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # --------------------------------------------------------------------------- @@ -687,12 +620,94 @@ def test_bsx_dispatcher_fallback_bad_shapes(): assert not ok(logits, pre_idx, sl_row, out, top_k, 2, 4, None, None) +@skip_not_sm100 +def test_bsx_accepts_order_row(tie_aware_check): + """``order_row`` (the LJF hint dsa.py computes for every batch with + num_rows >= 2 * num_sms) must NOT turn bsx off: the guard accepts it and + the tiers ignore it, so the per-row index SET is identical with and + without the permutation (emission order is unordered by contract). + Reuses the (4096, 256, 512) tp cell compiled by ``test_bsx_tp`` (no + extra JIT).""" + npad, bs, top_k = 4096, 256, 512 + _skip_if_cluster_capped(bs, npad, top_k) + logits, pre_idx, seq_lens = _make_bsx_inputs( + bs, npad, top_k, seed=npad + 7 * bs, kind="randn", varlen=True + ) + order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() + out_dummy = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + assert bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out_dummy, top_k, NEXT_N, CR, order_row, None + ), "guard must accept order_row (scheduling hint, ignored by the tiers)" + out_plain = _run_op(logits, pre_idx, seq_lens, top_k) + out_ordered = torch.empty_like(out_plain) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_ordered, + top_k=top_k, + next_n=NEXT_N, + compress_ratio=CR, + order_row=order_row, + ) + torch.cuda.synchronize() + assert torch.equal(out_plain.sort(-1).values, out_ordered.sort(-1).values), ( + "bsx per-row index set must be independent of the ignored order_row permutation" + ) + tie_aware_check(out_ordered, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) + # Production shape sanity (guard-only, no launch): num_rows >= 2*num_sms + # — the batch band where dsa.py always sends order_row — stays accepted. + big_bs = 2 * torch.cuda.get_device_properties().multi_processor_count + 8 + logits_b = torch.randn(big_bs, npad, dtype=torch.float32, device="cuda") + pre_b = torch.zeros(big_bs, top_k, dtype=torch.int32, device="cuda") + sl_b = torch.full((big_bs,), npad * CR, dtype=torch.int32, device="cuda") + out_b = torch.empty(big_bs, top_k, dtype=torch.int32, device="cuda") + order_b = torch.argsort(sl_b.long(), descending=True).int().contiguous() + assert bsx_dispatch.is_bsx_supported( + logits_b, pre_b, sl_b, out_b, top_k, NEXT_N, CR, order_b, None + ), "large-batch (num_rows >= 2*num_sms) calls with order_row must stay on bsx" + + +@skip_not_sm100 +def test_bsx_disable_kill_switch(monkeypatch): + """TRTLLM_BSX_DISABLE rejects everything at the guard (host-only).""" + bs, npad, top_k = 4, 65536, 512 + logits = torch.randn(bs, npad, dtype=torch.float32, device="cuda") + seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") + pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + ok = bsx_dispatch.is_bsx_supported + assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + monkeypatch.setenv("TRTLLM_BSX_DISABLE", "1") + bsx_dispatch._reset_env_cache() + assert not ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + monkeypatch.delenv("TRTLLM_BSX_DISABLE", raising=False) + bsx_dispatch._reset_env_cache() + assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) + + +def test_bsx_env_malformed_soft_fail(monkeypatch): + """Malformed tuning-knob values fail soft (warn + baked default), never + raise on the decode path.""" + baseline = bsx_dispatch.route(256, 20480, 512) + for bad in ("true", " 16x", "8.5"): + monkeypatch.setenv("TRTLLM_BSX_TP_BS", bad) + bsx_dispatch._reset_env_cache() + assert bsx_dispatch.route(256, 20480, 512) == baseline + # whitespace-padded but valid values still parse. + monkeypatch.setenv("TRTLLM_BSX_TP_BS", " 0 ") + bsx_dispatch._reset_env_cache() + assert bsx_dispatch.route(256, 20480, 512).startswith(("reg", "direct")) + monkeypatch.delenv("TRTLLM_BSX_TP_BS", raising=False) + bsx_dispatch._reset_env_cache() + + # --------------------------------------------------------------------------- # MTP (next_n > 1) + cr in {1, 4}: exactness on all three tiers, checked # against BOTH the torch.topk host N_eff/offset simulation -# (``_tie_aware_check``) and a differential in-tree arm — the SAME inputs -# through the in-tree kernel (forced via ``order_row``, which the bsx guard -# rejects), per-row value-multiset equality. Ragged: per-request varlen +# (``tie_aware_check``) and a differential in-tree arm — the SAME inputs +# through the in-tree kernel (forced via ``TRTLLM_BSX_DISABLE``; the guard +# accepts+ignores order_row), per-row value-multiset equality. Ragged: # seq_lens make N_eff differ across requests AND (via row % next_n) across # the MTP rows of one request; tails are poisoned with +1e30. # --------------------------------------------------------------------------- @@ -751,8 +766,10 @@ def _make_mtp_inputs( def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): - """bsx arm + in-tree differential arm (order_row forces the in-tree - sort path; the bsx guard rejects order_row). Returns (out_bsx, out_ref).""" + """bsx arm + in-tree differential arm (the ``TRTLLM_BSX_DISABLE`` kill + switch forces the in-tree path; ``order_row`` no longer does — the guard + accepts and ignores it). The ref call still passes ``order_row`` so the + in-tree sort-indirect path stays exercised. Returns (out_bsx, out_ref).""" num_rows = logits.shape[0] out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") assert bsx_dispatch.is_bsx_supported( @@ -763,16 +780,22 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): ) order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() out_ref = torch.empty_like(out) - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - logits, - pre_idx, - seq_lens, - out_ref, - top_k=top_k, - next_n=next_n, - compress_ratio=cr, - order_row=order_row, - ) + os.environ["TRTLLM_BSX_DISABLE"] = "1" + bsx_dispatch._reset_env_cache() + try: + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_ref, + top_k=top_k, + next_n=next_n, + compress_ratio=cr, + order_row=order_row, + ) + finally: + del os.environ["TRTLLM_BSX_DISABLE"] + bsx_dispatch._reset_env_cache() torch.cuda.synchronize() return out, out_ref @@ -817,7 +840,7 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): "tier,next_n,cr", _MTP_COMBOS, ids=[f"{t}-{n}-{c}" for t, n, c in _MTP_COMBOS] ) @pytest.mark.parametrize("kind,preidx", _MTP_KINDS, ids=[f"{k}-{p}" for k, p in _MTP_KINDS]) -def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx): +def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx, tie_aware_check): npad, bs_req, top_k = _MTP_TIER_CELLS[tier] num_rows = bs_req * next_n if tier == "tp" and num_rows < 16: @@ -830,7 +853,7 @@ def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx): ) out, out_ref = _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr) # Independent torch.topk + host N_eff/offset simulation reference. - _tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) # Differential oracle: per-row value multiset equal to the in-tree arm. sel = torch.gather(logits, -1, out.long()).sort(-1, descending=True).values sel_ref = torch.gather(logits, -1, out_ref.long()).sort(-1, descending=True).values @@ -842,7 +865,7 @@ def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx): @skip_not_sm100 @pytest.mark.parametrize("next_n,cr", [(2, 1), (3, 4)]) @pytest.mark.parametrize("preidx", ["zeros", "random"]) -def test_bsx_mtp_preidx_hardening(next_n, cr, preidx): +def test_bsx_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): """bsx-only MTP robustness: hint sets that VIOLATE the op's argmax- slot-0 contract (pure zeros / out-of-range garbage) must neither fault nor break exactness on the bsx tiers. No differential arm here: the @@ -872,13 +895,13 @@ def test_bsx_mtp_preidx_hardening(next_n, cr, preidx): logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr ) torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) + tie_aware_check(out, logits, seq_lens, top_k, next_n=next_n, compress_ratio=cr) @skip_not_sm100 @pytest.mark.parametrize("next_n", [2, 3]) @pytest.mark.parametrize("cr", [1, 4]) -def test_bsx_mtp_degenerate_rows(next_n, cr): +def test_bsx_mtp_degenerate_rows(next_n, cr, tie_aware_check): """Degenerate MTP requests (N_eff <= K for some or all of the next_n rows): identity emit [0..N_eff-1] + -1 pad must hold per ROW, with the per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr.""" @@ -912,7 +935,7 @@ def test_bsx_mtp_degenerate_rows(next_n, cr): # into an equivalent next_n=1 seq_lens so the shared checker # scans exactly N_eff(r) columns. sl_row = seq_lens[r // next_n : r // next_n + 1] - next_n + (r % next_n) + 1 - _tie_aware_check( + tie_aware_check( out[r : r + 1], logits[r : r + 1], sl_row, top_k, next_n=1, compress_ratio=cr ) nd = (n_eff > top_k).nonzero().flatten() 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 06b89b0382b8..392d6a4e953a 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 @@ -133,6 +133,19 @@ def _make_inputs_impl( _ref_vals_cache: dict = {} +def _gvr_check(check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1): + """Delegate to the shared conftest checker (``tie_aware_check`` fixture), + attaching this module's reference-values memo only when (logits, + seq_lens) are pinned for the process lifetime by ``_inputs_cache`` + (id-keyed caching on transient tensors would alias after GC).""" + ref = ( + _ref_vals_cache + if any(logits is v[0] and seq_lens is v[2] for v in _inputs_cache.values()) + else None + ) + check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio, ref_vals_cache=ref) + + def _make_inputs( num_rows: int, N: int, @@ -179,108 +192,6 @@ def _make_inputs( return _inputs_cache[key] -def _tie_aware_check( - out_indices: torch.Tensor, - logits: torch.Tensor, - seq_lens: torch.Tensor, - top_k: int, - next_n: int, - compress_ratio: int = 1, -) -> None: - """Vectorized multi-row tie-aware correctness check with strict sort+allclose. - - Per row r: scan range is ``logits[r, :N_eff(r)]`` where N_eff mirrors - the kernel's exact formula (see ``GvrTopKKernel.gvr_topk_kernel``): - - actual_kv_len = seq_lens[r // next_n] - next_n + (r % next_n) + 1 - N_eff = actual_kv_len // compress_ratio # cr=1 is identity - - Reference ``torch.topk`` is masked to this range so the reference and - kernel scan exactly the same columns under any (next_n, cr) combo - (including next_n>=3 + cr>=2 where the floor-division makes per-row - N_eff vary within a group). - - All checks (out-of-range, duplicates, n_below, sort+allclose) run as - batched GPU ops; only assertion-failure diagnostics fall back to host. - """ - num_rows, top_k_out = out_indices.shape - assert top_k_out == top_k - device = logits.device - logits_f32 = logits.to(torch.float32) - N = logits.shape[1] - - # Per-row N_eff mirroring the kernel formula. seq_lens is per-group - # (length num_rows // next_n); broadcast across the next_n rows of - # each group before computing actual_kv_len // cr. - row_idx = torch.arange(num_rows, device=device) - group_idx = row_idx // next_n - ofs = row_idx % next_n - seq_lens_per_row = seq_lens.to(device=device, dtype=torch.long)[group_idx] - actual_kv_len = seq_lens_per_row - next_n + ofs + 1 - N_eff = actual_kv_len // compress_ratio # [num_rows] - - # Reference per-row top-K, sorted descending, over logits masked beyond - # per-row N_eff. Memoized when (logits, seq_lens) come from the pinned - # ``_inputs_cache`` (identity match), since the reference only depends on - # (logits, seq_lens, top_k, next_n, compress_ratio) — not on the launch - # variant (cluster_size / order_row / ...) the test is exercising. - ref_key = None - if any(logits is v[0] and seq_lens is v[2] for v in _inputs_cache.values()): - ref_key = (id(logits), id(seq_lens), top_k, next_n, compress_ratio) - if ref_key is not None and ref_key in _ref_vals_cache: - ref_vals = _ref_vals_cache[ref_key] - else: - col_idx = torch.arange(N, device=device) - in_range_mask = col_idx[None, :] < N_eff[:, None] # [num_rows, N] - masked_logits = torch.where(in_range_mask, logits_f32, float("-inf")) - ref_vals, _ = torch.topk(masked_logits, k=top_k, largest=True, sorted=True, dim=-1) - if ref_key is not None: - _ref_vals_cache[ref_key] = ref_vals - - # ---- 1. Out-of-range / -1 placeholder check (single fused mask) ---- - out_of_range = (out_indices < 0) | (out_indices >= N_eff[:, None]) - if bool(out_of_range.any().item()): - bad_row = int(out_of_range.any(dim=1).int().argmax().item()) - bad_indices = out_indices[bad_row].cpu().tolist() - raise AssertionError( - f"row={bad_row}: kernel returned out-of-range index " - f"(N_eff={int(N_eff[bad_row].item())}, indices={bad_indices})" - ) - - # ---- 2. Duplicate-index check (sort each row, scan consecutive eq) ---- - sorted_idx, _ = out_indices.sort(dim=-1) - has_dup = (sorted_idx[:, 1:] == sorted_idx[:, :-1]).any(dim=-1) - if bool(has_dup.any().item()): - bad_row = int(has_dup.int().argmax().item()) - raise AssertionError( - f"row={bad_row}: kernel returned duplicate indices: " - f"{out_indices[bad_row].cpu().tolist()}" - ) - - # ---- 3. Gather selected values (safe — already in range) ---- - sel_vals = torch.gather(logits_f32, dim=-1, index=out_indices.long()) - - # ---- 4. n_below check vs per-row K-th value ---- - kth_vals = ref_vals[:, -1:] # [num_rows, 1] - n_below_per_row = (sel_vals < kth_vals).sum(dim=-1) - if bool((n_below_per_row > 0).any().item()): - bad_row = int(n_below_per_row.argmax().item()) - n_below = int(n_below_per_row[bad_row].item()) - kth = float(kth_vals[bad_row, 0].item()) - raise AssertionError( - f"row={bad_row}: {n_below} selected values < Kth-rank value ({kth:.6f})" - ) - - # ---- 5. Strict: sorted-value multiset == torch.topk reference ---- - sel_sorted, _ = sel_vals.sort(dim=-1, descending=True) - diff = (sel_sorted - ref_vals).abs() - if not bool(torch.allclose(sel_sorted, ref_vals, rtol=1e-5, atol=1e-5)): - per_row_max = diff.max(dim=-1).values - bad_row = int(per_row_max.argmax().item()) - max_diff = float(per_row_max[bad_row].item()) - raise AssertionError(f"row={bad_row}: sorted-value mismatch — max diff {max_diff:.4e}") - - @skip_not_sm100 @pytest.mark.parametrize( "dtype,top_k", @@ -312,6 +223,7 @@ def test_cute_dsl_gvr_topk_decode( compress_ratio, preidx_hit_rate, cluster_size, + tie_aware_check, ): """Compare custom op output against torch.topk reference (tie-aware). @@ -359,7 +271,9 @@ def test_cute_dsl_gvr_topk_decode( ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) @skip_not_sm100 @@ -380,7 +294,7 @@ def test_cute_dsl_gvr_topk_decode( ], ) def test_cute_dsl_gvr_topk_decode_seqlen_sorted( - dtype, top_k, N, batch_size, varlen, next_n, compress_ratio, cluster_size + dtype, top_k, N, batch_size, varlen, next_n, compress_ratio, cluster_size, tie_aware_check ): """LJF host-side dispatch order: ``order_row`` = descending argsort of ``seq_lens`` passed through the custom op. @@ -421,7 +335,9 @@ def test_cute_dsl_gvr_topk_decode_seqlen_sorted( ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) # =========================================================================== @@ -441,7 +357,9 @@ def test_cute_dsl_gvr_topk_decode_seqlen_sorted( "dtype,top_k", [(torch.bfloat16, 512), (torch.float32, 2048)], ) -def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, cluster_size): +def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary( + dtype, top_k, cluster_size, tie_aware_check +): """GVR top-K multi-CTA short-row degrade: correctness at the cluster transition boundary. When ``cluster_size > 1``, each row is dispatched to a cluster of @@ -512,8 +430,14 @@ def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, clu cluster_size=cluster_size, ) torch.cuda.synchronize() - _tie_aware_check( - out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + _gvr_check( + tie_aware_check, + out_indices, + logits, + seq_lens, + top_k, + next_n, + compress_ratio=compress_ratio, ) # Mixed batch: alternating degrade (even rows) and co-op (odd rows). @@ -548,7 +472,9 @@ def test_cute_dsl_gvr_topk_multi_cta_shortrow_degrade_boundary(dtype, top_k, clu cluster_size=cluster_size, ) torch.cuda.synchronize() - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio) + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=compress_ratio + ) # =========================================================================== @@ -638,7 +564,9 @@ def test_lb_prepare_partition(B, ratio): ) @pytest.mark.parametrize("batch_size", [4, 32]) @pytest.mark.parametrize("next_n", [1, 2]) -def test_lb_main_branches(dtype, top_k, scenario, N, seq_lens_mode, batch_size, next_n): +def test_lb_main_branches( + dtype, top_k, scenario, N, seq_lens_mode, batch_size, next_n, tie_aware_check +): """Each LB branch (all_long / all_short / mixed) produces correct top-K. For ``mixed_half`` half the rows are forced to be short (seq_len < threshold) @@ -703,7 +631,7 @@ def test_lb_main_branches(dtype, top_k, scenario, N, seq_lens_mode, batch_size, elif scenario == "mixed_half": assert n_long == batch_size - batch_size // 2 - _tie_aware_check(out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1) + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, next_n, compress_ratio=1) @skip_not_sm100 @@ -733,6 +661,7 @@ def test_lb_vs_reference( batch_size, compress_ratio, preidx_hit_rate, + tie_aware_check, ): """LB kernel output matches torch.topk tie-aware reference across the same param sweep used by the single-CTA UT.""" @@ -779,7 +708,8 @@ def test_lb_vs_reference( max_batch_size=max_batch_size, ) torch.cuda.synchronize() - _tie_aware_check( + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, @@ -934,7 +864,9 @@ def _make_r0_pre_idx(logits, top_k, hint, seed): @pytest.mark.parametrize("batch_size", [1, 16]) @pytest.mark.parametrize("hint", ["real", "rand"]) @pytest.mark.parametrize("cluster_size", [1, 4, 8]) -def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hint, cluster_size): +def test_cute_dsl_gvr_topk_decode_r0_equivalence( + dtype, top_k, N, batch_size, hint, cluster_size, tie_aware_check +): """R0 admission (``enable_r0=True``, the new default) selects the same top-K as the secant baseline (``enable_r0=False``), by index set. @@ -965,8 +897,8 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hi ) # 1. Both arms independently produce a valid top-K (tie-aware value set). - _tie_aware_check(out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) - _tie_aware_check(out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) # 2. Equivalence. fp32 logits are ALMOST tie-free, so R0 and base must # return the identical index set (order-independent) — but randn @@ -997,7 +929,7 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence(dtype, top_k, N, batch_size, hi ) @pytest.mark.parametrize("hint", ["real", "rand"]) def test_cute_dsl_gvr_topk_decode_r0_equivalence_bigbs( - dtype, top_k, N, batch_size, hint, cluster_size + dtype, top_k, N, batch_size, hint, cluster_size, tie_aware_check ): """Big-batch R0-vs-secant equivalence: multi-wave grids only. @@ -1017,8 +949,8 @@ def test_cute_dsl_gvr_topk_decode_r0_equivalence_bigbs( out_r0 = _run_gvr_direct( logits, pre_idx, seq_lens, top_k, enable_r0=True, cluster_size=cluster_size ) - _tie_aware_check(out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) - _tie_aware_check(out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_base, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_r0, logits, seq_lens, top_k, next_n=1, compress_ratio=1) if dtype == torch.float32: _assert_index_sets_equal_tie_aware(out_base, out_r0, logits) @@ -1111,7 +1043,7 @@ def test_cute_dsl_gvr_topk_decode_pick_config_policy(): (torch.bfloat16, 1024, 65536, 256), # cs=1 multi-wave big-BS ], ) -def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size): +def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size, tie_aware_check): """``GvrTopKKernel.launch`` (pick_config + variant cache) produces a valid top-K at every launch-shape regime the policy can pick, including cluster_size=8. Direct-drive users get production-equivalent shapes.""" @@ -1125,14 +1057,14 @@ def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size) _GvrTopKKernel.launch(logits, pre_idx, seq_lens, out, top_k) torch.cuda.synchronize() - _tie_aware_check(out, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out, logits, seq_lens, top_k, next_n=1, compress_ratio=1) # Override path: forcing the secant arm through launch() must also be a # valid top-K and (fp32, tie-free) the identical index set. out_sec = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") _GvrTopKKernel.launch(logits, pre_idx, seq_lens, out_sec, top_k, enable_r0=False) torch.cuda.synchronize() - _tie_aware_check(out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1) + _gvr_check(tie_aware_check, out_sec, logits, seq_lens, top_k, next_n=1, compress_ratio=1) if dtype == torch.float32: assert torch.equal(out.sort(dim=-1).values, out_sec.sort(dim=-1).values) @@ -1143,8 +1075,12 @@ def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): """16-bit exact-tail adversarial: two DISTINCT half-precision values in ONE fine bin straddling the K boundary (window-relative binning cannot separate them under a wide Phase-2 bracket — e.g. fp16 1.0 vs 1.25 - under a [0, 65504]-scale bracket). With p4_exact_tail now default-on - for all dtypes the tail radix must keep every 1.25 above every 1.0.""" + under a [0, 65504]-scale bracket). The 16-bit default stays OFF (the + ambiguity gate fires on virtually every 16-bit input, measured gm + 1.29-1.36x envelope cost, while typical 16-bit inputs are value-exact + without the tail); this covers the explicit OPT-IN: with + ``p4_exact_tail=True`` the tail radix must keep every 1.25 above + every 1.0.""" torch.manual_seed(11) top_k, n = 1024, 32768 bs = 2 @@ -1165,7 +1101,7 @@ def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): pre[:, 0] = lo.float().argmax(dim=-1).int() pre[:, 1:] = torch.arange(1, top_k, dtype=torch.int32, device="cuda") out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") - _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1) + _GvrTopKKernel.launch(lo, pre, seq_lens, out, top_k, compress_ratio=1, p4_exact_tail=True) torch.cuda.synchronize() sel = torch.gather(lo.float(), -1, out.long()) # every 1.25 must be selected before any 1.0 fills the remainder From fd97411067e256aeb94723720eaea33a572b82b1 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:24:39 +0000 Subject: [PATCH 16/19] [None][chore] conftest: explicit optional type for ref_vals_cache (RUF013) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Co-Authored-By: Claude Fable 5 --- tests/unittest/_torch/attention/sparse/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/attention/sparse/conftest.py b/tests/unittest/_torch/attention/sparse/conftest.py index 60e6b73f470b..85d0aa06c3e2 100644 --- a/tests/unittest/_torch/attention/sparse/conftest.py +++ b/tests/unittest/_torch/attention/sparse/conftest.py @@ -32,7 +32,7 @@ def _tie_aware_check_impl( top_k: int, next_n: int, compress_ratio: int = 1, - ref_vals_cache: dict = None, + ref_vals_cache: dict | None = None, ) -> None: """Vectorized multi-row tie-aware correctness check with strict sort+allclose. From abee8b8c22eaee8e5b0a20c48ea84b988a55076a Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:01:52 +0000 Subject: [PATCH 17/19] [None][chore] bsx: memoize cluster-cap verdict, doc fallback-bands knob, scrub leftover codenames; pin gvr tests to the in-tree path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (PR #16877): - dispatch: is_bsx_supported re-ran route()+_parse_reg (~2.5-3us host) on every eager forward — the cost _DISPATCH_CACHE exists to avoid. Memoize the cluster-cap verdict per (bs, npad, K); cleared by _reset_env_cache since routes depend on the env thresholds. Covered by test_bsx_cluster_cap_verdict_memoized. - dispatch: document TRTLLM_BSX_FALLBACK_BANDS in the module-docstring env-knob list (it was introduced with the band table but missing from the knob family reference). - bsx_tp: scrub three internal development codenames that survived the codename-scrub commit (two pre-existing, one introduced by the SM-count follow-up); neutral phrasing, load-bearing facts kept. - gvr tests: autouse fixture pins TRTLLM_BSX_DISABLE=1 — this module tests the in-tree kernel contract, but its op-level fp32 cases fall inside the bsx envelope and were silently routed to the bsx tiers (which also ignore the explicit cluster_size they parametrize over). Mirrors (inverted) the bands-off fixture in the bsx test module. Tested: bsx dispatcher/guard subset 9 passed; gvr fp32/K2048 op-level subset 96 passed / 32 skipped (pre-existing skips) on B200. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_bsx_dispatch.py | 20 +++++++++--- .../blackwell/top_k/gvr_topk_decode_bsx_tp.py | 10 +++--- .../sparse/test_cute_dsl_bsx_topk_decode.py | 31 +++++++++++++++++++ .../sparse/test_cute_dsl_gvr_topk_decode.py | 19 ++++++++++++ 4 files changed, 71 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py index 771e4a6b7d62..7c14f4589fab 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py @@ -48,6 +48,9 @@ TRTLLM_BSX_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); else the bs threshold at which the tp tier takes over. TRTLLM_BSX_DENSE_BS same, for the dense (tb=1024) reg tiers. + TRTLLM_BSX_FALLBACK_BANDS + 0 -> disable the measured fallback-band table (bsx + serves every guarded shape); unset/other -> active. TRTLLM_BSX_DISABLE any value other than unset/""/"0" -> the guard rejects everything (kill switch: every call takes the in-tree kernel path). @@ -105,6 +108,7 @@ def _env_flag(name): def _reset_env_cache(): _ENV.clear() _DISPATCH_CACHE.clear() # routes depend on the env thresholds + _CAP_OK_CACHE.clear() # ditto (verdict embeds the routed tier) def _thresholds(npad): @@ -236,6 +240,7 @@ def route_cluster_size(bs: int, npad: int, K: int) -> int: # [env thresholds are cached at first use, like the CUDA static locals], so # bind it once per key to a closure. _DISPATCH_CACHE = {} # (bs, npad, K, next_n, cr) -> callable(logits, pre, seq_lens, out) +_CAP_OK_CACHE = {} # (bs, npad, K) -> routed tier's cluster size within hw max def _bind(bs, npad, K, next_n, cr): @@ -316,10 +321,17 @@ def is_bsx_supported( ): return False # Cluster cap: fall back to the in-tree kernel (dispatcher-level) rather - # than silently degrading the tier's cluster shape. - if route_cluster_size(bs, npad, top_k) > _query_max_cluster_size(): - return False - return True + # than silently degrading the tier's cluster shape. The verdict is pure + # in (bs, npad, K) once the env thresholds are cached, and computing it + # re-runs route() + the _parse_reg string parse (the ~2.5-3us host cost + # _DISPATCH_CACHE exists to avoid) — so memoize it the same way. + key = (bs, npad, top_k) + ok = _CAP_OK_CACHE.get(key) + if ok is None: + ok = _CAP_OK_CACHE[key] = ( + route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() + ) + return ok def bsx_topk( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py index f2d090462268..64b4482293fc 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py @@ -35,7 +35,7 @@ P1 : hint gather + minmax + two-stage 64-bin histogram -> CCDF rung ladder P2a : uniform every-32nd-float4 sampled multi-rung count (+ per-rung float4 occupancy for a clustering-aware sigma) -> 4-stage pivot: - stage 0 is the hint-ladder ADMISSION (in-tree R0 op#26 parity) — + stage 0 is the hint-ladder ADMISSION (in-tree R0 parity) — tightest rung whose sampled CI sits inside [K, 0.6*kC] (the legacy pivot-band hi) — stage 0b prefers the legacy band pick when it is strictly leaner and safe by its clustering-aware @@ -1197,8 +1197,8 @@ def gvr_tp_kernel( tgt = cutlass.Int32(tgt_py) best = cutlass.Int32(AR - 1) bestd = cutlass.Int32(0x7FFFFFFF) - # Stage 0 — hint-ladder ADMISSION (in-tree R0 parity, - # op#26 lineage): accept the TIGHTEST ladder rung + # Stage 0 — hint-ladder ADMISSION (in-tree R0 + # parity): accept the TIGHTEST ladder rung # (highest threshold = smallest admitted-candidate set) # whose sampled-count confidence interval lies inside # the [K, kC] acceptance window, mirroring the R0 rule @@ -1838,8 +1838,8 @@ def _two_waves_rows() -> int: here — the import runs the other way). NOTE: only this term scales with the device; every other constant in ``tp_cluster_size`` and the dispatch band tables is frozen B200 calibration (the kernel family - measured HW-invariant across B200/B300 in the op9/op17 cross-arch - A/Bs, so the bands are kept device-independent on purpose). + measured HW-invariant in earlier B200/B300 cross-arch A/Bs, so the + bands are kept device-independent on purpose). """ if not hasattr(_two_waves_rows, "_value"): _two_waves_rows._value = 2 * torch.cuda.get_device_properties().multi_processor_count diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py index d80a2bc80dfa..81329f73106d 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py @@ -288,6 +288,37 @@ def test_bsx_route_env_knobs(monkeypatch): bsx_dispatch._reset_env_cache() +@skip_not_sm100 +def test_bsx_cluster_cap_verdict_memoized(): + """The guard's cluster-cap verdict is memoized per (bs, npad, K) — + the per-call route()+_parse_reg host cost the _DISPATCH_CACHE comment + motivates must not be re-paid by is_bsx_supported on every eager + forward — and _reset_env_cache clears it (routes depend on the env + thresholds).""" + bs, npad, top_k = 4, 65536, 512 + logits, pre_idx, seq_lens = _make_bsx_inputs(bs, npad, top_k, seed=0) + out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") + + bsx_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in bsx_dispatch._CAP_OK_CACHE + ok = bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ) + expected = bsx_dispatch.route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() + assert ok == expected + assert bsx_dispatch._CAP_OK_CACHE[(bs, npad, top_k)] == expected + # Second call must hit the memo (same verdict, no recompute observable + # beyond the cache entry staying put). + assert ( + bsx_dispatch.is_bsx_supported( + logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None + ) + == expected + ) + bsx_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in bsx_dispatch._CAP_OK_CACHE + + # --------------------------------------------------------------------------- # reg tier: every launch-table instance's ROUTE is asserted (host-only, # free); a LAUNCH (JIT compile ~5-7s each) runs only for the 6 instances 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 392d6a4e953a..1c00e7146bc8 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py @@ -22,6 +22,9 @@ from cutlass.cute import runtime as _crt import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_bsx_dispatch as _bsx_dispatch, +) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import ( GvrTopKKernel as _GvrTopKKernel, ) @@ -33,6 +36,22 @@ ) +@pytest.fixture(autouse=True) +def _bsx_off(monkeypatch): + """This module tests the IN-TREE kernel contract. Its op-level cases + (fp32 shapes within the bsx envelope) would otherwise be routed to the + bsx tiers by the dispatcher — which also ignores the explicit + ``cluster_size`` these tests parametrize over — so pin the op to the + in-tree path. The bsx tiers are covered by + ``test_cute_dsl_bsx_topk_decode.py`` (which flips the routing the other + way: fallback bands off so every test reaches a bsx tier).""" + monkeypatch.setenv("TRTLLM_BSX_DISABLE", "1") + _bsx_dispatch._reset_env_cache() + yield + monkeypatch.delenv("TRTLLM_BSX_DISABLE", raising=False) + _bsx_dispatch._reset_env_cache() + + def _make_inputs_impl( num_rows: int, N: int, From 4e4abbd164c7d4c86b7650810354da768d9eec73 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:41:37 +0000 Subject: [PATCH 18/19] [None][chore] drop the BSX working name: the tiers are GVR optimizations, not a different algorithm Review follow-up (@limin2021): pure rename, zero logic change. The three tiers run the same Guess-Verify-Refine algorithm as the existing kernel and differ only in where the row's logits are resident (direct / reg / tp), so the development codename goes away in favor of GVR-family names: - files: gvr_topk_decode_bsx_{dispatch,direct,reg,tp}.py -> gvr_topk_decode_{dispatch,direct,reg,tp}.py (git mv); test_cute_dsl_bsx_topk_decode.py -> test_cute_dsl_gvr_topk_tiers.py (l0_b300.yml updated) - symbols: bsx_topk -> tiered_topk, is_bsx_supported -> is_tiered_topk_supported - env knobs: TRTLLM_BSX_{DISABLE,FALLBACK_BANDS,TP_BS,DENSE_BS} -> TRTLLM_GVR_{TIERS_DISABLE,FALLBACK_BANDS,TP_BS,DENSE_BS} - the unreachable 'gvr(cs=16,tb=512)' route label becomes 'cluster(cs=16,tb=512)' so 'GVR' unambiguously names the algorithm family, not one tier - prose: 'bsx' -> 'the (GVR) tiers'; 'in-tree kernel' still names the pre-existing single kernel (gvr_topk_decode.py) Tested post-rename on B200: tiers suite 95 passed / 8 skipped; GVR suite 684 passed / 144 skipped (4-way GPU-sharded, counts identical to the pre-rename baseline). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 16 +- .../blackwell/top_k/__init__.py | 12 +- ...sx_direct.py => gvr_topk_decode_direct.py} | 10 +- ...ispatch.py => gvr_topk_decode_dispatch.py} | 81 +++-- ...code_bsx_reg.py => gvr_topk_decode_reg.py} | 8 +- ...decode_bsx_tp.py => gvr_topk_decode_tp.py} | 4 +- .../test_lists/test-db/l0_b300.yml | 4 +- .../_torch/attention/sparse/conftest.py | 2 +- .../sparse/test_cute_dsl_gvr_topk_decode.py | 22 +- ...ode.py => test_cute_dsl_gvr_topk_tiers.py} | 282 +++++++++--------- 10 files changed, 220 insertions(+), 221 deletions(-) rename tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/{gvr_topk_decode_bsx_direct.py => gvr_topk_decode_direct.py} (98%) rename tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/{gvr_topk_decode_bsx_dispatch.py => gvr_topk_decode_dispatch.py} (84%) rename tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/{gvr_topk_decode_bsx_reg.py => gvr_topk_decode_reg.py} (99%) rename tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/{gvr_topk_decode_bsx_tp.py => gvr_topk_decode_tp.py} (99%) rename tests/unittest/_torch/attention/sparse/{test_cute_dsl_bsx_topk_decode.py => test_cute_dsl_gvr_topk_tiers.py} (82%) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 7e55cc241b6b..95c3f61e3e00 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7271,10 +7271,10 @@ def warmup_cute_dsl_radix_topk_decode( # ------------------------------------------------------------------ # from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import \ GvrTopKKernel as _GvrTopKKernel - from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_bsx_dispatch import \ - bsx_topk as _bsx_topk - from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_bsx_dispatch import \ - is_bsx_supported as _is_bsx_supported + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ + tiered_topk as _tiered_topk + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ + is_tiered_topk_supported as _is_tiered_topk_supported class CuteDSLGvrTopKDecodeRunner: """Runner for the GVR Top-K cuTe DSL kernel (Blackwell SM100). @@ -7533,20 +7533,20 @@ def forward( ``counters`` without ``order_row`` is rejected. """ - # BSX tier fast path: fp32 / next_n >= 1 (MTP) / + # Tiered-GVR fast path: fp32 / next_n >= 1 (MTP) / # cr in {1, 4} / npad <= 262144 decode rows route to the # direct/reg/tp CuTe DSL tiers; everything else (half-prec, LB, # oversize npad, hw cluster cap) falls through to the in-tree # kernel below. ``order_row`` (the LJF hint dsa.py computes for - # num_rows >= 2 * num_sms) is accepted and ignored: the bsx + # num_rows >= 2 * num_sms) is accepted and ignored: the GVR # tiers launch per-row CTAs and do not consume the permutation. # Host-only guard — no device sync. The op signature and output # contract are unchanged (unordered int32 indices, -1 pad only # for degenerate rows). - if _is_bsx_supported(logits, pre_idx, seq_lens, output_indices, + if _is_tiered_topk_supported(logits, pre_idx, seq_lens, output_indices, top_k, next_n, compress_ratio, order_row, counters): - _bsx_topk(logits, pre_idx, seq_lens, output_indices, top_k, + _tiered_topk(logits, pre_idx, seq_lens, output_indices, top_k, next_n, compress_ratio) return diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index 042e8e70fe2f..a28e9a62fcfe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -17,10 +17,10 @@ from .filtered_top_k_decode_varlen import FilteredTopKKernelVarlenDecode from .filtered_top_k_varlen_util import FilteredTopKKernelVarlen from .gvr_topk_decode import GvrParams, GvrTopKKernel -from .gvr_topk_decode_bsx_direct import DirectTopKKernel -from .gvr_topk_decode_bsx_dispatch import bsx_topk, is_bsx_supported -from .gvr_topk_decode_bsx_reg import GvrRegKernel -from .gvr_topk_decode_bsx_tp import GvrTpKernel +from .gvr_topk_decode_direct import DirectTopKKernel +from .gvr_topk_decode_dispatch import tiered_topk, is_tiered_topk_supported +from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel __all__ = [ @@ -32,6 +32,6 @@ "GvrTpKernel", "GvrRegKernel", "DirectTopKKernel", - "bsx_topk", - "is_bsx_supported", + "tiered_topk", + "is_tiered_topk_supported", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py similarity index 98% rename from tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py rename to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py index 098a110cc1fc..5dbe90ced305 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_direct.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_direct.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""BSX direct (short-row) top-K tier — CuTe DSL, Blackwell SM100. +"""Direct (short-row) GVR top-K tier — CuTe DSL, Blackwell SM100. CuTe DSL translation of the original CUDA ``direct_topk_kernel``, adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract (ragged N via a device ``seq_lens`` tensor + per-row degenerate identity emit; see the module -docstring of ``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory). +docstring of ``gvr_topk_decode_tp`` for the shared adaptation inventory). Exact top-K indices for short padded rows (npad <= DKCMAX = 12288), one CTA per row, TB = 1024 threads. No threshold solve: the whole row is collected @@ -46,7 +46,7 @@ from cutlass.utils.smem_allocator import SmemAllocator FLT_MAX = 3.4028234663852886e38 -DKCMAX = 12288 # direct-path candidate capacity (mirrors gvr_bsx.cu) +DKCMAX = 12288 # direct-path candidate capacity (mirrors the CUDA development arm) @cute.jit @@ -69,7 +69,7 @@ class DirectTopKKernel: (TB, default 1024), ``next_n`` / ``compress_ratio`` for the per-row N_eff arithmetic (constexpr; the direct tier reads no hints, so MTP support is the N_eff formula alone). SMEM layout - mirrors gvr_bsx.cu's DSmem with the same packed (key << 32 | idx) + mirrors the CUDA development arm's DSmem with the same packed (key << 32 | idx) u64 layout (measured: split key/idx arrays cost 2x smem transactions in collect + emits vs CUDA's single STS.64 / LDS.64 per candidate): @@ -98,7 +98,7 @@ def __init__( # ------------------------------------------------------------------ # Per-row valid length (ragged N) — mirrors the in-tree run_one_row - # arithmetic exactly (see gvr_topk_decode_bsx_tp._row_n_eff). + # arithmetic exactly (see gvr_topk_decode_tp._row_n_eff). # ------------------------------------------------------------------ @cute.jit def _row_n_eff(self, seq_lens: cute.Tensor, row): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py similarity index 84% rename from tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py rename to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py index 7c14f4589fab..57dd2e3f03de 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py @@ -12,15 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""BSX top-K tier dispatcher — routes ``cute_dsl_gvr_topk_decode`` calls to -the BSX CuTe DSL tiers (direct / reg / tp). +"""Tiered-GVR top-K dispatcher — routes ``cute_dsl_gvr_topk_decode`` calls +to the CuTe DSL GVR tiers (direct / reg / tp). Exact transcription of the original CUDA implementation's dispatch (``gvr_topk_launch_batched``) (``gvr_topk_launch_batched`` + ``launch_dense`` + ``launch_tp<512>``). The gvr streaming tier (cs=16 fallback for npad > 262144) is intentionally NOT ported: it is unreachable inside the deployment envelope -(npad <= 262144), which :func:`is_bsx_supported` enforces. +(npad <= 262144), which :func:`is_tiered_topk_supported` enforces. Dispatcher guard (anything else falls back to the in-tree ``GvrTopKKernel`` path in ``CuteDSLGvrTopKDecodeRunner.forward``): @@ -31,10 +31,10 @@ queried hardware max (never silently degrade the cluster shape). ``order_row`` is accepted and IGNORED: it is the LJF scheduling hint the caller (``dsa.py``) computes for the in-tree persistent kernel whenever -num_rows >= 2 * num_sms, and the bsx tiers launch per-row CTAs that the +num_rows >= 2 * num_sms, and the GVR tiers launch per-row CTAs that the hardware schedules directly — the permutation affects neither their correctness nor their measured performance (the tier mesh was benched in -natural row order). Rejecting it would silently turn bsx off for every +natural row order). Rejecting it would silently turn the tiers off for every large batch in production. ``counters`` (LB mode) still falls back: the LB partition contract belongs to the in-tree kernel. pre_idx / seq_lens are request-level ([num_rows // next_n, K] / @@ -42,16 +42,15 @@ (N_eff <= K) and ragged N are handled INSIDE the tiers, so the guard needs no device sync. -Env knobs (identical semantics to the CUDA arm's ``GVR_BSX_*`` static locals, -renamed for the production tree; cached at first use like the CUDA static -locals; :func:`_reset_env_cache` re-reads for tests): - TRTLLM_BSX_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); +Env knobs (identical semantics to the CUDA development arm's static locals; +cached at first use like those; :func:`_reset_env_cache` re-reads for tests): + TRTLLM_GVR_TP_BS unset/-1 -> baked per-npad bands; 0 -> disable (2^30); else the bs threshold at which the tp tier takes over. - TRTLLM_BSX_DENSE_BS same, for the dense (tb=1024) reg tiers. - TRTLLM_BSX_FALLBACK_BANDS - 0 -> disable the measured fallback-band table (bsx - serves every guarded shape); unset/other -> active. - TRTLLM_BSX_DISABLE any value other than unset/""/"0" -> the guard + TRTLLM_GVR_DENSE_BS same, for the dense (tb=1024) reg tiers. + TRTLLM_GVR_FALLBACK_BANDS + 0 -> disable the measured fallback-band table (the + tiers serve every guarded shape); unset/other -> active. + TRTLLM_GVR_TIERS_DISABLE any value other than unset/""/"0" -> the guard rejects everything (kill switch: every call takes the in-tree kernel path). Malformed numeric values fail soft: a warning is logged once and the knob @@ -64,9 +63,9 @@ from tensorrt_llm.logger import logger -from .gvr_topk_decode_bsx_direct import DKCMAX, direct_topk -from .gvr_topk_decode_bsx_reg import reg_topk -from .gvr_topk_decode_bsx_tp import tp_cluster_size, tp_topk +from .gvr_topk_decode_direct import DKCMAX, direct_topk +from .gvr_topk_decode_reg import reg_topk +from .gvr_topk_decode_tp import tp_cluster_size, tp_topk from .single_pass_multi_cta_radix_topk_cluster import _query_max_cluster_size _BIG = 1 << 30 @@ -112,17 +111,17 @@ def _reset_env_cache(): def _thresholds(npad): - tpb = _env_threshold("TRTLLM_BSX_TP_BS") + tpb = _env_threshold("TRTLLM_GVR_TP_BS") if tpb < 0: tpb = 256 if npad <= 20480 else (128 if npad < 32768 else 16) - dnb = _env_threshold("TRTLLM_BSX_DENSE_BS") + dnb = _env_threshold("TRTLLM_GVR_DENSE_BS") if dnb < 0: dnb = 8 if npad >= 163840 else (64 if npad < 32768 else _BIG) return tpb, dnb # Measured fallback band table (full-grid calibration, 2026-07-28): -# (npad, bs) buckets where the in-tree kernel is faster than every bsx +# (npad, bs) buckets where the in-tree kernel is faster than every GVR # tier by >1.10x on at least one production layer (865 real decode cells # x 11 BS, same-rep cold-L2 nsys pairs; to recalibrate, re-run that # paired sweep and route every (npad, bs) bucket whose per-case floor @@ -130,10 +129,10 @@ def _thresholds(npad): # in-tree exact-count ladder admits a leaner candidate set and its # row-slice cluster split keeps every CTA busy through P4. Routing them # to the in-tree kernel caps the worst case at 1.10x while keeping the -# bsx win elsewhere (full-grid gm 1.40 vs the in-tree head). +# tier win elsewhere (full-grid gm 1.40 vs the in-tree head). # Keys are the nearest power-of-two of npad; values are inclusive bs -# ranges. TRTLLM_BSX_FALLBACK_BANDS=0 disables the table (bsx serves -# every guarded shape). +# ranges. TRTLLM_GVR_FALLBACK_BANDS=0 disables the table (the tiers +# serve every guarded shape). # 2026-07-29 recalibration: the 131072 band was extended down to bs=8 for # TRUE 128K shapes only (see _BAND_LOW_BS_NPAD_MAX). At bs=8 the reg tier # dips to 0.82-0.91x vs the in-tree kernel on 5 production layers at @@ -151,13 +150,13 @@ def _thresholds(npad): # The bs=8 end of the 131072 band applies only up to this npad: shapes above # it (e.g. npad 163776) round INTO the 131072 bucket but behave like the next -# tier band, where the bsx reg tier is far ahead. They keep the calibrated +# tier band, where the reg tier is far ahead. They keep the calibrated # bs>=16 routing. _BAND_LOW_BS_NPAD_MAX = {131072: 147456} def _in_fallback_band(bs: int, npad: int) -> bool: - if _env_threshold("TRTLLM_BSX_FALLBACK_BANDS") == _BIG: # "0" -> off + if _env_threshold("TRTLLM_GVR_FALLBACK_BANDS") == _BIG: # "0" -> off return False up = 1 << max(npad - 1, 1).bit_length() # pow2 >= npad np2 = up if 4 * npad >= 3 * up else up >> 1 # arithmetic-midpoint nearest @@ -174,8 +173,8 @@ def _in_fallback_band(bs: int, npad: int) -> bool: # tier name -> (kind, params). reg params = (cs, tb, maxv, ar). def route(bs: int, npad: int, K: int) -> str: """Tier the CUDA gvr_topk_launch_batched would take. Returns - 'tp' | 'direct' | 'reg(cs=C,tb=T,maxv=M,ar=A)' | 'gvr(cs=16,tb=512)' - ('gvr' is unreachable through :func:`bsx_topk` — see module docstring).""" + 'tp' | 'direct' | 'reg(cs=C,tb=T,maxv=M,ar=A)' | 'cluster(cs=16,tb=512)' + ('cluster' is unreachable through :func:`tiered_topk` — see module docstring).""" tpb, dnb = _thresholds(npad) if bs >= tpb: return "tp" @@ -191,7 +190,7 @@ def route(bs: int, npad: int, K: int) -> str: return "reg(cs=4,tb=1024,maxv=8,ar=8)" if npad <= 262144: return "reg(cs=8,tb=1024,maxv=8,ar=8)" - return "gvr(cs=16,tb=512)" + return "cluster(cs=16,tb=512)" # latency ladder if npad <= DKCMAX: return "direct" @@ -213,7 +212,7 @@ def route(bs: int, npad: int, K: int) -> str: if K == 2048: return "reg(cs=16,tb=512,maxv=8,ar=8)" return "reg(cs=16,tb=512,maxv=8,ar=6)" - return "gvr(cs=16,tb=512)" + return "cluster(cs=16,tb=512)" def _parse_reg(tier): @@ -232,7 +231,7 @@ def route_cluster_size(bs: int, npad: int, K: int) -> int: return 1 if tier.startswith("reg"): return _parse_reg(tier)[0] - return 16 # gvr + return 16 # cluster # Per-call route()+_parse_reg() string work costs ~2.5-3us host-submit wall @@ -253,10 +252,10 @@ def fn(lg, pre, sl, out): def fn(lg, pre, sl, out): direct_topk(lg, sl, out, K, next_n, cr) - elif tier.startswith("gvr"): + elif tier.startswith("cluster"): raise ValueError( - f"bsx gvr tier is not ported (npad beyond the deployment " - f"envelope); is_bsx_supported must gate this out (bs={bs}, " + f"the cluster tier is not ported (npad beyond the deployment " + f"envelope); is_tiered_topk_supported must gate this out (bs={bs}, " f"npad={npad}, K={K})" ) else: @@ -268,7 +267,7 @@ def fn(lg, pre, sl, out): return fn -def is_bsx_supported( +def is_tiered_topk_supported( logits: torch.Tensor, pre_idx: torch.Tensor, seq_lens: torch.Tensor, @@ -279,9 +278,9 @@ def is_bsx_supported( order_row, counters, ) -> bool: - """Host-only guard for the bsx tiers (no device sync; see module + """Host-only guard for the GVR tiers (no device sync; see module docstring). Returns False -> caller uses the in-tree kernel.""" - if _env_flag("TRTLLM_BSX_DISABLE"): + if _env_flag("TRTLLM_GVR_TIERS_DISABLE"): return False if logits.dtype != torch.float32: return False @@ -334,7 +333,7 @@ def is_bsx_supported( return ok -def bsx_topk( +def tiered_topk( logits: torch.Tensor, pre_idx: torch.Tensor, seq_lens: torch.Tensor, @@ -343,14 +342,14 @@ def bsx_topk( next_n: int = 1, compress_ratio: int = 4, ) -> None: - """Unified bsx tier dispatch, replicating gvr_topk_launch_batched. + """Unified GVR tier dispatch, replicating gvr_topk_launch_batched. logits [BS, npad] fp32 (BS = num_requests * next_n; npad multiple of 64; per-row tail beyond N_eff may be garbage — masked in-kernel), pre_idx [BS // next_n, K] int32 (request-level), seq_lens [BS // next_n] int32 (request-level, uncompressed-token space), output_indices [BS, K] int32. Caller must have passed - :func:`is_bsx_supported`. + :func:`is_tiered_topk_supported`. """ bs, npad = logits.shape key = (bs, npad, top_k, next_n, compress_ratio) @@ -361,8 +360,8 @@ def bsx_topk( __all__ = [ - "bsx_topk", - "is_bsx_supported", + "tiered_topk", + "is_tiered_topk_supported", "route", "route_cluster_size", "_reset_env_cache", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py rename to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py index c5dc1637bce3..472edfb0ef3d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_reg.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_reg.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""BSX register-resident (reg) GVR Top-K tier — CuTe DSL, Blackwell SM100. +"""Register-resident (reg) GVR Top-K tier — CuTe DSL, Blackwell SM100. CuTe DSL translation of the CUDA ``gvr_topk_reg`` register-resident GVR top-K kernel (tuned CUDA head), adapted for the production ``trtllm::cute_dsl_gvr_topk_decode`` contract (see the module docstring of -``gvr_topk_decode_bsx_tp`` for the shared adaptation inventory: ragged-N +``gvr_topk_decode_tp`` for the shared adaptation inventory: ragged-N masking, pre_idx clamping, per-row degenerate identity emit). The row is loaded ONCE from GMEM into per-thread registers (MAXV float4s per @@ -29,7 +29,7 @@ global traffic pinned at 4*npad bytes/row. Phase skeleton (shared with the tp tier — helpers REUSED from -``gvr_topk_decode_bsx_tp``): +``gvr_topk_decode_tp``): P1 : hint gather + two-stage 64-bin histogram -> rung ladder (tp.phase1) P2 : full AR-rung register count + cluster exchange, secant refine / max-below plateau descent (count_reg here) @@ -71,7 +71,7 @@ from cutlass.cutlass_dsl import dsl_user_op from cutlass.utils.smem_allocator import SmemAllocator -from .gvr_topk_decode_bsx_tp import ( +from .gvr_topk_decode_tp import ( FLT_MAX, INF, MAXPASS, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py rename to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py index 64b4482293fc..9c194a556546 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_tp.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""BSX throughput (tp) GVR Top-K tier — CuTe DSL, Blackwell SM100. +"""Throughput (tp) GVR Top-K tier — CuTe DSL, Blackwell SM100. CuTe DSL translation of the CUDA ``gvr_topk_tp`` throughput GVR top-K kernel (tuned CUDA head), @@ -469,7 +469,7 @@ def count_pass( # compound-Poisson sigma cnt/sqrt(occ) (equal to the classic # sqrt(cnt) Poisson sigma on IID data where occ == cnt). Packing # keeps registers, SMEM and the exchange at their pre-admission - # sizes; no field overflow: the bsx envelope pins npad <= 262144, so + # sizes; no field overflow: the tier envelope pins npad <= 262144, so # cluster-total cnt <= npad/32 = 8192 < 2^16 and occ <= npad/128. # ------------------------------------------------------------------ @cute.jit diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index b9db233d6d19..747679751e5b 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -16,11 +16,11 @@ l0_b300: backend: pytorch tests: # ------------- PyTorch tests --------------- - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py - - unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/thop/serial - unittest/_torch/executor # 250s diff --git a/tests/unittest/_torch/attention/sparse/conftest.py b/tests/unittest/_torch/attention/sparse/conftest.py index 85d0aa06c3e2..02af87693b6f 100644 --- a/tests/unittest/_torch/attention/sparse/conftest.py +++ b/tests/unittest/_torch/attention/sparse/conftest.py @@ -18,7 +18,7 @@ via the ``tie_aware_check`` fixture: conftest is collected by pytest regardless of rootdir/package-resolution style, so this sidesteps the cross-module import problem that previously forced a keep-in-sync duplicate -of the checker in ``test_cute_dsl_bsx_topk_decode.py``. +of the checker in ``test_cute_dsl_gvr_topk_tiers.py``. """ import pytest 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 1c00e7146bc8..789749bff5b4 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py @@ -23,7 +23,7 @@ import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( - gvr_topk_decode_bsx_dispatch as _bsx_dispatch, + gvr_topk_decode_dispatch as _tier_dispatch, ) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import ( GvrTopKKernel as _GvrTopKKernel, @@ -37,19 +37,19 @@ @pytest.fixture(autouse=True) -def _bsx_off(monkeypatch): +def _tiers_off(monkeypatch): """This module tests the IN-TREE kernel contract. Its op-level cases - (fp32 shapes within the bsx envelope) would otherwise be routed to the - bsx tiers by the dispatcher — which also ignores the explicit + (fp32 shapes within the tier envelope) would otherwise be routed to the + GVR tiers by the dispatcher — which also ignores the explicit ``cluster_size`` these tests parametrize over — so pin the op to the - in-tree path. The bsx tiers are covered by - ``test_cute_dsl_bsx_topk_decode.py`` (which flips the routing the other - way: fallback bands off so every test reaches a bsx tier).""" - monkeypatch.setenv("TRTLLM_BSX_DISABLE", "1") - _bsx_dispatch._reset_env_cache() + in-tree path. The GVR tiers are covered by + ``test_cute_dsl_gvr_topk_tiers.py`` (which flips the routing the other + way: fallback bands off so every test reaches a GVR tier).""" + monkeypatch.setenv("TRTLLM_GVR_TIERS_DISABLE", "1") + _tier_dispatch._reset_env_cache() yield - monkeypatch.delenv("TRTLLM_BSX_DISABLE", raising=False) - _bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_TIERS_DISABLE", raising=False) + _tier_dispatch._reset_env_cache() def _make_inputs_impl( diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py similarity index 82% rename from tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py rename to tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py index 81329f73106d..a49be7eb0ddc 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""BSX (direct/reg/tp CuTe DSL tiers) top-K decode tests. +"""Tiered GVR (direct/reg/tp CuTe DSL tiers) top-K decode tests. CI-sized exactness grid for the guarded fp32 fast path inside ``trtllm::cute_dsl_gvr_topk_decode`` (next_n >= 1, cr in {1, 4}): every reg @@ -25,7 +25,7 @@ x all three tiers, each case checked against BOTH the torch.topk host N_eff/offset simulation (the shared ``tie_aware_check`` conftest fixture) and a differential in-tree arm (same inputs through the in-tree kernel via -the ``TRTLLM_BSX_DISABLE`` kill switch, per-row value-multiset equality). +the ``TRTLLM_GVR_TIERS_DISABLE`` kill switch, per-row value-multiset equality). """ import os @@ -35,7 +35,7 @@ import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( - gvr_topk_decode_bsx_dispatch as bsx_dispatch, + gvr_topk_decode_dispatch as tier_dispatch, ) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.single_pass_multi_cta_radix_topk_cluster import ( # noqa: E501 _query_max_cluster_size, @@ -44,7 +44,7 @@ skip_not_sm100 = pytest.mark.skipif( get_sm_version() not in (100, 103), - reason=f"CuTe DSL BSX Top-K only supports SM 100/103, got SM {get_sm_version()}", + reason=f"CuTe DSL GVR tiers only support SM 100/103, got SM {get_sm_version()}", ) CR = 4 # default axis of the legacy (pre-MTP) cases: compress_ratio == 4 (DSv4) @@ -52,33 +52,33 @@ @pytest.fixture(autouse=True) -def _bsx_bands_off(monkeypatch): - """Kernel-contract tests must reach the bsx tiers: disable the measured +def _tiers_bands_off(monkeypatch): + """Kernel-contract tests must reach the GVR tiers: disable the measured fallback band table (it legitimately routes several tested (npad, bs) shapes to the in-tree kernel in production). The table itself is covered - by ``test_bsx_fallback_band_table``.""" - monkeypatch.setenv("TRTLLM_BSX_FALLBACK_BANDS", "0") - bsx_dispatch._reset_env_cache() + by ``test_tiers_fallback_band_table``.""" + monkeypatch.setenv("TRTLLM_GVR_FALLBACK_BANDS", "0") + tier_dispatch._reset_env_cache() yield - monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() @skip_not_sm100 @pytest.mark.parametrize("bands", ["on", "off"]) -def test_bsx_cuda_graph_capture_replay(monkeypatch, bands, tie_aware_check): +def test_tiers_cuda_graph_capture_replay(monkeypatch, bands, tie_aware_check): """The op must be CUDA-graph capturable and replay-consistent on both sides of the fallback band table. Shape (npad=32768, bs=64) sits inside a fallback bucket: bands=on captures the in-tree branch, bands=off the - bsx tp branch — the dispatch decision is host-side per shape, so it + tp tier branch — the dispatch decision is host-side per shape, so it bakes into the graph at capture time and must hold across replays, including replays over in-place rewritten inputs.""" if bands == "on": # the autouse fixture pinned the table off; restore the default - monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() bs, npad, top_k = 64, 32768, 2048 - logits, pre_idx, seq_lens = _make_bsx_inputs(bs, npad, top_k, seed=1234) + logits, pre_idx, seq_lens = _make_tier_inputs(bs, npad, top_k, seed=1234) out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") def call(): @@ -103,7 +103,7 @@ def call(): tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # replay over in-place rewritten inputs (fresh logits + fresh hints) - logits2, pre_idx2, seq_lens2 = _make_bsx_inputs(bs, npad, top_k, seed=4321) + logits2, pre_idx2, seq_lens2 = _make_tier_inputs(bs, npad, top_k, seed=4321) logits.copy_(logits2) pre_idx.copy_(pre_idx2) seq_lens.copy_(seq_lens2) @@ -112,14 +112,14 @@ def call(): tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) -def test_bsx_fallback_band_table(monkeypatch): +def test_tiers_fallback_band_table(monkeypatch): """The measured (npad, bs) fallback buckets route to the in-tree kernel - by default; the kill-switch restores bsx service; neighbours are not + by default; the kill-switch restores tier service; neighbours are not over-routed. Bands: full-grid calibration (2026-07-28), 131072 lower bound recalibrated to 8 (2026-07-29).""" - monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) - bsx_dispatch._reset_env_cache() - inb = bsx_dispatch._in_fallback_band + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() + inb = tier_dispatch._in_fallback_band # routed buckets (bucket floor <0.909 vs the in-tree head) assert inb(256, 8192) and inb(1024, 8192) assert inb(16, 32768) and inb(1024, 65536) @@ -131,20 +131,20 @@ def test_bsx_fallback_band_table(monkeypatch): assert not inb(8, 163776) assert inb(16, 163776) and inb(255, 163776) and not inb(256, 163776) assert inb(48, 40960) # off-grid npad resolves to the nearest pow2 - # neighbours stay on bsx + # neighbours stay on the tiers assert not inb(128, 8192) # direct/tp win band assert not inb(8, 32768) # latency reg band assert not inb(256, 131072) and not inb(128, 262144) # large-N tp win assert not inb(1024, 4096) # small-npad tp win # kill-switch - monkeypatch.setenv("TRTLLM_BSX_FALLBACK_BANDS", "0") - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_FALLBACK_BANDS", "0") + tier_dispatch._reset_env_cache() assert not inb(64, 32768) - monkeypatch.delenv("TRTLLM_BSX_FALLBACK_BANDS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_FALLBACK_BANDS", raising=False) + tier_dispatch._reset_env_cache() -def _make_bsx_inputs( +def _make_tier_inputs( bs: int, npad: int, top_k: int, @@ -154,7 +154,7 @@ def _make_bsx_inputs( preidx: str = "mixed", hit_rate: float = 0.5, ): - """Build (logits fp32, pre_idx int32, seq_lens int32) for the bsx path. + """Build (logits fp32, pre_idx int32, seq_lens int32) for the tier path. ``kind='ties'`` quantizes the logits to 0.25 steps so massive value plateaus straddle the K-th boundary (exercises the tie-ticket emits and @@ -215,20 +215,20 @@ def _run_op(logits, pre_idx, seq_lens, top_k): return out -def _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected_tier): +def _assert_tier_routes(logits, pre_idx, seq_lens, top_k, expected_tier): """Host-only: the dispatcher must accept this call and route it to ``expected_tier``.""" out = torch.empty(logits.shape[0], top_k, dtype=torch.int32, device="cuda") - assert bsx_dispatch.is_bsx_supported( + assert tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None - ), "expected the bsx fast path to accept this call" + ), "expected the tiered fast path to accept this call" bs, npad = logits.shape - tier = bsx_dispatch.route(bs, npad, top_k) + tier = tier_dispatch.route(bs, npad, top_k) assert tier == expected_tier, f"route({bs}, {npad}, {top_k}) = {tier} != {expected_tier}" def _skip_if_cluster_capped(bs, npad, top_k): - cs = bsx_dispatch.route_cluster_size(bs, npad, top_k) + cs = tier_dispatch.route_cluster_size(bs, npad, top_k) torch.zeros(1, device="cuda") # the driver-API query needs a live context hw_max = _query_max_cluster_size() if cs > hw_max: @@ -238,8 +238,8 @@ def _skip_if_cluster_capped(bs, npad, top_k): # --------------------------------------------------------------------------- # Host-only route-table asserts (mirror of the original CUDA dispatch). # --------------------------------------------------------------------------- -def test_bsx_route_table(): - r = bsx_dispatch.route +def test_tiers_route_table(): + r = tier_dispatch.route # latency ladder assert r(1, 4096, 512) == "direct" assert r(1, 12288, 2048) == "direct" @@ -262,61 +262,61 @@ def test_bsx_route_table(): assert r(16, 65536, 1024) == "tp" assert r(128, 262144, 2048) == "tp" # gvr tier only beyond the deployment envelope (guarded out upstream) - assert r(1, 262208, 512).startswith("gvr") + assert r(1, 262208, 512).startswith("cluster") -def test_bsx_route_env_knobs(monkeypatch): - """TRTLLM_BSX_TP_BS / TRTLLM_BSX_DENSE_BS keep the GVR_BSX_* semantics: +def test_tiers_route_env_knobs(monkeypatch): + """TRTLLM_GVR_TP_BS / TRTLLM_GVR_DENSE_BS keep the CUDA development arm's semantics: unset/-1 -> baked bands, 0 -> disable, else explicit bs threshold.""" - r = bsx_dispatch.route + r = tier_dispatch.route try: - monkeypatch.setenv("TRTLLM_BSX_DENSE_BS", "8") - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_DENSE_BS", "8") + tier_dispatch._reset_env_cache() assert r(8, 65536, 512) == "reg(cs=2,tb=1024,maxv=8,ar=8)" assert r(8, 131072, 1024) == "reg(cs=4,tb=1024,maxv=8,ar=8)" - monkeypatch.setenv("TRTLLM_BSX_TP_BS", "0") # 0 -> disable tp - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_TP_BS", "0") # 0 -> disable tp + tier_dispatch._reset_env_cache() assert r(1024, 65536, 512) != "tp" - monkeypatch.setenv("TRTLLM_BSX_TP_BS", "4") - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_TP_BS", "4") + tier_dispatch._reset_env_cache() assert r(4, 65536, 512) == "tp" finally: - monkeypatch.delenv("TRTLLM_BSX_TP_BS", raising=False) - monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_TP_BS", raising=False) + monkeypatch.delenv("TRTLLM_GVR_DENSE_BS", raising=False) + tier_dispatch._reset_env_cache() @skip_not_sm100 -def test_bsx_cluster_cap_verdict_memoized(): +def test_tiers_cluster_cap_verdict_memoized(): """The guard's cluster-cap verdict is memoized per (bs, npad, K) — the per-call route()+_parse_reg host cost the _DISPATCH_CACHE comment - motivates must not be re-paid by is_bsx_supported on every eager + motivates must not be re-paid by is_tiered_topk_supported on every eager forward — and _reset_env_cache clears it (routes depend on the env thresholds).""" bs, npad, top_k = 4, 65536, 512 - logits, pre_idx, seq_lens = _make_bsx_inputs(bs, npad, top_k, seed=0) + logits, pre_idx, seq_lens = _make_tier_inputs(bs, npad, top_k, seed=0) out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") - bsx_dispatch._reset_env_cache() - assert (bs, npad, top_k) not in bsx_dispatch._CAP_OK_CACHE - ok = bsx_dispatch.is_bsx_supported( + tier_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in tier_dispatch._CAP_OK_CACHE + ok = tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None ) - expected = bsx_dispatch.route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() + expected = tier_dispatch.route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() assert ok == expected - assert bsx_dispatch._CAP_OK_CACHE[(bs, npad, top_k)] == expected + assert tier_dispatch._CAP_OK_CACHE[(bs, npad, top_k)] == expected # Second call must hit the memo (same verdict, no recompute observable # beyond the cache entry staying put). assert ( - bsx_dispatch.is_bsx_supported( + tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None ) == expected ) - bsx_dispatch._reset_env_cache() - assert (bs, npad, top_k) not in bsx_dispatch._CAP_OK_CACHE + tier_dispatch._reset_env_cache() + assert (bs, npad, top_k) not in tier_dispatch._CAP_OK_CACHE # --------------------------------------------------------------------------- @@ -325,7 +325,7 @@ def test_bsx_cluster_cap_verdict_memoized(): # that together cover every codegen axis value (cs {1,2,8,16} x tb # {512,1024} x ar {6,8} x maxv {5,8} + the dense-knob-only path) — the # dropped launches differ only in axis combinations, not in code paths. -# The cs=2 dense instance is reachable only through the TRTLLM_BSX_DENSE_BS +# The cs=2 dense instance is reachable only through the TRTLLM_GVR_DENSE_BS # knob (default bands route around it), doubling as the env-knob dispatch # test on a live launch. # --------------------------------------------------------------------------- @@ -355,7 +355,7 @@ def test_bsx_cluster_cap_verdict_memoized(): ids=[t[4] + f"_n{t[0]}_bs{t[1]}_k{t[2]}" for t in _REG_INSTANCES], ) @pytest.mark.parametrize("kind", ["randn", "ties"]) -def test_bsx_reg_launch_table( +def test_tiers_reg_launch_table( npad, bs, top_k, dense_env, expected, launch, kind, monkeypatch, tie_aware_check ): if not launch and kind == "ties": @@ -363,19 +363,19 @@ def test_bsx_reg_launch_table( _skip_if_cluster_capped(bs, npad, top_k) try: if dense_env is not None: - monkeypatch.setenv("TRTLLM_BSX_DENSE_BS", dense_env) - bsx_dispatch._reset_env_cache() - logits, pre_idx, seq_lens = _make_bsx_inputs( + monkeypatch.setenv("TRTLLM_GVR_DENSE_BS", dense_env) + tier_dispatch._reset_env_cache() + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=npad + bs + top_k, kind=kind, varlen=True ) - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, expected) + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, expected) if launch: out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) finally: if dense_env is not None: - monkeypatch.delenv("TRTLLM_BSX_DENSE_BS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_DENSE_BS", raising=False) + tier_dispatch._reset_env_cache() # --------------------------------------------------------------------------- @@ -392,11 +392,11 @@ def test_bsx_reg_launch_table( (4096, 8, 512, "randn", False), # uniform N_eff == npad (no tail) ], ) -def test_bsx_direct(npad, bs, top_k, kind, varlen, tie_aware_check): - logits, pre_idx, seq_lens = _make_bsx_inputs( +def test_tiers_direct(npad, bs, top_k, kind, varlen, tie_aware_check): + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=npad * 3 + bs, kind=kind, varlen=varlen ) - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "direct") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "direct") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @@ -414,12 +414,12 @@ def test_bsx_direct(npad, bs, top_k, kind, varlen, tie_aware_check): (262144, 128, 2048, "randn"), # cs=1, uf=8 streaming path ], ) -def test_bsx_tp(npad, bs, top_k, kind, tie_aware_check): +def test_tiers_tp(npad, bs, top_k, kind, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=npad + 7 * bs, kind=kind, varlen=True ) - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @@ -437,9 +437,9 @@ def test_bsx_tp(npad, bs, top_k, kind, tie_aware_check): (65536, 16, 1024, "tp"), ], ) -def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): +def test_tiers_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=11, kind="randn", varlen=True ) # Rows 0..3: degenerate (N_eff = 0 / 1 / K-1 / K); the rest keep their @@ -451,7 +451,7 @@ def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): tail = col[None, :] >= n_eff[:, None] logits = torch.where(tail, torch.full_like(logits, 1e30), logits) - tier = bsx_dispatch.route(bs, npad, top_k) + tier = tier_dispatch.route(bs, npad, top_k) assert tier.startswith(expected_kind) or tier == expected_kind out = _run_op(logits, pre_idx, seq_lens, top_k) @@ -488,9 +488,9 @@ def test_bsx_degenerate_rows(npad, bs, top_k, expected_kind, tie_aware_check): (65536, 16, 1024), # tp tier ], ) -def test_bsx_preidx_hardening(npad, bs, top_k, preidx, tie_aware_check): +def test_tiers_preidx_hardening(npad, bs, top_k, preidx, tie_aware_check): _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=13, kind="randn", varlen=True, preidx=preidx ) out = _run_op(logits, pre_idx, seq_lens, top_k) @@ -520,20 +520,20 @@ def test_bsx_preidx_hardening(npad, bs, top_k, preidx, tie_aware_check): (131136, 128, 512), # cs=1 streaming tp (flash 512k production shape) ], ) -def test_bsx_tp_admission_hitrate(npad, bs, top_k, hit_rate, tie_aware_check): +def test_tiers_tp_admission_hitrate(npad, bs, top_k, hit_rate, tie_aware_check): """High-hr rows must admit (1-pass) and low-hr rows must stay exact via the pivot/secant fallback; both paths must produce exact top-K.""" _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=npad + bs + int(hit_rate * 100), varlen=True, hit_rate=hit_rate ) - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 -def test_bsx_tp_admission_tie_plateau(tie_aware_check): +def test_tiers_tp_admission_tie_plateau(tie_aware_check): """Tie plateau AT the admission threshold: K/2 distinct high values + a 3K-wide exact-tie plateau straddling the K-th rank. The hint set is the true top-K, so the ladder rungs land ON the plateau value and the @@ -552,14 +552,14 @@ def test_bsx_tp_admission_tie_plateau(tie_aware_check): logits[r, plateau] = 1.0 # exact ties straddling rank K seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 @pytest.mark.parametrize("top_k,n_plateau", [(1024, 8000), (2048, 12000)]) -def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau, tie_aware_check): +def test_tiers_tp_admission_overflow_fallback(top_k, n_plateau, tie_aware_check): """count(>= any admissible threshold) > kC: K-1 distinct high values, then an n_plateau-wide exact-tie plateau (n_plateau > kC = 6144/8192). Every threshold at or below the plateau overflows the candidate buffer @@ -580,24 +580,24 @@ def test_bsx_tp_admission_overflow_fallback(top_k, n_plateau, tie_aware_check): logits[r, plateau] = 1.0 seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") pre_idx = logits.topk(top_k, dim=-1).indices.int().contiguous() - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @skip_not_sm100 -def test_bsx_tp_admission_mixed_batch(tie_aware_check): +def test_tiers_tp_admission_mixed_batch(tie_aware_check): """Ragged batch mixing rows that admit (true-top-K hints) with rows that must fall back (all-zero cold-start hints): per-row admission escape is independent, so every row must stay exact regardless of which path its cluster takes.""" bs, npad, top_k = 16, 65536, 1024 # cs=8 tp _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=41, varlen=True, hit_rate=0.9 ) pre_idx[1::2] = 0 # odd rows: cold-start (degenerate ladder -> fallback) - _assert_bsx_routes(logits, pre_idx, seq_lens, top_k, "tp") + _assert_tier_routes(logits, pre_idx, seq_lens, top_k, "tp") out = _run_op(logits, pre_idx, seq_lens, top_k) tie_aware_check(out, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) @@ -607,7 +607,7 @@ def test_bsx_tp_admission_mixed_batch(tie_aware_check): # and still produce its results (op contract unchanged). # --------------------------------------------------------------------------- @skip_not_sm100 -def test_bsx_dispatcher_fallback_bf16(): +def test_tiers_dispatcher_fallback_bf16(): bs, npad, top_k = 4, 65536, 512 torch.manual_seed(21) torch.cuda.manual_seed(21) @@ -618,9 +618,9 @@ def test_bsx_dispatcher_fallback_bf16(): pre_idx[:, 0] = logits.float().argmax(dim=-1).int() out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") - assert not bsx_dispatch.is_bsx_supported( + assert not tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None - ), "bf16 must NOT take the bsx fast path" + ), "bf16 must NOT take the tiered fast path" # Guard-only on purpose (host, no JIT): the bf16 in-tree execution the # op falls back to is exhaustively covered by # test_cute_dsl_gvr_topk_decode.py; compiling that variant here again @@ -628,15 +628,15 @@ def test_bsx_dispatcher_fallback_bf16(): @skip_not_sm100 -def test_bsx_dispatcher_fallback_bad_shapes(): +def test_tiers_dispatcher_fallback_bad_shapes(): """Host-only: MTP contract violations must fall back to the in-tree - kernel (which asserts on them) rather than mis-launch a bsx tier.""" + kernel (which asserts on them) rather than mis-launch a tier.""" bs, npad, top_k = 4, 65536, 512 logits = torch.randn(bs * 2, npad, dtype=torch.float32, device="cuda") seq_lens = torch.full((bs,), npad * 4, dtype=torch.int32, device="cuda") pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") out = torch.empty(bs * 2, top_k, dtype=torch.int32, device="cuda") - ok = bsx_dispatch.is_bsx_supported + ok = tier_dispatch.is_tiered_topk_supported # next_n=2 with request-level pre_idx/seq_lens: accepted. assert ok(logits, pre_idx, seq_lens, out, top_k, 2, 4, None, None) # cr outside {1, 4}: rejected. @@ -652,21 +652,21 @@ def test_bsx_dispatcher_fallback_bad_shapes(): @skip_not_sm100 -def test_bsx_accepts_order_row(tie_aware_check): +def test_tiers_accepts_order_row(tie_aware_check): """``order_row`` (the LJF hint dsa.py computes for every batch with - num_rows >= 2 * num_sms) must NOT turn bsx off: the guard accepts it and + num_rows >= 2 * num_sms) must NOT turn the tiers off: the guard accepts it and the tiers ignore it, so the per-row index SET is identical with and without the permutation (emission order is unordered by contract). - Reuses the (4096, 256, 512) tp cell compiled by ``test_bsx_tp`` (no + Reuses the (4096, 256, 512) tp cell compiled by ``test_tiers_tp`` (no extra JIT).""" npad, bs, top_k = 4096, 256, 512 _skip_if_cluster_capped(bs, npad, top_k) - logits, pre_idx, seq_lens = _make_bsx_inputs( + logits, pre_idx, seq_lens = _make_tier_inputs( bs, npad, top_k, seed=npad + 7 * bs, kind="randn", varlen=True ) order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() out_dummy = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") - assert bsx_dispatch.is_bsx_supported( + assert tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out_dummy, top_k, NEXT_N, CR, order_row, None ), "guard must accept order_row (scheduling hint, ignored by the tiers)" out_plain = _run_op(logits, pre_idx, seq_lens, top_k) @@ -683,7 +683,7 @@ def test_bsx_accepts_order_row(tie_aware_check): ) torch.cuda.synchronize() assert torch.equal(out_plain.sort(-1).values, out_ordered.sort(-1).values), ( - "bsx per-row index set must be independent of the ignored order_row permutation" + "tier per-row index set must be independent of the ignored order_row permutation" ) tie_aware_check(out_ordered, logits, seq_lens, top_k, next_n=NEXT_N, compress_ratio=CR) # Production shape sanity (guard-only, no launch): num_rows >= 2*num_sms @@ -694,50 +694,50 @@ def test_bsx_accepts_order_row(tie_aware_check): sl_b = torch.full((big_bs,), npad * CR, dtype=torch.int32, device="cuda") out_b = torch.empty(big_bs, top_k, dtype=torch.int32, device="cuda") order_b = torch.argsort(sl_b.long(), descending=True).int().contiguous() - assert bsx_dispatch.is_bsx_supported( + assert tier_dispatch.is_tiered_topk_supported( logits_b, pre_b, sl_b, out_b, top_k, NEXT_N, CR, order_b, None - ), "large-batch (num_rows >= 2*num_sms) calls with order_row must stay on bsx" + ), "large-batch (num_rows >= 2*num_sms) calls with order_row must stay on the tiers" @skip_not_sm100 -def test_bsx_disable_kill_switch(monkeypatch): - """TRTLLM_BSX_DISABLE rejects everything at the guard (host-only).""" +def test_tiers_disable_kill_switch(monkeypatch): + """TRTLLM_GVR_TIERS_DISABLE rejects everything at the guard (host-only).""" bs, npad, top_k = 4, 65536, 512 logits = torch.randn(bs, npad, dtype=torch.float32, device="cuda") seq_lens = torch.full((bs,), npad * CR, dtype=torch.int32, device="cuda") pre_idx = torch.zeros(bs, top_k, dtype=torch.int32, device="cuda") out = torch.empty(bs, top_k, dtype=torch.int32, device="cuda") - ok = bsx_dispatch.is_bsx_supported + ok = tier_dispatch.is_tiered_topk_supported assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) - monkeypatch.setenv("TRTLLM_BSX_DISABLE", "1") - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_TIERS_DISABLE", "1") + tier_dispatch._reset_env_cache() assert not ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) - monkeypatch.delenv("TRTLLM_BSX_DISABLE", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.delenv("TRTLLM_GVR_TIERS_DISABLE", raising=False) + tier_dispatch._reset_env_cache() assert ok(logits, pre_idx, seq_lens, out, top_k, NEXT_N, CR, None, None) -def test_bsx_env_malformed_soft_fail(monkeypatch): +def test_tiers_env_malformed_soft_fail(monkeypatch): """Malformed tuning-knob values fail soft (warn + baked default), never raise on the decode path.""" - baseline = bsx_dispatch.route(256, 20480, 512) + baseline = tier_dispatch.route(256, 20480, 512) for bad in ("true", " 16x", "8.5"): - monkeypatch.setenv("TRTLLM_BSX_TP_BS", bad) - bsx_dispatch._reset_env_cache() - assert bsx_dispatch.route(256, 20480, 512) == baseline + monkeypatch.setenv("TRTLLM_GVR_TP_BS", bad) + tier_dispatch._reset_env_cache() + assert tier_dispatch.route(256, 20480, 512) == baseline # whitespace-padded but valid values still parse. - monkeypatch.setenv("TRTLLM_BSX_TP_BS", " 0 ") - bsx_dispatch._reset_env_cache() - assert bsx_dispatch.route(256, 20480, 512).startswith(("reg", "direct")) - monkeypatch.delenv("TRTLLM_BSX_TP_BS", raising=False) - bsx_dispatch._reset_env_cache() + monkeypatch.setenv("TRTLLM_GVR_TP_BS", " 0 ") + tier_dispatch._reset_env_cache() + assert tier_dispatch.route(256, 20480, 512).startswith(("reg", "direct")) + monkeypatch.delenv("TRTLLM_GVR_TP_BS", raising=False) + tier_dispatch._reset_env_cache() # --------------------------------------------------------------------------- # MTP (next_n > 1) + cr in {1, 4}: exactness on all three tiers, checked # against BOTH the torch.topk host N_eff/offset simulation # (``tie_aware_check``) and a differential in-tree arm — the SAME inputs -# through the in-tree kernel (forced via ``TRTLLM_BSX_DISABLE``; the guard +# through the in-tree kernel (forced via ``TRTLLM_GVR_TIERS_DISABLE``; the guard # accepts+ignores order_row), per-row value-multiset equality. Ragged: # seq_lens make N_eff differ across requests AND (via row % next_n) across # the MTP rows of one request; tails are poisoned with +1e30. @@ -761,8 +761,8 @@ def _make_mtp_inputs( ``pre_idx[..., 0] = per-group argmax`` (over the min-N_eff window, mirroring the in-tree test suite). The in-tree kernel REQUIRES this invariant for exactness — the differential arm is only valid with it. - The bsx tiers do not require it (clamp hardening); pass False for - bsx-only robustness cases.""" + The GVR tiers do not require it (clamp hardening); pass False for + tier-only robustness cases.""" torch.manual_seed(seed) torch.cuda.manual_seed(seed) num_rows = bs_req * next_n @@ -797,22 +797,22 @@ def _make_mtp_inputs( def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): - """bsx arm + in-tree differential arm (the ``TRTLLM_BSX_DISABLE`` kill + """tier arm + in-tree differential arm (the ``TRTLLM_GVR_TIERS_DISABLE`` kill switch forces the in-tree path; ``order_row`` no longer does — the guard accepts and ignores it). The ref call still passes ``order_row`` so the - in-tree sort-indirect path stays exercised. Returns (out_bsx, out_ref).""" + in-tree sort-indirect path stays exercised. Returns (out_tier, out_ref).""" num_rows = logits.shape[0] out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") - assert bsx_dispatch.is_bsx_supported( + assert tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None - ), "expected the bsx fast path to accept this MTP call" + ), "expected the tiered fast path to accept this MTP call" torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits, pre_idx, seq_lens, out, top_k=top_k, next_n=next_n, compress_ratio=cr ) order_row = torch.argsort(seq_lens.long(), descending=True).int().contiguous() out_ref = torch.empty_like(out) - os.environ["TRTLLM_BSX_DISABLE"] = "1" - bsx_dispatch._reset_env_cache() + os.environ["TRTLLM_GVR_TIERS_DISABLE"] = "1" + tier_dispatch._reset_env_cache() try: torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits, @@ -825,8 +825,8 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): order_row=order_row, ) finally: - del os.environ["TRTLLM_BSX_DISABLE"] - bsx_dispatch._reset_env_cache() + del os.environ["TRTLLM_GVR_TIERS_DISABLE"] + tier_dispatch._reset_env_cache() torch.cuda.synchronize() return out, out_ref @@ -871,14 +871,14 @@ def _run_mtp_both_arms(logits, pre_idx, seq_lens, top_k, next_n, cr): "tier,next_n,cr", _MTP_COMBOS, ids=[f"{t}-{n}-{c}" for t, n, c in _MTP_COMBOS] ) @pytest.mark.parametrize("kind,preidx", _MTP_KINDS, ids=[f"{k}-{p}" for k, p in _MTP_KINDS]) -def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx, tie_aware_check): +def test_tiers_mtp_exactness(tier, next_n, cr, kind, preidx, tie_aware_check): npad, bs_req, top_k = _MTP_TIER_CELLS[tier] num_rows = bs_req * next_n if tier == "tp" and num_rows < 16: bs_req = (16 + next_n - 1) // next_n num_rows = bs_req * next_n _skip_if_cluster_capped(num_rows, npad, top_k) - assert bsx_dispatch.route(num_rows, npad, top_k).startswith(tier) + assert tier_dispatch.route(num_rows, npad, top_k).startswith(tier) logits, pre_idx, seq_lens = _make_mtp_inputs( bs_req, next_n, cr, npad, top_k, seed=npad + 13 * next_n + cr, kind=kind, preidx=preidx ) @@ -889,20 +889,20 @@ def test_bsx_mtp_exactness(tier, next_n, cr, kind, preidx, tie_aware_check): sel = torch.gather(logits, -1, out.long()).sort(-1, descending=True).values sel_ref = torch.gather(logits, -1, out_ref.long()).sort(-1, descending=True).values assert torch.equal(sel, sel_ref), ( - f"bsx vs in-tree value-multiset mismatch (next_n={next_n}, cr={cr}, tier={tier})" + f"tier vs in-tree value-multiset mismatch (next_n={next_n}, cr={cr}, tier={tier})" ) @skip_not_sm100 @pytest.mark.parametrize("next_n,cr", [(2, 1), (3, 4)]) @pytest.mark.parametrize("preidx", ["zeros", "random"]) -def test_bsx_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): - """bsx-only MTP robustness: hint sets that VIOLATE the op's argmax- +def test_tiers_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): + """Tier-only MTP robustness: hint sets that VIOLATE the op's argmax- slot-0 contract (pure zeros / out-of-range garbage) must neither fault - nor break exactness on the bsx tiers. No differential arm here: the - in-tree kernel requires the argmax invariant for exactness, the bsx + nor break exactness on the GVR tiers. No differential arm here: the + in-tree kernel requires the argmax invariant for exactness, the GVR tiers deliberately do not (clamp hardening). Cell and (next_n, cr) - pairs match the reg combos of ``test_bsx_mtp_exactness`` so the JIT + pairs match the reg combos of ``test_tiers_mtp_exactness`` so the JIT variants are reused (zero extra compiles).""" bs_req, npad, top_k = 1, 24576, 2048 # reg tier at num_rows 2/3 num_rows = bs_req * next_n @@ -919,7 +919,7 @@ def test_bsx_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): argmax_slot0=False, ) out = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") - assert bsx_dispatch.is_bsx_supported( + assert tier_dispatch.is_tiered_topk_supported( logits, pre_idx, seq_lens, out, top_k, next_n, cr, None, None ) torch.ops.trtllm.cute_dsl_gvr_topk_decode( @@ -932,7 +932,7 @@ def test_bsx_mtp_preidx_hardening(next_n, cr, preidx, tie_aware_check): @skip_not_sm100 @pytest.mark.parametrize("next_n", [2, 3]) @pytest.mark.parametrize("cr", [1, 4]) -def test_bsx_mtp_degenerate_rows(next_n, cr, tie_aware_check): +def test_tiers_mtp_degenerate_rows(next_n, cr, tie_aware_check): """Degenerate MTP requests (N_eff <= K for some or all of the next_n rows): identity emit [0..N_eff-1] + -1 pad must hold per ROW, with the per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr.""" From 9dffd7162518676c33fc5f509e8f1ea069d3df1f Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 12 Aug 2026 00:30:00 +0000 Subject: [PATCH 19/19] [None][chore] apply pre-commit formatting missed in the rename commit yapf/ruff-only: continuation indent of the is_tiered_topk_supported call, alphabetical import order for the renamed symbols, and one collapsed assignment in the dispatcher. No behavior change. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 12 ++++++------ .../cute_dsl_kernels/blackwell/top_k/__init__.py | 2 +- .../blackwell/top_k/gvr_topk_decode_dispatch.py | 4 +--- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 95c3f61e3e00..a067f2fca7a5 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7271,10 +7271,10 @@ def warmup_cute_dsl_radix_topk_decode( # ------------------------------------------------------------------ # from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode import \ GvrTopKKernel as _GvrTopKKernel - from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ - tiered_topk as _tiered_topk from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ is_tiered_topk_supported as _is_tiered_topk_supported + from ..cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_dispatch import \ + tiered_topk as _tiered_topk class CuteDSLGvrTopKDecodeRunner: """Runner for the GVR Top-K cuTe DSL kernel (Blackwell SM100). @@ -7543,11 +7543,11 @@ def forward( # Host-only guard — no device sync. The op signature and output # contract are unchanged (unordered int32 indices, -1 pad only # for degenerate rows). - if _is_tiered_topk_supported(logits, pre_idx, seq_lens, output_indices, - top_k, next_n, compress_ratio, order_row, - counters): + if _is_tiered_topk_supported(logits, pre_idx, seq_lens, + output_indices, top_k, next_n, + compress_ratio, order_row, counters): _tiered_topk(logits, pre_idx, seq_lens, output_indices, top_k, - next_n, compress_ratio) + next_n, compress_ratio) return cute_dtype = _TORCH_TO_CUTLASS_DTYPE[logits.dtype] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index a28e9a62fcfe..68e8f7bf8a45 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -18,7 +18,7 @@ from .filtered_top_k_varlen_util import FilteredTopKKernelVarlen from .gvr_topk_decode import GvrParams, GvrTopKKernel from .gvr_topk_decode_direct import DirectTopKKernel -from .gvr_topk_decode_dispatch import tiered_topk, is_tiered_topk_supported +from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py index 57dd2e3f03de..06bfb73c17c4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_dispatch.py @@ -327,9 +327,7 @@ def is_tiered_topk_supported( key = (bs, npad, top_k) ok = _CAP_OK_CACHE.get(key) if ok is None: - ok = _CAP_OK_CACHE[key] = ( - route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() - ) + ok = _CAP_OK_CACHE[key] = route_cluster_size(bs, npad, top_k) <= _query_max_cluster_size() return ok