From 1fb56b8decc0d73d3cc011e884147fd2ad440d6a Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:24:25 +0000 Subject: [PATCH 01/40] [None][perf] Add self-sampling GVR top-K decode kernels (CuTeDSL, standalone) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a self-sampling variant of the GVR (Guess-Verify-Refine) heuristic top-K decode for the DSA indexer, translated to CuTeDSL from the optimized CUDA line (fork branch GVR-selfsampling-CuTeDSL), as two standalone modules under cute_dsl_kernels/blackwell/top_k: - gvr_topk_decode_self_sampling.py: merged device module — four kernel families (sampling-ladder main / register-resident reg / cluster clus / cluster-register reg_clus), lazily JIT-compiled per constexpr tuple. - gvr_topk_decode_self_sampling_host.py: host companion — pure-function dispatch route(b, n, npad, k) (bit-exact transcription of the CUDA host dispatch, cross-checked by a 1,159,168-case boundary+fuzz sweep plus a 300k-case parity fuzz of this merged form), per-device workspace slab, and run/run_ws DPS entries with the CUDA binding's hardening battery. Contract (documented in the host module): batch-uniform host-int n_valid in compressed index space; fp32; K in {512, 1024, 2048}; 64-element-multiple row stride. Exact (tie-interchangeable) top-K. NOT wired into the decode path: the production engine reads per-request seq_lens on-device with per-row MTP offsets (heuristicTopKDecode.cu); adopting that per-row contract inside these kernels is follow-up work, so this module must not substitute for the tiered path under continuous batching, MTP, or CUDA-graph capture. Evidence (B200): 886-cell x 11-BS real-decode-capture grid = 9,746 cases vs the production CUDA arm: 0 INEXACT, geomean ratio 0.974. Unit tests: tie-aware exactness (signed-zero normalized, poisoned-pad immunity) across gate-edge and envelope shapes, run_ws with caller workspace, guard predicates, dispatch totality; sm_100-gated, picked up by the existing unittest/_torch/thop/parallel sweep. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/__init__.py | 6 + .../top_k/gvr_topk_decode_self_sampling.py | 5862 +++++++++++++++++ .../gvr_topk_decode_self_sampling_host.py | 811 +++ .../parallel/test_gvr_selfsampling_topk.py | 143 + 4 files changed, 6822 insertions(+) create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py create mode 100644 tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py 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 68e8f7bf8a45..507662eaa0c0 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 @@ -20,6 +20,9 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_self_sampling_host import route as selfsampling_topk_route +from .gvr_topk_decode_self_sampling_host import run as selfsampling_topk_run +from .gvr_topk_decode_self_sampling_host import run_ws as selfsampling_topk_run_ws from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -34,4 +37,7 @@ "DirectTopKKernel", "tiered_topk", "is_tiered_topk_supported", + "selfsampling_topk_route", + "selfsampling_topk_run", + "selfsampling_topk_run_ws", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py new file mode 100644 index 000000000000..a9f6af2c2e0a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -0,0 +1,5862 @@ +# 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. + +"""Self-sampling GVR top-K decode kernels (CuTe DSL, Blackwell sm_100a). + +Sample-calibrated threshold ladders for exact single-pass top-K: the kernel +derives its selection threshold from an in-kernel sample of the row itself +(one sample histogram yields a bracketed ladder of candidate thresholds, +resolved and harvested in a single streaming pass); exactness is guaranteed +by count-crossing invariants, never by the estimate; the temporal hint +(pre_idx) survives only as a degenerate-case anchor. + +High-fidelity translation of the CUDA ``GVR-selfsampling-CUDA`` branch +(layer-10 head d76f4387c1): four kernel families -- sampling-ladder +slab/streaming (main), register-resident (reg), cluster streaming (clus), +clustered register-resident (regclus) -- merged into this single device +module for in-tree convention parity. The development/audit form (one +module per family, per-era commit history, SASS audit trail) lives on the +contributor fork branch ``GVR-selfsampling-CuTeDSL``. + +This file is GENERATED by a rename-only merge (collision symbols get a +``__`` suffix; no semantic edits) and is verified against the +per-family source by SASS bit-equality over every reachable compile tuple. +Do not hand-edit kernel bodies here without re-running that gate. + +Validation record (2026-08-17): 886-cell x 11-BS real-capture grid +(9,746 cases) vs the CUDA source arm -- 0 INEXACT, face-value gm 0.9743. +Known family-level exception (filed 2026-08-15): v32_8k + flash_32k at +BS >= 256, isolated-retest upper bound ~1.065 vs the 1.02 per-case bar +(compiler barrier-arrival raggedness; tracked for re-check on every +nvidia-cutlass-dsl version bump). +""" + +import contextlib +import sys + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm, nvvm +from cutlass._mlir.dialects import llvm as mlir_llvm +from cutlass._mlir.dialects import math as mlir_math +from cutlass.cute import runtime as _crt +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +# `C.` references throughout resolve against this module itself +# (the per-family sources imported the shared-units sibling as C). +C = sys.modules[__name__] + + +# =========================================================================== +# ==== ct_common.py — shared device units ====================== +# =========================================================================== +"""ct_common.py — op46 shared device-helper library (CuTeDSL port). + +Translated ONCE from the frozen CUDA source +`op46_selfsampling_cutedsl/src_cuda/kernel.cu` L15-376 (helpers shared by +gvr_main / gvr_topk_reg / gvr_clus / gvr_reg_clus) per TRANSLATION_SPEC.md §5 +head, with every spelling pinned by probes/PROBE_RESULTS.md (P1..P15, all +BINDING) and op43 prior-art idioms (ct_tp/ct_gvr/ct_reg/ct_direct). + +Conventions for kernel translators +---------------------------------- +* Crossing-scan helpers write their scalar outputs into an Int32 smem tensor + `s_res` using the slot map RES_B=0, RES_M=1, RES_ABOVE=2, RES_TOT=3, + RES_B2=4, RES_B3=5 (mirror of the CUDA `int *s_B,*s_m,*s_above,*s_tot` + out-params + the TWO/THREE extra pins). Slots are written ONLY on a pin, + exactly like the CUDA. +* Histograms are Int32 smem tensors (CUDA uint32_t): adds/scans are + bit-identical mod 2^32 and every compare the CUDA does against `target` + is already `(int)`-cast there. Totals < 2^31 by dispatch domain. +* Warp-0-only helpers (find_cross / scan_cross0 / merge_scan0) contain NO + barrier (probe P14); scan_cross and scan_cross_w contain EXACTLY ONE + internal barrier (kernel.cu L199 / L306); gather_hint contains EXACTLY TWO + (L349 / L357). Do not add or drop any (op43 lesson L5). +* All warp collectives use the full mask FULLM = 0xffffffff (kernel.cu L15). +""" + + +# --------------------------------------------------------------------------- +# constants (kernel.cu L15-30) +# --------------------------------------------------------------------------- +FULLM = 0xFFFFFFFF +NB = 1024 # register-family base bin count (L16) +SNB = 256 # streaming-path bin count (L170-177) — MUST stay 256 +MAXC = 160 # multi-CTA SPLIT row cap (L17) +GCAP = 16384 # per-row slab capacity in int2 (L18) +QUADC = 96 # O(mc^2) rank gate, streaming/reg (L21) +QUADC_CLUS = 288 # clus + gvr_main gate (L28, spec §2 conflict resolution) +IDXB = 22 # packed candidate index bits (L29) +IDXM = (1 << IDXB) - 1 +GVR_WS_OFF_OFF = MAXC * 8 # workspace g_off byte offset (L42) +GVR_WS_BUF_OFF = 2048 # workspace g_buf byte offset (L43) + +# degenerate-hint sentinels (kernel.cu L629, L356, L1014 exact-equality flag) +SENT_LO = -3.0e38 +SENT_HI = 3.0e38 + +# s_res slot map (see module docstring) +RES_B = 0 +RES_M = 1 +RES_ABOVE = 2 +RES_TOT = 3 +RES_B2 = 4 +RES_B3 = 5 + + +# --------------------------------------------------------------------------- +# float <-> u32 bitcasts (op43 ct_tp.py:145-151) +# --------------------------------------------------------------------------- +def u32_of_f32(v): + """Raw fp32 bits as Uint32 (bit-cast, no conversion).""" + return cutlass.Uint32(llvm.bitcast(cutlass.Uint32.mlir_type, v.ir_value())) + + +def f32_of_u32(u): + """Uint32 bit pattern as Float32 (bit-cast).""" + return cutlass.Float32(llvm.bitcast(cutlass.Float32.mlir_type, u.ir_value())) + + +def f32_of_i32(i): + return cutlass.Float32(llvm.bitcast(cutlass.Float32.mlir_type, i.ir_value())) + + +def i32_of_f32(v): + return cutlass.Int32(llvm.bitcast(cutlass.Int32.mlir_type, v.ir_value())) + + +# --------------------------------------------------------------------------- +# fkey / invkey (kernel.cu L64-71) — order-preserving float->u32 radix key. +# fkey: u ^ (((int32)u >> 31) | 0x80000000) [arithmetic-shift sign trick, +# spelled 0 - (u >> 31) on Uint32 per op43 ct_tp.py:171-175] +# invkey: (K & 0x80000000) ? K ^ 0x80000000 : ~K [exact inverse] +# Monotone over all finite floats and +-inf; min identity 0xffffffff, max 0. +# --------------------------------------------------------------------------- +def fkey_bits(u): + """fkey on raw fp32 bits already held as Uint32.""" + neg = cutlass.Uint32(0) - (u >> cutlass.Uint32(31)) # 0 or 0xFFFFFFFF + return u ^ (neg | cutlass.Uint32(0x80000000)) + + +def fkey(x): + """CUDA fkey(float) (L68-71). x: dynamic Float32 -> Uint32 key.""" + return fkey_bits(u32_of_f32(x)) + + +def invkey_bits(K): + """CUDA invkey (L64-67) without the final bitcast: key -> fp32 bits.""" + s = K >> cutlass.Uint32(31) # 1 iff key top bit set + m = (s - cutlass.Uint32(1)) | cutlass.Uint32(0x80000000) + # s==1 -> m=0x80000000 (K^0x80000000); s==0 -> m=0xFFFFFFFF (~K) + return K ^ m + + +def invkey(K): + """CUDA invkey(uint32) (L64-67). K: dynamic Uint32 key -> Float32.""" + return f32_of_u32(invkey_bits(K)) + + +# --------------------------------------------------------------------------- +# warp redux wrappers (kernel.cu L59-62; scan_cross_w L307-315; probe P1/P13). +# Values passed to the u32 forms MUST be genuine cutlass.Uint32 — an Int32 +# silently lowers to redux.sync.{min,max}.s32 (probe P1 note). +# --------------------------------------------------------------------------- +def warp_min_u32(v): + """__reduce_min_sync(FULLM, v) -> redux.sync.min.u32 (single inst).""" + return cute.arch.warp_redux_sync(v, "min") + + +def warp_max_u32(v): + """__reduce_max_sync(FULLM, v) -> redux.sync.max.u32.""" + return cute.arch.warp_redux_sync(v, "max") + + +def warp_add_u32(v): + """__reduce_add_sync(FULLM, v) -> redux.sync.add.s32 (bit-identical u32).""" + return cute.arch.warp_redux_sync(v, "add") + + +def warp_add_i32(v): + """__reduce_add_sync on Int32 (scan_cross_w two-redux stage L314-315).""" + return cute.arch.warp_redux_sync(v, "add") + + +def fmin_f32(a, b): + """fminf -> native min.f32 (probe P13; op43 inline-PTX no longer needed).""" + return cute.arch.fmin(a, b) + + +def fmax_f32(a, b): + """fmaxf -> max.f32.""" + return cute.arch.fmax(a, b) + + +# --------------------------------------------------------------------------- +# ballot / popc / clz / ffs wrappers +# --------------------------------------------------------------------------- +def ballot(pred): + """__ballot_sync(FULLM, pred) -> Int32 mask.""" + return cute.arch.vote_ballot_sync(pred) + + +def popc(x): + return cute.arch.popc(x) + + +def clz_i32(x): + """__clz as Int32.""" + return cutlass.Int32(cute.arch.clz(x)) + + +def ffs_m1(x): + """__ffs(x) - 1 for x != 0 (bit index of lowest set bit). + + Spelled popc((x & -x) - 1) per op43 ct_direct.py:210-221. Caller must + guarantee x != 0 (every kernel.cu use is inside a mask-walk loop). + """ + return cutlass.Int32(cute.arch.popc((x & (cutlass.Int32(0) - x)) - cutlass.Int32(1))) + + +@cute.jit +def hi_bit_or_zero(msk): + """CUDA `msk ? (31 - __clz(msk)) : 0` (find_cross L92/L102).""" + r = cutlass.Int32(0) + if msk != cutlass.Int32(0): + r = cutlass.Int32(31) - clz_i32(msk) + return r + + +# --------------------------------------------------------------------------- +# warp shfl scans (op43 ct_tp.py:186-206 + the TWO-interleaved variant that +# gvr_topk_reg L1669-1673 needs) +# --------------------------------------------------------------------------- +@cute.jit +def _shfl_up_add(val, lane, offset: cutlass.Constexpr): + """Inclusive-scan step: val += shfl_up(val, offset) gated lane >= offset. + + Native shfl.sync.up (mask_and_clamp=0, the __shfl_up_sync lowering): + hardware clamps the source lane, deleting the VIMNMX+VIADD software + clamp of the previous idx-kind spelling. Lanes < offset receive an + undefined-but-discarded value (the gate keeps the result identical). + """ + other = cute.arch.shuffle_sync_up(val, offset, mask_and_clamp=0) + 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 lane+offset < 32. + + Native shfl.sync.down (mask_and_clamp=31 = __shfl_down_sync lowering); + hardware clamps, gate discards out-of-range lanes as before. + """ + other = cute.arch.shuffle_sync_down(val, offset, mask_and_clamp=31) + if lane + cutlass.Int32(offset) < cutlass.Int32(32): + val = val + other + return val + + +@cute.jit +def warp_incl_scan_add(val, lane): + """5-step inclusive __shfl_up_sync add scan (e.g. L144-146, L843-849).""" + for o in [1, 2, 4, 8, 16]: + val = _shfl_up_add(val, lane, o) + return val + + +@cute.jit +def warp_incl_scan_add2(v1, v2, lane): + """TWO interleaved inclusive shfl_up scans (gvr_topk_reg L1669-1673). + + Per step o: shfl(v1); gated add; shfl(v2); gated add — the two dependency + chains interleave so the second scan hides under the first's shfl latency + exactly as the CUDA dual-scan loop does. + """ + for o in [1, 2, 4, 8, 16]: + z1 = cute.arch.shuffle_sync_up(v1, o, mask_and_clamp=0) + if lane >= cutlass.Int32(o): + v1 = v1 + z1 + z2 = cute.arch.shuffle_sync_up(v2, o, mask_and_clamp=0) + if lane >= cutlass.Int32(o): + v2 = v2 + z2 + return v1, v2 + + +@cute.jit +def warp_suffix_scan_add(val, lane): + """5-step __shfl_down_sync suffix add scan (find_cross L86-89, L97-100).""" + for o in [1, 2, 4, 8, 16]: + val = _shfl_down_add(val, lane, o) + return val + + +# --------------------------------------------------------------------------- +# CTA-scope shared-memory atomics (probe P7: returns OLD value, ATOMS not RED, +# warp-aggregation preserved; never sys-scope utils.distributed.atomicAdd) +# --------------------------------------------------------------------------- +def atomic_add_cta(ptr, val): + """shared atomicAdd returning old value. ptr: cute Pointer + + (e.g. `s_hist.iterator + bin_idx`), val: Int32. + """ + return cutlass.Int32(cute.arch.atomic_add(ptr, val, sem="relaxed", scope="cta")) + + +def atomic_min_cta(ptr, val): + """shared atomicMin (s_kmin seeds, L1094-1101). Unsigned iff val is Uint32.""" + return cute.arch.atomic_min(ptr, val, sem="relaxed", scope="cta") + + +def atomic_max_cta(ptr, val): + """shared atomicMax (s_kmax seeds).""" + return cute.arch.atomic_max(ptr, val, sem="relaxed", scope="cta") + + +def atomic_or_cta(ptr, val): + """shared atomicOr (gvr_topk_reg bitmap path).""" + return cute.arch.atomic_or(ptr, val, sem="relaxed", scope="cta") + + +# --------------------------------------------------------------------------- +# gpu-scope fences + global u64 atomicAdd (SPLIT slab protocol, probe P8) +# --------------------------------------------------------------------------- +def threadfence_gpu(): + """__threadfence() == fence.acq_rel.gpu — use at BOTH L959 and L968.""" + cute.arch.fence_acq_rel_gpu() + + +def atomic_add_u64_gpu(ptr, val): + """atom.global.add.u64 returning the OLD value (L960 arrival RMW). + + ptr: cute Pointer to an Int64 gmem word; val: cutlass.Int64. + Packed arrival word: `cutlass.Int64(1 << 32) + cutlass.Int64(myn)`. + """ + return cutlass.Int64(cute.arch.atomic_add(ptr, val)) + + +# --------------------------------------------------------------------------- +# saturating converts (probe P4: native ctors emit cvt.rzi.{u32,s32}.f32) +# --------------------------------------------------------------------------- +def f2u_rz(v): + """__float2uint_rz: saturating (neg/-inf -> 0, huge -> 0xffffffff, NaN -> 0). + + Native ctor spelling verified on this exact toolchain (P4). Dynamic values + only — host constants raise OverflowError on inf. + """ + return cutlass.Uint32(v) + + +def f2s_rz(v): + """__float2int_rz: saturating (-inf -> INT_MIN, huge -> INT_MAX, NaN -> 0).""" + return cutlass.Int32(v) + + +# --------------------------------------------------------------------------- +# L2 prefetch escape hatch (op43 ct_gvr.py:42-54; kernel.cu sites L576/592/612) +# --------------------------------------------------------------------------- +@dsl_user_op +def _prefetch_l2(gaddr, *, loc=None, ip=None): + """prefetch.global.L2 [gaddr]; gaddr is a byte address (Int64).""" + llvm.inline_asm( + res=None, + operands_=[gaddr.ir_value(loc=loc, ip=ip)], + asm_string="prefetch.global.L2 [$0];", + constraints="l", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# global loads: 128-bit ldg (read-only) / plain, scalar forms, and __ldcg +# (L2-direct) vector forms for the slab consume (L929/945/978/1032/1172/1194) +# --------------------------------------------------------------------------- +def g2r_atom_f32(bits: int, invariant: bool = True): + """CopyG2ROp atom: bits=128 -> LDG.E.128[.CONSTANT], bits=32 -> scalar.""" + return cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=bits, invariant=invariant + ) + + +def g2r_atom_i32(bits: int, invariant: bool = False): + return cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=bits, invariant=invariant + ) + + +def ld_g_f32x4(copy_atom, base_addr, v_idx, frag): + """Load float4 #v_idx (16B units) from gmem byte base into frag[0..3]. + + op43 ct_tp.py:236-245 idiom. base_addr: Int64 byte address; frag: (4,) f32 + fragment. Issue ALL batch members before consuming any (op43 lesson L1). + """ + p = cute.make_ptr( + cutlass.Float32, + base_addr + cutlass.Int64(v_idx) * cutlass.Int64(16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + cute.copy(copy_atom, cute.make_tensor(p, cute.make_layout((4,))), frag) + + +def ldg_f32(base_addr, idx): + """__ldg(X + idx): scalar read-only 4B gather (gather_hint L343).""" + atom = g2r_atom_f32(32, invariant=True) + p = cute.make_ptr( + cutlass.Float32, + base_addr + cutlass.Int64(idx) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + frag = cute.make_fragment((1,), cutlass.Float32) + cute.copy(atom, cute.make_tensor(p, cute.make_layout((1,))), frag) + return frag[0] + + +def ld_g_i32(base_addr, idx): + """plain P[idx] scalar int32 load (gather_hint L340).""" + p = cute.make_ptr( + cutlass.Int32, + base_addr + cutlass.Int64(idx) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + return cutlass.Int32(cute.arch.load(p, cutlass.Int32)) + + +@dsl_user_op +def _ldcg_v2_i32(gaddr, *, loc=None, ip=None): + """__ldcg on an int2 (8B slab word): ld.global.cg.v2.u32 -> (x, y). + + x = value bits, y = index (workspace g_buf layout, kernel.cu L38-41). + gaddr: Int64 byte address, 8B-aligned. + """ + ret = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32()]), + [gaddr.ir_value(loc=loc, ip=ip)], + "ld.global.cg.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [0])), + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [1])), + ) + + +@dsl_user_op +def _ldcg_v4_i32(gaddr, *, loc=None, ip=None): + """ld.global.cg.v4.b32 (16B L2-direct load), returns 4 Int32.""" + ret = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [gaddr.ir_value(loc=loc, ip=ip)], + "ld.global.cg.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return tuple(cutlass.Int32(llvm.extractvalue(T.i32(), ret, [i])) for i in range(4)) + + +# --------------------------------------------------------------------------- +# 128-bit shared-memory ld/st (probe P5a copy-atom spelling) + ulonglong2 read +# --------------------------------------------------------------------------- +def smem_atom_i32_128(): + """CopyUniversalOp atom for ld/st.shared.v4.b32 on Int32 smem.""" + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) + + +def _smem_v4_tensor(base_addr, byte_off): + """4-elt Int32 smem tensor at 16B-aligned base_addr+byte_off (Int32 addr).""" + p = cute.make_ptr(cutlass.Int32, base_addr + byte_off, cute.AddressSpace.smem, assumed_align=16) + return cute.make_tensor(p, cute.make_layout((4,))) + + +def lds128_i32(copy_atom, base_addr, byte_off, frag): + """ld.shared.v4.b32 -> frag(4, Int32).""" + cute.copy(copy_atom, _smem_v4_tensor(base_addr, byte_off), frag) + + +def sts128_i32(copy_atom, frag, base_addr, byte_off): + """st.shared.v4.b32 <- frag(4, Int32).""" + cute.copy(copy_atom, frag, _smem_v4_tensor(base_addr, byte_off)) + + +@dsl_user_op +def _lds_v2_u64(saddr, *, loc=None, ip=None): + """ulonglong2 16B smem read (quad-rank L1080-1089, L2241): (lo, hi).""" + ret = llvm.inline_asm( + llvm.StructType.get_literal([T.i64(), T.i64()]), + [saddr.ir_value(loc=loc, ip=ip)], + "ld.shared.v2.u64 {$0, $1}, [$2];", + "=l,=l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + cutlass.Uint64(llvm.extractvalue(T.i64(), ret, [0])), + cutlass.Uint64(llvm.extractvalue(T.i64(), ret, [1])), + ) + + +# --------------------------------------------------------------------------- +# DSMEM op set (op43 ct_tp.py:44-125 + probe P5b 128-bit remote load). +# mapa returns a byte-addressed Int32 in the PEER's shared window; offset +# arithmetic after one mapa-per-rank is the proven op43/P5b form. +# --------------------------------------------------------------------------- +@dsl_user_op +def _mapa_shared_cluster(smem_ptr, peer_rank, *, loc=None, ip=None): + """mapa.shared::cluster of a local smem Pointer -> Int32 peer byte addr.""" + 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 _mapa_shared_cluster_addr(addr_i32, peer_rank, *, loc=None, ip=None): + """mapa of a raw Int32 shared-window byte address (already .toint()'d).""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [addr_i32.ir_value(loc=loc, ip=ip), 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 _ld_shared_cluster_v4_u32(mapped_addr, *, loc=None, ip=None): + """Single-shot remote 16B DSMEM load (probe P5b; merge_scan0 L136-137).""" + ret = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [0])), + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [1])), + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [2])), + cutlass.Int32(llvm.extractvalue(T.i32(), ret, [3])), + ) + + +@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_f32(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.f32 [$0], $1;", + constraints="r,f", + 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 packed 8B DSMEM candidate push (op43 lesson L3; kernel.cu L2185). + + val = (Uint64(key) << 32) | Uint64(idx_bits). + """ + 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): + """Remote CTA smem atomicAdd (cluster scope), returns old value.""" + 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, + ) + ) + + +# --------------------------------------------------------------------------- +# aligned cluster barrier (op43 ct_reg.py:100-113; cg::cluster.sync() == +# barrier.cluster.{arrive,wait}.aligned). Writers use the FULL (releasing) +# arrive — cluster_arrive_relaxed has NO release and races DSMEM (known +# lesson). Never substitute the non-aligned cute.arch forms. +# --------------------------------------------------------------------------- +@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(): + """cg::cluster_group::sync() (kernel.cu L2016/2135/2226/2474/2530/2648).""" + _cluster_arrive_aligned() + _cluster_wait_aligned() + + +# =========================================================================== +# find_cross (kernel.cu L73-109) +# highest bin B with sum_{j>=B} hist[j] >= target; also total, m = hist[B], +# above = sum_{j>B}. Warp-parallel (warp 0 only), bank-conflict free via the +# rotated indexing hist[lane*BPL + ((j+lane) & (BPL-1))] (L83 — DO NOT drop). +# Non-destructive. NO barrier inside. +# Writes s_res[RES_B/RES_M/RES_ABOVE] from the single pinning lane and +# s_res[RES_TOT] from lane 0. +# =========================================================================== +@cute.jit +def find_cross(s_hist, target, tidx, s_res, nb: cutlass.Constexpr): + BPL = nb // 32 # python int at trace time + if tidx < cutlass.Int32(32): + lane = tidx + # per-lane span sum with rotated bank-skew indexing (L82-84) + part = cutlass.Int32(0) + for j in cutlass.range_constexpr(BPL): + idx = lane * cutlass.Int32(BPL) + ((cutlass.Int32(j) + lane) & cutlass.Int32(BPL - 1)) + part = part + s_hist[idx] + # 5-step suffix scan (L85-89): v = sum of part over lanes >= lane + v = warp_suffix_scan_add(part, lane) + if lane == cutlass.Int32(0): + s_res[RES_TOT] = v + # level 1: highest lane whose suffix still reaches target (L91-92) + msk = ballot(v >= target) + L = hi_bit_or_zero(msk) + aboveL = cute.arch.shuffle_sync(v - part, L) # L93 + # level 2: one bin per lane inside lane L's span (L94-100) + h = cutlass.Int32(0) + if lane < cutlass.Int32(BPL): + h = s_hist[L * cutlass.Int32(BPL) + lane] + w = warp_suffix_scan_add(h, lane) + msk2 = ballot((aboveL + w) >= target) + J = hi_bit_or_zero(msk2) + if lane == J: # L103-107 + s_res[RES_B] = L * cutlass.Int32(BPL) + J + s_res[RES_M] = h + s_res[RES_ABOVE] = aboveL + (w - h) + + +# =========================================================================== +# scan_cross0 (kernel.cu L218-286) +# Warp-0-only single-barrier vectorized suffix scan (streaming workhorse, +# NB_=256 at every production call site). Contains NO barrier — the caller +# pays exactly one after it. Leaves hist[j] = per-bin OUTPUT CURSOR +# (count strictly above bin j), or ZEROS when zero=True (folds the next +# phase's histogram clear). two/three pin extra crossing bins for +# target2/target3 into RES_B2/RES_B3. addf folds the per-rank bin-offset +# vector s_addv into the cursors (L279-282). +# HOLD register guard (L236-240): NV<=2 holds the span in regs across the +# scan; wider instantiations re-READ their span (no barrier needed — each +# lane only touches its own span). +# =========================================================================== +@cute.jit +def scan_cross0( + s_hist, + target, + tidx, + s_res, + target2, + target3, + s_addv, + nb: cutlass.Constexpr, + zero: cutlass.Constexpr, + two: cutlass.Constexpr = False, + three: cutlass.Constexpr = False, + addf: cutlass.Constexpr = False, +): + BPT = nb // 32 # bins per lane (trace-time int) + NV = BPT // 4 # 16B vectors per lane + HOLD = NV <= 2 # register-pressure guard (L240) + if tidx < cutlass.Int32(32): + lane = tidx + atom = smem_atom_i32_128() + hbase = s_hist.iterator.toint() + # pass 1: span sum via NV uint4 LDS.128 (L243-251) + frags = [cute.make_fragment((4,), cutlass.Int32) for _ in range(NV)] + sm = cutlass.Int32(0) + for q in cutlass.range_constexpr(NV): + boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + lds128_i32(atom, hbase, boff, frags[q]) + sm = sm + frags[q][0] + frags[q][1] + frags[q][2] + frags[q][3] + # 5-step inclusive shfl_up scan (L252-254) + w = warp_incl_scan_add(sm, lane) + tot = cute.arch.shuffle_sync(w, cutlass.Int32(31)) + after = tot - w # bins strictly above my span + if lane == cutlass.Int32(0): + s_res[RES_TOT] = tot + base = lane * cutlass.Int32(BPT) + # pass 2: descending vector walk (L258-284) + for q in cutlass.range_constexpr(NV - 1, -1, -1): + if cutlass.const_expr(HOLD): + vv = frags[q] + else: + vv = cute.make_fragment((4,), cutlass.Int32) # re-read span + boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + lds128_i32(atom, hbase, boff, vv) + o4 = cute.make_fragment((4,), cutlass.Int32) + for j in cutlass.range_constexpr(3, -1, -1): + cq = vv[j] + if cutlass.const_expr(zero): + o4[j] = cutlass.Int32(0) + else: + o4[j] = after + gb = base + cutlass.Int32(4 * q + j) + cross = cutlass.Int32(0) + if after < target: + if (after + cq) >= target: + cross = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross = cutlass.Int32(1) + if cross != cutlass.Int32(0): + s_res[RES_B] = gb + s_res[RES_ABOVE] = after + s_res[RES_M] = cq + if cutlass.const_expr(two): + cross2 = cutlass.Int32(0) + if after < target2: + if (after + cq) >= target2: + cross2 = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross2 = cutlass.Int32(1) + if cross2 != cutlass.Int32(0): + s_res[RES_B2] = gb + if cutlass.const_expr(three): + cross3 = cutlass.Int32(0) + if after < target3: + if (after + cq) >= target3: + cross3 = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross3 = cutlass.Int32(1) + if cross3 != cutlass.Int32(0): + s_res[RES_B3] = gb + after = after + cq + if cutlass.const_expr(addf): # fold per-rank bin offset (L279-282) + av = cute.make_fragment((4,), cutlass.Int32) + aoff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + lds128_i32(atom, s_addv.iterator.toint(), aoff, av) + for j in cutlass.range_constexpr(4): + o4[j] = o4[j] + av[j] + boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + sts128_i32(atom, o4, hbase, boff) + + +# =========================================================================== +# scan_cross (kernel.cu L179-216) +# Block-parallel suffix scan over NB_ (<= BLK) bins. Leaves hist[j] = OUTPUT +# CURSOR (count in bins > j) and pins the crossing bin. Warps that hold no +# bin skip the body. EXACTLY ONE internal barrier (L199); the caller pays its +# usual publish barrier after. Used by gvr_clus whole-row degen (L2326). +# =========================================================================== +@cute.jit +def scan_cross( + s_hist, + s_ws, + target, + tidx, + s_res, + target2, + blk: cutlass.Constexpr, + nb: cutlass.Constexpr, + two: cutlass.Constexpr = False, +): + NWU = nb // 32 + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + c = cutlass.Int32(0) + w = cutlass.Int32(0) + if tidx < cutlass.Int32(nb): # L189 + c = s_hist[tidx] + w = warp_incl_scan_add(c, lane) # L191-193 + if lane == cutlass.Int32(31): + s_ws[wid] = w # L194 + cute.arch.barrier() # L199 — the ONE barrier + if tidx < cutlass.Int32(nb): # L200 + v2 = cutlass.Int32(0) + if lane < cutlass.Int32(NWU): + v2 = s_ws[lane] + pre = warp_incl_scan_add(v2, lane) # L202-204 + tot = cute.arch.shuffle_sync(pre, cutlass.Int32(31)) + off = cute.arch.shuffle_sync(pre - v2, wid) + after = tot - (off + w) + if tidx == cutlass.Int32(0): + s_res[RES_TOT] = tot + s_hist[tidx] = after # output cursor + cross = cutlass.Int32(0) + if after < target: + if (after + c) >= target: + cross = cutlass.Int32(1) + if tidx == cutlass.Int32(0): + cross = cutlass.Int32(1) + if cross != cutlass.Int32(0): + s_res[RES_B] = tidx + s_res[RES_ABOVE] = after + s_res[RES_M] = c + if cutlass.const_expr(two): # L212-214 + cross2 = cutlass.Int32(0) + if after < target2: + if (after + c) >= target2: + cross2 = cutlass.Int32(1) + if tidx == cutlass.Int32(0): + cross2 = cutlass.Int32(1) + if cross2 != cutlass.Int32(0): + s_res[RES_B2] = tidx + + +# =========================================================================== +# scan_cross_w (kernel.cu L288-327) +# Register-path block-parallel suffix scan for NB_ >= BLK: every thread owns +# a private contiguous BPT = NB_/BLK span, so its read->write needs no +# barrier. EXACTLY ONE internal barrier (L306). The second stage is TWO +# REDUCTIONS, not a scan (L307-315): tot = redux_add(vv), off = redux_add +# ((lane < wid) ? vv : 0) — wid is warp-uniform so the masked operand stays +# convergent. +# =========================================================================== +@cute.jit +def scan_cross_w(s_hist, s_ws, target, tidx, s_res, blk: cutlass.Constexpr, nb: cutlass.Constexpr): + BPT = nb // blk + NW = blk // 32 + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + loc = cute.make_fragment((BPT,), cutlass.Int32) + base = tidx * cutlass.Int32(BPT) + sm = cutlass.Int32(0) + for i in cutlass.range_constexpr(BPT): # L297-300 (#pragma unroll) + loc[i] = s_hist[base + cutlass.Int32(i)] + sm = sm + loc[i] + w = warp_incl_scan_add(sm, lane) # L301-304 + if lane == cutlass.Int32(31): + s_ws[wid] = w # L305 + cute.arch.barrier() # L306 — the ONE barrier + vv = cutlass.Int32(0) + if lane < cutlass.Int32(NW): + vv = s_ws[lane] # L313 + tot = cutlass.Int32(warp_add_i32(vv)) # L314 + sel = cutlass.Int32(0) + if lane < wid: + sel = vv + off = cutlass.Int32(warp_add_i32(sel)) # L315 + after = tot - (off + w) # L316 + if tidx == cutlass.Int32(0): + s_res[RES_TOT] = tot + for i in cutlass.range_constexpr(BPT - 1, -1, -1): # L318-326 + cq = loc[i] + s_hist[base + cutlass.Int32(i)] = after # per-bin OUTPUT CURSOR + gb = base + cutlass.Int32(i) + cross = cutlass.Int32(0) + if after < target: + if (after + cq) >= target: + cross = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross = cutlass.Int32(1) + if cross != cutlass.Int32(0): + s_res[RES_B] = gb + s_res[RES_ABOVE] = after + s_res[RES_M] = cq + after = after + cq + + +# =========================================================================== +# merge_scan0 (kernel.cu L111-168) +# Warp-0-fused cluster merge + suffix scan: each lane reads its BPT-bin span +# from EVERY rank's hist via 16B DSMEM loads (probe P5b), sums the cluster +# totals (and the r= target: + cross = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross = cutlass.Int32(1) + if cross != cutlass.Int32(0): + s_res[RES_B] = gb + s_res[RES_ABOVE] = after + s_res[RES_M] = cq + after = after + cq + boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + sts128_i32(atom, o4, s_mrg.iterator.toint(), boff) + + +# =========================================================================== +# gather_hint == GVR_GATHER_HINT(GM_, GX_, KPTV) (kernel.cu L329-358) +# LAZY block-wide (min,max) of logits[pre_idx[j]] over all k hint slots, in +# fkey space, returned as floats. Off the hot path by design: two dependent +# memory round trips (k coalesced P[j] words, then k scattered __ldg 4B +# gathers). Contains EXACTLY 2 barriers — call sites must be block-uniform. +# Outputs are block-uniform (every thread computes them). +# NaN-safe degeneracy guard (L356): if !(GM < GX) both become sentinels. +# +# x_addr / p_addr: Int64 byte base addresses of THIS ROW of logits/pre_idx +# (pass `t.iterator.toint() + row * stride_bytes`). s_wmn/s_wmx: Uint32 smem +# tensors of >= blk//32 slots. Returns (gm, gx) Float32. +# op43 lessons L1/L2: both round trips are issued as predicated flat batches. +# =========================================================================== +@cute.jit +def gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk: cutlass.Constexpr, kpt: cutlass.Constexpr +): + NW = blk // 32 + lane = tidx & cutlass.Int32(31) + # batch A: KPT coalesced pre_idx loads, predicated flat (L340) + pvs = [] + for t in cutlass.range_constexpr(kpt): + pv = cutlass.Int32(-1) + j = tidx + cutlass.Int32(t * blk) + if j < k: + pv = ld_g_i32(p_addr, j) + pvs.append(pv) + # batch B: KPT scattered read-only gathers, predicated flat (L341-343) + xs = [] + for t in cutlass.range_constexpr(kpt): + xv = cutlass.Float32(0.0) + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): # (unsigned)p < (unsigned)n + xv = ldg_f32(x_addr, pvs[t]) + xs.append(xv) + # fold (L344-346) + glmin = cutlass.Uint32(0xFFFFFFFF) + glmax = cutlass.Uint32(0) + for t in cutlass.range_constexpr(kpt): + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): + u2 = fkey(xs[t]) + if u2 < glmin: + glmin = u2 + if u2 > glmax: + glmax = u2 + # warp redux + staging (L347-348) + glmin = warp_min_u32(glmin) + glmax = warp_max_u32(glmax) + if lane == cutlass.Int32(0): + s_wmn[tidx >> cutlass.Int32(5)] = glmin + s_wmx[tidx >> cutlass.Int32(5)] = glmax + cute.arch.barrier() # L349 (barrier 1/2) + # cross-warp redux by EVERY thread — block-uniform outputs (L350-355) + a2 = cutlass.Uint32(0xFFFFFFFF) + c2 = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + a2 = s_wmn[lane] + c2 = s_wmx[lane] + gm = invkey(warp_min_u32(a2)) + gx = invkey(warp_max_u32(c2)) + # NaN-safe degeneracy guard (L356): !(GM < GX) — NaN compares false + ok = cutlass.Int32(0) + if gm < gx: + ok = cutlass.Int32(1) + if ok == cutlass.Int32(0): + gm = cutlass.Float32(SENT_LO) + gx = cutlass.Float32(SENT_HI) + cute.arch.barrier() # L357 (barrier 2/2) + return gm, gx + + +# =========================================================================== +# ==== ct_main.py — family: main =============================== +# =========================================================================== +"""ct_main.py — op46 gvr_main (streaming self-sampling GVR) CuTeDSL port. + +Ground truth: src_cuda/kernel.cu L377-1265 (frozen); phase contract, smem map +and barrier inventory per TRANSLATION_SPEC.md §5.1; every DSL spelling pinned +by probes/PROBE_RESULTS.md (BINDING) and op43 idioms (L1..L5 lessons). + +Ctor knobs (compile-time, mirror of the CUDA template params, spec §4d): + BLK ∈ {1024, 512, 256}, U ∈ {1,2,4,8}, MINB ∈ {1,2,4}, NBS = 256, + KPT ∈ {1,2,4,8}, SPLIT ∈ {True, False} +Derived constexprs (kernel.cu L394-523, bit-identical): + HB=NBS; KBIG=(KPT>=2 && KPT*BLK>=2048); SCPB=(BLK>=1024)?(SPLIT?8192:16384) + :(KBIG?8192:4096); CMPB=(BLK>=1024)?(KBIG?4096:2048):1024; SHD=!SPLIT; + VSTG=SPLIT||BLK>=512; PFD=(MINB<=2)?min(U,4):0; PF=PFD>0; NATT=SPLIT?1:3. + +Signature (ABI parity with kernel.cu L379-382 incl. dead SCAP_/CMP_): + run(logits[b,npad] f32, pre_idx[b,k] i32, out[b,k] i32, + n, npad, k, SCAP_, CMP_, R, SMP, TGT, Q, SS2, TGT2, ws) +Grid dim3(R, b) native 2-D (probe P11); block BLK; min_blocks_per_mp=MINB is +the 64-register wall (probes P2/P15); smem via one SmemAllocator blob (all +extents compile-time), dynamic-equivalent region byte-identical to the host +formula L3149: (SCPB+4)*(VSTG?8:4) + (CMPB+1)*8. + +int2 staging convention: an int2 (x=value bits, y=index) is ONE little-endian +Uint64 = (idx << 32) | value_bits, so cbuf2 / g_buf traffic is single u64 +ld/st (mirrors the CUDA ST.64/LD.64; __ldcg = ct_common._ldcg_v2_i32). + +Barrier inventory implemented (kernel.cu line cites, checklist per op43 L5): + L555, L655, L664, [retry: L753], L909, + SPLIT: L924, [overflow: L935], L956, L961, [last: L983, L985] + non-split: L992, [ladder: L1008 | L1017] + P6: L1075, L1097, L1100, per-level L1116/L1118, L1125 + degen A: L1162, per-level L1177/L1179, L1186 + degen B: L1221, per-level L1234/L1236, L1243 + + exactly 2 inside each gather_hint expansion (L349/L357, ct_common). +scan_cross0 contains NO barrier (probe P14 protocol). +""" + + +MAXC__main = C.MAXC +GCAP__main = C.GCAP +IDXB__main = C.IDXB +IDXM__main = C.IDXM +QUADC_CLUS__main = C.QUADC_CLUS +WS_BYTES = C.GVR_WS_BUF_OFF + MAXC__main * GCAP__main * 8 # 20,973,568 (kernel.cu L44-46) + +_NEG_INF = float("-inf") + + +# --------------------------------------------------------------------------- +# single-rounding fma.rn.f32 (probe P6 emit spelling; ct_common has no fma). +# Used at every CUDA fmaf() site: T/Tk/T3 rung math (L665/L689/L722), HIC +# (L708), window terms. (x-T)*SC classify shapes stay plain sub+mul (P6: +# structurally uncontractible). +# --------------------------------------------------------------------------- +@dsl_user_op +def _fmaf(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + mlir_math.fma( + a.ir_value(loc=loc, ip=ip), + b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), + fastmath=mlir_arith.FastMathFlags.none, + loc=loc, + ip=ip, + ) + ) + + +def _st_g_u64(addr_i64, val_u64): + """plain st.global.u64 (slab publish L929/945/952, g_don restore L969).""" + p = cute.make_ptr(cutlass.Uint64, addr_i64, cute.AddressSpace.gmem, assumed_align=8) + t = cute.make_tensor(p, cute.make_layout((1,))) + t[0] = val_u64 + + +def _st_g_u32(addr_i64, val_i32): + """plain st.global.u32 (g_off restore L969).""" + p = cute.make_ptr(cutlass.Int32, addr_i64, cute.AddressSpace.gmem, assumed_align=4) + t = cute.make_tensor(p, cute.make_layout((1,))) + t[0] = val_i32 + + +@dsl_user_op +def _st_s_v2_u32(saddr_i32, lo_u32, hi_u32, *, loc=None, ip=None): + """st.shared.v2.u32 [saddr], {lo, hi} — the CUDA make_int2 STS.64 spelling + (kernel.cu L1018-1019). Byte-identical to the little-endian u64 pack + ((hi << 32) | lo) but keeps the two words as independent 32-bit registers, + so ptxas can coalesce the emission bit-walk's loop-carried (xv, idx) pair + straight into the store pair (drops 2 IMAD.MOV/iter; op46 SASS diff).""" + mlir_llvm.inline_asm( + res=None, + operands_=[ + saddr_i32.ir_value(loc=loc, ip=ip), + lo_u32.ir_value(loc=loc, ip=ip), + hi_u32.ir_value(loc=loc, ip=ip), + ], + asm_string="st.shared.v2.u32 [$0], {$1, $2};", + constraints="r,r,r", + has_side_effects=True, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _pin_i64(v, *, loc=None, ip=None): + """Opaque identity mov.b64: pins a loop-invariant Int64 so NVVM cannot + rematerialize its defining chain (param ld.const + %ctaid reads + mul/add) + into every scf region body (PTX $L__BB0_123 evidence, op46 SASS diff).""" + return cutlass.Int64( + mlir_llvm.inline_asm( + T.i64(), + [v.ir_value(loc=loc, ip=ip)], + "mov.b64 $0, $1;", + "=l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _pin_i32(v, *, loc=None, ip=None): + """Opaque identity mov.b32 (Int32 twin of _pin_i64).""" + return cutlass.Int32( + mlir_llvm.inline_asm( + T.i32(), + [v.ir_value(loc=loc, ip=ip)], + "mov.b32 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _ldg_f32_rs(base_addr, idx, sc4): + """__ldg(X + idx) with the byte stride riding a register (fix-3 P1d). + + Identical to ct_common.ldg_f32 except `* 4` multiplies a caller-held + Int32: the row base is uniformized into URx by ptxas, IMAD.WIDE cannot + encode an immediate stride next to a UR addend, and a constant stride + register gets re-materialized INSIDE the survivor walk (its 26th + instruction; CUDA parity = 25). The caller loads the 4 from smem + (LDS results are opaque to ptxas value-tracking — asm movs and shfl + are NOT, see FIX2_CPU_VERIFY.md P1c falsification), so the register + stays live and the remat disappears.""" + atom = C.g2r_atom_f32(32, invariant=True) + p = cute.make_ptr( + cutlass.Float32, + base_addr + cutlass.Int64(idx) * cutlass.Int64(sc4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + frag = cute.make_fragment((1,), cutlass.Float32) + cute.copy(atom, cute.make_tensor(p, cute.make_layout((1,))), frag) + return frag[0] + + +@dsl_user_op +def _smem_addr_reg(addr, *, loc=None, ip=None): + """Pin a CTA-shared 32-bit byte address in ONE register (ct_reg A1 donor). + + Identity `mov` behind an asm boundary: without it LLVM re-folds the + `mov.b32 %r, __dynamic_shmem__0` symbol materialisation into EVERY use + site inside the divergent emission bit-walk — the residual +1 IMAD.MOV + per survivor at BS256 (VIOLATION_TRIAGE_20260814 §3). The asm result is + not duplicable, so the shared window is materialised exactly once. + Value-identical: a plain register copy.""" + return cutlass.Int32( + mlir_llvm.inline_asm( + T.i32(), + [addr.ir_value(loc=loc, ip=ip)], + "mov.u32 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_shared_add1(addr, *, loc=None, ip=None): + """CUDA `atomicAdd(&hist[bin], 1u)` with the result unused (ct_reg A1). + + `red` (not `atom`) is the result-less spelling — ptxas lowers it to the + same ATOMS.POPC.INC.32 RZ the CUDA arm emits. Same ordering contract as + ct_common.atomic_add_cta (.relaxed scope .cta). Takes the final shared + byte address as a plain Int32 so ptxas fuses the shl+add into one LEA + against the pinned `_smem_addr_reg` base.""" + mlir_llvm.inline_asm( + res=None, + operands_=[addr.ir_value(loc=loc, ip=ip)], + asm_string="red.relaxed.cta.shared.add.u32 [$0], 1;", + constraints="r", + has_side_effects=True, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class GvrMainKernel: + """CuTeDSL port of gvr_main (kernel.cu L377).""" + + def __init__( + self, blk: int, u: int, minb: int, nbs: int, kpt: int, split: bool, tshg: bool = False + ): + assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177, measured)" + assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) + assert kpt in (1, 2, 4, 8) and minb in (1, 2, 4) + self.blk = blk + self.u = u + self.minb = minb + self.nbs = nbs + self.kpt = kpt + self.split = bool(split) + # knife5 (layer 7): TSH-floor staging arm. SPLIT-only compile-time + # key; the CUDA form is a grid-uniform runtime gate over the same + # predicate (b > 15 && k <= 1024 && n4 <= 32768). + self.tshg = bool(tshg) and bool(split) + # derived constexprs (kernel.cu L394-523) + self.hb = nbs # L394 + self.kbig = (kpt >= 2) and (kpt * blk >= 2048) # L413 + self.scpb = ( + (8192 if split else 16384) if blk >= 1024 else (8192 if self.kbig else 4096) + ) # L422-423 + self.cmpb = (4096 if self.kbig else 2048) if blk >= 1024 else 1024 # L424 + self.shd = not split # L438 + self.vstg = split or blk >= 512 # L445 + self.pfd = (u if u < 4 else 4) if minb <= 2 else 0 # L522 + self.pf = self.pfd > 0 # L523 + self.natt = 1 if split else 3 # L733 + # smem blob byte map (kernel.cu L447-456): cbuf/cbuf2 alias @0, + # ck64 @ 4*(VSTG ? 2*(SCPB+4) : SCPB+4), size (CMPB+1)*8 + self.ck_off = 4 * ((2 * (self.scpb + 4)) if self.vstg else (self.scpb + 4)) + assert self.ck_off % 16 == 0, "ck64 must stay 16B aligned (ulonglong2)" + self.dyn_bytes = self.ck_off + (self.cmpb + 1) * 8 + self.lb = self.nbs.bit_length() - 1 # log2(NBS)=8 + + # ------------------------------------------------------------------ + # GVR_EMITC (kernel.cu L869-883): classify+stage one survivor. + # Returns pos+1. Branchless trash slot min(pos, SCPB) (L866-868). + # ------------------------------------------------------------------ + @cute.jit + def _emitc(self, xv, idx, pos, TF, SC, hb, cb2, s_hist, s_cbuf, s_cbuf2): + SCPB = self.scpb + NBS = self.nbs + if cutlass.const_expr(not self.split): + bn_u = C.f2u_rz((xv - TF) * SC) # saturating cvt.rzi (P4) + if bn_u > cutlass.Uint32(NBS - 1): + bn_u = cutlass.Uint32(NBS - 1) + bn = cutlass.Int32(bn_u) + if cutlass.const_expr(self.vstg): + # fix-2 P1: result unused -> resultless red off the pinned + # hist base (ct_reg A1 idiom) — no per-site smem-base refold + _red_shared_add1(hb + (bn << cutlass.Int32(2))) + else: + # VSTG=False tuples sit at the 64-reg wall (fix-1 spill-5 + # lesson): keep the original spelling, no pinned base here + C.atomic_add_cta(s_hist.iterator + bn, cutlass.Int32(1)) + if cutlass.const_expr(not self.vstg): + ps = pos + if ps > cutlass.Int32(SCPB): + ps = cutlass.Int32(SCPB) # trash slot (IMNMX) + s_cbuf[ps] = cutlass.Int32( + (bn_u << cutlass.Uint32(IDXB__main)) | cutlass.Uint32(idx) + ) + if cutlass.const_expr(self.vstg): + ps = pos + if ps > cutlass.Int32(SCPB): + ps = cutlass.Int32(SCPB) + # int2 {value bits, idx} via st.shared.v2.u32 — same bytes as the + # former (idx << 32) | bits u64 pack (+0=bits, +4=idx), but no i64 + # materialization inside the bit-walk (kernel.cu L1018-1019 parity) + # fix-2 P1: address = one LEA off the pinned cb2 base + _st_s_v2_u32(cb2 + ps * cutlass.Int32(8), C.u32_of_f32(xv), cutlass.Uint32(idx)) + return pos + cutlass.Int32(1) + + # ------------------------------------------------------------------ + # two-predicate warp-ballot emit step (shared by P6 L1124-1146 and both + # degen emits L1185-1210 / L1242-1263): q1 winners to out[base1+p] p> cutlass.Int32(2) # L477 + c0 = cutlass.Int32(0) + c1 = n4 + if cutlass.const_expr(self.split): # L478-479 + c0 = part * Q + c1 = c0 + Q + if c1 > n4: + c1 = n4 + tail0 = n4 << cutlass.Int32(2) # L480 + tailn = cutlass.Int32(0) + if part == cutlass.Int32(0): # L481 + tailn = n - tail0 + + if tidx == cutlass.Int32(0): # L483 + s_scal[0] = cutlass.Int32(0) # s_bufn + s_res[C.RES_B2] = cutlass.Int32(-1) + s_res[C.RES_B3] = cutlass.Int32(-1) + if tidx < cutlass.Int32(self.hb): # L484-487 (HB<=BLK always) + s_hist[tidx] = cutlass.Int32(0) + + # ============ P1: sample prefetch (hint gather LAZY, L489-529) ======= + atom128 = C.g2r_atom_f32(128, invariant=True) + fsa = cute.make_fragment((4,), cutlass.Float32) + fsb = cute.make_fragment((4,), cutlass.Float32) + shas = cutlass.Int32(0) + if tidx < SMP: + shas = cutlass.Int32(1) + if shas != cutlass.Int32(0): # L502-504 + p4 = tidx * SS2 * cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fsa) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) + + # ============ P2: quantile rung from the sample (L531-727) =========== + smn = cutlass.Float32(float("inf")) # L538 + smx = cutlass.Float32(float("-inf")) + if shas != cutlass.Int32(0): # L539-543 + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsa[t]) + smx = C.fmax_f32(smx, fsa[t]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsb[t]) + smx = C.fmax_f32(smx, fsb[t]) + fma_ = cute.make_fragment((4,), cutlass.Float32) # strided-tail pair bufs + fmb_ = cute.make_fragment((4,), cutlass.Float32) + j = tidx + cutlass.Int32(BLK) # L544-550 strided tail + while j < SMP: + p4 = j * SS2 * cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fma_) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fma_[t]) + smx = C.fmax_f32(smx, fma_[t]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fmb_[t]) + smx = C.fmax_f32(smx, fmb_[t]) + j = j + cutlass.Int32(BLK) + a0 = C.warp_min_u32(C.fkey(smn)) # L551-554 + c0m = C.warp_max_u32(C.fkey(smx)) + if lane == cutlass.Int32(0): + s_wmn[tidx >> cutlass.Int32(5)] = a0 + s_wmx[tidx >> cutlass.Int32(5)] = c0m + cute.arch.barrier() # ---- barrier L555 ---- + + # PRIME-LATE prefetch block (L556-616): strictly after the barrier. + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L524 + pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(max(PFD, 1))] + if cutlass.const_expr(self.pf): + fullsl = cutlass.Int32(0) + if (c1 - c0) >= cutlass.Int32(BLK * U): + fullsl = cutlass.Int32(1) + if fullsl != cutlass.Int32(0): # L557-559 prime, full slice + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4(atom128, x_addr, c0 + tidx + cutlass.Int32(uu * BLK), pf[uu]) + else: # L561-562 clamped prime + for uu in cutlass.range_constexpr(PFD): + i_ = c0 + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= c1: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + # asm prefetch site #1 (L573-577): gate (c1-c0)>=2*BLK*U && SMP>=160 + g1 = cutlass.Int32(0) + if (c1 - c0) >= cutlass.Int32(2 * BLK * U): + if SMP >= cutlass.Int32(160): + g1 = cutlass.Int32(1) + if g1 != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD, U): + C._prefetch_l2( + x_addr + + cutlass.Int64(c0 + tidx + cutlass.Int32(uu * BLK)) * cutlass.Int64(16) + ) + if cutlass.const_expr((not self.pf) and (not self.split)): + fullsl = cutlass.Int32(0) + if (c1 - c0) >= cutlass.Int32(BLK * U): + fullsl = cutlass.Int32(1) + if fullsl != cutlass.Int32(0): # site #2 (L589-592) + for uu in cutlass.range_constexpr(U): + C._prefetch_l2( + x_addr + + cutlass.Int64(c0 + tidx + cutlass.Int32(uu * BLK)) * cutlass.Int64(16) + ) + else: + if SMP > cutlass.Int32(0): # site #3 knife4-L1 (L593-615) + for uu in cutlass.range_constexpr(U): + i_ = c0 + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= c1: + ic = lim4 + C._prefetch_l2(x_addr + cutlass.Int64(ic) * cutlass.Int64(16)) + + # cross-warp sample reduce (L617-623) + av = cutlass.Uint32(0xFFFFFFFF) + cv = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + av = s_wmn[lane] + cv = s_wmx[lane] + SMIN = C.invkey(C.warp_min_u32(av)) + SMAX = C.invkey(C.warp_max_u32(cv)) + + GMIN = cutlass.Float32(C.SENT_LO) # L629 sentinels + GMAX = cutlass.Float32(C.SENT_HI) + T = cutlass.Float32(_NEG_INF) + HIC = cutlass.Float32(_NEG_INF) + w = cutlass.Float32(0.0) + sok = cutlass.Int32(0) # L633 + if SMP > cutlass.Int32(0): + if SMAX > SMIN: + sok = cutlass.Int32(1) + if sok != cutlass.Int32(0): # L635-654 sample histogram + w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) + # rcp.approx.ftz.f32 = the CUDA arm's --use_fast_math 1.0f/w + # (bare MUFU.RCP, no Newton refinement) — bitwise-aligned scale + sc_s = cute.arch.rcp_approx(w) + if shas != cutlass.Int32(0): + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsa[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsb[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = tidx + cutlass.Int32(BLK) # L646-653 tail re-loads + while j < SMP: + p4 = j * SS2 * cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fma_) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fma_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fmb_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = j + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L655 ---- + # triple-target ZERO scan (L659-660): TGT / TGT2 / 2*TGT + # (THREE = SHD || gated-SPLIT, knife5 layer 7) + C.scan_cross0( + s_hist, + TGT, + tidx, + s_res, + TGT2, + TGT * cutlass.Int32(2), + s_hist, + nb=NBS, + zero=True, + two=True, + three=(self.shd or self.tshg), + ) + cute.arch.barrier() # ---- barrier L664 ---- + + tot0 = s_res[C.RES_TOT] + b1v = s_res[C.RES_B] + if sok != cutlass.Int32(0): # L665 + if tot0 >= TGT: + T = _fmaf(cutlass.Float32(b1v), w, SMIN) + Trung = T # L666 snapshot + needg = cutlass.Int32(1) # L667-675 degenerate sample + if T > cutlass.Float32(_NEG_INF): + needg = cutlass.Int32(0) + if needg != cutlass.Int32(0): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) # 2 barriers inside + T = GMIN + if sok != cutlass.Int32(0): # L688-709 HIC tighten + if tot0 >= TGT: + b2v = s_res[C.RES_B2] + if b2v >= cutlass.Int32(0): + Tk = _fmaf(cutlass.Float32(b2v), w, SMIN) + anch = T + if cutlass.const_expr(not self.split): + anch = C.fmin_f32(T, Trung) # L703 + d_ = C.fmax_f32(Tk - anch, cutlass.Float32(0.0)) + HIC = C.fmax_f32( + _fmaf(cutlass.Float32(4.0), d_, T), _fmaf(cutlass.Float32(8.0), w, T) + ) # L708 + if cutlass.const_expr(self.shd or self.tshg): # TSH floor (knife5: +gated SPLIT) + if tidx == cutlass.Int32(0): + t5 = cutlass.Float32(_NEG_INF) + if sok != cutlass.Int32(0): + if tot0 >= TGT * cutlass.Int32(2): + b3v = s_res[C.RES_B3] + if b3v >= cutlass.Int32(0): + if T > GMIN: + T3 = _fmaf(cutlass.Float32(b3v), w, SMIN) + if T3 < T: + t5 = T3 + s_tsh[0] = t5 + + if cutlass.const_expr(self.tshg): + # knife5 (layer 7) TSH-FLOOR STAGING: SPLIT has no retry ladder, + # so a rung overshoot (count(>=T) < k) used to hand the LAST CTA + # a single-CTA whole-row narrowing. Stage at the sample's + # rank-(2*TGT) floor instead: staged population ~aim -> ~2*aim, + # and the merged histogram contains the k-crossing whenever + # count(>=TSH) >= k. TSH miss falls to GMIN/degen unchanged. + cute.arch.barrier() + t5s = s_tsh[0] + if t5s > cutlass.Float32(_NEG_INF): + if t5s < T: + T = t5s + + # ============ attempt loop (L729-1019) — MUST NOT unroll ============ + listN = cutlass.Int32(0) + above = cutlass.Int32(0) + m = cutlass.Int32(0) + need = cutlass.Int32(0) + B = cutlass.Int32(0) + SC = cutlass.Float32(1.0) + TF = T + complete = cutlass.Int32(0) + valid = cutlass.Int32(0) + fromg = cutlass.Int32(0) + alive = cutlass.Int32(1) + + fr = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(max(U - PFD, 1)) + ] # explicit batch (op43 L1) + att = cutlass.Int32(0) + running = cutlass.Int32(1) + while running != cutlass.Int32(0): + if cutlass.const_expr(not self.split): # SPLIT never retries (NATT=1) + if att > cutlass.Int32(0): # retry reset (L737-754) + if cutlass.const_expr(self.pf): + # EXACTNESS: re-prime pf[] (stale roll data, L738-749) + fullsl = cutlass.Int32(0) + if (c1 - c0) >= cutlass.Int32(BLK * U): + fullsl = cutlass.Int32(1) + if fullsl != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4( + atom128, x_addr, c0 + tidx + cutlass.Int32(uu * BLK), pf[uu] + ) + else: + for uu in cutlass.range_constexpr(PFD): + i_ = c0 + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= c1: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + if tidx < cutlass.Int32(NBS): # L750-751 + s_hist[tidx] = cutlass.Int32(0) + if tidx == cutlass.Int32(0): # L752 + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L753 ---- + + TF = T # window (L756-761) + hi = C.fmax_f32(GMAX, T) + if HIC > T: + if HIC < hi: + hi = HIC + WD = (hi - T) * cutlass.Float32(1.0 / 256.0) + wdok = cutlass.Int32(0) + if WD > cutlass.Float32(0.0): + wdok = cutlass.Int32(1) + if wdok == cutlass.Int32(0): + WD = cutlass.Float32(1e-30) + # fix-4: CUDA compiles its own `1.0f / WD` here to a bare + # MUFU.RCP (approximate) — div.rn's dependent rcp+Newton+CALL + # chain serializes the attempt prologue and is the last + # CALL-class divergence (era-3 ncu: DSL executes ~4% FEWER + # instructions yet runs ~5% slower = latency-shaped, so the + # cost is the chain, not the instruction count). blk==512 + # ONLY: the fix-1 spill-5 revert was bisected on (256,8,4,·,4) + # — that family keeps the original div.rn spelling below. + if cutlass.const_expr(self.blk == 512): + SC = cute.arch.rcp_approx(WD) + else: + SC = cutlass.Float32(1.0) / WD + + # ---- P3 row pass (L763-908) ---- + span = c1 - c0 + step = cutlass.Int32(BLK * U) + nFull = cutlass.Int32(0) + rem = cutlass.Int32(0) + if span > cutlass.Int32(0): # L776-779 peel + nFull = span // step + rem = span - nFull * step + # _pin_i32: the isfull peel predicate reads nFull every tile iter; + # unpinned, NVVM re-derives the whole ld.param+shr/sel div chain + # at the loop head (v3 SASS evidence) + nFull = _pin_i32(nFull) + nIt = nFull + if rem > cutlass.Int32(0): + nIt = nIt + cutlass.Int32(1) + # _pin_i32: stop NVVM re-deriving the ceil-div bound (ld.param n + + # shr/sel chain) inside the tile-loop condition region per iter + nIt = _pin_i32(nIt) + + it = cutlass.Int32(0) + while it < nIt: + i0 = c0 + it * step + tidx + M = cutlass.Int32(0) + isfull = cutlass.Int32(0) + if it < nFull: + isfull = cutlass.Int32(1) + if isfull != cutlass.Int32(0): # full body (L783-795) + for uu in cutlass.range_constexpr(PFD, U): + C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + else: # partial body (L796-815) + for uu in cutlass.range_constexpr(PFD, U): + i_ = i0 + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= c1: + ic = lim4 # clamped address (P10) + C.ld_g_f32x4(atom128, x_addr, ic, fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + i_ = i0 + cutlass.Int32(uu * BLK) + okq = cutlass.Int32(0) + if i_ < c1: + okq = cutlass.Int32(1) + if okq != cutlass.Int32(0): # ok-gated (+inf pad fix L804-813) + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + # prefetch roll-forward BEFORE reservation/walk (L816-842) + if cutlass.const_expr(self.pf): + hasnext = cutlass.Int32(0) + if it + cutlass.Int32(1) < nIt: + hasnext = cutlass.Int32(1) + if hasnext != cutlass.Int32(0): + j0 = i0 + step + infull = cutlass.Int32(0) # warp-uniform peel L835 + if it + cutlass.Int32(1) < nFull: + infull = cutlass.Int32(1) + if infull != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4(atom128, x_addr, j0 + cutlass.Int32(uu * BLK), pf[uu]) + else: + for uu in cutlass.range_constexpr(PFD): + j_ = j0 + cutlass.Int32(uu * BLK) + jc = j_ + if jc >= c1: + jc = lim4 + C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) + # warp-aggregated reservation (L843-854) + cnt = cutlass.Int32(C.popc(M)) + inc = C.warp_incl_scan_add(cnt, lane) + bpos = cutlass.Int32(0) + if lane == cutlass.Int32(31): + if inc != cutlass.Int32(0): + bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) + pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) + # survivor bit-walk, software-pipelined ONE deep (L884-898); + # reload X[idx] — do NOT hold the U float4s (+18% spill L855-859) + if M != cutlass.Int32(0): + bp = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx = ( + (i0 + (bp >> cutlass.Int32(2)) * cutlass.Int32(BLK)) << cutlass.Int32(2) + ) + (bp & cutlass.Int32(3)) + if cutlass.const_expr(self.vstg and self.blk == 512): + xv = _ldg_f32_rs(x_addr, idx, x4_pin) + else: + xv = C.ldg_f32(x_addr, idx) + while M != cutlass.Int32(0): + bp2 = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx2 = ( + (i0 + (bp2 >> cutlass.Int32(2)) * cutlass.Int32(BLK)) + << cutlass.Int32(2) + ) + (bp2 & cutlass.Int32(3)) + if cutlass.const_expr(self.vstg and self.blk == 512): + xv2 = _ldg_f32_rs(x_addr, idx2, x4_pin) + else: + xv2 = C.ldg_f32(x_addr, idx2) + pos = self._emitc( + xv, idx, pos, TF, SC, hb_pin, cb2_pin, s_hist, s_cbuf, s_cbuf2 + ) + idx = idx2 + xv = xv2 + pos = self._emitc( + xv, idx, pos, TF, SC, hb_pin, cb2_pin, s_hist, s_cbuf, s_cbuf2 + ) + it = it + cutlass.Int32(1) + # scalar tail, part 0 only (L900-906) + i = tidx + while i < tailn: + x = C.ldg_f32(x_addr, tail0 + i) + if x >= TF: + post = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) + post = self._emitc( + x, tail0 + i, post, TF, SC, hb_pin, cb2_pin, s_hist, s_cbuf, s_cbuf2 + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L909 ---- + myn = s_scal[0] # L911 + + if cutlass.const_expr(self.split): + # ---- SLAB HAND-OFF (L913-989); exactly ONE attempt ---- + if tidx == cutlass.Int32(0): # L923 + pgo = cute.make_ptr( + cutlass.Int32, + goff_addr + row64 * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + s_scal[3] = cutlass.Int32(cute.arch.atomic_add(pgo, myn)) + cute.arch.barrier() # ---- barrier L924 ---- + base = s_scal[3] + if myn <= cutlass.Int32(SCPB): # L926-930 coalesced publish + i = tidx + while i < myn: + p = base + i + if p < cutlass.Int32(GCAP__main): + _st_g_u64(gbuf_row + cutlass.Int64(p) * cutlass.Int64(8), s_cbuf2[i]) + i = i + cutlass.Int32(BLK) + else: # L931-955 overflow re-sweep + if tidx == cutlass.Int32(0): + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L935 ---- + lo2 = c0 << cutlass.Int32(2) + hi2 = c1 << cutlass.Int32(2) + i = lo2 + tidx + while i < hi2: + x = C.ldg_f32(x_addr, i) + if x >= TF: + pq = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) + p = base + pq + if p < cutlass.Int32(GCAP__main): + _st_g_u64( + gbuf_row + cutlass.Int64(p) * cutlass.Int64(8), + (cutlass.Uint64(cutlass.Uint32(i)) << cutlass.Uint64(32)) + | cutlass.Uint64(C.u32_of_f32(x)), + ) + i = i + cutlass.Int32(BLK) + i = tidx # true tail (L948-954) + while i < tailn: + x = C.ldg_f32(x_addr, tail0 + i) + if x >= TF: + pq = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) + p = base + pq + if p < cutlass.Int32(GCAP__main): + _st_g_u64( + gbuf_row + cutlass.Int64(p) * cutlass.Int64(8), + ( + cutlass.Uint64(cutlass.Uint32(tail0 + i)) + << cutlass.Uint64(32) + ) + | cutlass.Uint64(C.u32_of_f32(x)), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L956 ---- + if tidx == cutlass.Int32(0): # L959-960 release + RMW + C.threadfence_gpu() + pdon = cute.make_ptr( + cutlass.Int64, + gdon_addr + row64 * cutlass.Int64(8), + cute.AddressSpace.gmem, + assumed_align=8, + ) + s_pk[0] = C.atomic_add_u64_gpu( + pdon, cutlass.Int64(1 << 32) + cutlass.Int64(myn) + ) + cute.arch.barrier() # ---- barrier L961 ---- + pk = s_pk[0] + alive = cutlass.Int32(0) # L963 last-CTA test + if cutlass.Int32(pk >> cutlass.Int64(32)) == R - cutlass.Int32(1): + alive = cutlass.Int32(1) + if alive != cutlass.Int32(0): + C.threadfence_gpu() # L968 acquire + if tidx == cutlass.Int32(0): # L969 ZERO-RESTORE + _st_g_u32(goff_addr + row64 * cutlass.Int64(4), cutlass.Int32(0)) + _st_g_u64(gdon_addr + row64 * cutlass.Int64(8), cutlass.Uint64(0)) + total = cutlass.Int32(pk & cutlass.Int64(0xFFFFFFFF)) + myn + if total <= cutlass.Int32(GCAP__main): # L971-988 one-pass consume + listN = total + if total > cutlass.Int32(SCPB): + fromg = cutlass.Int32(1) + i = tidx + while i < listN: + gvx, gvy = C._ldcg_v2_i32( + gbuf_row + cutlass.Int64(i) * cutlass.Int64(8) + ) + if fromg == cutlass.Int32(0): + s_cbuf2[i] = ( + cutlass.Uint64(cutlass.Uint32(gvy)) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(gvx)) + bq = C.f2s_rz((C.f32_of_i32(gvx) - TF) * SC) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + # fix-2 P2: resultless red off the P1 hist pin + _red_shared_add1(hb_pin + (bq << cutlass.Int32(2))) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L983 ---- + C.scan_cross0( + s_hist, + k, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=False, + ) + cute.arch.barrier() # ---- barrier L985 ---- + if s_res[C.RES_TOT] >= k: # L986-987 + valid = cutlass.Int32(1) + complete = cutlass.Int32(1) + above = s_res[C.RES_ABOVE] + m = s_res[C.RES_M] + need = k - above + B = s_res[C.RES_B] + running = cutlass.Int32(0) # L989 break (NATT==1) + else: + # ---- non-split verify + rung ladder (L990-1018) ---- + C.scan_cross0( + s_hist, + k, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=False, + ) + cute.arch.barrier() # ---- barrier L992 ---- + tot = s_res[C.RES_TOT] + acc = cutlass.Int32(0) + if tot >= k: + acc = cutlass.Int32(1) + if acc != cutlass.Int32(0): # L994-998 accept + valid = cutlass.Int32(1) + complete = cutlass.Int32(0) + if myn <= cutlass.Int32(SCPB): + complete = cutlass.Int32(1) + listN = myn + above = s_res[C.RES_ABOVE] + m = s_res[C.RES_M] + need = k - above + B = s_res[C.RES_B] + running = cutlass.Int32(0) + else: + if att == cutlass.Int32(NATT - 1): # L999 ladder exhausted + running = cutlass.Int32(0) + else: + tshtaken = cutlass.Int32(0) # L1005-1010 TSH retry + if cutlass.const_expr(self.shd): + if att == cutlass.Int32(0): + T5 = s_tsh[0] + if T5 > cutlass.Float32(_NEG_INF): + if T5 < TF: + T = T5 + tshtaken = cutlass.Int32(1) + if tshtaken != cutlass.Int32(0): + cute.arch.barrier() # ---- barrier L1008 ---- + else: + # LAZY GATHER (L1014, sentinel equality flag) + if GMIN == cutlass.Float32(C.SENT_LO): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT + ) + floorhit = cutlass.Int32(1) # L1015 + if T > GMIN: + floorhit = cutlass.Int32(0) + if floorhit != cutlass.Int32(0): + running = cutlass.Int32(0) + else: + T = GMIN # L1016 + cute.arch.barrier() # ---- barrier L1017 ---- + att = att + cutlass.Int32(1) + + # ============ classification (L1021-1024) ============ + if alive != cutlass.Int32(0): + whole = cutlass.Int32(0) + if valid != cutlass.Int32(0): + if need >= m: + whole = cutlass.Int32(1) + lim1 = above + if whole != cutlass.Int32(0): + lim1 = above + m + degen = cutlass.Int32(0) + if valid == cutlass.Int32(0): + degen = cutlass.Int32(1) + if m > cutlass.Int32(CMPB): + degen = cutlass.Int32(1) + mc = cutlass.Int32(0) + if degen == cutlass.Int32(0): + mc = m + + if degen == cutlass.Int32(0): + # ---- P5 cursor emit (L1026-1071) ---- + if complete != cutlass.Int32(0): # L1028-1048 + i = tidx + while i < listN: + idv = cutlass.Int32(0) + bq = cutlass.Int32(0) + xv = cutlass.Float32(0.0) + if cutlass.const_expr(self.vstg): + vx = cutlass.Int32(0) + vy = cutlass.Int32(0) + if cutlass.const_expr(self.split): + if fromg != cutlass.Int32(0): + vx, vy = C._ldcg_v2_i32( + gbuf_row + cutlass.Int64(i) * cutlass.Int64(8) + ) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + vy = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + vy = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + xv = C.f32_of_i32(vx) + idv = vy + bq = C.f2s_rz((xv - TF) * SC) # L1034 + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + else: + wpk = cutlass.Uint32(s_cbuf[i]) # L1036-1037 + idv = cutlass.Int32(wpk & cutlass.Uint32(IDXM__main)) + bq = cutlass.Int32(wpk >> cutlass.Uint32(IDXB__main)) + if bq >= B: # L1039-1047 + p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + if p < lim1: + out_row[p] = idv + else: + if whole == cutlass.Int32(0): + q2 = p - above + if q2 < cutlass.Int32(CMPB): + if cutlass.const_expr(self.vstg): + kk = C.fkey(xv) + else: + kk = C.fkey(C.ldg_f32(x_addr, idv)) + s_ck64[q2] = ( + cutlass.Uint64(kk) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(idv)) + i = i + cutlass.Int32(BLK) + else: + # collect overflow: scalar re-sweep, exact tail remap + # (L1049-1070) — zero extra live registers by design + lo2 = c0 << cutlass.Int32(2) + hi2 = c1 << cutlass.Int32(2) + i0_ = lo2 + tidx + while i0_ < hi2 + tailn: + i_ = i0_ + if i0_ >= hi2: + i_ = tail0 + (i0_ - hi2) # L1057-1058 + x = C.ldg_f32(x_addr, i_) + if x >= TF: + bq = C.f2s_rz((x - TF) * SC) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + if bq >= B: + p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + if p < lim1: + out_row[p] = i_ + else: + if whole == cutlass.Int32(0): + q2 = p - above + if q2 < cutlass.Int32(CMPB): + s_ck64[q2] = ( + cutlass.Uint64(C.fkey(x)) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(i_)) + i0_ = i0_ + cutlass.Int32(BLK) + + # ---- P6 refine (L1073-1147) ---- + if whole == cutlass.Int32(0): + cute.arch.barrier() # ---- barrier L1075 ---- + if mc <= cutlass.Int32(QUADC_CLUS__main): # L1077-1092 O(mc^2) rank + mc2 = mc & cutlass.Int32(~1) + i = tidx + while i < mc: + # NOTE: values crossing a dynamic-while region are + # re-wrapped SIGNED by the DSL — every u64 compare + # must re-assert Uint64 at the USE site (found via + # odd-tail rank corruption; see notes G1). + u64v = s_ck64[i] + r_ = cutlass.Int32(0) + jq = cutlass.Int32(0) + while jq < mc2: # ulonglong2 16B reads + vlo, vhi = C._lds_v2_u64(ck_addr + jq * cutlass.Int32(8)) + r_ = ( + r_ + + cutlass.Int32(vlo > cutlass.Uint64(u64v)) + + cutlass.Int32(vhi > cutlass.Uint64(u64v)) + ) + jq = jq + cutlass.Int32(2) + if mc2 < mc: # odd tail L1089 + r_ = r_ + cutlass.Int32( + cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) + ) + if r_ < need: + out_row[above + r_] = cutlass.Int32( + cutlass.Uint32( + cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) + ) + ) + i = i + cutlass.Int32(BLK) + else: + # key-space narrowing over ck64 (L1094-1123) + if tidx == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + if tidx < cutlass.Int32(NBS): # cleared ONCE (L1096) + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1097 ---- + i = tidx + while i < mc: + kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + C.atomic_min_cta(s_kmm.iterator + 0, kk) + C.atomic_max_cta(s_kmm.iterator + 1, kk) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1100 ---- + rlo = s_kmm[0] + rhi = s_kmm[1] + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L1103-1123 (<=6 levels) + if needC == mm: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(6): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + i = tidx + while i < mc: # L1111-1115 re-bin + uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), + cutlass.Int32(1), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1116 ---- + C.scan_cross0( + s_hist, + needC, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=True, + ) + cute.arch.barrier() # ---- barrier L1118 ---- + aboveC = aboveC + s_res[C.RES_ABOVE] + needC = needC - s_res[C.RES_ABOVE] + mm = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): # L1121 + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L1124 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1125 ---- + it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L1127-1146 ballot emit + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + if i < mc: + w64 = s_ck64[i] + iu = cutlass.Int64(cutlass.Uint32(w64 >> cutlass.Uint64(32))) + idv = cutlass.Int32( + cutlass.Uint32(w64 & cutlass.Uint64(0xFFFFFFFF)) + ) + if iu > ethr: + p1 = cutlass.Int32(1) + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, + p2, + idv, + above, + aboveC, + above + aboveC, + needC, + out_row, + s_scal, + lane, + ) + it = it + cutlass.Int32(1) + else: + dga = cutlass.Int32(0) # L1156 gate: valid && complete + if valid != cutlass.Int32(0): + if complete != cutlass.Int32(0): + dga = cutlass.Int32(1) + if dga != cutlass.Int32(0): + # ---- degen A: narrowing over STAGED candidates + # (L1150-1212) ---- + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + above2 = cutlass.Int32(0) + need2 = k + m2 = listN + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) + if tidx < cutlass.Int32(NBS): # L1161 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1162 ---- + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L1163-1184 (<=8 levels) + if need2 == m2: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + above2 = above2 + m2 + need2 = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(8): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + i = tidx + while i < listN: # L1170-1176 + uq = cutlass.Uint32(0) + if cutlass.const_expr(self.vstg): + vx = cutlass.Int32(0) + vy = cutlass.Int32(0) + if cutlass.const_expr(self.split): + if fromg != cutlass.Int32(0): + vx, vy = C._ldcg_v2_i32( + gbuf_row + cutlass.Int64(i) * cutlass.Int64(8) + ) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + uq = C.fkey_bits(cutlass.Uint32(vx)) + else: + id0 = cutlass.Int32( + cutlass.Uint32(s_cbuf[i]) & cutlass.Uint32(IDXM__main) + ) + uq = C.fkey(C.ldg_f32(x_addr, id0)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1177 ---- + C.scan_cross0( + s_hist, + need2, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=True, + ) + cute.arch.barrier() # ---- barrier L1179 ---- + above2 = above2 + s_res[C.RES_ABOVE] + need2 = need2 - s_res[C.RES_ABOVE] + m2 = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L1185 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1186 ---- + nA = k # L1187 + nT = cutlass.Int32(0) + if tie_m != cutlass.Int32(0): + nA = above2 + nT = need2 + it2 = (listN + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L1189-1210 + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + if i < listN: + uq = cutlass.Uint32(0) + if cutlass.const_expr(self.vstg): + vx = cutlass.Int32(0) + vy = cutlass.Int32(0) + if cutlass.const_expr(self.split): + if fromg != cutlass.Int32(0): + vx, vy = C._ldcg_v2_i32( + gbuf_row + cutlass.Int64(i) * cutlass.Int64(8) + ) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + vy = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + else: + pk64 = s_cbuf2[i] + vx = cutlass.Int32( + cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF)) + ) + vy = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + uq = C.fkey_bits(cutlass.Uint32(vx)) + idv = vy + else: + idv = cutlass.Int32( + cutlass.Uint32(s_cbuf[i]) & cutlass.Uint32(IDXM__main) + ) + uq = C.fkey(C.ldg_f32(x_addr, idv)) + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, p2, idv, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane + ) + it = it + cutlass.Int32(1) + else: + # ---- degen B: whole-row narrowing (L1214-1264) ---- + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + above2 = cutlass.Int32(0) + need2 = k + m2 = n + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) + if tidx < cutlass.Int32(NBS): # L1220 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1221 ---- + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L1222-1241 (<=8 levels) + if need2 == m2: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + above2 = above2 + m2 + need2 = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(8): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + i = tidx + while i < n: # L1229-1233 whole row + uq = C.fkey(C.ldg_f32(x_addr, i)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1234 ---- + C.scan_cross0( + s_hist, + need2, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=True, + ) + cute.arch.barrier() # ---- barrier L1236 ---- + above2 = above2 + s_res[C.RES_ABOVE] + need2 = need2 - s_res[C.RES_ABOVE] + m2 = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L1242 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L1243 ---- + nA = k + nT = cutlass.Int32(0) + if tie_m != cutlass.Int32(0): + nA = above2 + nT = need2 + it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L1246-1263 + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + if i < n: + uq = C.fkey(C.ldg_f32(x_addr, i)) + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane + ) + it = it + cutlass.Int32(1) + + # ------------------------------------------------------------------ + # host launcher (grid dim3(R, b) L2750; MINB wall via min_blocks_per_mp) + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + out: cute.Tensor, + ws: cute.Tensor, + n: cutlass.Int32, + npad: cutlass.Int32, + k: cutlass.Int32, + scap_dead: cutlass.Int32, + cmp_dead: cutlass.Int32, + R: cutlass.Int32, + SMP: cutlass.Int32, + TGT: cutlass.Int32, + Q: cutlass.Int32, + SS2: cutlass.Int32, + TGT2: cutlass.Int32, + stream, + ): + b = logits.shape[0] + self.kern( + logits, pre_idx, out, ws, n, npad, k, scap_dead, cmp_dead, R, SMP, TGT, Q, SS2, TGT2 + ).launch(grid=(R, b, 1), block=(self.blk, 1, 1), stream=stream, min_blocks_per_mp=self.minb) + + +# --------------------------------------------------------------------------- +# compile cache + torch-facing entry +# --------------------------------------------------------------------------- +_COMPILE_CACHE = {} + + +def get_compiled(tpl, options_extra: str = ""): + """Compile (or fetch) the gvr_main variant for constexpr tuple + tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG).""" + key = (tuple(tpl), options_extra) + hit = _COMPILE_CACHE.get(key) + if hit is not None: + return hit + blk, u, minb, nbs, kpt, split, tshg = tpl + kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg)) + r0, c0 = cute.sym_int(), cute.sym_int() + r1, c1 = cute.sym_int(), cute.sym_int() + r2, c2 = cute.sym_int(), cute.sym_int() + w0 = cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (r0, c0), stride_order=(1, 0), assumed_align=16 + ) + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 + ) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r2, c2), stride_order=(1, 0), assumed_align=16 + ) + ws_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (w0,), stride_order=(0,), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kern, + logits_fake, + pre_fake, + out_fake, + ws_fake, + *([cutlass.Int32(0)] * 11), + stream=fake_stream, + options=("--enable-tvm-ffi " + options_extra).strip(), + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +def workspace_bytes() -> int: + return WS_BYTES + + +def run(logits, pre_idx, n: int, out, ws): + """torch-facing single-call entry: routes (b, n, k) through ct_dispatch, + asserts the shape lands on gvr_main, launches the matching variant. + ws: zero-initialised >=20,973,568-B CUDA buffer (reused across launches; + the kernel restores the zeros it consumes).""" + try: + from . import gvr_topk_decode_self_sampling_host as ct_dispatch + except ImportError: + import gvr_topk_decode_self_sampling_host as ct_dispatch + b, npad = logits.shape + k = pre_idx.shape[1] + r = ct_dispatch.route(b, int(n), npad, k) + assert r["kernel"] == "main", f"shape routes to {r['kernel']}, not gvr_main" + assert ws.numel() * ws.element_size() >= WS_BYTES + rt = r["rt"] + fn = get_compiled(tuple(r["tpl"])) + fn( + logits, + pre_idx, + out, + ws, + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP_"], + rt["CMP_"], + rt["R"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + return r + + +# =========================================================================== +# ==== ct_reg.py — family: reg ================================= +# =========================================================================== +"""ct_reg.py — op46 gvr_topk_reg CuTeDSL translation. + +CUDA ground truth: src_cuda/kernel.cu L1269-1776 (register-resident exact +top-K, one CTA per row, histogram bins in FLOAT space). Contract: +TRANSLATION_SPEC.md §5.2. All probe verdicts (probes/PROBE_RESULTS.md) are +binding; op43 lessons L1-L5 applied throughout. + +Template knobs (CUDA `gvr_topk_reg`): +ctor args of :class:`GvrTopkRegKernel`, read as `cutlass.const_expr(self.x)` +inside the kernel. Runtime args mirror the CUDA `(n, npad, k, CMP, IMGOFF, +QC)` — npad/k come from tensor shapes, IMGOFF is dropped (dispatch pins +IMGOFF == NBSEL == NBH at every site, asserted in the host wrapper). + +Shared-memory map (single dynamic window, word offsets; the CUDA static +__shared__ block is folded into the first 512 B so occupancy accounting +matches nvcc's static+dynamic sum): + + [0..5] s_res (ct_common slot map RES_B/M/ABOVE/TOT/B2/B3) + [6..7] s_cnt (s_o1, s_oc) + [8..9] s_kmm (s_kmin, s_kmax — Uint32) + [10..11] s_e12 (s_e1, s_e2) + [16..16+NW) ws (scan_cross_w workspace) + [48..48+NW) wmn (Uint32 warp min partials) + [80..80+NW) wmx (Uint32 warp max partials) + [128..128+NBH) hist (kernel.cu L1296) + [128+NBH..128+NBH+CMP) ck (Uint32 crossing keys; CMP dynamic) + [128+NBH+CMP..+2CMP) ci (Int32 crossing indices) + img/bm alias ck at word 128+NBH (kernel.cu L1299/L1409; IMGOFF==NBH) + +Launch smem = 512 + dispatch_smem_bytes (dynamic Int32). + +TOOLCHAIN GOTCHA (documented in notes/ct_reg_NOTES.md): dynamic launch smem +with min_blocks_per_mp>1 crashes cutlass_dsl._build_kernel_attrs (host ceil() +on a dynamic value while computing the PREFERRED_SHARED_MEMORY_CARVEOUT +hint). `_no_carveout()` scopes a monkeypatch around cute.compile dropping +ONLY that hint (CUDA __launch_bounds__ sets no carveout either); +`.reqntid`/`.minnctapersm` (the register wall) are unaffected — verified. +""" + + +NB__reg = 1024 # kernel.cu L16 (NBH default) +STATIC_WORDS = 128 # DSL smem prelude (static-__shared__ mirror) +STATIC_BYTES = STATIC_WORDS * 4 +_NEG_INF__reg = float("-inf") +_POS_INF = float("inf") + + +# --------------------------------------------------------------------------- +# module-local FP spellings (probe P6) +# --------------------------------------------------------------------------- +@dsl_user_op +def _fmaf__reg(a, b, c, *, loc=None, ip=None): + """CUDA fmaf: single fma.rn.f32 (P6 emit spelling).""" + return cutlass.Float32( + mlir_math.fma( + a.ir_value(loc=loc, ip=ip), + b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), + fastmath=mlir_arith.FastMathFlags.none, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _submul_asm(v, t, sc, *, loc=None, ip=None): + """(v - t) * sc with two roundings, opaque to CSE/contraction (P6). + + Used at the !BRL classify site (kernel.cu L1521) so no sub-expression is + shared with the emit's `fmaf(v - T, SC, OFF)` — the CUDA deliberately + spells the two sites differently to stop nvcc holding all S q's live + across the barrier (L1506-1509). + """ + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [v.ir_value(loc=loc, ip=ip), t.ir_value(loc=loc, ip=ip), sc.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 rtmp;\n\tsub.rn.f32 rtmp, $1, $2;\n\tmul.rn.f32 $0, rtmp, $3;\n\t}", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smem_addr_reg__reg(addr, *, loc=None, ip=None): + """Pin a CTA-shared 32-bit byte address in ONE register (SASS audit fix). + + Identity `mov` behind an asm boundary: without it LLVM re-folds the + `mov.b32 %r, __dynamic_shmem__0` symbol materialisation into EVERY use + site, and ptxas then re-derives the CGA shared window (S2UR SR_CgaCtaId + + UMOV + ULEA, 3 instructions) inside each divergent classify block — + measured +24 warp-instructions/warp vs the CUDA arm, which keeps the + base in one UR. The asm result is not duplicable, so the window is + materialised exactly once. Value-identical: a plain register copy. + """ + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [addr.ir_value(loc=loc, ip=ip)], + "mov.u32 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_shared_add1__reg(addr, *, loc=None, ip=None): + """CUDA classify `atomicAdd(&hist[bin], 1u)` with the result unused. + + `red` (not `atom`) is the result-less spelling — ptxas lowers it to the + same ATOMS.POPC.INC.32 RZ the CUDA arm emits (kernel.cu L1521-1526). + Same ordering contract as atomic_add_cta: .relaxed scope .cta. Takes the + final shared byte address as a plain Int32 so the address datapath stays + ordinary IR (ptxas fuses the shl+add into one LEA against the pinned + `_smem_addr_reg__reg` base). + """ + llvm.inline_asm( + res=None, + operands_=[addr.ir_value(loc=loc, ip=ip)], + asm_string="red.relaxed.cta.shared.add.u32 [$0], 1;", + constraints="r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _umin_u32(a, b): + """unsigned min(a, b) — CUDA min() on the bin clamp (IMNMX).""" + r = a + if b < a: + r = b + return r + + +@cute.jit +def _fabsf(x): + """|x| via sign-bit clear (exact, matches fabsf).""" + return f32_of_u32(u32_of_f32(x) & cutlass.Uint32(0x7FFFFFFF)) + + +def _f32_smem_atom(): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128) + + +def _sts128_f32(atom, frag, base_addr, byte_off): + p = cute.make_ptr( + cutlass.Float32, base_addr + byte_off, cute.AddressSpace.smem, assumed_align=16 + ) + cute.copy(atom, frag, cute.make_tensor(p, cute.make_layout((4,)))) + + +def _smem_view(dtype, sbase, word_off: int, length: int, align: int = 16): + """Typed tensor view at a constexpr word offset into the smem window.""" + p = cute.make_ptr( + dtype, sbase + cutlass.Int32(word_off * 4), cute.AddressSpace.smem, assumed_align=align + ) + return cute.make_tensor(p, cute.make_layout((length,))) + + +def _val(frags, s: int): + """val[s] accessor over the float4[VPT] register batch (constexpr s).""" + return frags[s // 4][s % 4] + + +@contextlib.contextmanager +def _no_carveout(): + """Scoped: drop the DSL's carveout hint (see module docstring).""" + import cutlass.cutlass_dsl.cutlass as _cdsl + + orig = _cdsl._build_kernel_attrs + _cdsl._build_kernel_attrs = lambda config: {} + try: + yield + finally: + _cdsl._build_kernel_attrs = orig + + +class GvrTopkRegKernel: + """gvr_topk_reg (kernel.cu L1287).""" + + def __init__( + self, + blk: int, + vpt: int, + minb: int, + kpt: int, + cur: bool, + deg: bool, + img: bool, + nbh: int = NB__reg, + pdl: bool = False, + ): + assert blk in (256, 512, 1024) and vpt in (1, 2, 4) + assert nbh in (256, 512, 1024, 2048) + assert nbh % blk == 0 or blk % nbh == 0 + self.blk = blk + self.vpt = vpt + self.minb = minb + self.kpt = kpt + self.cur = bool(cur) + self.deg = bool(deg) + self.img = bool(img) + self.nbh = nbh + self.pdl = bool(pdl) + # derived compile-time constants (kernel.cu L1292-1294, L1359, L1389, L1485) + self.S = vpt * 4 + self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) + self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 + self.use_img = img and vpt == 1 + self.brl = (minb * blk <= 1024) or (vpt == 1) + + # ------------------------------------------------------------------ + @cute.kernel + def kern( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + cmp_: cutlass.Int32, + qc: cutlass.Int32, + smem_bytes: cutlass.Int32, + ): + BLK = cutlass.const_expr(self.blk) + VPT = cutlass.const_expr(self.vpt) + KPT = cutlass.const_expr(self.kpt) + NBH = cutlass.const_expr(self.nbh) # noqa: F841 + S = cutlass.const_expr(self.S) + LNBH = cutlass.const_expr(self.lnbh) + NW = cutlass.const_expr(self.blk // 32) + + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_wait() # L1291 (knob default off) + + tid, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + lane = tid & cutlass.Int32(31) + + # ------------------------------------------------------------------ + # Predeclarations: the DSL AST transformer requires every scalar that + # is (re)assigned under a dynamic if/while region to pre-exist with a + # stable type at every enclosing region level. Constant inits are + # sunk/dead-coded by LLVM, so this costs no registers. + # ------------------------------------------------------------------ + i = cutlass.Int32(0) + j = cutlass.Int32(0) + r = cutlass.Int32(0) + tinc = cutlass.Int32(0) + cnt = cutlass.Int32(0) + bit = cutlass.Int32(0) + abv = cutlass.Int32(0) + nA = cutlass.Int32(0) + nT = cutlass.Int32(0) + n1 = cutlass.Int32(0) + n2 = cutlass.Int32(0) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + p1e = cutlass.Int32(0) + p2e = cutlass.Int32(0) + lml = cutlass.Int32(0) + aboveC = cutlass.Int32(0) + needC = cutlass.Int32(0) + mm = cutlass.Int32(0) + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + b2w = cutlass.Int32(0) + sh2 = cutlass.Int32(0) + it = cutlass.Int32(0) + it2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + q1e = cutlass.Int32(0) + q2e = cutlass.Int32(0) + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + b_lv = cutlass.Int32(0) + mc = cutlass.Int32(0) + quad = cutlass.Int32(0) + lim1 = cutlass.Int32(0) + p = cutlass.Int32(0) + q2i = cutlass.Int32(0) + idx = cutlass.Int32(0) + m1 = cutlass.Int32(0) + m2 = cutlass.Int32(0) + t1 = cutlass.Int32(0) + t2 = cutlass.Int32(0) + c1 = cutlass.Int32(0) + c2 = cutlass.Int32(0) + s1 = cutlass.Int32(0) + s2 = cutlass.Int32(0) + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + wm = cutlass.Int32(0) + sdyn = cutlass.Int32(0) + nbw = cutlass.Int32(0) + uq = cutlass.Uint32(0) + vq = cutlass.Uint32(0) + kt = cutlass.Uint32(0) + klo = cutlass.Uint32(0) + kv = cutlass.Uint32(0) + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0) + d2 = cutlass.Uint32(0) + unar = cutlass.Uint32(0) + bnn = cutlass.Uint32(0) + nlo = cutlass.Uint32(0) + uke = cutlass.Uint32(0) + uk = cutlass.Uint32(0) + bn = cutlass.Uint32(0) + w = cutlass.Uint32(0) + wt = cutlass.Uint32(0) + ethr = cutlass.Int64(0) + u64 = cutlass.Int64(0) + LOQ = cutlass.Float32(0.0) + HIf = cutlass.Float32(0.0) + LOf = cutlass.Float32(0.0) + qt2 = cutlass.Float32(0.0) + qt3 = cutlass.Float32(0.0) + + npad = cutlass.Int32(logits.shape[1]) # noqa: F841 + k = cutlass.Int32(pre_idx.shape[1]) + out_row = out[row, None] + x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base + p_addr = pre_idx[row, None].iterator.toint() + + # ---- shared-memory window (map in module docstring) ---- + sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) + sbase = sptr.toint() # Int32 shared addr + + s_res = _smem_view(cutlass.Int32, sbase, 0, 6) + s_cnt = _smem_view(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_oc + s_kmm = _smem_view(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax + s_e12 = _smem_view(cutlass.Int32, sbase, 10, 2) # [0]=s_e1 [1]=s_e2 + s_ws = _smem_view(cutlass.Int32, sbase, 16, 32) + s_wmn = _smem_view(cutlass.Uint32, sbase, 48, 32) + s_wmx = _smem_view(cutlass.Uint32, sbase, 80, 32) + s_hist = _smem_view(cutlass.Int32, sbase, STATIC_WORDS, self.nbh) + ck_base = sbase + cutlass.Int32((STATIC_WORDS + self.nbh) * 4) + ck = cute.make_tensor( + cute.make_ptr(cutlass.Uint32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) # typed view, no bound + ci = cute.make_tensor( + cute.make_ptr( + cutlass.Int32, + ck_base + cmp_ * cutlass.Int32(4), + cute.AddressSpace.smem, + assumed_align=4, + ), + cute.make_layout((65536,)), + ) + img_f = cute.make_tensor( # aliases ck/ci (L1299) + cute.make_ptr(cutlass.Float32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) + bm = cute.make_tensor( # aliases ck (L1409) + cute.make_ptr(cutlass.Int32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) + + n4 = n >> cutlass.Int32(2) + ntail = n - (n4 << cutlass.Int32(2)) + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L1351 + + # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any dependent + # gather (L1314-1322); compiled out under DEG. + pvs = [] + if cutlass.const_expr(not self.deg): + for t in cutlass.range_constexpr(KPT): + pv = cutlass.Int32(-1) + j = tid + cutlass.Int32(t * self.blk) + if j < k: + pv = ld_g_i32(p_addr, j) + pvs.append(pv) + + # ---- row load: exact-fit peel + float4[VPT] register batch (L1327-1350) + atom128 = g2r_atom_f32(128, invariant=True) + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + if n4 >= cutlass.Int32(self.blk * self.vpt): # block-uniform peel + for u in cutlass.range_constexpr(VPT): + ld_g_f32x4(atom128, x_addr, tid + cutlass.Int32(u * self.blk), frags[u]) + else: # predicated flat batch + for u in cutlass.range_constexpr(VPT): + i = tid + cutlass.Int32(u * self.blk) + if i < n4: + ld_g_f32x4(atom128, x_addr, i, frags[u]) + for u in cutlass.range_constexpr(VPT): + i = tid + cutlass.Int32(u * self.blk) + if i >= n4: # -INFINITY fill L1346 + for q in cutlass.range_constexpr(4): + frags[u][q] = cutlass.Float32(_NEG_INF__reg) + + tval = cutlass.Float32(_NEG_INF__reg) + if tid < ntail: + tval = ldg_f32(x_addr, tix) # L1352 + + # ---- init (L1391-1392) + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + s_cnt[1] = cutlass.Int32(0) + for z in cutlass.range_constexpr(self.nbh // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + + # ---- bracket: 4 mutually exclusive compile-time arms (L1393-1451) + lmin = cutlass.Uint32(0xFFFFFFFF) + lmax = cutlass.Uint32(0) + if cutlass.const_expr(self.use_img): + fatom = _f32_smem_atom() + for u in cutlass.range_constexpr(VPT): # VPT == 1 here + i = tid + cutlass.Int32(u * self.blk) + if i < n4: + _sts128_f32(fatom, frags[u], ck_base, i * cutlass.Int32(16)) + if tid < ntail: + img_f[tix] = tval + cute.arch.barrier() # L1400 + for t in cutlass.range_constexpr(KPT): + p = pvs[t] + if cutlass.Uint32(p) < cutlass.Uint32(n): + uk = fkey(img_f[p]) + if uk < lmin: + lmin = uk + if uk > lmax: + lmax = uk + cute.arch.barrier() # L1406 (img dies) + elif cutlass.const_expr(self.use_bm): + nbw = (n + cutlass.Int32(31)) >> cutlass.Int32(5) + i = tid + while i < nbw: # bitmap clear (L1410) + bm[i] = cutlass.Int32(0) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1411 + for t in cutlass.range_constexpr(KPT): + p = pvs[t] + if cutlass.Uint32(p) < cutlass.Uint32(n): + atomic_or_cta( + bm.iterator + (p >> cutlass.Int32(5)), + cutlass.Int32(1) << (p & cutlass.Int32(31)), + ) + cute.arch.barrier() # L1417 + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for u in cutlass.range_constexpr(VPT): + base = (tid + cutlass.Int32(u * self.blk)) << cutlass.Int32(2) + w = cutlass.Uint32(0) + if cutlass.Uint32(base) < cutlass.Uint32(n): + w = cutlass.Uint32(bm[base >> cutlass.Int32(5)]) >> cutlass.Uint32( + base & cutlass.Int32(31) + ) + for cbit in cutlass.range_constexpr(4): + if (w & cutlass.Uint32(1 << cbit)) != cutlass.Uint32(0): + lmn = fmin_f32(lmn, _val(frags, 4 * u + cbit)) + lmx = fmax_f32(lmx, _val(frags, 4 * u + cbit)) + if tid < ntail: + wt = cutlass.Uint32(bm[tix >> cutlass.Int32(5)]) >> cutlass.Uint32( + tix & cutlass.Int32(31) + ) + if (wt & cutlass.Uint32(1)) != cutlass.Uint32(0): + lmn = fmin_f32(lmn, tval) + lmx = fmax_f32(lmx, tval) + lmin = fkey(lmn) + lmax = fkey(lmx) # monotone (L1428) + cute.arch.barrier() # L1429 (bm dies) + elif cutlass.const_expr(self.deg): + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for s in cutlass.range_constexpr(S): # L1436-1439 + v = _val(frags, s) + if v > cutlass.Float32(_NEG_INF__reg): + lmn = fmin_f32(lmn, v) + lmx = fmax_f32(lmx, v) + if tid < ntail: + lmn = fmin_f32(lmn, tval) + lmx = fmax_f32(lmx, tval) + lmin = fkey(lmn) + lmax = fkey(lmx) + else: + # default: KPT scattered fkey ldg gathers, batch-then-fold (L1443-1450) + xs = [] + for t in cutlass.range_constexpr(KPT): + xv = cutlass.Float32(0.0) + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): + xv = ldg_f32(x_addr, pvs[t]) + xs.append(xv) + for t in cutlass.range_constexpr(KPT): + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): + uk = fkey(xs[t]) + if uk < lmin: + lmin = uk + if uk > lmax: + lmax = uk + + # ---- block min/max in ONE barrier (L1452-1462); publishes hist clear + lmin = warp_min_u32(lmin) + lmax = warp_max_u32(lmax) + if lane == cutlass.Int32(0): + s_wmn[tid >> cutlass.Int32(5)] = lmin + s_wmx[tid >> cutlass.Int32(5)] = lmax + cute.arch.barrier() # L1456 + a = cutlass.Uint32(0xFFFFFFFF) + c = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + a = cutlass.Uint32(s_wmn[lane]) + c = cutlass.Uint32(s_wmx[lane]) + lmin = warp_min_u32(a) + lmax = warp_max_u32(c) + Tv = invkey(lmin) + GMAX = invkey(lmax) + + # ---- collapse guard, NaN-safe (L1464-1471) + okc = cutlass.Int32(0) + if Tv < GMAX: + if (GMAX - Tv) > cutlass.Float32(1e-30): + okc = cutlass.Int32(1) + if okc == cutlass.Int32(0): + Tv = cutlass.Float32(SENT_LO) + GMAX = cutlass.Float32(SENT_HI) + + # ---- bin transform constants (L1485-1511) + BRL = cutlass.const_expr(self.brl) # noqa: F841 + OFFf = cutlass.Float32(1.0 if self.brl else 0.0) + recip = 1.0 / float(self.nbh - (2 if self.brl else 0)) + WD = (GMAX - Tv) * cutlass.Float32(recip) + wsel = cutlass.Float32(1e-30) + if WD > cutlass.Float32(0.0): + wsel = WD + # rcp.approx (single MUFU.RCP) — the CUDA arm's exact lowering of + # `1.0f / wsel`; the previous `1.0 / wsel` spelling emitted the IEEE + # div.rn Newton triple + slowpath CALL on the barrier-bounded chain + # feeding all S classify FMULs. Output exactness is SC-invariant + # (any SC > 0 preserves the sign/monotonicity invariants, L1485-1511) + # and the WD > 0 arm is now bit-identical to CUDA's MUFU.RCP. + SC = cute.arch.rcp_approx(wsel) + QCAPf = cutlass.Float32(float(self.nbh - 1)) + CQ0 = OFFf - Tv * SC + CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf(CQ0) + cutlass.Float32(1.0)) + + # ---- histogram (L1513-1526) + if cutlass.const_expr(self.brl): + # fix-2 P4 (GATED, removable as one hunk): A1 ported to the BRL + # classify arm — hist base pinned ONCE via the same + # _smem_addr_reg__reg identity-mov used in the !BRL arm below, and the + # result-discarded classify atomics spelled as resultless + # red.shared (_red_shared_add1__reg). Value-identical: same +1 to the + # same byte address (hb + 4*bn == &s_hist[bn]), same .relaxed.cta + # ordering; the q/bn computations are untouched so classify/emit + # bit-identity (BRL requirement) is preserved. Emit-path hist + # atomics (results used, L1630+) are NOT touched. + hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) + for s in cutlass.range_constexpr(S): + q = _fmaf__reg(_val(frags, s), SC, CQ) + bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) + _red_shared_add1__reg(hb + (cutlass.Int32(bn) << cutlass.Int32(2))) + qt = _fmaf__reg(tval, SC, CQ) # unconditional (L1517) + bnt = _umin_u32(f2u_rz(qt), cutlass.Uint32(self.nbh - 1)) + _red_shared_add1__reg(hb + (cutlass.Int32(bnt) << cutlass.Int32(2))) + else: + # hist base pinned ONCE (byte addr, +STATIC_BYTES = word 128 map); + # each site below is then LEA + ATOMS exactly like the CUDA arm + # instead of re-deriving the shared window per divergent block. + hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) + for s in cutlass.range_constexpr(S): + q = _submul_asm(_val(frags, s), Tv, SC) # anti-CSE classify + if q >= cutlass.Float32(0.0): + _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(q, QCAPf)) << cutlass.Int32(2))) + qt = _submul_asm(tval, Tv, SC) + if qt >= cutlass.Float32(0.0): + _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(qt, QCAPf)) << cutlass.Int32(2))) + cute.arch.barrier() # L1527 + + # ---- crossing-bin find (L1528-1538) + if cutlass.const_expr(self.cur or self.nbh > 1024): + scan_cross_w(s_hist, s_ws, k, tid, s_res, blk=self.blk, nb=self.nbh) + else: + find_cross(s_hist, k, tid, s_res, nb=self.nbh) + cute.arch.barrier() # L1535 + above = s_res[RES_ABOVE] + m = s_res[RES_M] + Bv = s_res[RES_B] + need = k - above + whole = cutlass.Int32(0) + if need >= m: + whole = cutlass.Int32(1) + + # ---- prod-fix ESCAPE (L1540-1617): 32-step key-space bisection + esc = cutlass.Int32(0) + if whole == cutlass.Int32(0): + if m > cmp_: + esc = cutlass.Int32(1) + if esc == cutlass.Int32(1): + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + s_cnt[1] = cutlass.Int32(0) + # DEVIATION (race fix, see notes): the CUDA zeroes s_o1/s_oc + # again between the nA read (L1581) and the emit (L1584) with + # only ONE barrier pair around both — a read/write race that + # nvcc's schedule happens to win and ptxas' does not (observed + # ~8% row corruption under CTA co-residency). We instead emit + # through the path-exclusive s_e1/s_e2 slots, zeroed HERE under + # the existing L1550 barrier; the racy mid-emit rezero is + # dropped. Barrier count unchanged. + s_e12[0] = cutlass.Int32(0) + s_e12[1] = cutlass.Int32(0) + cute.arch.barrier() # L1550 + klo = cutlass.Uint32(0) + bit = cutlass.Int32(31) + while bit >= cutlass.Int32(0): + kt = klo | (cutlass.Uint32(1) << cutlass.Uint32(bit)) + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(S): + ix = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if ix < n: + if fkey(_val(frags, s)) >= kt: + cnt = cnt + cutlass.Int32(1) + if tid < ntail: + if fkey(tval) >= kt: + cnt = cnt + cutlass.Int32(1) + cnt = cutlass.Int32(warp_add_i32(cnt)) + if lane == cutlass.Int32(0): + if cnt != cutlass.Int32(0): + atomic_add_cta(s_cnt.iterator, cnt) + cute.arch.barrier() # L1563 + if s_cnt[0] >= k: + klo = kt + cute.arch.barrier() # L1565 + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + cute.arch.barrier() # L1567 + bit = bit - cutlass.Int32(1) + ethr = cutlass.Int64(klo) # k-th largest key + abv = cutlass.Int32(0) + for s in cutlass.range_constexpr(S): + ix = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if ix < n: + if cutlass.Int64(fkey(_val(frags, s))) > ethr: + abv = abv + cutlass.Int32(1) + if tid < ntail: + if cutlass.Int64(fkey(tval)) > ethr: + abv = abv + cutlass.Int32(1) + abv = cutlass.Int32(warp_add_i32(abv)) + if lane == cutlass.Int32(0): + if abv != cutlass.Int32(0): + atomic_add_cta(s_cnt.iterator + 1, abv) + cute.arch.barrier() # L1580 + nA = s_cnt[1] + nT = k - nA + # (rezero dropped — emit counters live in s_e12, see race-fix note) + cute.arch.barrier() # L1583 + lml = cutlass.Int32(cute.arch.lanemask_lt()) + for s in cutlass.range_constexpr(S): + ixv = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + u64 = cutlass.Int64(-1) + if ixv < n: + u64 = cutlass.Int64(fkey(_val(frags, s))) + q1e = cutlass.Int32(0) + q2e = cutlass.Int32(0) + if u64 > ethr: + q1e = cutlass.Int32(1) + if u64 == ethr: + q2e = cutlass.Int32(1) + n1 = ballot(q1e == cutlass.Int32(1)) + n2 = ballot(q2e == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_e12.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1e == cutlass.Int32(1): + if p1e < nA: + out_row[p1e] = ixv + if q2e == cutlass.Int32(1): + if p2e < nT: + out_row[nA + p2e] = ixv + # tail element (L1601-1615) + u64 = cutlass.Int64(-1) + if tid < ntail: + u64 = cutlass.Int64(fkey(tval)) + q1e = cutlass.Int32(0) + q2e = cutlass.Int32(0) + if u64 > ethr: + q1e = cutlass.Int32(1) + if u64 == ethr: + q2e = cutlass.Int32(1) + n1 = ballot(q1e == cutlass.Int32(1)) + n2 = ballot(q2e == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_e12.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1e == cutlass.Int32(1): + if p1e < nA: + out_row[p1e] = tix + if q2e == cutlass.Int32(1): + if p2e < nT: + out_row[nA + p2e] = tix + # (CUDA returns here — everything below is the else-arm) + else: + # ---- emit (L1619-1701) + if cutlass.const_expr(self.cur): + LOQ = cutlass.Float32(Bv) # int->float cvt (L1626) + lim1 = above + if whole == cutlass.Int32(1): + lim1 = above + m + for s in cutlass.range_constexpr(S): + if cutlass.const_expr(self.brl): + q = _fmaf__reg(_val(frags, s), SC, CQ) # bit-identical to classify + else: + q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) # L1630 emit spelling + idx = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + p = cutlass.Int32(0) + if q >= LOQ: + bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) + p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + if p < lim1: + out_row[p] = idx + else: + if whole == cutlass.Int32(0): + q2i = p - above + if q2i < cmp_: # escape-made-safe guard + ck[q2i] = fkey(_val(frags, s)) + ci[q2i] = idx + # tail (L1640-1647) + if cutlass.const_expr(self.brl): + qt2 = _fmaf__reg(tval, SC, CQ) + else: + qt2 = _fmaf__reg(tval - Tv, SC, OFFf) + p = cutlass.Int32(0) + if qt2 >= LOQ: + bn = _umin_u32(f2u_rz(qt2), cutlass.Uint32(self.nbh - 1)) + p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + if p < lim1: + out_row[p] = tix + else: + if whole == cutlass.Int32(0): + q2i = p - above + if q2i < cmp_: + ck[q2i] = fkey(tval) + ci[q2i] = tix + else: + # two-mask ballot emit (L1650-1701) + HIf = cutlass.Float32(_POS_INF) + LOf = cutlass.Float32(_POS_INF) + if whole == cutlass.Int32(1): + HIf = cutlass.Float32(Bv) + else: + if Bv < cutlass.Int32(self.nbh - 1): + HIf = cutlass.Float32(Bv + cutlass.Int32(1)) + LOf = cutlass.Float32(Bv) + m1 = cutlass.Int32(0) + m2 = cutlass.Int32(0) + for s in cutlass.range_constexpr(S): + if cutlass.const_expr(self.brl): + q = _fmaf__reg(_val(frags, s), SC, CQ) + else: + q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) + if q >= HIf: + m1 = m1 | cutlass.Int32(1 << s) + else: + if q >= LOf: + m2 = m2 | cutlass.Int32(1 << s) + if cutlass.const_expr(self.brl): + qt3 = _fmaf__reg(tval, SC, CQ) + else: + qt3 = _fmaf__reg(tval - Tv, SC, OFFf) + t1 = cutlass.Int32(0) + t2 = cutlass.Int32(0) + if qt3 >= HIf: + t1 = cutlass.Int32(1) + else: + if qt3 >= LOf: + t2 = cutlass.Int32(1) + c1 = popc(m1) + t1 + c2 = popc(m2) + t2 + s1, s2 = warp_incl_scan_add2(c1, c2, lane) # L1669-1673 + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(31): + b1 = atomic_add_cta(s_cnt.iterator, s1) + b2 = atomic_add_cta(s_cnt.iterator + 1, s2) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(31)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(31)) + p1 = b1 + (s1 - c1) + p2 = b2 + (s2 - c2) + lim1 = above + if whole == cutlass.Int32(1): + lim1 = k + wm = m1 # sparse set-bit walk + while wm != cutlass.Int32(0): + sdyn = ffs_m1(wm) + idx = ( + (tid + (sdyn >> cutlass.Int32(2)) * cutlass.Int32(self.blk)) + << cutlass.Int32(2) + ) + (sdyn & cutlass.Int32(3)) + if p1 < lim1: + out_row[p1] = idx + p1 = p1 + cutlass.Int32(1) + wm = wm & (wm - cutlass.Int32(1)) + if t1 == cutlass.Int32(1): + if p1 < lim1: + out_row[p1] = tix + p1 = p1 + cutlass.Int32(1) + if m2 != cutlass.Int32(0): # static-unrolled (L1689) + for s in cutlass.range_constexpr(S): + if (m2 & cutlass.Int32(1 << s)) != cutlass.Int32(0): + idx = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if p2 < cmp_: + ck[p2] = fkey(_val(frags, s)) + ci[p2] = idx + p2 = p2 + cutlass.Int32(1) + if t2 == cutlass.Int32(1): + if p2 < cmp_: + ck[p2] = fkey(tval) + ci[p2] = tix + p2 = p2 + cutlass.Int32(1) + + # ---- refine (skipped when whole — CUDA returned inside emit) + if whole == cutlass.Int32(0): + cute.arch.barrier() # L1703 + if cutlass.const_expr(self.cur): + mc = m + if mc > cmp_: + mc = cmp_ + else: + mc = s_cnt[1] + if mc > cmp_: + mc = cmp_ + quad = cutlass.Int32(0) + if mc >= m: + if mc <= qc: + quad = cutlass.Int32(1) + if quad == cutlass.Int32(1): + # O(mc^2) index-tie-broken rank (L1706-1718) + i = tid + while i < mc: + uq = cutlass.Uint32(ck[i]) + r = cutlass.Int32(0) + j = cutlass.Int32(0) + while j < mc: + vq = cutlass.Uint32(ck[j]) + tinc = cutlass.Int32(0) + if vq > uq: + tinc = cutlass.Int32(1) + if vq == uq: + if j < i: + tinc = cutlass.Int32(1) + r = r + tinc + j = j + cutlass.Int32(1) + if r < need: + out_row[above + r] = ci[i] + i = i + cutlass.Int32(BLK) + else: + # ---- fallback: exact key-space narrowing (L1720-1775) + if tid == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + cute.arch.barrier() # L1724 + i = tid + while i < mc: + kv = cutlass.Uint32(ck[i]) + atomic_min_cta(s_kmm.iterator, kv) + atomic_max_cta(s_kmm.iterator + 1, kv) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1726 + rlo = cutlass.Uint32(s_kmm[0]) + rhi = cutlass.Uint32(s_kmm[1]) + ethr = cutlass.Int64(rlo) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + while done == cutlass.Int32(0): + if needC == mm: # L1730 + ethr = cutlass.Int64(rlo) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + if rlo >= rhi: # L1731 + ethr = cutlass.Int64(rlo) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(6): # L1732 + ethr = cutlass.Int64(rlo) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + d2 = rhi - rlo + b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + sh2 = cutlass.Int32(0) + if b2w > cutlass.Int32(LNBH): + sh2 = b2w - cutlass.Int32(LNBH) + for z in cutlass.range_constexpr(self.nbh // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + cute.arch.barrier() # L1737 + i = tid + while i < mc: + unar = cutlass.Uint32(ck[i]) + if unar >= rlo: + if unar <= rhi: + bnn = (unar - rlo) >> cutlass.Uint32(sh2) + bnn = _umin_u32(bnn, cutlass.Uint32(self.nbh - 1)) + atomic_add_cta( + s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1743 + if cutlass.const_expr(self.nbh > 1024): + scan_cross_w( + s_hist, s_ws, needC, tid, s_res, blk=self.blk, nb=self.nbh + ) + else: + find_cross(s_hist, needC, tid, s_res, nb=self.nbh) + cute.arch.barrier() # L1746 + aboveC = aboveC + s_res[RES_ABOVE] + needC = needC - s_res[RES_ABOVE] + mm = s_res[RES_M] + b_lv = s_res[RES_B] + nlo = rlo + (cutlass.Uint32(b_lv) << cutlass.Uint32(sh2)) + if b_lv != cutlass.Int32(self.nbh - 1): + rhi = nlo + ( + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + # final two-predicate ballot emit (L1752-1775) + if tid == cutlass.Int32(0): + s_e12[0] = cutlass.Int32(0) + s_e12[1] = cutlass.Int32(0) + cute.arch.barrier() # L1753 + lml = cutlass.Int32(cute.arch.lanemask_lt()) + it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it = cutlass.Int32(0) + while it < it2: + i = it * cutlass.Int32(BLK) + tid + uke = cutlass.Uint32(0) + idv = cutlass.Int32(0) + if i < mc: + uke = cutlass.Uint32(ck[i]) + idv = ci[i] + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + if i < mc: + if cutlass.Int64(uke) > ethr: + q1f = cutlass.Int32(1) + if cutlass.Int64(uke) == ethr: + q2f = cutlass.Int32(1) + n1 = ballot(q1f == cutlass.Int32(1)) + n2 = ballot(q2f == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_e12.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1f == cutlass.Int32(1): + if p1e < aboveC: + out_row[above + p1e] = idv + if q2f == cutlass.Int32(1): + if p2e < needC: + out_row[above + aboveC + p2e] = idv + it = it + cutlass.Int32(1) + + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + cmp_: cutlass.Int32, + qc: cutlass.Int32, + smem_bytes: cutlass.Int32, + stream, + ): + b = logits.shape[0] + self.kern(logits, pre_idx, out, n, cmp_, qc, smem_bytes).launch( + grid=(b, 1, 1), + block=(self.blk, 1, 1), + stream=stream, + smem=smem_bytes, + min_blocks_per_mp=self.minb, + use_pdl=self.pdl, + ) + + +# --------------------------------------------------------------------------- +# host wrapper: compile cache + route()-driven entry +# --------------------------------------------------------------------------- +_COMPILE_CACHE__reg: dict = {} + + +def get_compiled__reg(tpl, dump_dir=None, pdl=False): + """Compile (or fetch) the variant for constexpr tuple + (BLK, VPT, MINB, KPT, CUR, DEG, IMG, NBH).""" + key = (tuple(tpl), bool(pdl)) + compiled = _COMPILE_CACHE__reg.get(key) + if compiled is None: + from cutlass.cute import runtime as _crt + + blk, vpt, minb, kpt, cur, deg, img, nbh = tpl + kernel = GvrTopkRegKernel(blk, vpt, minb, kpt, cur, deg, img, nbh, pdl=pdl) + nb_, nc_ = cute.sym_int(), cute.sym_int() + nb2_, nc2_ = cute.sym_int(), cute.sym_int() + nb3_, nc3_ = cute.sym_int(), cute.sym_int() + lg_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (nb_, nc_), stride_order=(1, 0), assumed_align=16 + ) + pi_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (nb2_, nc2_), stride_order=(1, 0), assumed_align=16 + ) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (nb3_, nc3_), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + opts = "--enable-tvm-ffi" + if dump_dir: + opts += f" --keep-ptx --keep-cubin --dump-dir {dump_dir}" + with _no_carveout(): + compiled = cute.compile( + kernel, + lg_fake, + pi_fake, + out_fake, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + stream=fake_stream, + options=opts, + ) + _COMPILE_CACHE__reg[key] = compiled + return compiled + + +def reg_topk(logits, pre_idx, n, out, rd=None): + """torch-facing entry for the register family. + + logits [b, npad] f32, pre_idx [b, k] i32, out [b, >=k] i32, n = valid len. + rd: optional pre-computed ct_dispatch.route() dict (must be reg/regimg). + """ + if rd is None: + try: + from .gvr_topk_decode_self_sampling_host import route + except ImportError: + from gvr_topk_decode_self_sampling_host import route + rd = route(logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1]) + assert rd["kernel"] in ("reg", "regimg"), rd["kernel"] + tpl = rd["tpl"] + rt = rd["rt"] + assert rt["IMGOFF"] == tpl[7], (rt["IMGOFF"], tpl[7]) # IMGOFF == NBH + compiled = get_compiled__reg(tpl) + smem = STATIC_BYTES + rd["smem"] + compiled(logits, pre_idx, out, int(n), rt["CMP"], rt["QC"], smem) + return out + + +# =========================================================================== +# ==== ct_clus.py — family: clus =============================== +# =========================================================================== +"""ct_clus.py — op46 gvr_clus (clustered streaming GVR) CuTeDSL port. + +Ground truth: src_cuda/kernel.cu L1793-2356 (frozen); phase contract, smem +map and barrier inventory per TRANSLATION_SPEC.md §5.3; DSL spellings pinned +by probes/PROBE_RESULTS.md (BINDING) + op43 lessons L1..L5; per-CTA stream +mirrors gvr_main (idioms reused from the proven src/ct_main.py). + +Ctor knobs (compile-time, mirror of the CUDA template params, spec §4c): + BLK = 1024, U ∈ {1,2,4,8}, MINB = 1, NBS = 256, CS ∈ {2,4,8} + (+ scap/cmp smem-extent knobs: every reachable route has 8192/2048 — + SCAP/CMP stay LIVE runtime args for all value logic, ABI parity). +Derived: HB=NBS, STEPC=BLK*U, PFD=min(U,4) (kernel.cu L1905). + +Signature (ABI parity with kernel.cu L1797-1799; Q is dead in-kernel): + run__clus(logits[b,npad] f32, pre_idx[b,k] i32, out[b,k] i32) via + kern(..., n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2) +Grid dim3(CS, b) native 2-D + cluster (CS,1,1) (probe P11); block 1024; +min_blocks_per_mp=1 (64-reg wall, probes P2/P15); smem one SmemAllocator +blob mirroring the CUDA dynamic map hist|cbuf|ck64c|mrg (L1802-1814), +dyn-equivalent bytes == host smc formula L3130 (asserted in run__clus()). + +int2 staging convention (same as ct_main): int2(value bits, index) is ONE +little-endian Uint64 = (idx << 32) | value_bits — single u64 smem ld/st. + +Barrier / cluster-op inventory implemented (kernel.cu line cites, op43 L5): + L1902 (sample redux publish), L1948 (sample hist), L1956 (scan publish), + [degenerate sample: 2 inside gather_hint], + retry preamble: clus.sync L2016 + __syncthreads L2024, + clus.sync L2135 (merge), __syncthreads L2143 (merge publish), + [ladder gather: 2 inside gather_hint], + clus.sync L2226 (EXIT RENDEZVOUS — the only one; rank!=0 falls through), + narrowing: L2255, L2257, per-level L2270/L2273, L2281, + degen: per-level L2321/L2325/L2327 (+1 INSIDE scan_cross), L2337. + NO loop-tail ladder barriers (gvr_clus has none — unlike gvr_main). + All clus.sync = releasing aligned arrive+wait (risk R3, never relaxed). +Cluster ops: merge = _merge_scan0_local, a LOCAL patched copy of the frozen + ct_common.merge_scan0 that rematerializes mapa per (q, r) like the CUDA + L135-137 (register-pressure fix, see notes; DSMEM v4 spelling = probe + P5b via ct_common ops); ONE packed u64 st.shared::cluster candidate push + to rank-0 ck64c (L2185/L2205/L2219 — never split 4B, op43 L3); mapa of + ck64c to rank 0. PTX inventory audited: 3 arrive+3 wait (aligned), 3 + st.shared::cluster.u64 sites, 2*CS ld.shared::cluster.v4 sites, 4 + prefetch.global.L2 sites on U=8 only, zero griddepcontrol. +Every rung/ladder decision is cluster-uniform by construction (identical +sample locations on every rank; merged tot; block-uniform gather) — the +conditional retry clus.sync at L2016 cannot deadlock (spec §5.3). +""" + + +QUADC_CLUS__clus = C.QUADC_CLUS + +_NEG_INF__clus = float("-inf") + + +# --------------------------------------------------------------------------- +# single-rounding fma.rn.f32 (probe P6 emit spelling; local — ct_common has +# no fma export). Sites: T (L1957), Tk/T3 (L1966/1977/1992), HIC (L1981). +# (x-TF)*SC classify shapes stay plain sub+mul (P6: uncontractible). +# --------------------------------------------------------------------------- +@dsl_user_op +def _fmaf__clus(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + mlir_math.fma( + a.ir_value(loc=loc, ip=ip), + b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), + fastmath=mlir_arith.FastMathFlags.none, + loc=loc, + ip=ip, + ) + ) + + +class GvrClusKernel: + """CuTeDSL port of gvr_clus (kernel.cu L1795).""" + + def __init__( + self, blk: int, u: int, minb: int, nbs: int, cs: int, scap: int = 8192, cmp_: int = 2048 + ): + assert blk == 1024, "gvr_clus is always BLK=1024 (dispatch L3132)" + assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1) (L1796)" + assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177)" + assert u in (1, 2, 4, 8) and cs in (2, 4, 8) + self.blk = blk + self.u = u + self.minb = minb + self.nbs = nbs + self.cs = cs + self.scap = scap # smem extents only — + self.cmp = cmp_ # value logic uses rt args + self.hb = nbs # L1801 + self.stepc = blk * u # L1836 + self.pfd = u if u < 4 else 4 # L1905 PFD=min(U,4) + self.lb = nbs.bit_length() - 1 # log2(NBS)=8 + # dynamic-region byte map (L1802-1814): hist | cbuf(int2) | ck64c | mrg + self.cbuf_bytes = (scap + 4) * 8 + assert self.cbuf_bytes % 16 == 0 + self.ck_off = self.cbuf_bytes # inside the blob + self.dyn_bytes = nbs * 4 + self.cbuf_bytes + cmp_ * 8 + nbs * 4 + # == host smc = SNB*8 + (SCAP+4)*8 + CMP*8 (L3130) + + # ------------------------------------------------------------------ + # GVR_EMITK (kernel.cu L2096-2104): classify+stage one survivor. + # bn via UNSIGNED saturating convert (f2u_rz, P4); staging store is ONE + # u64; branchless trash slot min(pos, SCAP) (runtime SCAP). Returns pos+1. + # ------------------------------------------------------------------ + @cute.jit + def _emitk(self, xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2): + NBS = self.nbs + bn_u = C.f2u_rz((xv - TF) * SC) + if bn_u > cutlass.Uint32(NBS - 1): + bn_u = cutlass.Uint32(NBS - 1) + bn = cutlass.Int32(bn_u) + C.atomic_add_cta(s_hist.iterator + bn, cutlass.Int32(1)) + ps = pos + if ps > SCAP: + ps = SCAP # trash slot (IMNMX) + s_cbuf2[ps] = (cutlass.Uint64(cutlass.Uint32(idx)) << cutlass.Uint64(32)) | cutlass.Uint64( + C.u32_of_f32(xv) + ) + return pos + cutlass.Int32(1) + + # ------------------------------------------------------------------ + # P5 emit step (kernel.cu L2178-2186 == L2196-2204 == L2211-2219): + # bn via SIGNED rz convert (__float2int_rz, L2180); bn>=B gate; LOCAL + # mrg atomicAdd whose result is a CLUSTER-GLOBAL position (prefix-biased + # cursors from merge_scan0); overflow -> ONE packed u64 DSMEM store to + # rank-0 ck64c (probe P5b, op43 L3 — never split 4B). + # ------------------------------------------------------------------ + @cute.jit + def _p5_emit(self, xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64): + NBS = self.nbs + bn = C.f2s_rz((xv - TF) * SC) + if bn > cutlass.Int32(NBS - 1): + bn = cutlass.Int32(NBS - 1) + if bn >= B: + p = C.atomic_add_cta(s_mrg.iterator + bn, cutlass.Int32(1)) + if p < lim1: + out_row[p] = idv + else: + if whole == cutlass.Int32(0): + q2 = p - above + if q2 < CMP: + C._st_shared_cluster_u64( + rk64 + q2 * cutlass.Int32(8), + (cutlass.Uint64(C.fkey(xv)) << cutlass.Uint64(32)) + | cutlass.Uint64(cutlass.Uint32(idv)), + ) + + # ------------------------------------------------------------------ + # LOCAL patched copy of ct_common.merge_scan0 (frozen sibling — not + # edited): rematerializes mapa per (q, r) exactly like the CUDA L135-137 + # instead of holding CS mapped base addresses across the whole merge. + # The hoisted-array form costs CS extra long-lived registers; with the + # U>=4 sixteen-register pf prime batch it tips ptxas into spilling the + # batch across the rung phase (reg-audit finding, notes). Semantics, + # the DSMEM v4 load spelling (probe P5b), the register accumulation and + # the prefix-biased STS.128 cursor write are IDENTICAL to ct_common. + # NO barrier inside (caller pays L2143). + # ------------------------------------------------------------------ + @cute.jit + def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): + NBS = self.nbs + CS = self.cs + BPT = NBS // 32 + NV = BPT // 4 + if tidx < cutlass.Int32(32): + lane = tidx + atom = C.smem_atom_i32_128() + hbase = s_hist.iterator.toint() + # pass 1 (L131-143): remote v4 accumulation of tot/pre per vector + tot_r = [] + pre_r = [] + sm = cutlass.Int32(0) + for q in cutlass.range_constexpr(NV): + boff = (lane * cutlass.Int32(BPT) + cutlass.Int32(4 * q)) * cutlass.Int32(4) + t = [cutlass.Int32(0)] * 4 + p = [cutlass.Int32(0)] * 4 + for r in cutlass.range_constexpr(CS): + mapped = C._mapa_shared_cluster_addr( + hbase + boff, cutlass.Int32(r) + ) # per-use mapa (L136) + v0, v1, v2, v3 = C._ld_shared_cluster_v4_u32(mapped) + t[0] = t[0] + v0 + t[1] = t[1] + v1 + t[2] = t[2] + v2 + t[3] = t[3] + v3 + if cutlass.Int32(r) < rank: # L140 predicated adds + p[0] = p[0] + v0 + p[1] = p[1] + v1 + p[2] = p[2] + v2 + p[3] = p[3] + v3 + tot_r.append(t) + pre_r.append(p) + sm = sm + t[0] + t[1] + t[2] + t[3] + # inclusive scan + totals (L144-148) + w = C.warp_incl_scan_add(sm, lane) + tt = cute.arch.shuffle_sync(w, cutlass.Int32(31)) + after = tt - w + if lane == cutlass.Int32(0): + s_res[C.RES_TOT] = tt + base = lane * cutlass.Int32(BPT) + # descending walk: crossing pin + prefix-biased cursors into mrg + for q in cutlass.range_constexpr(NV - 1, -1, -1): # L151-165 + o4 = cute.make_fragment((4,), cutlass.Int32) + for j in cutlass.range_constexpr(3, -1, -1): + cq = tot_r[q][j] + o4[j] = after + pre_r[q][j] + gb = base + cutlass.Int32(4 * q + j) + cross = cutlass.Int32(0) + if after < target: + if (after + cq) >= target: + cross = cutlass.Int32(1) + if gb == cutlass.Int32(0): + cross = cutlass.Int32(1) + if cross != cutlass.Int32(0): + s_res[C.RES_B] = gb + s_res[C.RES_ABOVE] = after + s_res[C.RES_M] = cq + after = after + cq + boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) + C.sts128_i32(atom, o4, s_mrg.iterator.toint(), boff) + + # ------------------------------------------------------------------ + # two-predicate warp-ballot emit step (narrowing L2279-2301 and degen + # L2338-2354) — same helper as ct_main. s_scal[1]=s_o1, s_scal[2]=s_o2. + # ------------------------------------------------------------------ + @cute.jit + def _ballot_pair_emit(self, p1, p2, idv, base1, cap1, base2, cap2, out_row, s_scal, lane): + n1 = C.ballot(p1 != cutlass.Int32(0)) + n2 = C.ballot(p2 != cutlass.Int32(0)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = C.atomic_add_cta(s_scal.iterator + 1, cutlass.Int32(C.popc(n1))) + if n2 != cutlass.Int32(0): + b2 = C.atomic_add_cta(s_scal.iterator + 2, cutlass.Int32(C.popc(n2))) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + lm = cutlass.Int32(cute.arch.lanemask_lt()) + if p1 != cutlass.Int32(0): + p = b1 + cutlass.Int32(C.popc(n1 & lm)) + if p < cap1: + out_row[base1 + p] = idv + if p2 != cutlass.Int32(0): + p = b2 + cutlass.Int32(C.popc(n2 & lm)) + if p < cap2: + out_row[base2 + p] = idv + + # ------------------------------------------------------------------ + # kernel + # ------------------------------------------------------------------ + @cute.kernel + def kern( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + npad: cutlass.Int32, + k: cutlass.Int32, + SCAP: cutlass.Int32, + CMP: cutlass.Int32, + SMP: cutlass.Int32, + TGT: cutlass.Int32, + Q: cutlass.Int32, + SS2: cutlass.Int32, + TGT2: cutlass.Int32, + ): + BLK = self.blk + U = self.u + NBS = self.nbs + CS = self.cs + PFD = self.pfd + STEPC = self.stepc + NW = BLK // 32 + + tidx, _, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() # (rank, row) L1824-1825 + rank = bx + row = by + lane = tidx & cutlass.Int32(31) + + # ---- shared memory (CUDA dynamic map order L1802-1814, then static allocs) ---- + smem = SmemAllocator() + s_hist = smem.allocate_tensor( # hist[NBS] @ blob start + cutlass.Int32, cute.make_ordered_layout((self.hb,), order=(0,)), byte_alignment=128 + ) + blob = smem.allocate_tensor( # cbuf(int2) | ck64c + cutlass.Int8, + cute.make_ordered_layout((self.cbuf_bytes + self.cmp * 8,), order=(0,)), + byte_alignment=16, + ) + s_mrg = smem.allocate_tensor( # mrg[NBS] (L1814) + cutlass.Int32, cute.make_ordered_layout((self.nbs,), order=(0,)), byte_alignment=16 + ) + s_ws = smem.allocate_tensor( # L1817 (degen scan only) + cutlass.Int32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 + ) # Int32: ct_common + # scan_cross predeclares its second-stage partial as Int32 and reads + # s_ws inside a dynamic if — a Uint32 ws tensor trips the DSL type- + # stability check (frozen sibling; counts < 2^31 so Int32 is exact). + s_wmn = smem.allocate_tensor( + cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 + ) + s_wmx = smem.allocate_tensor( + cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 + ) + s_res = smem.allocate_tensor( # ct_common slot map + cutlass.Int32, cute.make_ordered_layout((8,), order=(0,)), byte_alignment=16 + ) + # scalar block: [0]=s_bufn [1]=s_o1 [2]=s_o2 (L1818) + s_scal = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 + ) + s_tsh = smem.allocate_tensor( # L1820 + cutlass.Float32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 + ) + s_kmm = smem.allocate_tensor( # L1821 [0]=kmin [1]=kmax + cutlass.Uint32, cute.make_ordered_layout((2,), order=(0,)), byte_alignment=8 + ) + sbase = blob.iterator.toint() + s_cbuf2 = cute.make_tensor( # int2 staged as u64 + cute.make_ptr(cutlass.Uint64, sbase, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((self.scap + 4,)), + ) + ck_addr = sbase + cutlass.Int32(self.ck_off) + s_ck64 = cute.make_tensor( + cute.make_ptr(cutlass.Uint64, ck_addr, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((self.cmp,)), + ) + + # ---- row bases (L1830-1833) ---- + row64 = cutlass.Int64(row) + x_addr = logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) + p_addr = pre_idx.iterator.toint() + row64 * cutlass.Int64(k) * cutlass.Int64(4) + out_row = out[row, None] + + # ---- interleaved chunk ownership (L1835-1846) ---- + n4 = n >> cutlass.Int32(2) + nCh = (n4 + cutlass.Int32(STEPC - 1)) // cutlass.Int32(STEPC) + nFullG = n4 // cutlass.Int32(STEPC) + tail0 = n4 << cutlass.Int32(2) + tailn = cutlass.Int32(0) + if rank == cutlass.Int32(0): + tailn = n - tail0 + + if tidx == cutlass.Int32(0): # L1848 + s_res[C.RES_B2] = cutlass.Int32(-1) + s_res[C.RES_B3] = cutlass.Int32(-1) + s_scal[0] = cutlass.Int32(0) # s_bufn + if tidx < cutlass.Int32(self.hb): # L1849 (HB<=BLK) + s_hist[tidx] = cutlass.Int32(0) + + # ============ P1: QUAD sample (hint gather LAZY, L1851-1878) ======== + # one 64B line = 4 float4 per location, TWO threads: tid takes the + # lower pair at p4, tid+SMP the upper pair at p4+2 (L1868-1869). + atom128 = C.g2r_atom_f32(128, invariant=True) + fsa = cute.make_fragment((4,), cutlass.Float32) + fsb = cute.make_fragment((4,), cutlass.Float32) + smp2 = SMP * cutlass.Int32(2) + shas = cutlass.Int32(0) + if tidx < smp2: + shas = cutlass.Int32(1) + if shas != cutlass.Int32(0): + p4 = tidx * SS2 * cutlass.Int32(4) + if tidx >= SMP: + p4 = (tidx - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fsa) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) + + # ============ P2: quantile rung, redundant per CTA (L1883-1997) ===== + smn = cutlass.Float32(float("inf")) + smx = cutlass.Float32(float("-inf")) + if shas != cutlass.Int32(0): # L1928-1932 + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsa[t]) + smx = C.fmax_f32(smx, fsa[t]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsb[t]) + smx = C.fmax_f32(smx, fsb[t]) + fma_ = cute.make_fragment((4,), cutlass.Float32) # mop-up pair bufs + fmb_ = cute.make_fragment((4,), cutlass.Float32) + j = tidx + cutlass.Int32(BLK) # mop-up L1891-1897 + while j < smp2: + p4 = j * SS2 * cutlass.Int32(4) + if j >= SMP: + p4 = (j - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fma_) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fma_[t]) + smx = C.fmax_f32(smx, fma_[t]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fmb_[t]) + smx = C.fmax_f32(smx, fmb_[t]) + j = j + cutlass.Int32(BLK) + a0 = C.warp_min_u32(C.fkey(smn)) # L1898-1901 + c0m = C.warp_max_u32(C.fkey(smx)) + if lane == cutlass.Int32(0): + s_wmn[tidx >> cutlass.Int32(5)] = a0 + s_wmx[tidx >> cutlass.Int32(5)] = c0m + cute.arch.barrier() # ---- barrier L1902 ---- + + # PRIME-LATE (L1903-1916): every rank's sample has landed; prime NOW. + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L1904 + pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(PFD)] + for uu in cutlass.range_constexpr(PFD): # clamped prime L1906 + i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + # asm prefetch gate (L1912-1916): DEEP rows only; empty for U<=PFD + if cutlass.const_expr(U > PFD): + gpp = cutlass.Int32(0) + if n4 >= cutlass.Int32(32768): + if (rank + cutlass.Int32(1)) * cutlass.Int32(STEPC) <= n4: + gpp = cutlass.Int32(1) + if gpp != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD, U): + C._prefetch_l2( + x_addr + + cutlass.Int64( + rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ) + * cutlass.Int64(16) + ) + + # cross-warp sample reduce (L1917-1923) + av = cutlass.Uint32(0xFFFFFFFF) + cv = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + av = s_wmn[lane] + cv = s_wmx[lane] + SMIN = C.invkey(C.warp_min_u32(av)) + SMAX = C.invkey(C.warp_max_u32(cv)) + + GMIN = cutlass.Float32(C.SENT_LO) # L1925-1926 + GMAX = cutlass.Float32(C.SENT_HI) + T = cutlass.Float32(_NEG_INF__clus) + HIC = cutlass.Float32(_NEG_INF__clus) + w = cutlass.Float32(0.0) + sok = cutlass.Int32(0) # L1930 + if SMP > cutlass.Int32(0): + if SMAX > SMIN: + sok = cutlass.Int32(1) + if sok != cutlass.Int32(0): # L1932-1947 sample hist + w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) + sc_s = cutlass.Float32(1.0) / w + if shas != cutlass.Int32(0): + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsa[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsb[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = tidx + cutlass.Int32(BLK) # mop-up reloads + while j < smp2: + p4 = j * SS2 * cutlass.Int32(4) + if j >= SMP: + p4 = (j - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fma_) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fma_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fmb_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = j + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1948 ---- + # triple-target ZERO scan (L1952-1955): TGT / TGT2 / 2*TGT + C.scan_cross0( + s_hist, + TGT, + tidx, + s_res, + TGT2, + TGT * cutlass.Int32(2), + s_hist, + nb=NBS, + zero=True, + two=True, + three=True, + ) + cute.arch.barrier() # ---- barrier L1956 ---- + + tot0 = s_res[C.RES_TOT] + b1v = s_res[C.RES_B] + if sok != cutlass.Int32(0): # L1957 + if tot0 >= TGT: + T = _fmaf__clus(cutlass.Float32(b1v), w, SMIN) + needg = cutlass.Int32(1) # L1958-1963 + if T > cutlass.Float32(_NEG_INF__clus): + needg = cutlass.Int32(0) + if needg != cutlass.Int32(0): + # degenerate sample: identical on every rank of the cluster + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers + T = GMIN + if sok != cutlass.Int32(0): # L1964-1982 HIC + if tot0 >= TGT: + b2v = s_res[C.RES_B2] + if b2v >= cutlass.Int32(0): + Tk = _fmaf__clus(cutlass.Float32(b2v), w, SMIN) + up = C.fmax_f32(Tk - T, cutlass.Float32(0.0)) + # heavy-tail cap by T - T3 (rank-TGT..rank-2TGT distance) + if tot0 >= TGT * cutlass.Int32(2): + b3v = s_res[C.RES_B3] + if b3v >= cutlass.Int32(0): + T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) + if T > T3: + up = C.fmin_f32(up, cutlass.Float32(2.0) * (T - T3)) + HIC = C.fmax_f32( + _fmaf__clus(cutlass.Float32(4.0), up, T), + _fmaf__clus(cutlass.Float32(8.0), w, T), + ) + # r4 (a000) ladder floor in SHARED (64-reg wall, L1983-1996) + if tidx == cutlass.Int32(0): + t5 = cutlass.Float32(_NEG_INF__clus) + if sok != cutlass.Int32(0): + if tot0 >= TGT * cutlass.Int32(2): + b3v = s_res[C.RES_B3] + if b3v >= cutlass.Int32(0): + if T > GMIN: + T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) + if T3 < T: + t5 = T3 + s_tsh[0] = t5 + + # ============ attempt loop (L1999-2162) — MUST NOT unroll =========== + listN = cutlass.Int32(0) + above = cutlass.Int32(0) + m = cutlass.Int32(0) + need = cutlass.Int32(0) + B = cutlass.Int32(0) + SC = cutlass.Float32(1.0) + TF = T + complete = cutlass.Int32(0) + valid = cutlass.Int32(0) + + fr = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(U - PFD) + ] # explicit batch (op43 L1) + # (empty for U<=PFD — every row-pass float4 then comes from pf[]) + att = cutlass.Int32(0) + running = cutlass.Int32(1) + while running != cutlass.Int32(0): + if att > cutlass.Int32(0): # retry preamble L2005-2024 + # EXACTNESS: re-prime pf[] (stale roll data, L2006-2015) + if rank < nFullG: + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4( + atom128, + x_addr, + rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK), + pf[uu], + ) + else: + for uu in cutlass.range_constexpr(PFD): + i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + C._cluster_sync_aligned() # ==== clus.sync L2016 ==== + if tidx < cutlass.Int32(NBS): # L2021-2022 + s_hist[tidx] = cutlass.Int32(0) + if tidx == cutlass.Int32(0): + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2024 ---- + + TF = T # window L2026-2031 + hi = C.fmax_f32(GMAX, T) + if HIC > T: + if HIC < hi: + hi = HIC + WD = (hi - T) * cutlass.Float32(1.0 / 256.0) + wdok = cutlass.Int32(0) + if WD > cutlass.Float32(0.0): + wdok = cutlass.Int32(1) + if wdok == cutlass.Int32(0): + WD = cutlass.Float32(1e-30) + SC = cutlass.Float32(1.0) / WD + + # ---- P3 row pass over OWNED CHUNKS (L2033-2121) ---- + g = rank + cutlass.Int32(0) + while g < nCh: + i0 = g * cutlass.Int32(STEPC) + tidx + M = cutlass.Int32(0) + isfull = cutlass.Int32(0) + if g < nFullG: + isfull = cutlass.Int32(1) + if isfull != cutlass.Int32(0): # full body L2042-2049 + for uu in cutlass.range_constexpr(PFD, U): + C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + else: # partial body L2051-2064 + for uu in cutlass.range_constexpr(PFD, U): + i_ = i0 + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 # clamp in [n, npad) + C.ld_g_f32x4(atom128, x_addr, ic, fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + i_ = i0 + cutlass.Int32(uu * BLK) + okq = cutlass.Int32(0) + if i_ < n4: + okq = cutlass.Int32(1) + if okq != cutlass.Int32(0): # +inf-pad escape, ok-gated + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + # ROLL THE PREFETCH FORWARD (L2066-2081): next OWNED chunk, + # issued before the reservation and the survivor walk. + g2 = g + cutlass.Int32(CS) + if g2 < nCh: + j0 = g2 * cutlass.Int32(STEPC) + tidx + infull = cutlass.Int32(0) + if g2 < nFullG: + infull = cutlass.Int32(1) + if infull != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4(atom128, x_addr, j0 + cutlass.Int32(uu * BLK), pf[uu]) + else: + for uu in cutlass.range_constexpr(PFD): + j_ = j0 + cutlass.Int32(uu * BLK) + jc = j_ + if jc >= n4: + jc = lim4 + C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) + # warp-aggregated slot reservation (L2082-2095) + cnt = cutlass.Int32(C.popc(M)) + inc = C.warp_incl_scan_add(cnt, lane) + bpos = cutlass.Int32(0) + if lane == cutlass.Int32(31): + if inc != cutlass.Int32(0): + bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) + pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) + # survivor bit-walk, software-pipelined ONE deep (L2105-2119); + # reload X[idx] — never hold the U float4s across the walk + if M != cutlass.Int32(0): + bp = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx = ( + (i0 + (bp >> cutlass.Int32(2)) * cutlass.Int32(BLK)) << cutlass.Int32(2) + ) + (bp & cutlass.Int32(3)) + xv = C.ldg_f32(x_addr, idx) + while M != cutlass.Int32(0): + bp2 = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx2 = ( + (i0 + (bp2 >> cutlass.Int32(2)) * cutlass.Int32(BLK)) + << cutlass.Int32(2) + ) + (bp2 & cutlass.Int32(3)) + xv2 = C.ldg_f32(x_addr, idx2) + pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) + idx = idx2 + xv = xv2 + pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) + g = g + cutlass.Int32(CS) + # rank-0 scalar tail (L2122-2130): per-thread atomics, bound-check + i = tidx + while i < tailn: + x = C.ldg_f32(x_addr, tail0 + i) + if x >= TF: + bq = C.f2s_rz((x - TF) * SC) # signed form L2125 + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + post = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) + if post < SCAP: + s_cbuf2[post] = ( + cutlass.Uint64(cutlass.Uint32(tail0 + i)) << cutlass.Uint64(32) + ) | cutlass.Uint64(C.u32_of_f32(x)) + i = i + cutlass.Int32(BLK) + + # ---- cluster merge (L2132-2148) ---- + C._cluster_sync_aligned() # ==== clus.sync L2135 ==== + myn = s_scal[0] # L2140 + self._merge_scan0_local(s_hist, s_mrg, rank, k, tidx, s_res) + cute.arch.barrier() # ---- barrier L2143 ---- + tot = s_res[C.RES_TOT] + acc = cutlass.Int32(0) + if tot >= k: + acc = cutlass.Int32(1) + if acc != cutlass.Int32(0): # L2145-2148 accept + valid = cutlass.Int32(1) + complete = cutlass.Int32(0) + if myn <= SCAP: + complete = cutlass.Int32(1) + listN = myn + above = s_res[C.RES_ABOVE] + m = s_res[C.RES_M] + need = k - s_res[C.RES_ABOVE] + B = s_res[C.RES_B] + running = cutlass.Int32(0) + else: + if att == cutlass.Int32(2): # L2149 + running = cutlass.Int32(0) + else: + # rung ladder (L2150-2161) — cluster-uniform on every arm + tshtaken = cutlass.Int32(0) + if att == cutlass.Int32(0): + T5 = s_tsh[0] + if T5 > cutlass.Float32(_NEG_INF__clus): + if T5 < TF: + T = T5 + tshtaken = cutlass.Int32(1) + if tshtaken == cutlass.Int32(0): + # LAZY GATHER — every rank computes identical GMIN + if GMIN == cutlass.Float32(C.SENT_LO): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers inside + floorhit = cutlass.Int32(1) # L2160 + if T > GMIN: + floorhit = cutlass.Int32(0) + if floorhit != cutlass.Int32(0): + running = cutlass.Int32(0) + else: + T = GMIN # L2161 + att = att + cutlass.Int32(1) + + # ============ classification (L2165-2173) ============ + whole = cutlass.Int32(0) + if valid != cutlass.Int32(0): + if need >= m: + whole = cutlass.Int32(1) + lim1 = above + if whole != cutlass.Int32(0): + lim1 = above + m + degen = cutlass.Int32(0) + if valid == cutlass.Int32(0): + degen = cutlass.Int32(1) + if m > CMP: + degen = cutlass.Int32(1) + mc = cutlass.Int32(0) + if degen == cutlass.Int32(0): + mc = m + # crossing candidates land in RANK 0's ck64c via DSMEM (L2173) + rk64 = C._mapa_shared_cluster_addr(ck_addr, cutlass.Int32(0)) + + if degen == cutlass.Int32(0): + if complete != cutlass.Int32(0): + # ---- P5 emit from staged cbuf (L2176-2187) ---- + i = tidx + while i < listN: + pk64 = s_cbuf2[i] + vx = cutlass.Int32(cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF))) + idv = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + xv = C.f32_of_i32(vx) + self._p5_emit(xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64) + i = i + cutlass.Int32(BLK) + else: + # ---- EXACTNESS re-sweep: OWNED CHUNKS + rank-0 true tail + # (L2188-2222) ---- + g = rank + cutlass.Int32(0) + while g < nCh: + lo2 = (g * cutlass.Int32(STEPC)) << cutlass.Int32(2) + e4 = (g + cutlass.Int32(1)) * cutlass.Int32(STEPC) + if e4 > n4: + e4 = n4 + hi2 = e4 << cutlass.Int32(2) + i = lo2 + tidx + while i < hi2: + x = C.ldg_f32(x_addr, i) + if x >= TF: + self._p5_emit( + x, i, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 + ) + i = i + cutlass.Int32(BLK) + g = g + cutlass.Int32(CS) + t2 = tidx + while t2 < tailn: + ii = tail0 + t2 + x = C.ldg_f32(x_addr, ii) + if x >= TF: + self._p5_emit( + x, ii, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 + ) + t2 = t2 + cutlass.Int32(BLK) + + # ============ EXIT RENDEZVOUS (L2226) ============ + # all DSMEM traffic retired; the ONLY exit rendezvous. rank!=0 falls + # through to the kernel end (post-barrier asymmetric exit L2227); + # NO later cluster barrier. + C._cluster_sync_aligned() # ==== clus.sync L2226 ==== + + if rank == cutlass.Int32(0): + if degen == cutlass.Int32(0): + if whole == cutlass.Int32(0): + # ---- P6 rank-0 refine (L2229-2303) ---- + if mc <= cutlass.Int32(QUADC_CLUS__clus): # L2232-2247 O(mc^2) + mc2 = mc & cutlass.Int32(~1) + i = tidx + while i < mc: + # G1: re-assert Uint64 at every unsigned compare + # in/after dynamic loops (ct_main notes). + u64v = s_ck64[i] + r_ = cutlass.Int32(0) + jq = cutlass.Int32(0) + while jq < mc2: # ulonglong2 16B reads + vlo, vhi = C._lds_v2_u64(ck_addr + jq * cutlass.Int32(8)) + r_ = ( + r_ + + cutlass.Int32(vlo > cutlass.Uint64(u64v)) + + cutlass.Int32(vhi > cutlass.Uint64(u64v)) + ) + jq = jq + cutlass.Int32(2) + if mc2 < mc: # odd tail L2244 + r_ = r_ + cutlass.Int32( + cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) + ) + if r_ < need: + out_row[above + r_] = cutlass.Int32( + cutlass.Uint32( + cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) + ) + ) + i = i + cutlass.Int32(BLK) + else: + # key-space narrowing over ck64c (L2249-2278) + if tidx == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + if tidx < cutlass.Int32(NBS): # cleared ONCE L2251 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2255 ---- + i = tidx + while i < mc: + kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + C.atomic_min_cta(s_kmm.iterator + 0, kk) + C.atomic_max_cta(s_kmm.iterator + 1, kk) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2257 ---- + rlo = s_kmm[0] + rhi = s_kmm[1] + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L2258-2278 (<=6 lvl) + if needC == mm: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(6): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + i = tidx + while i < mc: # L2266-2269 re-bin + uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), + cutlass.Int32(1), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2270 ---- + C.scan_cross0( + s_hist, + needC, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=True, + ) + cute.arch.barrier() # ---- barrier L2273 ---- + aboveC = aboveC + s_res[C.RES_ABOVE] + needC = needC - s_res[C.RES_ABOVE] + mm = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): # L2276 + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L2279 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2281 ---- + it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L2282-2301 ballot emit + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + if i < mc: + w64 = s_ck64[i] + iu = cutlass.Int64(cutlass.Uint32(w64 >> cutlass.Uint64(32))) + idv = cutlass.Int32( + cutlass.Uint32(w64 & cutlass.Uint64(0xFFFFFFFF)) + ) + if iu > ethr: + p1 = cutlass.Int32(1) + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, + p2, + idv, + above, + aboveC, + above + aboveC, + needC, + out_row, + s_scal, + lane, + ) + it = it + cutlass.Int32(1) + else: + # ---- degen fallback: whole-row key-space narrowing + # (L2305-2355; per-level clear + scan_cross w/ ws) ---- + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + above2 = cutlass.Int32(0) + need2 = k + m2 = n + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L2312-2331 (<=8 lvl) + if need2 == m2: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + above2 = above2 + m2 + need2 = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(8): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + if tidx < cutlass.Int32(NBS): # per-level clear L2320 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2321 ---- + i = tidx + while i < n: # L2322-2324 whole row + uq = C.fkey(C.ldg_f32(x_addr, i)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2325 ---- + # block-parallel scan (ONE internal barrier; only use + # of ws in this kernel, L2326) + C.scan_cross( + s_hist, + s_ws, + need2, + tidx, + s_res, + cutlass.Int32(0), + blk=BLK, + nb=NBS, + two=False, + ) + cute.arch.barrier() # ---- barrier L2327 ---- + above2 = above2 + s_res[C.RES_ABOVE] + need2 = need2 - s_res[C.RES_ABOVE] + m2 = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L2336 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2337 ---- + nA = k # L2338 + nT = cutlass.Int32(0) + if tie_m != cutlass.Int32(0): + nA = above2 + nT = need2 + it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L2340-2354 + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + if i < n: + uq = C.fkey(C.ldg_f32(x_addr, i)) + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane + ) + it = it + cutlass.Int32(1) + + # ------------------------------------------------------------------ + # host launcher: grid dim3(CS, b) + cluster (CS,1,1) (probe P11); + # min_blocks_per_mp=1 == __launch_bounds__(1024, 1) 64-reg wall. + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + npad: cutlass.Int32, + k: cutlass.Int32, + SCAP: cutlass.Int32, + CMP: cutlass.Int32, + SMP: cutlass.Int32, + TGT: cutlass.Int32, + Q: cutlass.Int32, + SS2: cutlass.Int32, + TGT2: cutlass.Int32, + stream, + ): + b = logits.shape[0] + self.kern(logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2).launch( + grid=(self.cs, b, 1), + block=(self.blk, 1, 1), + cluster=(self.cs, 1, 1), + stream=stream, + min_blocks_per_mp=self.minb, + ) + + +# --------------------------------------------------------------------------- +# compile cache + torch-facing entry +# --------------------------------------------------------------------------- +_COMPILE_CACHE__clus = {} + + +def get_compiled__clus(tpl, scap: int = 8192, cmp_: int = 2048, options_extra: str = ""): + """Compile (or fetch) the gvr_clus variant for constexpr tuple + tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every + reachable route has 8192/2048 — asserted by run__clus()).""" + key = (tuple(tpl), scap, cmp_, options_extra) + hit = _COMPILE_CACHE__clus.get(key) + if hit is not None: + return hit + blk, u, minb, nbs, cs = tpl + kern = GvrClusKernel(blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_) + r0, c0 = cute.sym_int(), cute.sym_int() + r1, c1 = cute.sym_int(), cute.sym_int() + r2, c2 = cute.sym_int(), cute.sym_int() + logits_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (r0, c0), stride_order=(1, 0), assumed_align=16 + ) + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 + ) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r2, c2), 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_fake, + out_fake, + *([cutlass.Int32(0)] * 10), + stream=fake_stream, + options=("--enable-tvm-ffi " + options_extra).strip(), + ) + _COMPILE_CACHE__clus[key] = compiled + return compiled + + +def run__clus(logits, pre_idx, n: int, out): + """torch-facing single-call entry: routes (b, n, k) through ct_dispatch, + asserts the shape lands on gvr_clus, launches the matching variant. + gvr_clus takes NO workspace (spec §4c).""" + try: + from . import gvr_topk_decode_self_sampling_host as ct_dispatch + except ImportError: + import gvr_topk_decode_self_sampling_host as ct_dispatch + b, npad = logits.shape + k = pre_idx.shape[1] + r = ct_dispatch.route(b, int(n), npad, k) + assert r["kernel"] == "clus", f"shape routes to {r['kernel']}, not gvr_clus" + rt = r["rt"] + kobj = GvrClusKernel(*r["tpl"], scap=rt["SCAP"], cmp_=rt["CMP"]) + assert r["smem"] == kobj.dyn_bytes, (r["smem"], kobj.dyn_bytes) + fn = get_compiled__clus(tuple(r["tpl"]), scap=rt["SCAP"], cmp_=rt["CMP"]) + fn( + logits, + pre_idx, + out, + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP"], + rt["CMP"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + return r + + +def run_manual(logits, pre_idx, n: int, out, tpl, rt): + """Manual-lattice entry for route()-unreachable (U, CS) members: launches + tpl with caller-supplied runtime scalars (must be route()-consistent for + the same CS; U only changes the chunk geometry).""" + fn = get_compiled__clus(tuple(tpl), scap=rt["SCAP"], cmp_=rt["CMP"]) + fn( + logits, + pre_idx, + out, + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP"], + rt["CMP"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + + +# =========================================================================== +# ==== ct_regclus.py — family: regclus ========================= +# =========================================================================== +"""ct_regclus.py — op46 gvr_reg_clus CuTeDSL translation. + +CUDA ground truth: src_cuda/kernel.cu L2359-2648 (CLUSTERED register-resident +GVR: the register algorithm — T = GMIN directly, one float-space histogram, +one register sweep — run across a cluster of CS CTAs; per-CTA instruction +stream intentionally identical to the single-CTA reg path plus two hardware +cluster barriers and CS DSMEM reads per bin). Contract: TRANSLATION_SPEC.md +§5.4. Probe verdicts (probes/PROBE_RESULTS.md) binding; op43 lessons L1-L5 +applied; ct_main G1 signedness mitigation applied at every unsigned +compare/shift in/after dynamic loops. + +Template knobs (CUDA `gvr_reg_clus`, all instantiations +BLK=BLKC=1024): ctor args of :class:`GvrRegClusKernel`. Runtime args mirror +the CUDA `(n, npad, k)` — npad/k come from tensor shapes, so only `n` crosses +the ABI. Launch: grid=(CS, b), cluster=(CS,1,1), block=1024, dynamic smem +45,056 B (+512 B static-mirror prelude), `__launch_bounds__(1024,1)` == +min_blocks_per_mp=1 -> 64-register wall. + +Shared-memory map (single dynamic window, word offsets; the CUDA static +__shared__ block folded into the first 512 B — byte-identical layout in every +CTA, a mapa/DSMEM requirement): + + [0..5] s_res (ct_common slot map RES_B/M/ABOVE/TOT/B2/B3) + [6..7] s_cnt (s_o1, s_o2 — kernel.cu L2392) + [8..9] s_kmm (s_kmin, s_kmax — Uint32, L2393) + [16..16+32) ws (scan_cross_w workspace, L2390) + [48..48+32) wmn (Uint32 warp min partials) + [80..80+32) wmx (Uint32 warp max partials) + [128..1152) hist (this CTA's raw counts, L2384) + [1152..2176) mrg (cluster totals -> per-CTA global write cursors) + [2176..3200) hoff (this CTA's rank-exclusive bin offset) + [3200..7296) ck (crossing keys, Uint32, CMPC=4096 slots) + [7296..11392) ci (crossing indices, Int32, CMPC slots) + +Launch smem = 45,568 B (compile-time constant -> plain int at .launch(); +MINB==1 so the _build_kernel_attrs carveout path is not taken and ct_reg's +_no_carveout workaround is unnecessary here). +""" + + +# ---- constants (kernel.cu lines) ------------------------------------------- +NB__regclus = 1024 # L16 (histogram bins; == BLKC here) +LNB = 10 # L1267 log2(NB__regclus) — reg_clus narrowing shift +QUADC__regclus = 96 # L21 O(mc^2) rank gate (L2534) +CMPC = 4096 # L2372 crossing slots PER CTA (pow2) +LCMPC = 12 # L2373 log2(CMPC) +BLKC = 1024 # L2374 CTA size + +STATIC_WORDS__regclus = 128 # DSL smem prelude (static-__shared__ mirror) +STATIC_BYTES__regclus = STATIC_WORDS__regclus * 4 +DYN_SMEM_BYTES = (3 * NB__regclus + 2 * CMPC) * 4 # 45,056 (L2926) +SMEM_BYTES = STATIC_BYTES__regclus + DYN_SMEM_BYTES # 45,568 + +# word offsets into the shared window (module docstring) +W_HIST = STATIC_WORDS__regclus +W_MRG = STATIC_WORDS__regclus + NB__regclus +W_HOFF = STATIC_WORDS__regclus + 2 * NB__regclus +W_CK = STATIC_WORDS__regclus + 3 * NB__regclus +W_CI = STATIC_WORDS__regclus + 3 * NB__regclus + CMPC + +_NEG_INF__regclus = float("-inf") +_POS_INF__regclus = float("inf") + + +# --------------------------------------------------------------------------- +# module-local FP/util spellings (copied from frozen sibling ct_reg.py — +# probe P6 fma discipline; kept local so this module is self-contained) +# --------------------------------------------------------------------------- +@dsl_user_op +def _fmaf__regclus(a, b, c, *, loc=None, ip=None): + """CUDA fmaf: single fma.rn.f32 (P6 spelling; classify == emit bit-exact).""" + return cutlass.Float32( + mlir_math.fma( + a.ir_value(loc=loc, ip=ip), + b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), + fastmath=mlir_arith.FastMathFlags.none, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _umin_u32__regclus(a, b): + """unsigned min(a, b) — CUDA min() on the bin clamp (IMNMX).""" + r = a + if b < a: + r = b + return r + + +@cute.jit +def _fabsf__regclus(x): + """|x| via sign-bit clear (exact, matches fabsf).""" + return f32_of_u32(u32_of_f32(x) & cutlass.Uint32(0x7FFFFFFF)) + + +def _smem_view__regclus(dtype, sbase, word_off: int, length: int, align: int = 16): + """Typed tensor view at a constexpr word offset into the smem window.""" + p = cute.make_ptr( + dtype, sbase + cutlass.Int32(word_off * 4), cute.AddressSpace.smem, assumed_align=align + ) + return cute.make_tensor(p, cute.make_layout((length,))) + + +def _val__regclus(frags, s: int): + """val[s] accessor over the float4[VPT] register batch (constexpr s).""" + return frags[s // 4][s % 4] + + +class GvrRegClusKernel: + """gvr_reg_clus (kernel.cu L2376-2379).""" + + def __init__(self, blk: int, vpt: int, cs: int, pdl: bool = False): + assert blk == BLKC, "all instantiations BLK=BLKC=1024 (spec §4b)" + assert vpt in (1, 2, 4) and cs in (2, 4, 8) + self.blk = blk + self.vpt = vpt + self.cs = cs + self.pdl = bool(pdl) + self.S = vpt * 4 # L2381 + self.span = blk * vpt # L2382 (float4 per CTA) + + # ------------------------------------------------------------------ + @cute.kernel + def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + BLK = cutlass.const_expr(self.blk) + VPT = cutlass.const_expr(self.vpt) + CS = cutlass.const_expr(self.cs) + S = cutlass.const_expr(self.S) + NW = cutlass.const_expr(self.blk // 32) + + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_wait() # L2380 (knob default off) + + tid, _, _ = cute.arch.thread_idx() + rank, row, _ = cute.arch.block_idx() # L2396-2397 (P11: bx=rank) + lane = tid & cutlass.Int32(31) + + # ------------------------------------------------------------------ + # Predeclarations (DSL AST rule: every scalar (re)assigned under a + # dynamic if/while must pre-exist with a stable type; constant inits + # are dead-coded — ct_reg precedent, reg-audited free). + # ------------------------------------------------------------------ + i = cutlass.Int32(0) + j = cutlass.Int32(0) + rnk = cutlass.Int32(0) + tinc = cutlass.Int32(0) + mc = cutlass.Int32(0) + p = cutlass.Int32(0) + q2i = cutlass.Int32(0) + idx = cutlass.Int32(0) + lim1 = cutlass.Int32(0) + aboveC = cutlass.Int32(0) + needC = cutlass.Int32(0) + mm = cutlass.Int32(0) + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + b2w = cutlass.Int32(0) + sh2 = cutlass.Int32(0) + b_lv = cutlass.Int32(0) + it = cutlass.Int32(0) + it2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + n1 = cutlass.Int32(0) + n2 = cutlass.Int32(0) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + p1e = cutlass.Int32(0) + p2e = cutlass.Int32(0) + lml = cutlass.Int32(0) + nA = cutlass.Int32(0) + nT = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + pv0 = cutlass.Int32(-1) + okc = cutlass.Int32(0) + whole = cutlass.Int32(0) + degen = cutlass.Int32(0) + pre_a = cutlass.Int32(0) + tot_a = cutlass.Int32(0) + uk = cutlass.Uint32(0) + uq = cutlass.Uint32(0) + vq = cutlass.Uint32(0) + kv = cutlass.Uint32(0) + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0) + d2 = cutlass.Uint32(0) + unar = cutlass.Uint32(0) + bnn = cutlass.Uint32(0) + nlo = cutlass.Uint32(0) + uke = cutlass.Uint32(0) + bn = cutlass.Uint32(0) + ethr = cutlass.Int64(0) + tval = cutlass.Float32(_NEG_INF__regclus) + LOQ = cutlass.Float32(0.0) + qv = cutlass.Float32(0.0) + + npad = cutlass.Int32(logits.shape[1]) # noqa: F841 + k = cutlass.Int32(pre_idx.shape[1]) + out_row = out[row, None] + x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base + p_addr = pre_idx[row, None].iterator.toint() + + # ---- shared-memory window (map in module docstring) ---- + sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) + sbase = sptr.toint() # Int32 shared addr + + s_res = _smem_view__regclus(cutlass.Int32, sbase, 0, 6) + s_cnt = _smem_view__regclus(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_o2 + s_kmm = _smem_view__regclus(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax + s_ws = _smem_view__regclus(cutlass.Int32, sbase, 16, 32) + s_wmn = _smem_view__regclus(cutlass.Uint32, sbase, 48, 32) + s_wmx = _smem_view__regclus(cutlass.Uint32, sbase, 80, 32) + s_hist = _smem_view__regclus(cutlass.Int32, sbase, W_HIST, NB__regclus) + s_mrg = _smem_view__regclus(cutlass.Int32, sbase, W_MRG, NB__regclus) + s_hoff = _smem_view__regclus(cutlass.Int32, sbase, W_HOFF, NB__regclus) + s_ck = _smem_view__regclus(cutlass.Uint32, sbase, W_CK, CMPC) + s_ci = _smem_view__regclus(cutlass.Int32, sbase, W_CI, CMPC, align=4) + # raw byte bases for DSMEM (mapa) addressing + hist_addr = sbase + cutlass.Int32(W_HIST * 4) + ck_addr = sbase + cutlass.Int32(W_CK * 4) + ci_addr = sbase + cutlass.Int32(W_CI * 4) + + n4 = n >> cutlass.Int32(2) # L2405 + ntail = n - (n4 << cutlass.Int32(2)) # L2406 + base4 = rank * cutlass.Int32(self.span) # L2407 + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L2425 + + # ---- P0: redundant hint gather, EVERY CTA (L2410-2413; k<=BLK by + # dispatch gate L2897). One coalesced word per thread, NO cluster + # barrier — GMIN/GMAX identical everywhere by construction. + if tid < k: + pv0 = ld_g_i32(p_addr, tid) + + # ---- P1: row load — predicated flat float4[VPT] batch (L2415-2424; + # the CUDA has NO exact-fit peel here, guard is per-load). Issue all + # loads first (op43 L1), then -INFINITY-fill missed slots (op43 L2). + atom128 = g2r_atom_f32(128, invariant=True) + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + for u in cutlass.range_constexpr(VPT): + i = base4 + tid + cutlass.Int32(u * self.blk) + if i < n4: + ld_g_f32x4(atom128, x_addr, i, frags[u]) + for u in cutlass.range_constexpr(VPT): + i = base4 + tid + cutlass.Int32(u * self.blk) + if i >= n4: # -INFINITY fill L2421 + for z in cutlass.range_constexpr(4): + frags[u][z] = cutlass.Float32(_NEG_INF__regclus) + # tail element: rank 0 only (L2425-2426) + if rank == cutlass.Int32(0): + if tid < ntail: + tval = ldg_f32(x_addr, tix) + + # ---- P2: init (L2428-2429). NB__regclus == BLK -> single-pass hist clear. + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + s_cnt[1] = cutlass.Int32(0) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + + # ---- P3: GMIN/GMAX from the hint (L2431-2445), ONE barrier fold. + lmin = cutlass.Uint32(0xFFFFFFFF) + lmax = cutlass.Uint32(0) + if cutlass.Uint32(pv0) < cutlass.Uint32(n): + uk = fkey(ldg_f32(x_addr, pv0)) # __ldg(X+pv0) L2433 + lmin = uk + lmax = uk + lmin = warp_min_u32(lmin) + lmax = warp_max_u32(lmax) + if lane == cutlass.Int32(0): + s_wmn[tid >> cutlass.Int32(5)] = lmin + s_wmx[tid >> cutlass.Int32(5)] = lmax + cute.arch.barrier() # L2438 + a = cutlass.Uint32(0xFFFFFFFF) + c = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + a = cutlass.Uint32(s_wmn[lane]) + c = cutlass.Uint32(s_wmx[lane]) + lmin = warp_min_u32(a) + lmax = warp_max_u32(c) + Tv = invkey(lmin) + GMAX = invkey(lmax) + + # ---- collapse guard, NaN-safe (L2446-2453) + okc = cutlass.Int32(0) + if Tv < GMAX: + if (GMAX - Tv) > cutlass.Float32(1e-30): + okc = cutlass.Int32(1) + if okc == cutlass.Int32(0): + Tv = cutlass.Float32(SENT_LO) + GMAX = cutlass.Float32(SENT_HI) + + # ---- bin transform constants (L2454-2467): branchless trash bin. + WD = (GMAX - Tv) * cutlass.Float32(1.0 / float(NB__regclus - 2)) + wsel = cutlass.Float32(1e-30) + if WD > cutlass.Float32(0.0): + wsel = WD + SC = cutlass.Float32(1.0) / wsel + CQ0 = cutlass.Float32(1.0) - Tv * SC + CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf__regclus(CQ0) + cutlass.Float32(1.0)) + + # ---- P4: histogram (L2469-2472); tval add UNCONDITIONAL (trash bin + # swallows -INFINITY via the saturating cvt). + for s in cutlass.range_constexpr(S): + qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + qv = _fmaf__regclus(tval, SC, CQ) + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + + # ---- P5: cluster merge (L2474-2484) + _cluster_sync_aligned() # L2474 + for z in cutlass.range_constexpr(NB__regclus // self.blk): + i = tid + cutlass.Int32(z * self.blk) + # CS-unrolled remote u32 loads: batch-issue, then fold (#pragma + # unroll L2477; one mapa per (i, r) exactly like map_shared_rank) + hvals = [] + for r in cutlass.range_constexpr(CS): + ma = _mapa_shared_cluster_addr( + hist_addr + (i << cutlass.Int32(2)), cutlass.Int32(r) + ) + hvals.append(_ld_shared_cluster_i32(ma)) + tot_a = cutlass.Int32(0) + pre_a = cutlass.Int32(0) + for r in cutlass.range_constexpr(CS): + if cutlass.Int32(r) < rank: + pre_a = pre_a + hvals[r] # rank-exclusive + tot_a = tot_a + hvals[r] + s_mrg[i] = tot_a + s_hoff[i] = pre_a + + # ---- P6: scan (L2485-2492) + cute.arch.barrier() # L2485 + scan_cross_w(s_mrg, s_ws, k, tid, s_res, blk=self.blk, nb=NB__regclus) + cute.arch.barrier() # L2487 + above = s_res[RES_ABOVE] + m = s_res[RES_M] + Bv = s_res[RES_B] + need = k - above + whole = cutlass.Int32(0) + if need >= m: + whole = cutlass.Int32(1) + degen = cutlass.Int32(0) + if m > cutlass.Int32(CS * CMPC): + degen = cutlass.Int32(1) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + i = tid + cutlass.Int32(z * self.blk) + s_mrg[i] = s_mrg[i] + s_hoff[i] # L2491 global cursor + cute.arch.barrier() # L2492 + + # ---- P7: register sweep emit (L2494-2527, !degen) + if degen == cutlass.Int32(0): + LOQ = cutlass.Float32(Bv) # L2495 + lim1 = above + if whole == cutlass.Int32(1): + lim1 = above + m # L2496 + for s in cutlass.range_constexpr(S): + qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) # bit-identical L2499 + if qv >= LOQ: + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + idx = ( + (base4 + tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if p < lim1: + out_row[p] = idx + else: + if whole == cutlass.Int32(0): + # crossing overflow -> striped DSMEM slabs; TWO + # separate u32 remote stores (NOT packed, L2507-10) + q2i = p - above + rnk = q2i >> cutlass.Int32(LCMPC) + j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) + _st_shared_cluster_i32( + _mapa_shared_cluster_addr(ck_addr + j, rnk), + fkey(_val__regclus(frags, s)), + ) + _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), idx) + # tail element (L2514-2526): tval == -INF fails q>=LOQ elsewhere + qv = _fmaf__regclus(tval, SC, CQ) + if qv >= LOQ: + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + if p < lim1: + out_row[p] = tix + else: + if whole == cutlass.Int32(0): + q2i = p - above + rnk = q2i >> cutlass.Int32(LCMPC) + j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) + _st_shared_cluster_i32( + _mapa_shared_cluster_addr(ck_addr + j, rnk), fkey(tval) + ) + _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), tix) + + # ---- P8 (L2529-2530): release staging to rank 0 + cute.arch.barrier() # L2529 + _cluster_sync_aligned() # L2530 + + # ---- P9: rank-0 selection (L2532-2647) + if rank == cutlass.Int32(0): + if whole == cutlass.Int32(0): + mc = m + if degen == cutlass.Int32(1): + mc = cutlass.Int32(0) # L2533 + if degen == cutlass.Int32(0): + if mc <= cutlass.Int32(QUADC__regclus): + # (1) quad-96: all candidates LOCAL (96 < CMPC), + # O(mc^2) slot-order tie-broken rank (L2535-2543) + i = tid + while i < mc: + uq = cutlass.Uint32(s_ck[i]) + rnk = cutlass.Int32(0) + j = cutlass.Int32(0) + while j < mc: + vq = cutlass.Uint32(s_ck[j]) + tinc = cutlass.Int32(0) + if vq > uq: + tinc = cutlass.Int32(1) + if vq == uq: + if j < i: + tinc = cutlass.Int32(1) + rnk = rnk + tinc + j = j + cutlass.Int32(1) + if rnk < need: + out_row[above + rnk] = s_ci[i] + i = i + cutlass.Int32(BLK) + else: + # (2) key-space narrowing over striped DSMEM slabs + # (L2544-2596): slot = i & (CMPC-1), rank = i >> LCMPC + if tid == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + cute.arch.barrier() # L2546 + i = tid + while i < mc: + kv = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), + i >> cutlass.Int32(LCMPC), + ) + ) + ) + atomic_min_cta(s_kmm.iterator, kv) + atomic_max_cta(s_kmm.iterator + 1, kv) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L2551 + rlo = cutlass.Uint32(s_kmm[0]) + rhi = cutlass.Uint32(s_kmm[1]) + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + while done == cutlass.Int32(0): # <=6 levels L2553 + if needC == mm: # L2554 + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(6): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2w = cutlass.Int32(32) - clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = cutlass.Int32(0) + if b2w > cutlass.Int32(LNB): + sh2 = b2w - cutlass.Int32(LNB) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + cute.arch.barrier() # L2563 + i = tid + while i < mc: + unar = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ( + (i & cutlass.Int32(CMPC - 1)) + << cutlass.Int32(2) + ), + i >> cutlass.Int32(LCMPC), + ) + ) + ) + if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): + if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): + bnn = ( + cutlass.Uint32(unar) - cutlass.Uint32(rlo) + ) >> cutlass.Uint32(sh2) + bnn = _umin_u32__regclus( + bnn, cutlass.Uint32(NB__regclus - 1) + ) + atomic_add_cta( + s_hist.iterator + cutlass.Int32(bnn), + cutlass.Int32(1), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L2568 + find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) + cute.arch.barrier() # L2570 + aboveC = aboveC + s_res[RES_ABOVE] + needC = needC - s_res[RES_ABOVE] + mm = s_res[RES_M] + b_lv = s_res[RES_B] + nlo = cutlass.Uint32(rlo) + ( + cutlass.Uint32(b_lv) << cutlass.Uint32(sh2) + ) + if b_lv != cutlass.Int32(NB__regclus - 1): + rhi = nlo + ( + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) + - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + cute.arch.barrier() # L2576 + # two-predicate ballot emit over the striped slabs + lml = cutlass.Int32(cute.arch.lanemask_lt()) + it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it = cutlass.Int32(0) + while it < it2: + i = it * cutlass.Int32(BLK) + tid + uke = cutlass.Uint32(0) + idv = cutlass.Int32(0) + if i < mc: # predicated remote + uke = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), + i >> cutlass.Int32(LCMPC), + ) + ) + ) + idv = _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ci_addr + + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), + i >> cutlass.Int32(LCMPC), + ) + ) + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + if i < mc: + if cutlass.Int64(cutlass.Uint32(uke)) > ethr: + q1f = cutlass.Int32(1) + if cutlass.Int64(cutlass.Uint32(uke)) == ethr: + q2f = cutlass.Int32(1) + n1 = ballot(q1f == cutlass.Int32(1)) + n2 = ballot(q2f == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_cnt.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_cnt.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1f == cutlass.Int32(1): + if p1e < aboveC: + out_row[above + p1e] = idv + if q2f == cutlass.Int32(1): + if p2e < needC: + out_row[above + aboveC + p2e] = idv + it = it + cutlass.Int32(1) + else: + # (3) degen safety net (L2597-2645): crossing bin larger + # than the whole cluster buffer -> exact whole-row + # key-space narrowing by rank 0 alone, <=8 levels. + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + aboveC = cutlass.Int32(0) # above2 + needC = k # need2 + mm = n # m2 + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + while done == cutlass.Int32(0): + if needC == mm: # L2603 + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(8): # L2605 + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + sh2 = cutlass.Int32(0) + if b2w > cutlass.Int32(LNB): + sh2 = b2w - cutlass.Int32(LNB) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + cute.arch.barrier() # L2612 + i = tid + while i < n: # whole-row bin L2613 + unar = fkey(ldg_f32(x_addr, i)) + if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): + if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): + bnn = ( + cutlass.Uint32(unar) - cutlass.Uint32(rlo) + ) >> cutlass.Uint32(sh2) + bnn = _umin_u32__regclus( + bnn, cutlass.Uint32(NB__regclus - 1) + ) + atomic_add_cta( + s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L2618 + find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) + cute.arch.barrier() # L2620 + aboveC = aboveC + s_res[RES_ABOVE] + needC = needC - s_res[RES_ABOVE] + mm = s_res[RES_M] + b_lv = s_res[RES_B] + nlo = cutlass.Uint32(rlo) + ( + cutlass.Uint32(b_lv) << cutlass.Uint32(sh2) + ) + if b_lv != cutlass.Int32(NB__regclus - 1): + rhi = nlo + ( + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + cute.arch.barrier() # L2626 + nA = k # tie_m ? above2 : k + if tie_m == cutlass.Int32(1): + nA = aboveC + nT = cutlass.Int32(0) + if tie_m == cutlass.Int32(1): + nT = needC + lml = cutlass.Int32(cute.arch.lanemask_lt()) + it2 = (n + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it = cutlass.Int32(0) + while it < it2: # L2628-2645 + i = it * cutlass.Int32(BLK) + tid + uke = cutlass.Uint32(0) + if i < n: + uke = fkey(ldg_f32(x_addr, i)) + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + if i < n: + if cutlass.Int64(cutlass.Uint32(uke)) > ethr: + q1f = cutlass.Int32(1) + if tie_m == cutlass.Int32(1): + if cutlass.Int64(cutlass.Uint32(uke)) == ethr: + q2f = cutlass.Int32(1) + n1 = ballot(q1f == cutlass.Int32(1)) + n2 = ballot(q2f == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_cnt.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_cnt.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1f == cutlass.Int32(1): + if p1e < nA: + out_row[p1e] = i + if q2f == cutlass.Int32(1): + if p2e < nT: + out_row[nA + p2e] = i + it = it + cutlass.Int32(1) + + # ---- P10: FINAL cluster rendezvous (L2648) — ALL ranks reach it; + # keeps peers resident until rank 0 has read their ck/ci. + _cluster_sync_aligned() + + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, stream + ): + b = logits.shape[0] + self.kern(logits, pre_idx, out, n).launch( + grid=(self.cs, b, 1), + block=(self.blk, 1, 1), + cluster=(self.cs, 1, 1), + stream=stream, + smem=SMEM_BYTES, + min_blocks_per_mp=1, + use_pdl=self.pdl, + ) + + +# --------------------------------------------------------------------------- +# host wrapper: compile cache + route()-driven entry +# --------------------------------------------------------------------------- +_COMPILE_CACHE__regclus: dict = {} + + +def get_compiled__regclus(tpl, dump_dir=None, pdl=False): + """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" + key = (tuple(tpl), bool(pdl)) + compiled = _COMPILE_CACHE__regclus.get(key) + if compiled is None: + from cutlass.cute import runtime as _crt + + blk, vpt, cs = tpl + kernel = GvrRegClusKernel(blk, vpt, cs, pdl=pdl) + nb_, nc_ = cute.sym_int(), cute.sym_int() + nb2_, nc2_ = cute.sym_int(), cute.sym_int() + nb3_, nc3_ = cute.sym_int(), cute.sym_int() + lg_fake = _crt.make_fake_compact_tensor( + cutlass.Float32, (nb_, nc_), stride_order=(1, 0), assumed_align=16 + ) + pi_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (nb2_, nc2_), stride_order=(1, 0), assumed_align=16 + ) + out_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (nb3_, nc3_), stride_order=(1, 0), assumed_align=16 + ) + fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) + opts = "--enable-tvm-ffi" + if dump_dir: + opts += f" --keep-ptx --keep-cubin --dump-dir {dump_dir}" + compiled = cute.compile( + kernel, lg_fake, pi_fake, out_fake, cutlass.Int32(0), stream=fake_stream, options=opts + ) + _COMPILE_CACHE__regclus[key] = compiled + return compiled + + +def regclus_topk(logits, pre_idx, n, out, rd=None): + """torch-facing entry for the clustered register family. + + logits [b, npad] f32, pre_idx [b, k] i32, out [b, >=k] i32, n = valid len. + rd: optional pre-computed ct_dispatch.route() dict (must be reg_clus). + """ + if rd is None: + try: + from .gvr_topk_decode_self_sampling_host import route + except ImportError: + from gvr_topk_decode_self_sampling_host import route + rd = route(logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1]) + assert rd["kernel"] == "reg_clus", rd["kernel"] + tpl = tuple(rd["tpl"]) + assert pre_idx.shape[1] <= tpl[0], "k <= BLK enforced by dispatch L2897" + assert rd["smem"] == DYN_SMEM_BYTES + compiled = get_compiled__regclus(tpl) + compiled(logits, pre_idx, out, int(n)) + return out + + +__all__ = [ + "get_compiled", + "get_compiled__clus", + "get_compiled__reg", + "get_compiled__regclus", + "run", + "run__clus", +] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py new file mode 100644 index 000000000000..bb7aef06258c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -0,0 +1,811 @@ +# 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. + +"""Self-sampling GVR top-K decode — host side (dispatch, workspace, entry). + +Companion to ``gvr_topk_decode_self_sampling.py`` (the merged device +module). Three sections, each a rename-only merge of the per-family source +(fork branch ``GVR-selfsampling-CuTeDSL``): + +1. dispatch — bit-exact transcription of the CUDA host dispatch, a pure + function ``route(b, n, npad, k)`` (cross-checked against an independent + second transcription by a 1,159,168-case boundary+fuzz sweep); +2. workspace — one zero-initialised per-device slab (20,973,568 B) via the + torch caching allocator, keep-alive + double-checked locking semantics + mirrored from the CUDA binding; +3. operator entry — ``run(logits, pre_idx, n_valid, indices)`` / + ``run_ws(..., workspace)`` DPS forms with the CUDA binding's hardening + battery, BIND-ONCE launch cache keyed on ``(b, n, npad, k)``. + +OPERATOR CONTRACT (standalone; not wired into the decode path): ``n_valid`` +is one host python int for the whole batch — every row shares the same +valid prefix, in COMPRESSED index space (the caller applies any +``compressRatio`` division and, for DSv3.2 (cr=1), the +1 temporal shift of +``pre_idx``; see ``heuristicTopKDecode.cu``). The production decode engine +instead reads per-request ``seq_lens`` on-device with per-row MTP offsets +(sync-free, CUDA-graph-replay safe with growing KV); adopting that per-row +contract inside these kernels is tracked follow-up work. Until then this +module is exercised standalone (unit tests / benchmarking) and must not be +substituted for the tiered path under continuous batching, MTP +(``next_n > 1``), or CUDA-graph capture. +""" + +import math +import operator +import threading + +import torch + +_dev_mod = None + + +def _device(): + """Lazy import of the merged device module (first routed shape compiles; + a broken/absent device module only fails when actually reached).""" + global _dev_mod + if _dev_mod is None: + try: + from . import gvr_topk_decode_self_sampling as _m # in-tree + except ImportError: # standalone dir + import gvr_topk_decode_self_sampling as _m + _dev_mod = _m + return _dev_mod + + +# =========================================================================== +# ==== dispatch (ct_dispatch.py) ============================================ +# =========================================================================== +"""Pure-Python transcription of the frozen GVR CUDA dispatch (gvr_topk_launch). + +Source of truth: ../src_cuda/kernel.cu (3197 lines). route(b, n, npad, k) is a +PURE function of its four ints -- no env knobs, no GPU, stdlib only. + +Branch map (kernel.cu line citations): + constants NB L16, QUADC L21, SNB L170, CMPC L2372, BLKC L2374 + reg-block prologue L2757-2822: wide=(b<=148) L2757; n4=n>>2 L2759; + CMP=min(n,2560) L2764; QC=(b>148?1024:QUADC) L2768; + CURE L2775; DEGE L2788; DEG widens CMP to n L2791; + NBSEL L2820; IMGOFF=NBSEL L2821; smem=(NBSEL+2*CMP)*4 L2822 + LAUNCH_REG2/DEG/REG macros L2823-2847 (KPT ladder 1/2/4; DEG forces KPT=1, + CUR=CURE both places); IMGW/smi/IMGE L2852-2854; + LAUNCH_REGIMG L2861-2863 -> gvr_topk_reg<...,KPT=1,CUR=true, + DEG=false,IMG=true,NBH=2*NB> via launch_regimg L2672-2686 + n4 rungs n4<=256 L2864; n4<=512 L2865; n4<=1024 wide/img/else L2866-2884 + clustered reg path L2897-2940: gate n4>4096 && n4<=8*BLKC*4 && k<=BLKC L2897; + av/amax L2898-2899; two-pass cs=8 co-residency veto + (pass==0 && c==8 && b>15 -> skip) L2917-2923; 64-bit product + (long long)c*BLKC*v < n4 L2919; smc=(3*NB+2*CMPC)*4 L2926; + grid dim3(cs,b) L2666 + wide 4k fallback n4<=4096 && wide -> LAUNCH_REG(1024,4,1,2*NB) L2945-2947 + streaming R L2959-2975 (b<=32: R=min(148/b, ((n>>2)+1023)/1024)); + r11 shallow split b<=74 && n4>=16384 && k<=1024 -> R=2 L2985; + cluster clamp R->pow2, useclus, only if 2<=R<=8 && k<=1024 L2994 + big/SCAP/CMP big=(b*R<=148) L2995; SCAP L3009-3010; CMP L3011 + aim L3039-3040; sqrt floor r=int(0.5+sqrt(6LL*n)) L3041-3042; + SFAC L3072-3073; amin L3079-3080; clamps L3081-3082 + sample geometry small_dense gate L3091 ((k>1024)&&!big&&n<=SCAP&&n>2*k); + PAIR form (sel>>3, half=n4s>>1, SMP*8) L3092-3109; + clus QUAD override (sel>>4, quarter=n4s>>2, SMP*16), + gated n>SCAP only, L3115-3128 + Q Q=(n4s+R-1)/R L3110 + clus launch smc=SNB*8+(SCAP+4)*8+CMP*8 L3130; U ladder per=Q>>10 + L3132-3142; CS=R in {2,4,8} L3143-3145; grid dim3(CS,b) L2704 + main launch smem=(SCAP+4)*((R>1||b<=296)?8:4)+(CMP+1)*8 L3149; + KPT ladder 1/2/4/8 L3150-3169; big: per=Q>>10 U ladder, + SPLIT=(R>1), grid dim3(R,b) L3173-3185 + L2750; + b<=296 -> (512,2,8,false) L3193; else (256,4,8,false) L3194 + +rt carries the FULL runtime scalar list each kernel receives, in signature +order, always starting with (n, npad, k) -- every launch site passes them +(L2666-2667 reg_clus, L2684-2685 regimg, L2704-2705 clus, L2726-2727 reg, +L2750-2751 main). [dispatch x-check 2026-08-13: rt previously omitted the +leading n/npad/k; fixed for full-ABI parity with the independent spec +transcription.] + +Dead ABI-parity args: gvr_main's 7th/8th params are declared `int SCAP_, int CMP_` +(kernel.cu L381) and are NEVER read by the kernel body -- it recomputes SCPB/CMPB +as constexprs of (BLK, SPLIT, KBIG) (L413-424) that mirror the host formulas +bit-identically. They are kept in rt under their source names 'SCAP_'/'CMP_' +purely for ABI parity. gvr_clus's SCAP/CMP (L1798) are LIVE runtime args. +`aim` and `SFAC` are host-side intermediates only (never cross the ABI), so they +do not appear in rt. + +C-semantics notes encoded here: + * every `/` on ints is C truncating division -> Python `//` (all operands + are non-negative on every reachable path); + * `sel = (long long)SFAC * n / aim` and the TGT/TGT2 products are 64-bit in C; + Python ints are exact, so `//` reproduces them; + * `int r = (int)(0.5 + sqrt((double)(6LL*n)))` truncates toward zero after + the +0.5 -> `int(0.5 + math.sqrt(float(6*n)))`; + * `IMGW = (n + 3) & ~3` four-element float4 round-up; + * the reg-block CMP (possibly widened to n by DEGE) is scoped to the braces + at L2758-2949; the streaming path re-derives its own CMP. +""" + + +# ---- constants lifted from kernel.cu --------------------------------------- +NB = 1024 # L16 register-path histogram bins +QUADC = 96 # L21 crossing-bin O(mc^2) rank gate (streaming/reg paths) +SNB = 256 # L170 streaming-path bin count +CMPC = 4096 # L2372 crossing-bin slots per CTA, clustered register path +BLKC = 1024 # L2374 CTA size of the clustered register path + + +def route(b, n, npad, k): + """Mirror of gvr_topk_launch (kernel.cu L2754-3197). Pure. See module doc.""" + wide = b <= 148 # L2757 + + # ================= register-resident block (L2758-2949) ================= + n4 = n >> 2 # L2759 + CMP = n if n < 2560 else 2560 # L2764 + QC = 1024 if b > 148 else QUADC # L2768 + CURE = not (n < 2 * k and b > 148) # L2775 + DEGE = (n <= 3 * k) or (n <= 4 * k + 64) # L2788 + if DEGE and CMP < n: # L2791 + CMP = n + NBSEL = (2 * NB) if (n4 > 512 and not (n4 <= 1024 and not wide)) else NB # L2820 + IMGOFF = NBSEL # L2821 + smem_reg = (NBSEL + 2 * CMP) * 4 # L2822 + + def _reg(BLK, VPT, MINB, NBH): + # LAUNCH_REG (L2844-2847): DEG wins, else CUR flag; KPT ladder L2823-2834. + if DEGE: + tpl = (BLK, VPT, MINB, 1, CURE, True, False, NBH) # LAUNCH_DEG L2836-2843 + else: + kpt = 1 if k <= BLK else (2 if k <= 2 * BLK else 4) + tpl = (BLK, VPT, MINB, kpt, CURE, False, False, NBH) + return { + "kernel": "reg", + "tpl": tpl, + "rt": { + "n": n, + "npad": npad, + "k": k, # L2726-2727 full ABI + "CMP": CMP, + "IMGOFF": IMGOFF, + "QC": QC, + }, + "grid": (b, 1), + "cluster": 1, + "block": BLK, + "smem": smem_reg, + "ws": False, + } + + IMGW = (n + 3) & ~3 # L2852 + smi = (NBSEL + (2 * CMP if 2 * CMP > IMGW else IMGW)) * 4 # L2853 + IMGE = wide and (not DEGE) and k <= 1024 # L2854 + + if n4 <= 256: # L2864 + return _reg(256, 1, 8, NB) + if n4 <= 512: # L2865 + return _reg(512, 1, 4, NB) + if n4 <= 1024: # L2866-2884 + if wide: + if IMGE: # LAUNCH_REGIMG(1024,1,2) L2872 + # launch_regimg<1024,1,2,NBV=2*NB,KPTV=1> -> gvr_topk_reg + # <1024,1,2,1,true,false,true,2048> (L2672-2686) + return { + "kernel": "regimg", + "tpl": (1024, 1, 2, 1, True, False, True, 2 * NB), + "rt": { + "n": n, + "npad": npad, + "k": k, # L2684-2685 full ABI + "CMP": CMP, + "IMGOFF": IMGOFF, + "QC": QC, + }, + "grid": (b, 1), + "cluster": 1, + "block": 1024, + "smem": smi, + "ws": False, + } + return _reg(1024, 1, 2, 2 * NB) # L2872 else-arm + return _reg(512, 2, 4, NB) # L2883 + + # ---- clustered register-resident path (L2897-2940) ---- + if n4 > 4096 and n4 <= 8 * BLKC * 4 and k <= BLKC: # L2897 + av = 148 // (b if b > 0 else 1) # L2898 truncating + amax = 1 # L2899 + while (amax << 1) <= av and amax < 8: + amax <<= 1 + vsel = 0 + cs = 0 + if amax >= 2: # L2901 + # knife5 (layer 9): UNCONDITIONAL cs=8 co-residency veto -- + # the L2w pass-1 rescue is deleted; 512k b>15 falls through to + # streaming, made retry-safe by TSH-floor staging (S1) and the + # gvr_clus veto (S2). + for v in (1, 2, 4): + c = 1 # 64-bit product + while c * BLKC * v < n4: + c <<= 1 + if c == 8 and b > 15: # THE VETO + continue + if c <= amax: + vsel = v + cs = c + break + if vsel and cs >= 2: # L2925 + smc = (3 * NB + 2 * CMPC) * 4 # L2926 + return { + "kernel": "reg_clus", + "tpl": (BLKC, vsel, cs), + "rt": {"n": n, "npad": npad, "k": k}, # dims only, L2666-2667 + "grid": (cs, b), + "cluster": cs, + "block": BLKC, + "smem": smc, + "ws": False, + } + + if n4 <= 4096 and wide: # L2945-2947 + return _reg(1024, 4, 1, 2 * NB) + + # ================= streaming / collect path (L2950-3196) ================= + R = 1 # L2959 + if b <= 32: # L2960-2975 + r1 = 148 // b + if r1 < 1: + r1 = 1 + r2 = ((n >> 2) + 1023) // 1024 # L2972 + if r2 < 1: + r2 = 1 + R = r1 if r1 < r2 else r2 + if R < 1: + R = 1 + elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: # L2985 r11 split + R = 2 + + useclus = False # L2993-2994 + if 2 <= R <= 8 and k <= 1024: + p2 = 1 + while (p2 << 1) <= R: + p2 <<= 1 + # knife5 (layer 8): gvr_clus cs=8 hits the same GPC packing wall as + # the clustered register path; same veto, same b>15 threshold. + if p2 == 8 and b > 15: + p2 = 4 + R = p2 + useclus = True + + big = b * R <= 148 # L2995 + SCAP = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) # L3009-3010 + CMP = (4096 if k > 1024 else 2048) if big else 1024 # L3011 + + aim = ( + ((4 * k if k >= 1024 else 2 * k) if R == 1 else 2 * k) + if big + else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) # L3039-3040 + q = 6 * n # L3041: 6LL * n + r = int(0.5 + math.sqrt(float(q))) # L3041 C cast trunc + if r > aim: # L3042 + aim = r + SFAC = ( + (32 if R == 2 else (48 if k > 1024 else 16)) if R > 1 else (64 if k >= 1024 else 32) + ) # L3072-3073 + amin = 3 * k if R == 2 else (7 * k) // 2 # L3079 + if R > 1 and aim < amin: # L3080 + aim = amin + if aim > (SCAP >> 1): # L3081 + aim = SCAP >> 1 + if aim < k: # L3082 + aim = k + + n4s = n >> 2 # L3084 + SMP, SS2, TGT, TGT2 = 0, 1, 0, 0 # L3085 + small_dense = (k > 1024) and (not big) and n <= SCAP and n > 2 * k # L3091 + if (n > SCAP or small_dense) and n4s >= 4: # L3092: PAIR sample + sel = SFAC * n // aim # L3095 64-bit + if sel < 256: # L3096 + sel = 256 + if sel > n // 2: # L3097 + sel = n // 2 + pairs = sel >> 3 # L3098 + if pairs < 1: + pairs = 1 + half = n4s >> 1 # L3099 + if half < 1: + half = 1 + if pairs > half: # L3100 + pairs = half + SS2 = half // pairs # L3101 + if SS2 < 1: + SS2 = 1 + SMP = half // SS2 # L3102 + if SMP < 1: + SMP = 1 + TGT = (aim * (SMP * 8)) // n # L3103 64-bit + if TGT < 1: # L3104 + TGT = 1 + TGT2 = (k * (SMP * 8)) // n # L3107 64-bit + if TGT2 < 1: # L3108 + TGT2 = 1 + Q = (n4s + R - 1) // R # L3110 + + if useclus: # L3111-3147 + if n > SCAP and n4s >= 4: # L3115: QUAD override + sel = SFAC * n // aim # L3116 + if sel < 256: + sel = 256 + if sel > n // 2: + sel = n // 2 + quads = sel >> 4 # L3119 + if quads < 1: + quads = 1 + quarter = n4s >> 2 # L3120 + if quarter < 1: + quarter = 1 + if quads > quarter: # L3121 + quads = quarter + SS2 = quarter // quads # L3122 + if SS2 < 1: + SS2 = 1 + SMP = quarter // SS2 # L3123 + if SMP < 1: + SMP = 1 + TGT = (aim * (SMP * 16)) // n # L3124 + if TGT < 1: + TGT = 1 + TGT2 = (k * (SMP * 16)) // n # L3126 + if TGT2 < 1: + TGT2 = 1 + smc = SNB * 8 + (SCAP + 4) * 8 + CMP * 8 # L3130 + per = Q >> 10 # L3131 + U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) # L3134-3141 + CS = 2 if R == 2 else (4 if R == 4 else 8) # L3143-3145 + return { + "kernel": "clus", + "tpl": (1024, U, 1, SNB, CS), + "rt": { + "n": n, + "npad": npad, + "k": k, # L2704-2705 ABI (live) + "SCAP": SCAP, + "CMP": CMP, + "SMP": SMP, + "TGT": TGT, + "Q": Q, + "SS2": SS2, + "TGT2": TGT2, + }, + "grid": (CS, b), + "cluster": CS, + "block": 1024, + "smem": smc, + "ws": False, + } + + smem_main = (SCAP + 4) * (8 if (R > 1 or b <= 296) else 4) + (CMP + 1) * 8 # L3149 + + def _main(BLK, MINB, U, SPLIT): + # LAUNCH_MAIN KPT ladder 1/2/4/8 (L3150-3169); grid dim3(gx=R, gy=b) L2750. + kpt = 1 if k <= BLK else (2 if k <= 2 * BLK else (4 if k <= 4 * BLK else 8)) + # knife5 (layer 7) TSH-floor staging gate. CUDA form: grid-uniform + # RUNTIME gate gridDim.y > 15 && k <= 1024 && (n >> 2) <= 32768 with + # a dual scan-instantiation branch. Here: compile-time key -- the + # ungated variant IS the pre-knife5 kernel; per-launch semantics are + # identical because the gate is uniform over the grid. + tshg = bool(SPLIT) and b > 15 and k <= 1024 and (n >> 2) <= 32768 + return { + "kernel": "main", + "tpl": (BLK, U, MINB, SNB, kpt, SPLIT, tshg), + # SCAP_/CMP_ are DEAD ABI-parity args: gvr_main (L381) never reads + # them, it uses constexpr SCPB/CMPB (L413-424). Kept for ABI parity. + "rt": { + "n": n, + "npad": npad, + "k": k, # L2750-2751 full ABI + "SCAP_": SCAP, + "CMP_": CMP, + "R": R, + "SMP": SMP, + "TGT": TGT, + "Q": Q, + "SS2": SS2, + "TGT2": TGT2, + }, + "grid": (R, b), + "cluster": 1, + "block": BLK, + "smem": smem_main, + "ws": True, + } + + if big: # L3173-3185 + per = Q >> 10 # L3174 + U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) + return _main(1024, 1, U, R > 1) # SPLIT iff R>1 + if b <= 296: # L3193 + return _main(512, 2, 8, False) + return _main(256, 4, 8, False) # L3194 + + +if __name__ == "__main__": + smoke = [ + # (b, n, npad, k) expected family + (64, 1024, 1024, 512), # reg n4<=256 rung (DEG: n<=3k) + (64, 2048, 2048, 512), # reg n4<=512 rung + (1024, 4096, 4096, 1024), # reg n4<=1024, b>148 -> (512,2,4) + (64, 4096, 4096, 512), # regimg wide !DEGE k<=1024 + (64, 4096, 4096, 1024), # reg wide but DEGE (n<=4k+64) + (8, 65536, 65536, 1024), # reg_clus (vsel=2, cs=8; b<=15 no veto) + (16, 131072, 131072, 512), # knife5: veto fall-through -> SPLIT slab, tshg=True + (64, 16384, 16384, 1024), # reg wide 4k fallback (1024,4,1) + (64, 262144, 262144, 1024), # clus r11 R=2 shallow cluster split + (1, 1048576, 1048576, 1024), # main deep slab SPLIT R=148 + (20, 262144, 262144, 2048), # main k>1024 split (no useclus) + (512, 131072, 131072, 1024), # main b>296 BLK=256 + (256, 6144, 6144, 2048), # main small_dense sample gate + (256, 262144, 262144, 2048), # main v32 KBIG-domain, BLK=512 KPT=4 + ] + for shp in smoke: + print(shp, "->", route(*shp)) + +# =========================================================================== +# ==== workspace (ct_workspace.py) ========================================== +# =========================================================================== +"""op46 workspace mirror of src_cuda/main.cpp B2 (L15-37) + run_ws checks +(L107-114) and kernel.h workspace_bytes contract. + +B2 semantics mirrored exactly: + * ONE zero-initialised slab workspace per device, lazily allocated through + the torch caching allocator (main.cpp:32-33 `at::zeros(..., kByte)`); + * keep-alive store (`ws_keep[GVR_MAX_DEV]`) -> module dict `_ws_keep` + (tensor refcount = keep-alive, same as the C static array); + * double-checked locking: lock-free hot-path load (a GIL-atomic dict get + plays the `std::memory_order_acquire` load, main.cpp:26-27), slow path + re-checks under a mutex (main.cpp:28-31); + * device index bounds `0 <= d < GVR_MAX_DEV` (main.cpp:24-25) -- checked + BEFORE the CUDA-ness of the tensor, exactly like the C binding (run() + resolves the default workspace before run_impl's B1 checks, so a CPU + logits tensor dies here with "device index out of range: -1"). + +Concurrent STREAMS on one device that may both take the multi-CTA SPLIT path +must pass their own workspace via run_ws() (main.cpp:16-17). + +Size: gvr_topk_workspace_bytes() = GVR_WS_BUF_OFF + MAXC*GCAP*sizeof(int2) + = 2048 + 160*16384*8 = 20,973,568 B (kernel.cu L44-46). + +Kernel-facing view: ct_main's compiled signature takes the workspace as a +1-D contiguous int32 tensor (fake tensor dtype Int32, assumed_align=16 -- +torch caching-allocator bases are 256B-aligned so the default slab always +satisfies it). `kernel_view()` reproduces the C binding's raw +`workspace.data_ptr()` semantics for arbitrary user tensors by aliasing the +underlying storage at the tensor's byte offset. +""" + + +GVR_MAX_DEV = 64 # kernel.cu L19 / main.cpp:19 +_MAXC = 160 # kernel.cu L17 +_GCAP = 16384 # kernel.cu L18 +_GVR_WS_BUF_OFF = 2048 # kernel.cu L43 +WS_BYTES = _GVR_WS_BUF_OFF + _MAXC * _GCAP * 8 # 20,973,568 (kernel.cu L44-46) +assert WS_BYTES == 20_973_568 + +_mu = threading.Lock() # main.cpp:28 slow-path mutex +_ws_keep = {} # device index -> keep-alive int32 view + + +def workspace_bytes() -> int: + """kernel.h:12 gvr_topk_workspace_bytes().""" + return WS_BYTES + + +def default_workspace(ref: torch.Tensor) -> torch.Tensor: + """main.cpp:23-37 default_workspace(ref) -> per-device cached slab. + + Returns the kernel-facing 1-D int32 view (zero-initialised on first use; + the kernel restores the zeros it consumes, so one zeroing suffices for + the lifetime of the cache entry).""" + d = ref.get_device() + if not (0 <= d < GVR_MAX_DEV): + raise RuntimeError(f"device index out of range: {d}") + ws = _ws_keep.get(d) # hot path: one (GIL-atomic) load + if ws is not None: + return ws + with _mu: # slow path: double-checked + ws = _ws_keep.get(d) + if ws is not None: + return ws + # lazy zeros via the torch caching allocator (at::zeros kByte, + # main.cpp:32-33), viewed int32 for the DSL launch signature. + buf = torch.zeros(WS_BYTES, dtype=torch.uint8, device=ref.device) + ws = buf.view(torch.int32) + _ws_keep[d] = ws # keep-alive (ws_keep[d] = tensor) + return ws + + +def validate_run_ws(workspace: torch.Tensor, logits: torch.Tensor) -> None: + """main.cpp:107-114 run_ws() workspace hardening, same predicate order: + CUDA + same device as logits; numel*element_size >= workspace_bytes(); + base 8-byte aligned.""" + if not (workspace.is_cuda and workspace.get_device() == logits.get_device()): + raise RuntimeError("workspace must be a CUDA tensor on the same device") + if workspace.numel() * workspace.element_size() < WS_BYTES: + raise RuntimeError(f"workspace too small: need {WS_BYTES} bytes") + if workspace.data_ptr() & 7: + raise RuntimeError("workspace must be 8-byte aligned") + + +def kernel_view(workspace: torch.Tensor) -> torch.Tensor: + """Raw-pointer semantics of the C binding (main.cpp:115 passes + workspace.data_ptr() and nothing else): alias the first WS_BYTES bytes at + the tensor's data_ptr() as int32[WS_BYTES/4], ignoring dtype/shape. + + NOTE: the DSL-side fake tensor declares assumed_align=16; a workspace at + 8-but-not-16-byte alignment passes the C-contract check above but is + rejected by the DSL at conversion -- surfaced as a launch failure with + shape context by ct_op (documented in notes/ct_op_NOTES.md).""" + if ( + workspace.dtype is torch.int32 + and workspace.dim() == 1 + and workspace.is_contiguous() + and workspace.storage_offset() == 0 + and workspace.numel() == WS_BYTES // 4 + ): + return workspace # already the canonical view + off_bytes = workspace.storage_offset() * workspace.element_size() + if off_bytes & 3: + # unreachable past the 8B-alignment check for allocator-backed + # storages; kept as a hard error rather than silent misalias. + raise RuntimeError("workspace storage offset must be 4-byte aligned") + t = torch.empty(0, dtype=torch.int32, device=workspace.device) + t.set_(workspace.untyped_storage(), off_bytes // 4, (WS_BYTES // 4,)) + return t + + +def _reset_for_tests() -> None: + """Drop cached slabs (tests only; NOT part of the C contract).""" + with _mu: + _ws_keep.clear() + + +# =========================================================================== +# ==== operator entry (ct_op.py) ============================================ +# =========================================================================== +"""op46 operator entry: CuTeDSL mirror of src_cuda/main.cpp run()/run_ws()/ +workspace_bytes() (spec section 1). + +B1 hardening checks run in the SAME ORDER with the SAME PREDICATES as +main.cpp:43-88 (run_impl): + 1. all three tensors CUDA (main.cpp:43-44) + 2. dtypes: logits f32, pre_idx i32, indices i32 (45-47) + 3. all 2-D (48-49) + 4. all contiguous (50-51) + 5. n_valid unwrap (57-67): python-int fast path (strict integral cast, like + pybind cast); Tensor path checks + torch.cuda.is_current_stream_capturing() FIRST and fails loudly (B1d), + else .item() (the D2H sync) + 6. b/npad from logits, k = pre_idx.size(1) (68-70) + 7. b == 0 -> early no-op (71, B1f) + 8. npad % 4 == 0 (74-75, B1e float4 row loads) + 9. logits base 16-byte aligned (76-78) + 10. pre_idx/indices batch dims match (79-81) + 11. indices width >= k (84-85) + 12. n_valid >= 0 (86) + 13. n = min(nv, npad) clamped in unbounded ints BEFORE any narrowing (88) + +Dispatch: ct_dispatch.route(b, n, npad, k) -> compile-cache keyed on +(kernel family, constexpr tuple) inside each family module -> BIND-ONCE +launch cache keyed on the shape key (b, n, npad, k): caches the compiled +callable + the prebuilt runtime-scalar arg pack as plain Python ints (probe +P12: plain ints, never pre-wrapped cutlass.Int32; pre-binding removes only +route()/marshal-prep work -- the tvm-ffi per-argument cost is paid every +call). Hot enqueue target ~3-6 us (P12 arg-width tax); measured numbers in +notes/ct_op_NOTES.md. + +Error contract (spec 1.4): launch failures surface as exceptions WITH +(b, n, npad, k) context, mirroring main.cpp:94-95. + +All four family modules are imported LAZILY (first shape that routes to +them), so a missing/broken sibling only fails when actually routed to, with +(b, n, npad, k) context. Wired compiled ABIs (verified against each +module's __call__ signature): + ct_reg (logits, pre_idx, out, n, CMP, QC, smem_bytes) + ct_main (logits, pre_idx, out, ws, n, npad, k, SCAP_, CMP_, R, SMP, + TGT, Q, SS2, TGT2) [only family taking workspace] + ct_clus (logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, + SS2, TGT2) [get_compiled keyed +scap/cmp_] + ct_regclus (logits, pre_idx, out, n) +""" + + +# shape key (b, n, npad, k) -> (fn, args tuple of python ints, needs_ws) +_LAUNCH_CACHE = {} + +# hot-path local bindings (each torch. lookup costs ~0.1 us; the B1 +# battery runs on EVERY call — mirror of main.cpp's "sub-100ns predicted +# branches" intent within Python's reach; measured in notes/ct_op_NOTES.md) +_F32 = torch.float32 +_I32 = torch.int32 +_TENSOR = torch.Tensor +_is_capturing = torch.cuda.is_current_stream_capturing +_index = operator.index +_ws_hot = _ws_keep # shared dict object (hot-path load) +_GVR_MAX_DEV = GVR_MAX_DEV + + +# --------------------------------------------------------------------------- +# per-family launcher builders (cold path: once per distinct shape key) +# --------------------------------------------------------------------------- +def _build_launcher(b, n, npad, k): + rd = route(b, n, npad, k) + fam = rd["kernel"] + tpl = tuple(rd["tpl"]) + rt = rd["rt"] + if fam in ("reg", "regimg"): + dev = _device() + fn = dev.get_compiled__reg(tpl) + # compiled ABI: (logits, pre_idx, out, n, CMP, QC, smem_total) + args = (rt["n"], rt["CMP"], rt["QC"], dev.STATIC_BYTES + rd["smem"]) + return (fn, args, False) + if fam == "main": + dev = _device() + fn = dev.get_compiled(tpl) + # compiled ABI: (logits, pre_idx, out, ws, n, npad, k, SCAP_, CMP_, + # R, SMP, TGT, Q, SS2, TGT2) [SCAP_/CMP_ dead, ABI parity] + args = ( + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP_"], + rt["CMP_"], + rt["R"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + return (fn, args, True) + if fam == "clus": + dev = _device() + # compile key carries the smem-extent scalars (scap/cmp_); compiled + # ABI: (logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, + # SS2, TGT2) -- NO workspace (spec §4c) + fn = dev.get_compiled__clus(tpl, scap=rt["SCAP"], cmp_=rt["CMP"]) + args = ( + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP"], + rt["CMP"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + return (fn, args, False) + if fam == "reg_clus": + dev = _device() + # compiled ABI: (logits, pre_idx, out, n) -- smem/k derived in-module + fn = dev.get_compiled__regclus(tpl) + return (fn, (rt["n"],), False) + # unreachable: route() only emits the five families above + raise RuntimeError(f"unknown dispatch family {fam!r}") + + +# --------------------------------------------------------------------------- +# run_impl mirror (main.cpp:39-96) +# --------------------------------------------------------------------------- +def _run_impl(logits, pre_idx, n_valid, indices, ws): + if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if logits.dtype is not _F32: + raise RuntimeError("logits must be float32") + if pre_idx.dtype is not _I32: + raise RuntimeError("pre_idx must be int32") + if indices.dtype is not _I32: + raise RuntimeError("indices must be int32") + lsh, psh, ish = logits.shape, pre_idx.shape, indices.shape + if not (len(lsh) == 2 and len(psh) == 2 and len(ish) == 2): + raise RuntimeError("logits/pre_idx/indices must be 2-D") + if not (logits.is_contiguous() and pre_idx.is_contiguous() and indices.is_contiguous()): + raise RuntimeError("tensors must be contiguous") + + # n_valid unwrap (main.cpp:57-67): tensor path = D2H sync, illegal under + # CUDA graph capture -- fail loudly instead of crashing the capture (B1d). + if isinstance(n_valid, _TENSOR): + if _is_capturing(): + raise RuntimeError( + "tensor n_valid requires a D2H sync, illegal under CUDA " + "graph capture — pass n_valid as a python int" + ) + nv = int(n_valid.item()) + else: + # strict integral cast (pybind cast rejects floats/strings) + nv = _index(n_valid) + + b, npad = lsh + k = psh[1] + if b == 0: # empty batch: no-op (main.cpp:71, B1f) + return + if npad & 3: + raise RuntimeError(f"npad (logits stride) must be a multiple of 4, got {npad}") + if logits.data_ptr() & 15: + raise RuntimeError( + "logits base must be 16-byte aligned (storage-offset views break the float4 row loads)" + ) + if psh[0] != b or ish[0] != b: + raise RuntimeError(f"batch dims must match: logits {b} pre_idx {psh[0]} indices {ish[0]}") + if ish[1] < k: + raise RuntimeError(f"indices width {ish[1]} < k={k} (k is pre_idx.size(1))") + if nv < 0: + raise RuntimeError(f"n_valid must be non-negative, got {nv}") + # clamp BEFORE any narrowing (main.cpp:87-88; python ints are unbounded, + # so min() is the exact 64-bit clamp) + n = nv if nv < npad else npad + + # CUDA out-indexing mirror: every kernel derives O = out + row*k + # (kernel.cu L475/L1309 etc.) -- flat PACKED rows, ignoring the actual + # indices width. The DSL kernels index out[row, :] with the tensor's own + # row stride, so a wider `indices` must be re-viewed packed (pure view, + # no copy; contiguity already checked). + if ish[1] != k: + indices = indices.reshape(-1)[: b * k].view(b, k) + + key = (b, n, npad, k) + lc = _LAUNCH_CACHE.get(key) + try: + if lc is None: + lc = _build_launcher(b, n, npad, k) + _LAUNCH_CACHE[key] = lc + fn, args, needs_ws = lc + if needs_ws: + fn(logits, pre_idx, indices, ws, *args) + else: + fn(logits, pre_idx, indices, *args) + except Exception as e: + raise RuntimeError(f"gvr_topk launch failed (b={b} n={n} npad={npad} k={k}): {e}") from e + + +# --------------------------------------------------------------------------- +# exports (main.cpp:98-124) +# --------------------------------------------------------------------------- +def run(logits, pre_idx, n_valid, indices): + """Fast 4-arg form: signature-identical to the original candidate. + Default per-device slab workspace resolved FIRST (main.cpp:99-102 -- + a CPU logits tensor therefore dies with 'device index out of range'). + Hot path inlines the C binding's check + atomic-load + cache-hit + (main.cpp:24-27); the slow path allocates under ct_workspace's lock.""" + d = logits.get_device() + if not 0 <= d < _GVR_MAX_DEV: # main.cpp:25, EVERY call + raise RuntimeError(f"device index out of range: {d}") + ws = _ws_hot.get(d) + if ws is None: + ws = default_workspace(logits) + _run_impl(logits, pre_idx, n_valid, indices, ws) + + +def run_ws(logits, pre_idx, n_valid, indices, workspace): + """Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" + validate_run_ws(workspace, logits) + _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace)) + + +__all__ = [ + "route", + "run", + "run_ws", + "workspace_bytes", + "WS_BYTES", + "default_workspace", + "validate_run_ws", + "kernel_view", +] diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py new file mode 100644 index 000000000000..bebde7a058e3 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -0,0 +1,143 @@ +# 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. +"""Exactness tests for the standalone self-sampling GVR top-K decode kernels +(`gvr_topk_decode_self_sampling[_host].py`). + +Contract under test (see the host module docstring): batch-uniform host-int +``n_valid`` in compressed index space, fp32 logits with a 64-element-multiple +row stride, output = exact (tie-interchangeable) top-K indices of +``logits[:, :n_valid]`` per row. + +Checks per case: + - tie-aware exactness: the multiset of gathered output values equals the + ``torch.topk`` value multiset bitwise (signed zeros normalized); + - output indices are unique and within ``[0, n_valid)``; + - padding immunity: the padded tail ``[n_valid, npad)`` is filled with huge + values — any kernel read past ``n_valid`` fails the value comparison. +""" + +import pytest +import torch +from utils.util import getSMVersion + +import tensorrt_llm # noqa: F401 +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + +if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gvr_selfsampling_topk tests", allow_module_level=True) + +if not IS_CUTLASS_DSL_AVAILABLE: + pytest.skip("cutlass DSL is required for gvr_selfsampling_topk tests", allow_module_level=True) + +if getSMVersion() != 100: + pytest.skip("self-sampling GVR kernels target Blackwell sm_100", allow_module_level=True) + +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling_host as ss_host, +) + +_DEV = "cuda" + + +def _make_case(batch_size, n_valid, top_k, seed, hit_ratio=0.6): + """Decode-like fp32 logits + prev-step hint. The padded tail is poisoned + with +3e38 so any read past n_valid corrupts the top-K values.""" + gen = torch.Generator(device=_DEV).manual_seed(seed) + npad = (n_valid + 63) // 64 * 64 + logits = torch.randn((batch_size, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + logits[:, n_valid:] = 3e38 + + ref_vals, ref_idx = torch.topk(logits[:, :n_valid].float(), top_k, dim=1, largest=True) + # hint: argmax first (anchor), then a hit_ratio slice of the true top-K, + # the rest random valid indices — mirrors the decode-step temporal hint. + n_hits = int(top_k * hit_ratio) + rand_fill = torch.randint( + 0, n_valid, (batch_size, top_k), generator=gen, dtype=torch.int32, device=_DEV + ) + pre_idx = rand_fill.clone() + pre_idx[:, :n_hits] = ref_idx[:, :n_hits].to(torch.int32) + indices = torch.full((batch_size, top_k), -1, dtype=torch.int32, device=_DEV) + return logits, pre_idx, indices, ref_vals + + +def _check_exact(logits, indices, n_valid, ref_vals): + top_k = indices.shape[1] + idx64 = indices.to(torch.int64) + assert int(idx64.min()) >= 0, "negative output index" + assert int(idx64.max()) < n_valid, "output index past n_valid" + for row in range(indices.shape[0]): + assert int(torch.unique(idx64[row]).numel()) == top_k, f"row {row}: duplicate indices" + got = torch.gather(logits, 1, idx64) + # +0.0 maps -0.0 to +0.0 so signed zeros compare equal bitwise + got_sorted = torch.sort(got + 0.0, dim=1, descending=True).values + ref_sorted = torch.sort(ref_vals + 0.0, dim=1, descending=True).values + assert torch.equal(got_sorted, ref_sorted), ( + "top-K value multiset mismatch (inexact or padding read)" + ) + + +# (top_k, n_valid) — gate-edge (131075/131076 straddle the K=2048 hint-band +# gate), the small-N floor, the mid band, and the deployment-envelope top. +_CASES = [ + (512, 4099), + (512, 65536), + (512, 262143), + (1024, 16387), + (1024, 131072), + (2048, 4111), + (2048, 131075), + (2048, 131076), + (2048, 262144), +] + + +@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda b: f"bs{b}") +@pytest.mark.parametrize("top_k,n_valid", _CASES, ids=[f"k{k}_n{n}" for k, n in _CASES]) +def test_selfsampling_topk_exactness(batch_size, top_k, n_valid): + logits, pre_idx, indices, ref_vals = _make_case( + batch_size, n_valid, top_k, seed=n_valid * 31 + top_k + batch_size + ) + ss_host.run(logits, pre_idx, n_valid, indices) + torch.cuda.synchronize() + _check_exact(logits, indices, n_valid, ref_vals) + + +def test_selfsampling_topk_run_ws_explicit_workspace(): + """run_ws with a caller-owned workspace must agree with run().""" + top_k, n_valid = 1024, 65536 + logits, pre_idx, indices, ref_vals = _make_case(2, n_valid, top_k, seed=7) + ws = torch.zeros(ss_host.workspace_bytes(), dtype=torch.uint8, device=_DEV) + ss_host.run_ws(logits, pre_idx, n_valid, indices, ws) + torch.cuda.synchronize() + _check_exact(logits, indices, n_valid, ref_vals) + + +def test_selfsampling_topk_guards(): + logits, pre_idx, indices, _ = _make_case(1, 8192, 512, seed=3) + with pytest.raises(RuntimeError, match="float32"): + ss_host.run(logits.to(torch.bfloat16), pre_idx, 8192, indices) + with pytest.raises(RuntimeError, match="non-negative"): + ss_host.run(logits, pre_idx, -1, indices) + with pytest.raises(RuntimeError, match="batch dims"): + ss_host.run(logits, pre_idx[:0], 8192, indices) + + +def test_selfsampling_dispatch_is_pure_and_total(): + """route(b, n, npad, k) must return a plan for every in-envelope shape.""" + for k in (512, 1024, 2048): + for n in (k + 1, 4111, 65536, 131075, 131076, 262144): + npad = (n + 63) // 64 * 64 + r = ss_host.route(4, n, npad, k) + assert r["kernel"] in ("main", "reg", "clus", "reg_clus") + assert r["block"] >= 128 and r["grid"][0] >= 1 From 29354d11494c58891e7926930c5b4f6c2ace89c9 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:02:30 +0000 Subject: [PATCH 02/40] =?UTF-8?q?[TRTLLM-15293][perf]=20self-sampling=20to?= =?UTF-8?q?p-K:=20fix=20pre=5Fidx=20hint=20contract=20=E2=80=94=20raw=20in?= =?UTF-8?q?dices,=20uniform=20across=20DSv3.2/Flash/Pro?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The contract docstring previously instructed callers to apply the cr==1 +1 temporal shift to pre_idx (mirroring heuristicTopKDecode.cu). Drop it: hints only steer the sampling ladder — exactness never depends on them — and on real V3.2 decode captures raw prev-step hints overlap the current top-K at 0.773 vs 0.536 when +1-shifted (15 cells x 14 consecutive step-pairs, gap widening with ISL). One offset-free hint convention now serves all three models; the kernels already consume pre_idx as-is, so this is a contract-documentation fix only. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling_host.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index bb7aef06258c..399f9a0297e4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -31,8 +31,14 @@ OPERATOR CONTRACT (standalone; not wired into the decode path): ``n_valid`` is one host python int for the whole batch — every row shares the same valid prefix, in COMPRESSED index space (the caller applies any -``compressRatio`` division and, for DSv3.2 (cr=1), the +1 temporal shift of -``pre_idx``; see ``heuristicTopKDecode.cu``). The production decode engine +``compressRatio`` division). ``pre_idx`` is consumed AS-IS — raw prev-step +top-K indices, uniformly for DSv3.2 / DSv4 Flash / Pro. This deliberately +drops the +1 temporal shift ``heuristicTopKDecode.cu`` applies for cr==1: +hints only steer the sampling ladder (exactness never depends on them), and +on real V3.2 decode captures raw prev-step hints land on MORE of the current +top-K than +1-shifted ones (mean overlap 0.773 vs 0.536 across 15 cells x 14 +consecutive step-pairs, the gap widening with ISL), so one offset-free hint +convention serves all three models. The production decode engine instead reads per-request ``seq_lens`` on-device with per-row MTP offsets (sync-free, CUDA-graph-replay safe with growing KV); adopting that per-row contract inside these kernels is tracked follow-up work. Until then this From b1c5e674e67ff897c0bcc6f3ae69d497699b5ceb Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:04:51 +0000 Subject: [PATCH 03/40] [TRTLLM-15293][perf] self-sampling top-K: n <= topK short path (production pad convention) When the row has no more than topK valid entries every valid position is in the top-K: emit identity indices and pad the tail with -1, mirroring heuristicTopKDecode.cu:72-84. Host-level torch-op branch for the standalone module (the CUDA-graph-safe per-row rewrite will move it in-kernel, where a per-row fallback is impossible inside a graph). Closes the 'n <= topK unproven' gap from the integration audit. Tests: 8 boundary shapes (n in {64..2048}, k in {512,1024,2048}, n < k / n == k-1 / n == k) x bs {1,4} with poisoned padding, plus kernel-path regression just above the boundary (k512_n4099, k2048_n4111) verified exact on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 13 +++++++ .../parallel/test_gvr_selfsampling_topk.py | 39 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 399f9a0297e4..0a561f2528ce 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -766,6 +766,19 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws): if ish[1] != k: indices = indices.reshape(-1)[: b * k].view(b, k) + # ---- n <= k short path (heuristicTopKDecode.cu:72-84) ------------------- + # Every valid position is in the top-K: emit identity indices and pad the + # tail with -1 (the production pad convention; downstream treats -1 as + # invalid). Order is contract-irrelevant — exactness is tie-interchangeable + # SET semantics. Torch-op path for now; the CUDA-graph-safe per-row rewrite + # moves this branch in-kernel (it cannot fall back per row inside a graph). + if n <= k: + if n > 0: + indices[:, :n] = torch.arange(n, dtype=_I32, device=indices.device) + if n < k: + indices[:, n:] = -1 + return + key = (b, n, npad, k) lc = _LAUNCH_CACHE.get(key) try: diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index bebde7a058e3..784357f01e42 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -113,6 +113,45 @@ def test_selfsampling_topk_exactness(batch_size, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) +# n_valid <= top_k: every valid position is in the top-K. Production short +# path (heuristicTopKDecode.cu:72-84): identity indices + -1 tail padding. +_SHORT_CASES = [ + (512, 256), + (512, 511), + (512, 512), + (1024, 64), + (1024, 1024), + (2048, 1000), + (2048, 2047), + (2048, 2048), +] + + +@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda b: f"bs{b}") +@pytest.mark.parametrize( + "top_k,n_valid", _SHORT_CASES, ids=[f"k{k}_n{n}" for k, n in _SHORT_CASES] +) +def test_selfsampling_topk_short_path(batch_size, top_k, n_valid): + gen = torch.Generator(device=_DEV).manual_seed(top_k + n_valid) + npad = (n_valid + 63) // 64 * 64 + logits = torch.randn((batch_size, npad), generator=gen, dtype=torch.float32, device=_DEV) + logits[:, n_valid:] = 3e38 # poison pad: the short path must never read it + pre_idx = torch.randint( + 0, n_valid, (batch_size, top_k), generator=gen, dtype=torch.int32, device=_DEV + ) + indices = torch.full((batch_size, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run(logits, pre_idx, n_valid, indices) + torch.cuda.synchronize() + head = indices[:, :n_valid].to(torch.int64) + expect = torch.arange(n_valid, dtype=torch.int64, device=_DEV).expand(batch_size, n_valid) + assert torch.equal(torch.sort(head, dim=1).values, expect), "short-path head not {0..n-1}" + if n_valid < top_k: + assert torch.equal( + indices[:, n_valid:], + torch.full((batch_size, top_k - n_valid), -1, dtype=torch.int32, device=_DEV), + ), "short-path tail must be -1 padded" + + def test_selfsampling_topk_run_ws_explicit_workspace(): """run_ws with a caller-owned workspace must agree with run().""" top_k, n_valid = 1024, 65536 From 3953e906cf7123f5400e5e3078e019d94bce9dba Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:11:34 +0000 Subject: [PATCH 04/40] [TRTLLM-15293][perf] self-sampling top-K: opt-in values output (default off) Optional `values` DPS output on run()/run_ws(), production parity for the heuristicTopKDecode values writeback. Default None = OFF, matching dsa.py, which allocates the values scratch only for the non-CuTeDSL path. The indices are exact, so a gather epilogue reproduces the in-kernel writeback bit-for-bit at zero cost when disabled; the constexpr in-kernel form rides the CUDA-graph per-row rewrite. Short path pads values with -FLT_MAX (production convention). Tests: kernel path (values == gathered top-K == torch.topk multiset), short path (head copies logits, -FLT_MAX tail), wide-buffer packed re-view and dtype guard, verified on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 38 ++++++++++++++++--- .../parallel/test_gvr_selfsampling_topk.py | 36 ++++++++++++++++++ 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 0a561f2528ce..203a7acc500d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -710,7 +710,7 @@ def _build_launcher(b, n, npad, k): # --------------------------------------------------------------------------- # run_impl mirror (main.cpp:39-96) # --------------------------------------------------------------------------- -def _run_impl(logits, pre_idx, n_valid, indices, ws): +def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): raise RuntimeError("all tensors must be CUDA") if logits.dtype is not _F32: @@ -766,6 +766,26 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws): if ish[1] != k: indices = indices.reshape(-1)[: b * k].view(b, k) + # ---- optional values output (production parity, default OFF) ------------ + # dsa.py allocates the values scratch only for the non-CuTeDSL path, so + # values stay opt-in. The indices are exact, so a gather epilogue + # reproduces the in-kernel writeback bit-for-bit; the constexpr in-kernel + # form rides the CUDA-graph per-row rewrite. + if values is not None: + if not values.is_cuda: + raise RuntimeError("values must be CUDA") + if values.dtype is not _F32: + raise RuntimeError("values must be float32") + vsh = values.shape + if len(vsh) != 2 or not values.is_contiguous(): + raise RuntimeError("values must be 2-D contiguous") + if vsh[0] != b: + raise RuntimeError(f"batch dims must match: logits {b} values {vsh[0]}") + if vsh[1] < k: + raise RuntimeError(f"values width {vsh[1]} < k={k}") + if vsh[1] != k: + values = values.reshape(-1)[: b * k].view(b, k) + # ---- n <= k short path (heuristicTopKDecode.cu:72-84) ------------------- # Every valid position is in the top-K: emit identity indices and pad the # tail with -1 (the production pad convention; downstream treats -1 as @@ -775,8 +795,12 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws): if n <= k: if n > 0: indices[:, :n] = torch.arange(n, dtype=_I32, device=indices.device) + if values is not None: + values[:, :n] = logits[:, :n] if n < k: indices[:, n:] = -1 + if values is not None: + values[:, n:] = torch.finfo(_F32).min # -FLT_MAX pad return key = (b, n, npad, k) @@ -792,13 +816,17 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws): fn(logits, pre_idx, indices, *args) except Exception as e: raise RuntimeError(f"gvr_topk launch failed (b={b} n={n} npad={npad} k={k}): {e}") from e + if values is not None: + values.copy_(logits.gather(1, indices.to(torch.int64))) # --------------------------------------------------------------------------- # exports (main.cpp:98-124) # --------------------------------------------------------------------------- -def run(logits, pre_idx, n_valid, indices): +def run(logits, pre_idx, n_valid, indices, values=None): """Fast 4-arg form: signature-identical to the original candidate. + ``values`` (optional DPS output, default None = OFF) mirrors the + production values writeback; see _run_impl. Default per-device slab workspace resolved FIRST (main.cpp:99-102 -- a CPU logits tensor therefore dies with 'device index out of range'). Hot path inlines the C binding's check + atomic-load + cache-hit @@ -809,13 +837,13 @@ def run(logits, pre_idx, n_valid, indices): ws = _ws_hot.get(d) if ws is None: ws = default_workspace(logits) - _run_impl(logits, pre_idx, n_valid, indices, ws) + _run_impl(logits, pre_idx, n_valid, indices, ws, values) -def run_ws(logits, pre_idx, n_valid, indices, workspace): +def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): """Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" validate_run_ws(workspace, logits) - _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace)) + _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) __all__ = [ diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 784357f01e42..02ce53fca7b5 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -152,6 +152,42 @@ def test_selfsampling_topk_short_path(batch_size, top_k, n_valid): ), "short-path tail must be -1 padded" +def test_selfsampling_topk_values_output(): + """Opt-in values output (default None = off, matching dsa.py, which + allocates the values scratch only for the non-CuTeDSL path): must equal + the gathered top-K values and the torch.topk value multiset.""" + top_k, n_valid = 512, 8192 + logits, pre_idx, indices, ref_vals = _make_case(2, n_valid, top_k, seed=11) + values = torch.full((2, top_k), 7.0, dtype=torch.float32, device=_DEV) + ss_host.run(logits, pre_idx, n_valid, indices, values) + torch.cuda.synchronize() + _check_exact(logits, indices, n_valid, ref_vals) + assert torch.equal(values, torch.gather(logits, 1, indices.to(torch.int64))) + assert torch.equal( + torch.sort(values + 0.0, dim=1, descending=True).values, + torch.sort(ref_vals + 0.0, dim=1, descending=True).values, + ) + + +def test_selfsampling_topk_values_short_path(): + """Short path with values: head copies logits, tail pads with -FLT_MAX + (production heuristicTopKDecode pad convention).""" + top_k, n_valid, bs = 1024, 512, 4 + gen = torch.Generator(device=_DEV).manual_seed(42) + logits = torch.randn((bs, n_valid), generator=gen, dtype=torch.float32, device=_DEV) + pre_idx = torch.randint( + 0, n_valid, (bs, top_k), generator=gen, dtype=torch.int32, device=_DEV + ) + indices = torch.full((bs, top_k), -7, dtype=torch.int32, device=_DEV) + values = torch.full((bs, top_k), 7.0, dtype=torch.float32, device=_DEV) + ss_host.run(logits, pre_idx, n_valid, indices, values) + torch.cuda.synchronize() + assert torch.equal(values[:, :n_valid], logits) + fmin = torch.finfo(torch.float32).min + assert bool((values[:, n_valid:] == fmin).all()) + assert bool((indices[:, n_valid:] == -1).all()) + + def test_selfsampling_topk_run_ws_explicit_workspace(): """run_ws with a caller-owned workspace must agree with run().""" top_k, n_valid = 1024, 65536 From 5877b60c991f9145936877ca0ce62122cc5f0549 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:12:06 +0000 Subject: [PATCH 05/40] [TRTLLM-15293][test] self-sampling top-K: degenerate-hint hardening tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production can hand the kernel degenerate hint buffers: dsa.py initializes heuristic_prev_topk with zero_() (all-zero cold start; the prefill->decode seeding covers the common path but zero-init corners remain), and nothing forbids duplicated hints. Exactness must never depend on hint quality (hint-robustness bug class of PR #17550). Adds all-zero / all-same / all-max / half-duplicated hint cases on (k512, n8192) and the k2048 gate-edge n131075 — all verified exact on B200 (also probed on n131072/k1024 x bs{1,4} pre-commit, 24 configs, zero failures). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../parallel/test_gvr_selfsampling_topk.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 02ce53fca7b5..cbd1353e9a1e 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -188,6 +188,33 @@ def test_selfsampling_topk_values_short_path(): assert bool((indices[:, n_valid:] == -1).all()) +@pytest.mark.parametrize("hint_kind", ["all_zero", "all_same", "all_max", "half_dup"]) +@pytest.mark.parametrize( + "top_k,n_valid", [(512, 8192), (2048, 131075)], ids=["k512_n8192", "k2048_n131075"] +) +def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): + """Hints only steer the sampling ladder — exactness must survive the + degenerate hint buffers production can produce: the all-zero cold start + (dsa.py ``heuristic_prev_topk.zero_()`` init corners), fully duplicated + hints, and max-index hints (hint-robustness class of PR #17550).""" + logits, _, indices, ref_vals = _make_case(2, n_valid, top_k, seed=n_valid + top_k) + if hint_kind == "all_zero": + pre_idx = torch.zeros((2, top_k), dtype=torch.int32, device=_DEV) + elif hint_kind == "all_same": + pre_idx = torch.full((2, top_k), 1234, dtype=torch.int32, device=_DEV) + elif hint_kind == "all_max": + pre_idx = torch.full((2, top_k), n_valid - 1, dtype=torch.int32, device=_DEV) + else: + gen = torch.Generator(device=_DEV).manual_seed(1) + pre_idx = torch.randint( + 0, n_valid, (2, top_k), generator=gen, dtype=torch.int32, device=_DEV + ) + pre_idx[:, top_k // 2 :] = pre_idx[:, :1] + ss_host.run(logits, pre_idx, n_valid, indices) + torch.cuda.synchronize() + _check_exact(logits, indices, n_valid, ref_vals) + + def test_selfsampling_topk_run_ws_explicit_workspace(): """run_ws with a caller-owned workspace must agree with run().""" top_k, n_valid = 1024, 65536 From 61032328edad9c4e74dedd2fb766556963c068c2 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:23:30 +0000 Subject: [PATCH 06/40] [TRTLLM-15293][perf] self-sampling top-K: run_varlen production-contract entry (reference engine) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds run_varlen(logits, pre_idx, kv_lens, indices, next_n=, compress_ratio=, values=None) — the heuristicTopKDecode contract: per-request device kv_lens (TOTAL cache length, uncompressed token space), per-row n = (kv_len - next_n + row%next_n + 1)/compressRatio (the MTP window formula, cr 1 = DSv3.2 / 4 = DSv4 Flash+Pro), request-level raw pre_idx shared by a request's next_n rows, per-row n <= k short path. REFERENCE engine: one documented host read of kv_lens (raises under CUDA-graph capture), rows driven as b=1 launches through the batch-uniform engine. This pins the varlen/MTP contract and its test battery; the per-row in-kernel engine (device kv_lens reads, fixed-R thin-slicing, n-band parameter table) replaces the loop next without changing either. Tests: heterogeneous lengths with per-row poisoned padding, cr in {1,4} x next_n in {1,2,4} incl. in-request n variation and compressed-boundary-crossing rows, mixed short rows, values output, contract guards — all verified exact on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/__init__.py | 2 + .../gvr_topk_decode_self_sampling_host.py | 76 ++++++++++++++++++ .../parallel/test_gvr_selfsampling_topk.py | 77 +++++++++++++++++++ 3 files changed, 155 insertions(+) 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 507662eaa0c0..5e2cf2001e74 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 @@ -22,6 +22,7 @@ from .gvr_topk_decode_reg import GvrRegKernel from .gvr_topk_decode_self_sampling_host import route as selfsampling_topk_route from .gvr_topk_decode_self_sampling_host import run as selfsampling_topk_run +from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_self_sampling_host import run_ws as selfsampling_topk_run_ws from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -39,5 +40,6 @@ "is_tiered_topk_supported", "selfsampling_topk_route", "selfsampling_topk_run", + "selfsampling_topk_run_varlen", "selfsampling_topk_run_ws", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 203a7acc500d..f4e9b7c763a2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -846,10 +846,86 @@ def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) +def run_varlen(logits, pre_idx, kv_lens, indices, next_n=1, compress_ratio=1, values=None): + """Production-contract varlen entry — REFERENCE implementation. + + Row semantics (mirror of ``heuristicTopKDecode.cu`` and the in-tree + ``cute_dsl_gvr_topk_decode`` runner): + + ``num_rows = logits.shape[0]``, ``batch = num_rows // next_n``; + ``kv_lens`` int32 ``[batch]`` — per-request TOTAL cache length in + UNCOMPRESSED token space (dsa.py ``metadata.kv_lens_cuda_runtime``, + not new-token seq_lens); row ``r`` uses + ``n_r = (kv_lens[r // next_n] - next_n + (r % next_n) + 1) // + compress_ratio`` valid entries (cr 1 = DSv3.2, 4 = DSv4 Flash/Pro); + ``pre_idx`` ``[batch, k]`` is REQUEST-level raw prev-step top-K, + shared by a request's ``next_n`` rows (offset-free hint contract); + per-row ``n_r <= k`` takes the short path (identity + ``-1`` tail). + + REFERENCE ENGINE: one host read of ``kv_lens`` (a documented D2H sync — + raises under CUDA-graph capture), then each row is driven through the + batch-uniform engine as a b=1 launch. Correctness-first scaffolding: it + pins the varlen/MTP contract and its test battery; the per-row in-kernel + rewrite replaces the loop without changing either. + """ + if not (isinstance(kv_lens, _TENSOR) and kv_lens.is_cuda): + raise RuntimeError("kv_lens must be a CUDA tensor") + if kv_lens.dtype is not _I32: + raise RuntimeError("kv_lens must be int32") + if kv_lens.dim() != 1: + raise RuntimeError("kv_lens must be 1-D") + nn = _index(next_n) + cr = _index(compress_ratio) + if nn < 1: + raise RuntimeError(f"next_n must be >= 1, got {nn}") + if cr not in (1, 4): + raise RuntimeError(f"compress_ratio must be 1 (DSv3.2) or 4 (DSv4), got {cr}") + if len(logits.shape) != 2: + raise RuntimeError("logits must be 2-D") + num_rows = logits.shape[0] + if num_rows == 0: + return + if num_rows % nn: + raise RuntimeError(f"num_rows {num_rows} not divisible by next_n {nn}") + batch = num_rows // nn + if kv_lens.shape[0] != batch: + raise RuntimeError(f"kv_lens length {kv_lens.shape[0]} != num_rows/next_n = {batch}") + if len(pre_idx.shape) != 2 or pre_idx.shape[0] != batch: + raise RuntimeError( + f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" + ) + if _is_capturing(): + raise RuntimeError( + "run_varlen reference implementation reads kv_lens on host, " + "illegal under CUDA graph capture" + ) + d = logits.get_device() + if not 0 <= d < _GVR_MAX_DEV: + raise RuntimeError(f"device index out of range: {d}") + ws = _ws_hot.get(d) + if ws is None: + ws = default_workspace(logits) + kl = kv_lens.tolist() # the ONE documented D2H sync of this entry + for r in range(num_rows): + actual = kl[r // nn] - nn + (r % nn) + 1 + if actual < 0: + raise RuntimeError(f"row {r}: kv_len {kl[r // nn]} < next_n {nn}") + req = r // nn + _run_impl( + logits[r : r + 1], + pre_idx[req : req + 1], + actual // cr, + indices[r : r + 1], + ws, + None if values is None else values[r : r + 1], + ) + + __all__ = [ "route", "run", "run_ws", + "run_varlen", "workspace_bytes", "WS_BYTES", "default_workspace", diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index cbd1353e9a1e..a6b9849b23f0 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -215,6 +215,83 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) +def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): + """Build a per-row-poisoned varlen batch, run run_varlen, verify every + row against its own n_r (production formula) — short rows included.""" + batch, rows = len(kv), len(kv) * next_n + n_r = [(kv[r // next_n] - next_n + (r % next_n) + 1) // cr for r in range(rows)] + npad = (max(n_r) + 63) // 64 * 64 + gen = torch.Generator(device=_DEV).manual_seed(seed) + logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + for r in range(rows): + logits[r, n_r[r] :] = 3e38 # poison beyond each row's OWN n_r + pre_idx = torch.empty((batch, top_k), dtype=torch.int32, device=_DEV) + for q in range(batch): + nmin = max(min(n_r[q * next_n : (q + 1) * next_n]), 1) + pre_idx[q] = torch.randint(0, nmin, (top_k,), generator=gen, dtype=torch.int32, device=_DEV) + indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + values = ( + torch.full((rows, top_k), 7.0, dtype=torch.float32, device=_DEV) if with_values else None + ) + kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, pre_idx, kv_lens, indices, next_n=next_n, compress_ratio=cr, values=values) + torch.cuda.synchronize() + fmin = torch.finfo(torch.float32).min + for r in range(rows): + n = n_r[r] + if n <= top_k: + head = indices[r, :n].to(torch.int64) + assert torch.equal(torch.sort(head).values, torch.arange(n, device=_DEV)) + assert bool((indices[r, n:] == -1).all()) + if values is not None: + assert torch.equal(values[r, :n], logits[r, :n]) + assert bool((values[r, n:] == fmin).all()) + else: + idx = indices[r].to(torch.int64) + assert int(idx.min()) >= 0 and int(idx.max()) < n + assert int(torch.unique(idx).numel()) == top_k + ref = torch.topk(logits[r, :n], top_k).values + got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), f"row {r} inexact" + if values is not None: + assert torch.equal(values[r], torch.gather(logits[r], 0, idx)) + + +@pytest.mark.parametrize( + "kv,next_n,cr,top_k", + [ + ([33000, 8200, 300], 1, 1, 512), # v3.2-style heterogeneous + short row + ([131075, 32800, 2000], 1, 4, 512), # v4-style compressed index space + ([9000, 5001], 2, 1, 512), # MTP: n varies per row within a request + ([65540], 4, 4, 1024), # MTP: compressed-boundary-crossing rows + ], + ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4"], +) +def test_selfsampling_topk_varlen(kv, next_n, cr, top_k): + """run_varlen production contract: per-row n from device kv_lens with the + MTP window formula, request-level hints, per-row short path.""" + _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr) + + +def test_selfsampling_topk_varlen_values(): + _run_varlen_case([40000, 1900], 2, 4, 512, seed=5, with_values=True) + + +def test_selfsampling_topk_varlen_guards(): + logits = torch.randn((2, 8192), dtype=torch.float32, device=_DEV) + pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) + indices = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) + kv = torch.tensor([8192, 8192], dtype=torch.int32, device=_DEV) + with pytest.raises(RuntimeError, match="kv_lens length"): + ss_host.run_varlen(logits, pre_idx, kv[:1], indices) + with pytest.raises(RuntimeError, match="not divisible"): + ss_host.run_varlen(logits, pre_idx, kv, indices, next_n=3) + with pytest.raises(RuntimeError, match="compress_ratio"): + ss_host.run_varlen(logits, pre_idx, kv, indices, compress_ratio=2) + with pytest.raises(RuntimeError, match="CUDA tensor"): + ss_host.run_varlen(logits, pre_idx, kv.cpu(), indices) + + def test_selfsampling_topk_run_ws_explicit_workspace(): """run_ws with a caller-owned workspace must agree with run().""" top_k, n_valid = 1024, 65536 From 0349d1511bfa4fb2138de526fc4c31976dcd21ae Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:28:30 +0000 Subject: [PATCH 07/40] [TRTLLM-15293][perf] self-sampling top-K: two-time-scale dispatch split (route_static/route_dynamic) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Factors route(b, n, npad, k) into route_static — family, compile tuple, grid/cluster/block and the rt scalars that change only at discrete n-thresholds (freezable at CUDA-graph capture time, mirroring the in-tree runner's pick_tuning(graph_capture=...) pattern) — and route_dynamic — the n-continuous scalars a per-row kernel recomputes from its own row length (CMP and the reg smem footprint; the SMP/TGT/SS2/TGT2/Q sampling ladder for the streaming families). The device-side per-row engine will mirror exactly the route_dynamic formulas. Lossless by construction and by fuzz: recombining the halves reproduces route() bit-exactly on 163,755 (b, n, k) points (threshold windows, R-boundaries, prime-stride sweep, LCG random; full result-dict equality). route_bands() enumerates the static-constant n-intervals — the (b=8, k=1024, 262144) envelope collapses to 10 contiguous bands, so the eventual in-kernel band table is tiny. Host-side groundwork only; no kernel behavior change. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 155 ++++++++++++++++++ .../parallel/test_gvr_selfsampling_topk.py | 39 +++++ 2 files changed, 194 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index f4e9b7c763a2..1f063227f40c 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -462,6 +462,157 @@ def _main(BLK, MINB, U, SPLIT): for shp in smoke: print(shp, "->", route(*shp)) + +# --------------------------------------------------------------------------- +# two-time-scale dispatch split (per-row varlen / CUDA-graph groundwork) +# --------------------------------------------------------------------------- +# route(b, n, npad, k) factored into +# route_static(b, n, npad, k) — everything that must be frozen per launch: +# family, compile tuple, grid, cluster, block, and the rt scalars that +# change only at discrete n-thresholds; +# route_dynamic(static, n) — the n-continuous scalars a per-row kernel +# recomputes from its own row length (the device code will mirror these +# formulas): n, CMP (reg families), the sampling ladder +# SMP/TGT/SS2/TGT2/Q (streaming families), and the reg-family smem +# footprint. +# INVARIANT (fuzz-verified): merging route_dynamic back into route_static +# reproduces route() EXACTLY for every n. The capture-time policy of which n +# to freeze the static half at (e.g. max_seq_len) is a later, perf-only +# choice — this split only proves the factorization is lossless. + +_DYN_RT = { + "reg": ("n", "CMP"), + "regimg": ("n", "CMP"), + "reg_clus": ("n",), + "clus": ("n", "SMP", "TGT", "Q", "SS2", "TGT2"), + "main": ("n", "SMP", "TGT", "Q", "SS2", "TGT2"), +} +_DYN_SMEM = ("reg", "regimg") # smem depends on CMP/IMGW -> recomputed per n + + +def route_static(b, n, npad, k): + """route() with the n-continuous fields redacted (see _DYN_RT/_DYN_SMEM). + Constant on maximal n-intervals ("bands"); every redacted field is + reconstructible from (static, n) by route_dynamic.""" + plan = route(b, n, npad, k) + st = {key: (dict(val) if isinstance(val, dict) else val) for key, val in plan.items()} + for f in _DYN_RT[st["kernel"]]: + st["rt"].pop(f) + if st["kernel"] in _DYN_SMEM: + st.pop("smem") + return st + + +def route_dynamic(static, n): + """Recompute the redacted n-continuous scalars from (static, n). + Returns (rt_updates, smem). Transcribed independently from route() — + the factorization fuzz is the equivalence proof, and the device-side + per-row engine mirrors exactly these formulas.""" + fam = static["kernel"] + k = static["rt"]["k"] + if fam in ("reg", "regimg"): + dege = static["tpl"][5] + cmp_ = n if dege else (n if n < 2560 else 2560) + nbsel = static["rt"]["IMGOFF"] + if fam == "regimg": + imgw = (n + 3) & ~3 + smem = (nbsel + (2 * cmp_ if 2 * cmp_ > imgw else imgw)) * 4 + else: + smem = (nbsel + 2 * cmp_) * 4 + return {"n": n, "CMP": cmp_}, smem + if fam == "reg_clus": + return {"n": n}, static["smem"] + + # streaming families (main / clus): the sampling-ladder scalars + b = static["grid"][1] + if fam == "clus": + R = static["cluster"] + scap = static["rt"]["SCAP"] + else: + R = static["rt"]["R"] + scap = static["rt"]["SCAP_"] + big = b * R <= 148 + aim = ( + ((4 * k if k >= 1024 else 2 * k) if R == 1 else 2 * k) + if big + else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) + r_ = int(0.5 + math.sqrt(float(6 * n))) + if r_ > aim: + aim = r_ + sfac = (32 if R == 2 else (48 if k > 1024 else 16)) if R > 1 else (64 if k >= 1024 else 32) + amin = 3 * k if R == 2 else (7 * k) // 2 + if R > 1 and aim < amin: + aim = amin + if aim > (scap >> 1): + aim = scap >> 1 + if aim < k: + aim = k + + n4s = n >> 2 + smp, ss2, tgt, tgt2 = 0, 1, 0, 0 + small_dense = (k > 1024) and (not big) and n <= scap and n > 2 * k + if (n > scap or small_dense) and n4s >= 4: + sel = sfac * n // aim + sel = 256 if sel < 256 else sel + sel = n // 2 if sel > n // 2 else sel + pairs = max(sel >> 3, 1) + half = max(n4s >> 1, 1) + pairs = half if pairs > half else pairs + ss2 = max(half // pairs, 1) + smp = max(half // ss2, 1) + tgt = max((aim * (smp * 8)) // n, 1) + tgt2 = max((k * (smp * 8)) // n, 1) + q_ = (n4s + R - 1) // R + if fam == "clus" and n > scap and n4s >= 4: + sel = sfac * n // aim + sel = 256 if sel < 256 else sel + sel = n // 2 if sel > n // 2 else sel + quads = max(sel >> 4, 1) + quarter = max(n4s >> 2, 1) + quads = quarter if quads > quarter else quads + ss2 = max(quarter // quads, 1) + smp = max(quarter // ss2, 1) + tgt = max((aim * (smp * 16)) // n, 1) + tgt2 = max((k * (smp * 16)) // n, 1) + return ( + {"n": n, "SMP": smp, "TGT": tgt, "Q": q_, "SS2": ss2, "TGT2": tgt2}, + static["smem"], + ) + + +def route_split(b, n, npad, k): + """route_static + route_dynamic recombined — must equal route() exactly + (the factorization fuzz in the unit tests asserts this).""" + st = route_static(b, n, npad, k) + dyn, smem = route_dynamic(st, n) + plan = {key: (dict(val) if isinstance(val, dict) else val) for key, val in st.items()} + plan["rt"].update(dyn) + plan["smem"] = smem + return plan + + +def route_bands(b, npad, k, n_lo=None, n_hi=None): + """Enumerate maximal n-intervals on which route_static is constant. + Dense O(n_hi - n_lo) scan of the pure host dispatch — an offline / + engine-init tool (seconds for the 262144-token envelope), NOT a hot + path. Returns [(n_lo, n_hi, static_plan), ...].""" + lo = k + 1 if n_lo is None else max(n_lo, k + 1) + hi = npad if n_hi is None else min(n_hi, npad) + bands = [] + cur_key, cur_lo, cur_plan = None, lo, None + for n in range(lo, hi + 1): + st = route_static(b, n, npad, k) + key = repr(st) + if key != cur_key: + if cur_key is not None: + bands.append((cur_lo, n - 1, cur_plan)) + cur_key, cur_lo, cur_plan = key, n, st + if cur_key is not None: + bands.append((cur_lo, hi, cur_plan)) + return bands + + # =========================================================================== # ==== workspace (ct_workspace.py) ========================================== # =========================================================================== @@ -923,6 +1074,10 @@ def run_varlen(logits, pre_idx, kv_lens, indices, next_n=1, compress_ratio=1, va __all__ = [ "route", + "route_static", + "route_dynamic", + "route_split", + "route_bands", "run", "run_ws", "run_varlen", diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index a6b9849b23f0..3a2819dcb2a8 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -312,6 +312,45 @@ def test_selfsampling_topk_guards(): ss_host.run(logits, pre_idx[:0], 8192, indices) +def test_selfsampling_route_factorization(): + """Two-time-scale dispatch groundwork: route() must factor losslessly + into route_static (constant on n-bands, freezable at capture time) and + route_dynamic (the n-continuous scalars the per-row device engine will + recompute) — recombining them reproduces route() exactly. CPU-only.""" + npad = 1 << 20 + checked = 0 + for b in (1, 8, 16, 64, 148, 296, 1024): + for k in (512, 1024, 2048): + ns = set() + for c in ( + 2 * k, 3 * k, 4 * k + 64, 2560, 4096, 8192, 16384, + 4 * 1024, 4 * 4096, 4 * 32768, 65536, 131072, 262144, + ): + ns.update(v for v in range(c - 4, c + 5) if k < v <= npad) + ns.update(range(k + 1, npad + 1, 4999)) + s = 12345 + for _ in range(400): + s = (s * 1103515245 + 12345) % (1 << 31) + ns.add(k + 1 + s % (npad - k - 1)) + for n in sorted(ns): + assert ss_host.route_split(b, n, npad, k) == ss_host.route(b, n, npad, k), ( + b, n, k, + ) + checked += 1 + assert checked > 10_000 + + +def test_selfsampling_route_bands_contiguous(): + """route_bands must tile the envelope contiguously with n-free statics.""" + bands = ss_host.route_bands(8, 262144, 1024) + assert bands[0][0] == 1025 and bands[-1][1] == 262144 + for (_, h1, _), (l2, _, _) in zip(bands, bands[1:]): + assert l2 == h1 + 1 + for _, _, st in bands: + for f in ss_host._DYN_RT[st["kernel"]]: + assert f not in st["rt"] + + def test_selfsampling_dispatch_is_pure_and_total(): """route(b, n, npad, k) must return a plan for every in-envelope shape.""" for k in (512, 1024, 2048): From cb23e1d0b687b4331d3f6eb5a0c38d46fe4ed684 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:03:15 +0000 Subject: [PATCH 08/40] [TRTLLM-15293][perf] self-sampling top-K: per-row in-kernel varlen engine (gvr_main port) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gvr_main family gains a per-row varlen mode (production heuristicTopKDecode contract): each CTA reads its row's kv_len from a device kv_lens tensor and re-derives n plus the whole sampling ladder (SMP/SS2/TGT/TGT2/Q) via an exact device transcription of the route_dynamic() host formulas — integer round(sqrt(6n)) with an isqrt fixup (bit-parity with the host double form), Int64 TGT products (the CUDA host math is 64-bit), and constexpr next_n / cr_shift / r_const so every division strength-reduces. All values are pure functions of the row index, so the R split CTAs of a row stay grid-uniform by construction (workspace handshake unchanged). Design points: - No runtime return in CuTe DSL (in-tree gvr_topk_decode.py precedent): n <= k rows run the body as a zero-work pass (n=0, TGT=INT_MAX so no rung accepts, Q=0 empties the split slices) and an epilogue emits the production identity/-1-pad short path. - The TSH-floor staging gate becomes per-row runtime (tsh_en && per-row n4 <= 32768), exactly the CUDA original's grid-uniform runtime gate; legacy compiles keep bit-identical behavior (tsh_run == 1). - pre_idx is REQUEST-level in varlen mode (row // next_n mapping). - run_varlen(engine="auto") launches the batch in ONE kernel; with max_seq_len given (capture-stable engine constant) the call performs no host reads. engine="reference" keeps the b=1 loop as the differential oracle. Capture-time tuning comes from route_streaming() (the streaming half of route(), 110,003-point fuzz agreement). - Legacy batch-uniform ABI extended with dead trailing args (dummy kv_lens + five zeros) — one kern body, no duplication. Validated on B200: legacy regression 5/5 exact (extended ABI), engine vs reference differential 9/9 mixed-batch configs row-for-row equal — deep SPLIT (b=1, n=200k), 8-row cr=4 mix to n=225k, tsh band (b=16), MTP next_n {2,4}, all-short batches, 200-row small_dense k=2048, and the no-host-read max_seq_len path. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 224 +++++++++++++++- .../gvr_topk_decode_self_sampling_host.py | 248 ++++++++++++++++-- .../parallel/test_gvr_selfsampling_topk.py | 54 +++- 3 files changed, 495 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index a9f6af2c2e0a..d96d10a12bbb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -47,6 +47,7 @@ import cutlass import cutlass.cute as cute +import cutlass.cute.math as cmath from cutlass._mlir.dialects import arith as mlir_arith from cutlass._mlir.dialects import llvm, nvvm from cutlass._mlir.dialects import llvm as mlir_llvm @@ -1332,7 +1333,8 @@ class GvrMainKernel: """CuTeDSL port of gvr_main (kernel.cu L377).""" def __init__( - self, blk: int, u: int, minb: int, nbs: int, kpt: int, split: bool, tshg: bool = False + self, blk: int, u: int, minb: int, nbs: int, kpt: int, split: bool, tshg: bool = False, + varlen: bool = False, next_n: int = 1, cr_shift: int = 0, r_const: int = 1, ): assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177, measured)" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) @@ -1343,10 +1345,27 @@ def __init__( self.nbs = nbs self.kpt = kpt self.split = bool(split) + # ---- per-row varlen mode (production heuristicTopKDecode contract) -- + # n and the sampling-ladder scalars are re-derived PER ROW inside the + # kernel from a device kv_lens tensor (route_dynamic formula mirror); + # the scalar n/SMP/TGT/Q/SS2/TGT2 launch args become dead. next_n / + # cr_shift (log2 compressRatio: 0 = DSv3.2, 2 = DSv4) / r_const (the + # frozen grid.x) are compile-time so their divisions strength-reduce. + self.varlen = bool(varlen) + self.next_n = int(next_n) + self.cr_shift = int(cr_shift) + self.r_const = int(r_const) + if self.varlen: + assert self.next_n >= 1 and self.cr_shift in (0, 2) and self.r_const >= 1 # knife5 (layer 7): TSH-floor staging arm. SPLIT-only compile-time # key; the CUDA form is a grid-uniform runtime gate over the same - # predicate (b > 15 && k <= 1024 && n4 <= 32768). - self.tshg = bool(tshg) and bool(split) + # predicate (b > 15 && k <= 1024 && n4 <= 32768). varlen mode + # compiles the machinery in whenever SPLIT and gates it per row at + # runtime (tsh_en && n4 <= 32768) — mirroring the CUDA runtime gate. + if self.varlen: + self.tshg = bool(split) + else: + self.tshg = bool(tshg) and bool(split) # derived constexprs (kernel.cu L394-523) self.hb = nbs # L394 self.kbig = (kpt >= 2) and (kpt * blk >= 2048) # L413 @@ -1454,6 +1473,12 @@ def kern( Q: cutlass.Int32, SS2: cutlass.Int32, TGT2: cutlass.Int32, + kv_lens: cute.Tensor, + aim_base: cutlass.Int32, + sfac: cutlass.Int32, + amin: cutlass.Int32, + sd_en: cutlass.Int32, + tsh_en: cutlass.Int32, ): BLK = self.blk U = self.u @@ -1473,6 +1498,128 @@ def kern( part = bx lane = tidx & cutlass.Int32(31) + # ================= per-row varlen prologue (varlen mode only) ========= + # Production contract: row r serves request r // next_n with + # kv_len = kv_lens[r // next_n], n = (kv_len - next_n + r % next_n + 1) + # >> cr_shift. The sampling-ladder scalars are then re-derived from + # this row's n by the EXACT route_dynamic() host formulas (the scalar + # launch args are dead in this mode). n <= k rows have no runtime + # `return` in CuTe DSL: they run the body as a zero-work pass + # (n = 0, TGT = INT_MAX so no rung ever accepts) and the identity/pad + # emission happens in the epilogue at the end of the kernel. Every + # value below is a pure function of `row`, so all R split CTAs of a + # row (and all threads) compute identical scalars — grid-uniform per + # row by construction. + short = cutlass.Int32(0) + n_row = cutlass.Int32(0) + tsh_run = cutlass.Int32(1) + if cutlass.const_expr(self.varlen): + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > npad: + nv = npad + n_row = nv + if nv <= k: + short = cutlass.Int32(1) + n = cutlass.Int32(0) + SMP = cutlass.Int32(0) + SS2 = cutlass.Int32(1) + TGT = cutlass.Int32(0x7FFFFFFF) + TGT2 = cutlass.Int32(0x7FFFFFFF) + Q = cutlass.Int32(0) + if short == cutlass.Int32(0): + n = nv + # ---- aim ladder (route_dynamic mirror) ---- + # r6 = int(0.5 + sqrt(6n)) computed EXACTLY in integers: f32 + # sqrt seed, fixup to isqrt, then round-half-up via the + # (x - r*r > r) test (bit-parity with the host double form, + # fuzz-proven over the whole n domain). + x6 = cutlass.Int32(6) * n + ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) + while ri * ri > x6: + ri = ri - cutlass.Int32(1) + while (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: + ri = ri + cutlass.Int32(1) + r6 = ri + if x6 - ri * ri > ri: + r6 = ri + cutlass.Int32(1) + aim = aim_base + if r6 > aim: + aim = r6 + if cutlass.const_expr(self.r_const > 1): + if aim < amin: + aim = amin + scap_c = cutlass.Int32(SCPB) # SCAP == SCPB for gvr_main (proven identity) + if aim > (scap_c >> cutlass.Int32(1)): + aim = scap_c >> cutlass.Int32(1) + if aim < k: + aim = k + n4v = n >> cutlass.Int32(2) + SMP = cutlass.Int32(0) + SS2 = cutlass.Int32(1) + TGT = cutlass.Int32(0) + TGT2 = cutlass.Int32(0) + # pair-sample gate: (n > SCAP or small_dense) and n4 >= 4; + # small_dense = k > 1024 and not big and n <= SCAP and n > 2k + # (k/big folded into the launch-constant sd_en flag). + gate = cutlass.Int32(0) + if n > scap_c: + gate = cutlass.Int32(1) + if sd_en != cutlass.Int32(0): + if n <= scap_c: + if n > (k << cutlass.Int32(1)): + gate = cutlass.Int32(1) + if n4v < cutlass.Int32(4): + gate = cutlass.Int32(0) + if gate != cutlass.Int32(0): + sel = sfac * n // aim + if sel < cutlass.Int32(256): + sel = cutlass.Int32(256) + nh = n >> cutlass.Int32(1) + if sel > nh: + sel = nh + pairs = sel >> cutlass.Int32(3) + if pairs < cutlass.Int32(1): + pairs = cutlass.Int32(1) + half = n4v >> cutlass.Int32(1) + if half < cutlass.Int32(1): + half = cutlass.Int32(1) + if pairs > half: + pairs = half + SS2 = half // pairs + if SS2 < cutlass.Int32(1): + SS2 = cutlass.Int32(1) + SMP = half // SS2 + if SMP < cutlass.Int32(1): + SMP = cutlass.Int32(1) + # TGT/TGT2 are 64-bit products in the CUDA host (aim*SMP*8 + # overflows i32 at large n) — mirror with Int64. + smp8 = cutlass.Int64(SMP) * cutlass.Int64(8) + tgt64 = cutlass.Int64(aim) * smp8 // cutlass.Int64(n) + TGT = cutlass.Int32(tgt64) + if TGT < cutlass.Int32(1): + TGT = cutlass.Int32(1) + tgt264 = cutlass.Int64(k) * smp8 // cutlass.Int64(n) + TGT2 = cutlass.Int32(tgt264) + if TGT2 < cutlass.Int32(1): + TGT2 = cutlass.Int32(1) + if cutlass.const_expr(self.split): + Q = (n4v + cutlass.Int32(self.r_const - 1)) // cutlass.Int32(self.r_const) + else: + Q = cutlass.Int32(0) + # per-row TSH-floor runtime gate (CUDA parity: b>15 && k<=1024 in + # tsh_en, n4 <= 32768 per row) + tsh_run = cutlass.Int32(0) + if tsh_en != cutlass.Int32(0): + if (n >> cutlass.Int32(2)) <= cutlass.Int32(32768): + tsh_run = cutlass.Int32(1) + # ---- shared memory (one blob, compile-time offsets; spec §5.1 map) ---- smem = SmemAllocator() s_hist = smem.allocate_tensor( @@ -1559,7 +1706,13 @@ def kern( # _pin_i64: keep the row base a REGISTER across the attempt/tile scf # regions (NVVM otherwise re-derives ld.param+%ctaid.y+mul per region) x_addr = _pin_i64(logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4)) - p_addr = pre_idx.iterator.toint() + row64 * cutlass.Int64(k) * cutlass.Int64(4) + # varlen: pre_idx is REQUEST-level [num_rows/next_n, k] — a request's + # next_n rows share one hint row (production contract); legacy mode + # keeps the per-row mapping (next_n == 1 makes them identical). + prow64 = row64 + if cutlass.const_expr(self.varlen): + prow64 = cutlass.Int64(row // cutlass.Int32(self.next_n)) + p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) out_row = out[row, None] ws_addr = ws.iterator.toint() gdon_addr = ws_addr # L386-388 slab views @@ -1802,9 +1955,12 @@ def kern( # count(>=TSH) >= k. TSH miss falls to GMIN/degen unchanged. cute.arch.barrier() t5s = s_tsh[0] - if t5s > cutlass.Float32(_NEG_INF): - if t5s < T: - T = t5s + # varlen: per-row runtime gate (tsh_run == 1 always in legacy + # mode, so legacy codegen semantics are unchanged) + if tsh_run != cutlass.Int32(0): + if t5s > cutlass.Float32(_NEG_INF): + if t5s < T: + T = t5s # ============ attempt loop (L729-1019) — MUST NOT unroll ============ listN = cutlass.Int32(0) @@ -2676,6 +2832,23 @@ def kern( ) it = it + cutlass.Int32(1) + # ---- varlen short-row epilogue (production heuristicTopKDecode + # L72-84 convention): every valid position is in the top-K — emit + # identity indices and pad the tail with -1. The body above ran as + # a zero-work pass for these rows (n = 0, TGT = INT_MAX) so nothing + # was written; only part 0 of a SPLIT row emits. + if cutlass.const_expr(self.varlen): + if short != cutlass.Int32(0): + if part == cutlass.Int32(0): + i = tidx + while i < n_row: + out_row[i] = i + i = i + cutlass.Int32(BLK) + j = n_row + tidx + while j < k: + out_row[j] = cutlass.Int32(-1) + j = j + cutlass.Int32(BLK) + # ------------------------------------------------------------------ # host launcher (grid dim3(R, b) L2750; MINB wall via min_blocks_per_mp) # ------------------------------------------------------------------ @@ -2697,11 +2870,18 @@ def __call__( Q: cutlass.Int32, SS2: cutlass.Int32, TGT2: cutlass.Int32, + kv_lens: cute.Tensor, + aim_base: cutlass.Int32, + sfac: cutlass.Int32, + amin: cutlass.Int32, + sd_en: cutlass.Int32, + tsh_en: cutlass.Int32, stream, ): b = logits.shape[0] self.kern( - logits, pre_idx, out, ws, n, npad, k, scap_dead, cmp_dead, R, SMP, TGT, Q, SS2, TGT2 + logits, pre_idx, out, ws, n, npad, k, scap_dead, cmp_dead, R, SMP, TGT, Q, SS2, TGT2, + kv_lens, aim_base, sfac, amin, sd_en, tsh_en, ).launch(grid=(R, b, 1), block=(self.blk, 1, 1), stream=stream, min_blocks_per_mp=self.minb) @@ -2713,17 +2893,28 @@ def __call__( def get_compiled(tpl, options_extra: str = ""): """Compile (or fetch) the gvr_main variant for constexpr tuple - tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG).""" + tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG) — legacy, or + tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG, NEXT_N, CR_SHIFT, R_CONST) + — per-row varlen mode (TSHG slot is ignored: varlen compiles the TSH + machinery in whenever SPLIT and gates it per row at runtime).""" key = (tuple(tpl), options_extra) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit - blk, u, minb, nbs, kpt, split, tshg = tpl - kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg)) + if len(tpl) == 7: + blk, u, minb, nbs, kpt, split, tshg = tpl + kern = GvrMainKernel(blk, u, minb, nbs, kpt, bool(split), bool(tshg)) + else: + blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl + kern = GvrMainKernel( + blk, u, minb, nbs, kpt, bool(split), bool(tshg), + varlen=True, next_n=next_n, cr_shift=cr_shift, r_const=r_const, + ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() r2, c2 = cute.sym_int(), cute.sym_int() w0 = cute.sym_int() + v0 = cute.sym_int() logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (r0, c0), stride_order=(1, 0), assumed_align=16 ) @@ -2736,6 +2927,9 @@ def get_compiled(tpl, options_extra: str = ""): ws_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (w0,), stride_order=(0,), assumed_align=16 ) + kv_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (v0,), stride_order=(0,), assumed_align=4 + ) fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) compiled = cute.compile( kern, @@ -2744,6 +2938,8 @@ def get_compiled(tpl, options_extra: str = ""): out_fake, ws_fake, *([cutlass.Int32(0)] * 11), + kv_fake, + *([cutlass.Int32(0)] * 5), stream=fake_stream, options=("--enable-tvm-ffi " + options_extra).strip(), ) @@ -2787,6 +2983,12 @@ def run(logits, pre_idx, n: int, out, ws): rt["Q"], rt["SS2"], rt["TGT2"], + pre_idx.new_zeros(1), # dummy kv_lens (dead in legacy mode) + 0, + 0, + 0, + 0, + 0, ) return r diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 1f063227f40c..45c4cd3c4f7f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -592,6 +592,144 @@ def route_split(b, n, npad, k): return plan +def route_streaming(b, n, npad, k, force_main=False): + """route() restricted to its STREAMING half (main / clus) — the varlen + capture policy: per-row kernels must be picked from the families that are + correct for ANY row length, so the register-resident specialists are + skipped even when the envelope n would normally land on them. Where + route() itself lands on main/clus this is IDENTICAL to route() (fuzz: + 110,003/110,003 agreement). force_main additionally skips the clus + rounding (v1 varlen engine ships the gvr_main port first; the raw + min(r1, r2) R then matches the CUDA else-branch exactly).""" + R = 1 + if b <= 32: + r1 = max(148 // b, 1) + r2 = max(((n >> 2) + 1023) // 1024, 1) + R = max(min(r1, r2), 1) + elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: + R = 2 + useclus = False + if not force_main and 2 <= R <= 8 and k <= 1024: + p2 = 1 + while (p2 << 1) <= R: + p2 <<= 1 + if p2 == 8 and b > 15: + p2 = 4 + R = p2 + useclus = True + big = b * R <= 148 + scap = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) + cmp_ = (4096 if k > 1024 else 2048) if big else 1024 + aim = ( + ((4 * k if k >= 1024 else 2 * k) if R == 1 else 2 * k) + if big + else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) + r_ = int(0.5 + math.sqrt(float(6 * n))) + if r_ > aim: + aim = r_ + sfac = (32 if R == 2 else (48 if k > 1024 else 16)) if R > 1 else (64 if k >= 1024 else 32) + amin = 3 * k if R == 2 else (7 * k) // 2 + if R > 1 and aim < amin: + aim = amin + if aim > (scap >> 1): + aim = scap >> 1 + if aim < k: + aim = k + n4s = n >> 2 + smp, ss2, tgt, tgt2 = 0, 1, 0, 0 + small_dense = (k > 1024) and (not big) and n <= scap and n > 2 * k + if (n > scap or small_dense) and n4s >= 4: + sel = min(max(sfac * n // aim, 256), n // 2) + pairs = min(max(sel >> 3, 1), max(n4s >> 1, 1)) + half = max(n4s >> 1, 1) + ss2 = max(half // pairs, 1) + smp = max(half // ss2, 1) + tgt = max((aim * (smp * 8)) // n, 1) + tgt2 = max((k * (smp * 8)) // n, 1) + q_ = (n4s + R - 1) // R + if useclus: + if n > scap and n4s >= 4: + sel = min(max(sfac * n // aim, 256), n // 2) + quads = min(max(sel >> 4, 1), max(n4s >> 2, 1)) + quarter = max(n4s >> 2, 1) + ss2 = max(quarter // quads, 1) + smp = max(quarter // ss2, 1) + tgt = max((aim * (smp * 16)) // n, 1) + tgt2 = max((k * (smp * 16)) // n, 1) + smc = SNB * 8 + (scap + 4) * 8 + cmp_ * 8 + per = q_ >> 10 + u_ = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) + cs = 2 if R == 2 else (4 if R == 4 else 8) + return { + "kernel": "clus", + "tpl": (1024, u_, 1, SNB, cs), + "rt": {"n": n, "npad": npad, "k": k, "SCAP": scap, "CMP": cmp_, + "SMP": smp, "TGT": tgt, "Q": q_, "SS2": ss2, "TGT2": tgt2}, + "grid": (cs, b), "cluster": cs, "block": 1024, "smem": smc, "ws": False, + } + smem_main = (scap + 4) * (8 if (R > 1 or b <= 296) else 4) + (cmp_ + 1) * 8 + + def _main(blk_, minb_, u_, split_): + kpt = 1 if k <= blk_ else (2 if k <= 2 * blk_ else (4 if k <= 4 * blk_ else 8)) + tshg = bool(split_) and b > 15 and k <= 1024 and (n >> 2) <= 32768 + return { + "kernel": "main", + "tpl": (blk_, u_, minb_, SNB, kpt, split_, tshg), + "rt": {"n": n, "npad": npad, "k": k, "SCAP_": scap, "CMP_": cmp_, "R": R, + "SMP": smp, "TGT": tgt, "Q": q_, "SS2": ss2, "TGT2": tgt2}, + "grid": (R, b), "cluster": 1, "block": blk_, "smem": smem_main, "ws": True, + } + + if big: + per = q_ >> 10 + u_ = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) + return _main(1024, 1, u_, R > 1) + if b <= 296: + return _main(512, 2, 8, False) + return _main(256, 4, 8, False) + + +_VARLEN_CACHE = {} + + +def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): + """Capture-time varlen plan + compiled launcher (v1 = the gvr_main port, + universally correct across the envelope; the clus port follows). Every + choice here is a function of capture-stable quantities only — mirroring + the in-tree runner's pick_tuning(graph_capture=...) discipline.""" + key = (num_rows, npad, k, n_env, next_n, cr) + hit = _VARLEN_CACHE.get(key) + if hit is not None: + return hit + plan = route_streaming(num_rows, max(min(n_env, npad), k + 1), npad, k, force_main=True) + tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) + rt = plan["rt"] + r_const = rt["R"] + cr_shift = 0 if cr == 1 else 2 + dev = _device() + fn = dev.get_compiled(tpl + (next_n, cr_shift, r_const)) + big = num_rows * r_const <= 148 + aim_base = ( + ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) + if big + else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) + sfac = ( + (32 if r_const == 2 else (48 if k > 1024 else 16)) + if r_const > 1 + else (64 if k >= 1024 else 32) + ) + amin = 3 * k if r_const == 2 else (7 * k) // 2 + sd_en = 1 if (k > 1024 and not big) else 0 + tsh_en = 1 if (tpl[5] and num_rows > 15 and k <= 1024) else 0 + pre = (0, npad, k, rt["SCAP_"], rt["CMP_"], r_const, 0, 0, 0, 0, 0) + tail = (aim_base, sfac, amin, sd_en, tsh_en) + lc = (fn, pre, tail) + _VARLEN_CACHE[key] = lc + return lc + + def route_bands(b, npad, k, n_lo=None, n_hi=None): """Enumerate maximal n-intervals on which route_static is constant. Dense O(n_hi - n_lo) scan of the pure host dispatch — an offline / @@ -784,6 +922,17 @@ def _reset_for_tests() -> None: # shape key (b, n, npad, k) -> (fn, args tuple of python ints, needs_ws) _LAUNCH_CACHE = {} +_DUMMY_KV = {} + + +def _dummy_kv(dev_index, device): + """Cached 1-element int32 tensor per device — the dead kv_lens slot of + the extended gvr_main ABI in legacy (batch-uniform) mode.""" + t = _DUMMY_KV.get(dev_index) + if t is None: + t = torch.zeros(1, dtype=_I32, device=device) + _DUMMY_KV[dev_index] = t + return t # hot-path local bindings (each torch. lookup costs ~0.1 us; the B1 # battery runs on EVERY call — mirror of main.cpp's "sub-100ns predicted @@ -813,9 +962,15 @@ def _build_launcher(b, n, npad, k): return (fn, args, False) if fam == "main": dev = _device() - fn = dev.get_compiled(tpl) + raw = dev.get_compiled(tpl) # compiled ABI: (logits, pre_idx, out, ws, n, npad, k, SCAP_, CMP_, - # R, SMP, TGT, Q, SS2, TGT2) [SCAP_/CMP_ dead, ABI parity] + # R, SMP, TGT, Q, SS2, TGT2, + # kv_lens, aim_base, sfac, amin, sd_en, tsh_en) + # [SCAP_/CMP_ dead, ABI parity; the trailing varlen block is dead in + # legacy mode — a cached dummy kv_lens tensor + five zeros] + def fn(lg, pi, o, w, *a, _raw=raw): + _raw(lg, pi, o, w, *a, _dummy_kv(lg.get_device(), lg.device), 0, 0, 0, 0, 0) + args = ( rt["n"], rt["npad"], @@ -997,8 +1152,18 @@ def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) -def run_varlen(logits, pre_idx, kv_lens, indices, next_n=1, compress_ratio=1, values=None): - """Production-contract varlen entry — REFERENCE implementation. +def run_varlen( + logits, + pre_idx, + kv_lens, + indices, + next_n=1, + compress_ratio=1, + values=None, + max_seq_len=None, + engine="auto", +): + """Production-contract varlen entry (per-row device kv_lens). Row semantics (mirror of ``heuristicTopKDecode.cu`` and the in-tree ``cute_dsl_gvr_topk_decode`` runner): @@ -1013,11 +1178,15 @@ def run_varlen(logits, pre_idx, kv_lens, indices, next_n=1, compress_ratio=1, va shared by a request's ``next_n`` rows (offset-free hint contract); per-row ``n_r <= k`` takes the short path (identity + ``-1`` tail). - REFERENCE ENGINE: one host read of ``kv_lens`` (a documented D2H sync — - raises under CUDA-graph capture), then each row is driven through the - batch-uniform engine as a b=1 launch. Correctness-first scaffolding: it - pins the varlen/MTP contract and its test battery; the per-row in-kernel - rewrite replaces the loop without changing either. + ENGINES: ``engine="auto"`` (default) launches the per-row IN-KERNEL + gvr_main varlen port — ONE launch for the whole batch; each CTA reads its + row's kv_len on device and re-derives the sampling ladder (route_dynamic + formula mirror), so with ``max_seq_len`` given (a capture-stable engine + constant, e.g. dsa.py's ``indexer_max_seq_len``) the call performs NO + host reads. Without ``max_seq_len`` the envelope comes from ONE + ``kv_lens.max()`` host read (documented sync, refused under capture). + ``engine="reference"`` keeps the b=1 host-loop reference implementation — + the differential oracle the in-kernel engine is validated against. """ if not (isinstance(kv_lens, _TENSOR) and kv_lens.is_cuda): raise RuntimeError("kv_lens must be a CUDA tensor") @@ -1045,18 +1214,67 @@ def run_varlen(logits, pre_idx, kv_lens, indices, next_n=1, compress_ratio=1, va raise RuntimeError( f"pre_idx must be [batch={batch}, k] REQUEST-level, got {tuple(pre_idx.shape)}" ) - if _is_capturing(): - raise RuntimeError( - "run_varlen reference implementation reads kv_lens on host, " - "illegal under CUDA graph capture" - ) d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: raise RuntimeError(f"device index out of range: {d}") ws = _ws_hot.get(d) if ws is None: ws = default_workspace(logits) - kl = kv_lens.tolist() # the ONE documented D2H sync of this entry + + if engine == "auto": + # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- + npad = logits.shape[1] + k = pre_idx.shape[1] + cshift = 0 if cr == 1 else 2 + if max_seq_len is not None: + n_env = int(max_seq_len) >> cshift + else: + if _is_capturing(): + raise RuntimeError( + "run_varlen without max_seq_len reads kv_lens.max() on " + "host — pass max_seq_len (a capture-stable engine " + "constant) under CUDA graph capture" + ) + n_env = int(kv_lens.max().item()) >> cshift + n_env = min(max(n_env, 1), npad) + key = (num_rows, npad, k, n_env, nn, cr) + lc = _VARLEN_CACHE.get(key) + if lc is None: + if _is_capturing(): + raise RuntimeError( + "varlen launcher not compiled for this shape — warm up " + "before CUDA graph capture" + ) + lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) + fn, pre, tail = lc + if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: + raise RuntimeError("logits must be float32; pre_idx/indices int32") + if not (logits.is_contiguous() and pre_idx.is_contiguous() and indices.is_contiguous()): + raise RuntimeError("tensors must be contiguous") + if npad & 3: + raise RuntimeError(f"npad (logits stride) must be a multiple of 4, got {npad}") + idx = indices + if idx.shape[1] != k: + idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) + vals = values + if vals is not None and vals.shape[1] != k: + vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) + fn(logits, pre_idx, idx, ws, *pre, kv_lens, *tail) + if vals is not None: + idx64 = idx.to(torch.int64) + vals.copy_(logits.gather(1, idx64.clamp_min(0))) + vals.masked_fill_(idx < 0, torch.finfo(_F32).min) + return + if engine != "reference": + raise RuntimeError(f"engine must be 'auto' or 'reference', got {engine!r}") + + # ---- reference engine (differential oracle): b=1 host loop -------------- + if _is_capturing(): + raise RuntimeError( + "run_varlen reference engine reads kv_lens on host, " + "illegal under CUDA graph capture" + ) + kl = kv_lens.tolist() # the ONE documented D2H sync of this engine for r in range(num_rows): actual = kl[r // nn] - nn + (r % nn) + 1 if actual < 0: diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 3a2819dcb2a8..a037309e3e88 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -215,7 +215,7 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): _check_exact(logits, indices, n_valid, ref_vals) -def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): +def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="auto"): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" batch, rows = len(kv), len(kv) * next_n @@ -234,7 +234,10 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): torch.full((rows, top_k), 7.0, dtype=torch.float32, device=_DEV) if with_values else None ) kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) - ss_host.run_varlen(logits, pre_idx, kv_lens, indices, next_n=next_n, compress_ratio=cr, values=values) + ss_host.run_varlen( + logits, pre_idx, kv_lens, indices, + next_n=next_n, compress_ratio=cr, values=values, engine=engine, + ) torch.cuda.synchronize() fmin = torch.finfo(torch.float32).min for r in range(rows): @@ -257,6 +260,7 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): assert torch.equal(values[r], torch.gather(logits[r], 0, idx)) +@pytest.mark.parametrize("engine", ["auto", "reference"]) @pytest.mark.parametrize( "kv,next_n,cr,top_k", [ @@ -267,10 +271,50 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): ], ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4"], ) -def test_selfsampling_topk_varlen(kv, next_n, cr, top_k): +def test_selfsampling_topk_varlen(kv, next_n, cr, top_k, engine): """run_varlen production contract: per-row n from device kv_lens with the - MTP window formula, request-level hints, per-row short path.""" - _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr) + MTP window formula, request-level hints, per-row short path — on BOTH the + per-row in-kernel engine ("auto") and the b=1 reference loop.""" + _run_varlen_case(kv, next_n, cr, top_k, seed=sum(kv) + next_n + cr, engine=engine) + + +def test_selfsampling_topk_varlen_engine_matches_reference(): + """Differential: the in-kernel engine's per-row value multisets must + equal the reference loop's on a mixed batch (deep SPLIT rows, tsh band, + short rows, compressed space).""" + kv = [524288, 131075, 32800, 2000, 65540, 8192, 262144, 900] + top_k, next_n, cr = 1024, 1, 4 + rows = len(kv) + n_r = [(v - 1 + 1) // cr for v in kv] + npad = (max(n_r) + 63) // 64 * 64 + gen = torch.Generator(device=_DEV).manual_seed(77) + logits = torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + for r in range(rows): + logits[r, n_r[r] :] = 3e38 + pre_idx = torch.empty((rows, top_k), dtype=torch.int32, device=_DEV) + for q in range(rows): + pre_idx[q] = torch.randint( + 0, max(n_r[q], 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV + ) + kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) + out_a = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + out_r = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, pre_idx, kv_lens, out_a, compress_ratio=cr) + ss_host.run_varlen(logits, pre_idx, kv_lens, out_r, compress_ratio=cr, engine="reference") + torch.cuda.synchronize() + for r in range(rows): + if n_r[r] <= top_k: + assert torch.equal(out_a[r], out_r[r]) or torch.equal( + torch.sort(out_a[r]).values, torch.sort(out_r[r]).values + ) + else: + ga = torch.sort( + torch.gather(logits[r], 0, out_a[r].to(torch.int64)) + 0.0, descending=True + ).values + gr = torch.sort( + torch.gather(logits[r], 0, out_r[r].to(torch.int64)) + 0.0, descending=True + ).values + assert torch.equal(ga, gr), f"row {r}: engine != reference" def test_selfsampling_topk_varlen_values(): From 4ea46e5f65fb455b16cb438c6e7006214c150e0f Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:05:27 +0000 Subject: [PATCH 09/40] [TRTLLM-15293][test] self-sampling top-K: CUDA-graph capture/replay test for the varlen engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Warm up, capture one run_varlen(engine=auto, max_seq_len=...) launch — no host reads, no JIT inside capture — then replay while kv_lens grows in place: a row crossing the n <= topK short-path boundary INSIDE the graph, a row walking the 131072 band edge, and a 200k deep row. Every replay verified tie-aware exact on B200 (8-replay standalone run all green). This closes the CUDA-graph-safety design goal for the gvr_main varlen port: geometry and tuple are frozen from capture-stable quantities (route_streaming at max_seq_len), all N-dependence is per-row device arithmetic. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../parallel/test_gvr_selfsampling_topk.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index a037309e3e88..337996586259 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -336,6 +336,64 @@ def test_selfsampling_topk_varlen_guards(): ss_host.run_varlen(logits, pre_idx, kv.cpu(), indices) +def test_selfsampling_topk_varlen_cuda_graph(): + """CUDA-graph safety of the in-kernel varlen engine: warm up, capture one + launch with max_seq_len (capture-stable constant — no host reads, no JIT + inside capture), then replay while kv_lens grows in place — including a + row that crosses the n <= k short-path boundary INSIDE the graph and a + row walking over the 131072 band edge.""" + rows, top_k, msl = 4, 512, 262144 + npad = msl + logits = torch.randn((rows, npad), dtype=torch.float32, device=_DEV) - 2.0 + pre_idx = torch.zeros((rows, top_k), dtype=torch.int32, device=_DEV) + kv_lens = torch.tensor([100, 4099, 131070, 200000], dtype=torch.int32, device=_DEV) + indices = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + + def refresh(step): + kv = [100 + step * 137, 4099 + step * 977, 131070 + step, 200000 + step * 3] + kv_lens.copy_(torch.tensor(kv, dtype=torch.int32, device=_DEV)) + gen = torch.Generator(device=_DEV).manual_seed(1000 + step) + logits.copy_( + torch.randn((rows, npad), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + ) + for r in range(rows): + n = min(kv[r], npad) + logits[r, n:] = 3e38 + pre_idx[r] = torch.randint( + 0, max(n, 1), (top_k,), generator=gen, dtype=torch.int32, device=_DEV + ) + return kv + + refresh(0) + ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + ss_host.run_varlen(logits, pre_idx, kv_lens, indices, compress_ratio=1, max_seq_len=msl) + for step in range(1, 6): + kv = refresh(step) + indices.fill_(-7) + graph.replay() + torch.cuda.synchronize() + for r in range(rows): + n = min(kv[r], npad) + if n <= top_k: + head = torch.sort(indices[r, :n].to(torch.int64)).values + assert torch.equal(head, torch.arange(n, device=_DEV)) + assert bool((indices[r, n:] == -1).all()) + else: + idx = indices[r].to(torch.int64) + assert int(idx.min()) >= 0 and int(idx.max()) < n + assert int(torch.unique(idx).numel()) == top_k + ref = torch.topk(logits[r, :n], top_k).values + got = torch.sort( + torch.gather(logits[r], 0, idx) + 0.0, descending=True + ).values + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), ( + f"replay step {step} row {r} inexact (n={n})" + ) + + def test_selfsampling_topk_run_ws_explicit_workspace(): """run_ws with a caller-owned workspace must agree with run().""" top_k, n_valid = 1024, 65536 From 7986d32916885e8a3cb47e6412bffe5fef87a00c Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:08:01 +0000 Subject: [PATCH 10/40] [TRTLLM-15293][perf] self-sampling top-K: opt-in dsa.py decode seam (TRTLLM_GVR_SELF_SAMPLING=1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the varlen engine into the DSA indexer decode top-K dispatch as the highest-priority branch, env-gated (TRTLLM_GVR_SELF_SAMPLING=1) and contract-gated at init (cutlass DSL present, sm100+, index_topk in {512,1024,2048}, compress_ratio in {1,4}) — covering DSv3.2 (K=2048, cr=1), DSv4 Flash (K=512, cr=4) and DSv4 Pro (K=1024, cr=4). The call reuses the exact buffers of the existing cute_dsl_gvr_topk_decode branch (request-level heuristic_prev_topk, kv_lens_cuda_runtime, topk_indices_buffer, indexer_max_seq_len as the capture-stable tuning constant) — no new metadata plumbing. Lazy import behind the gate; contract violations raise loudly (explicit experiment flag, not a silent fallback). First warmup call pays the one-time DSL JIT. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index eae55ac98ae7..a729e8493e2e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -691,6 +691,28 @@ def __init__( self._enable_heuristic_topk = ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ) + # Opt-in self-sampling GVR top-K decode (standalone CuTeDSL modules, + # env-gated experimental path: TRTLLM_GVR_SELF_SAMPLING=1). Same + # operator contract as the tiered heuristic path (per-request device + # kv_lens, request-level raw prev-top-K hints, per-row MTP window, + # in-kernel n <= topK short path); tuning is frozen from + # indexer_max_seq_len at capture time, so the launch is + # CUDA-graph-replay safe. Contract violations raise loudly rather + # than silently falling back — this flag is an explicit experiment. + self._use_self_sampling_topk = ( + os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" + and IS_CUTLASS_DSL_AVAILABLE + and get_sm_version() >= 100 + and sparse_params.index_topk in (512, 1024, 2048) + and compress_ratio in (1, 4) + ) + self._selfsampling_run_varlen = None + if self._use_self_sampling_topk: + from ....cute_dsl_kernels.blackwell.top_k import ( + selfsampling_topk_run_varlen, + ) + + self._selfsampling_run_varlen = selfsampling_topk_run_varlen self.mtp_index_share = sparse_params.mtp_index_share if self._enable_heuristic_topk and layer_idx == 0: @@ -1813,7 +1835,26 @@ def sparse_attn_indexer( if not metadata.use_cute_dsl_topk: heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens] - if self.use_cute_dsl_topk and self._enable_heuristic_topk: + if ( + self._use_self_sampling_topk + and self._enable_heuristic_topk + and pre_idx is not None + ): + # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): + # one launch for the batch; per-row n from device kv_lens, + # capture-stable tuning from indexer_max_seq_len (no host + # reads — CUDA-graph safe). Hints are consumed raw for all + # of DSv3.2 / Flash / Pro (offset-free hint contract). + self._selfsampling_run_varlen( + logits_decode, + pre_idx, + gen_kv_lens_cuda, + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=indexer_max_seq_len, + ) + elif self.use_cute_dsl_topk and self._enable_heuristic_topk: torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits_decode, pre_idx, From b54eb2bb6a05afbee37e04ac277f9d0a917ff039 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:13:22 +0000 Subject: [PATCH 11/40] [TRTLLM-15293][fix] self-sampling top-K: short-row sentinel must survive the TGT*2 scan target MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The zero-work short pass used TGT = 0x7FFFFFFF, whose TGT*2 third scan target overflows Int32 to -2, flipping every tot0 >= TGT*2 gate on the all-zero histogram (benign downstream today — empty candidate sets emit nothing — but an unnecessary cliff). Use 2^30-1 so the doubled target stays positive; 'never accepts' semantics unchanged. Full differential battery (9/9 mixed-batch configs) + graph capture/replay (8 replays) re-verified green on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_self_sampling.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index d96d10a12bbb..996371083b05 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1530,8 +1530,11 @@ def kern( n = cutlass.Int32(0) SMP = cutlass.Int32(0) SS2 = cutlass.Int32(1) - TGT = cutlass.Int32(0x7FFFFFFF) - TGT2 = cutlass.Int32(0x7FFFFFFF) + # "never accepts" sentinels; 2^30-1 so the TGT*2 scan target + # stays positive (0x7FFFFFFF would overflow to -2 and flip + # every tot0 >= TGT*2 gate on the all-zero histogram) + TGT = cutlass.Int32(0x3FFFFFFF) + TGT2 = cutlass.Int32(0x3FFFFFFF) Q = cutlass.Int32(0) if short == cutlass.Int32(0): n = nv From 2efd66e29c33cc0030ad2d041e75cf82fd0d6e15 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:37:33 +0000 Subject: [PATCH 12/40] [TRTLLM-15293][fix] self-sampling top-K: adversarial-review hardening of the varlen engine and seam MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes from a 4-dimension adversarial review of the varlen stack: - CRITICAL seam unit bug: indexer_max_seq_len is COMPRESSED index space (metadata divides by compress_ratio) but run_varlen's max_seq_len is kv-token space — the seam shifted twice, freezing a 4x-too-small tuning envelope for cr=4. Seam now multiplies back (exact identity: n_env == indexer_max_seq_len). - Seam crash on legal configs: the DSL paged-MQA logits arena is a 256-aligned buffer column-sliced to max_seq_len — a NON-contiguous view the engine used to reject. The engine now accepts row-major views (inner stride 1) and widens them back to a compact [rows, row_stride] view over the same storage (as_strided, zero copy; the tail columns are never classified — per-row n gates all reads). A dispatch-site hardware-format gate (stride %4, 16B base) falls through to the existing branches for layouts the kernel cannot address (odd-npad DeepGEMM). - OOB-write guard (three reviewers converged): the engine never validated indices/values batch dims — a request-level-shaped buffer under MTP would be silently written past its end (grid comes from logits rows). Full B1-style battery now runs on the engine path (CUDA/dtype/2-D/contiguity/batch/width/alignment), kv_lens contiguity included. - Engine/reference convention alignment: wider-than-k buffers now follow the flat-packed contract identically on both engines; the reference clamps kv_len < next_n to the empty row (all -1) exactly like the kernel — padded/evicted graph slots are a legal input, and the differential oracle can now cover them. - Eager-mode compile churn: without max_seq_len the data-dependent envelope is quantized up to the next power of two (bounded plan set and _VARLEN_CACHE; a growing decode no longer recompiles at every R increment). - Multi-stream escape hatch: run_varlen(workspace=...) (run_ws parity) so concurrent streams do not share the SPLIT publish slab. - Legacy hot path: per-device cached dummy kv_lens (no per-call allocation for the dead ABI slot). New tests: b=16 SPLIT + per-row TSH runtime gate, b=200 BLK=512 non-split, zero-kv slot mixed with live MTP rows (both engines), wide-buffer flat-packed convention (both engines), num_rows / strided-kv_lens guards. Full battery re-verified on B200 (10/10 incl. a non-contiguous arena-view differential). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 20 +++- .../top_k/gvr_topk_decode_self_sampling.py | 16 ++- .../gvr_topk_decode_self_sampling_host.py | 99 +++++++++++++++---- .../parallel/test_gvr_selfsampling_topk.py | 88 +++++++++++++++++ 4 files changed, 201 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index a729e8493e2e..f3e8f219d673 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1839,12 +1839,24 @@ def sparse_attn_indexer( self._use_self_sampling_topk and self._enable_heuristic_topk and pre_idx is not None + # engine hardware-format gate (falls through otherwise): + # row-major logits with a float4-aligned row stride and a + # 16B-aligned base (the DSL paged-MQA arena view — column- + # sliced from a 256-aligned buffer — satisfies this; odd + # max_seq_len DeepGEMM layouts do not) + and logits_decode.stride(1) == 1 + and logits_decode.stride(0) % 4 == 0 + and logits_decode.data_ptr() % 16 == 0 ): # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): # one launch for the batch; per-row n from device kv_lens, - # capture-stable tuning from indexer_max_seq_len (no host - # reads — CUDA-graph safe). Hints are consumed raw for all - # of DSv3.2 / Flash / Pro (offset-free hint contract). + # capture-stable tuning from the max-seq-len engine + # constant (no host reads — CUDA-graph safe). Hints are + # consumed raw for all of DSv3.2 / Flash / Pro + # (offset-free hint contract). indexer_max_seq_len is in + # COMPRESSED index space (metadata divides by + # compress_ratio); run_varlen's max_seq_len is in kv-token + # space like kv_lens — multiply back. self._selfsampling_run_varlen( logits_decode, pre_idx, @@ -1852,7 +1864,7 @@ def sparse_attn_indexer( topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], next_n=next_n, compress_ratio=self.compress_ratio, - max_seq_len=indexer_max_seq_len, + max_seq_len=indexer_max_seq_len * self.compress_ratio, ) elif self.use_cute_dsl_topk and self._enable_heuristic_topk: torch.ops.trtllm.cute_dsl_gvr_topk_decode( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 996371083b05..32d96906832e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -2986,7 +2986,7 @@ def run(logits, pre_idx, n: int, out, ws): rt["Q"], rt["SS2"], rt["TGT2"], - pre_idx.new_zeros(1), # dummy kv_lens (dead in legacy mode) + _legacy_dummy_kv(pre_idx), # dummy kv_lens (dead in legacy mode) 0, 0, 0, @@ -2996,6 +2996,20 @@ def run(logits, pre_idx, n: int, out, ws): return r +_LEGACY_DUMMY_KV = {} + + +def _legacy_dummy_kv(ref): + """Cached 1-element int32 tensor per device for the dead kv_lens ABI slot + (avoids a per-call allocator round-trip on the legacy hot path).""" + d = ref.get_device() + t = _LEGACY_DUMMY_KV.get(d) + if t is None: + t = ref.new_zeros(1) + _LEGACY_DUMMY_KV[d] = t + return t + + # =========================================================================== # ==== ct_reg.py — family: reg ================================= # =========================================================================== diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 45c4cd3c4f7f..88ebd084cdb6 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1162,6 +1162,7 @@ def run_varlen( values=None, max_seq_len=None, engine="auto", + workspace=None, ): """Production-contract varlen entry (per-row device kv_lens). @@ -1217,14 +1218,66 @@ def run_varlen( d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: raise RuntimeError(f"device index out of range: {d}") - ws = _ws_hot.get(d) - if ws is None: - ws = default_workspace(logits) + if workspace is not None: + # multi-stream escape hatch (run_ws parity): concurrent varlen + # launches on one device must not share the SPLIT publish slab + validate_run_ws(workspace, logits) + ws = kernel_view(workspace) + else: + ws = _ws_hot.get(d) + if ws is None: + ws = default_workspace(logits) if engine == "auto": # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- - npad = logits.shape[1] + # Full B1-style validation battery (the engine bypasses _run_impl — + # every check the legacy path enforces is replayed here; the batch-dim + # check is CRITICAL: the kernel grid comes from logits.shape[0], so a + # short indices/values tensor would be written out of bounds). + if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: + raise RuntimeError("logits must be float32; pre_idx/indices int32") + if len(indices.shape) != 2 or indices.shape[0] != num_rows: + raise RuntimeError( + f"indices must be [num_rows={num_rows}, >=k], got {tuple(indices.shape)}" + ) k = pre_idx.shape[1] + if indices.shape[1] < k: + raise RuntimeError(f"indices width {indices.shape[1]} < k={k}") + if not (pre_idx.is_contiguous() and indices.is_contiguous() and kv_lens.is_contiguous()): + raise RuntimeError("pre_idx/indices/kv_lens must be contiguous") + # logits: accept row-major views with a wider row stride (the DSL + # paged-MQA logits arena is 256-aligned and column-sliced — a legal + # NON-contiguous view). The kernel only needs (base, row stride): + # widen back to a compact [rows, stride] view over the same storage; + # the tail columns are never classified (per-row n gates all reads). + if logits.stride(1) != 1: + raise RuntimeError("logits inner stride must be 1") + npad = logits.stride(0) if num_rows > 1 else logits.shape[1] + lg = logits + if not logits.is_contiguous(): + need = logits.storage_offset() + num_rows * npad + if logits.untyped_storage().size() // 4 < need: + raise RuntimeError("logits view storage too small to widen to its row stride") + lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) + if npad & 3: + raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") + if lg.data_ptr() & 15: + raise RuntimeError("logits base must be 16-byte aligned") + if values is not None: + if not values.is_cuda or values.dtype is not _F32: + raise RuntimeError("values must be CUDA float32") + if ( + len(values.shape) != 2 + or values.shape[0] != num_rows + or values.shape[1] < k + or not values.is_contiguous() + ): + raise RuntimeError( + f"values must be contiguous [num_rows={num_rows}, >=k], " + f"got {tuple(values.shape)}" + ) cshift = 0 if cr == 1 else 2 if max_seq_len is not None: n_env = int(max_seq_len) >> cshift @@ -1236,6 +1289,10 @@ def run_varlen( "constant) under CUDA graph capture" ) n_env = int(kv_lens.max().item()) >> cshift + # eager mode: quantize the data-dependent envelope up to the next + # power of two so a growing decode does not recompile at every + # R increment (bounded plans, bounded _VARLEN_CACHE) + n_env = 1 << max(n_env - 1, 1).bit_length() n_env = min(max(n_env, 1), npad) key = (num_rows, npad, k, n_env, nn, cr) lc = _VARLEN_CACHE.get(key) @@ -1247,22 +1304,16 @@ def run_varlen( ) lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) fn, pre, tail = lc - if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: - raise RuntimeError("logits must be float32; pre_idx/indices int32") - if not (logits.is_contiguous() and pre_idx.is_contiguous() and indices.is_contiguous()): - raise RuntimeError("tensors must be contiguous") - if npad & 3: - raise RuntimeError(f"npad (logits stride) must be a multiple of 4, got {npad}") idx = indices if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) vals = values if vals is not None and vals.shape[1] != k: vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - fn(logits, pre_idx, idx, ws, *pre, kv_lens, *tail) + fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) if vals is not None: idx64 = idx.to(torch.int64) - vals.copy_(logits.gather(1, idx64.clamp_min(0))) + vals.copy_(lg.gather(1, idx64.clamp_min(0))) vals.masked_fill_(idx < 0, torch.finfo(_F32).min) return if engine != "reference": @@ -1274,19 +1325,33 @@ def run_varlen( "run_varlen reference engine reads kv_lens on host, " "illegal under CUDA graph capture" ) + # match the engine's flat-packed output convention for wider-than-k + # buffers (pack ONCE from the tensor base, then slice per row) + k = pre_idx.shape[1] + idx = indices + if len(idx.shape) != 2 or idx.shape[0] != num_rows: + raise RuntimeError( + f"indices must be [num_rows={num_rows}, >=k], got {tuple(indices.shape)}" + ) + if idx.shape[1] != k: + idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) + vals = values + if vals is not None and vals.shape[1] != k: + vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) kl = kv_lens.tolist() # the ONE documented D2H sync of this engine for r in range(num_rows): - actual = kl[r // nn] - nn + (r % nn) + 1 - if actual < 0: - raise RuntimeError(f"row {r}: kv_len {kl[r // nn]} < next_n {nn}") + # production graph slots can carry kv_len < next_n (padded / evicted + # requests): clamp to the empty row, emitting all -1 — the same + # contract the in-kernel engine implements + actual = max(kl[r // nn] - nn + (r % nn) + 1, 0) req = r // nn _run_impl( logits[r : r + 1], pre_idx[req : req + 1], actual // cr, - indices[r : r + 1], + idx[r : r + 1], ws, - None if values is None else values[r : r + 1], + None if vals is None else vals[r : r + 1], ) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 337996586259..736407f00b5f 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -321,6 +321,87 @@ def test_selfsampling_topk_varlen_values(): _run_varlen_case([40000, 1900], 2, 4, 512, seed=5, with_values=True) +@pytest.mark.parametrize( + "rows,base,step,top_k", + [ + (16, 40000, 977, 1024), # R=9 SPLIT + per-row TSH runtime gate (tsh_en=1) + (200, 6000, 64, 512), # BLK=512, SPLIT=False, non-big launch mode + ], + ids=["b16_tsh_split", "b200_blk512_nonsplit"], +) +def test_selfsampling_topk_varlen_launch_modes(rows, base, step, top_k): + """Exercise the varlen kernel's other launch modes: the SPLIT + per-row + TSH-floor runtime gate domain (16 <= b <= 74, k <= 1024) and the + BLK=512 non-split wide-batch plan.""" + _run_varlen_case([base + step * i for i in range(rows)], 1, 1, top_k, seed=rows) + + +def test_selfsampling_topk_varlen_zero_kv_slot(): + """Padded / evicted CUDA-graph request slots can carry kv_len < next_n + (even 0): both engines must emit the empty short row (all -1), not raise — + mixed with live MTP rows of another request in the same launch.""" + for engine in ("auto", "reference"): + gen = torch.Generator(device=_DEV).manual_seed(9) + logits = torch.randn((8, 8192), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) + kv_lens = torch.tensor([0, 8192], dtype=torch.int32, device=_DEV) + indices = torch.full((8, 512), -7, dtype=torch.int32, device=_DEV) + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + logits[r, n:] = 3e38 + ss_host.run_varlen( + logits, pre_idx, kv_lens, indices, next_n=4, compress_ratio=1, engine=engine + ) + torch.cuda.synchronize() + assert bool((indices[:4] == -1).all()), f"{engine}: kv=0 rows must be all -1" + for r in range(4, 8): + n = 8192 - 4 + (r - 4) + 1 + idx = indices[r].to(torch.int64) + assert int(idx.min()) >= 0 and int(idx.max()) < n + ref = torch.topk(logits[r, :n], 512).values + got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values) + + +def test_selfsampling_topk_varlen_wide_buffers_flat_packed(): + """indices wider than k follow the CUDA flat-packed contract (rows at + stride k from the tensor base) IDENTICALLY on both engines.""" + kv = [9000, 300] + top_k, width = 512, 512 + 64 + gen = torch.Generator(device=_DEV).manual_seed(21) + logits = torch.randn((2, 9024), generator=gen, dtype=torch.float32, device=_DEV) - 2.0 + logits[0, 9000:] = 3e38 + logits[1, 300:] = 3e38 + pre_idx = torch.randint(0, 300, (2, top_k), generator=gen, dtype=torch.int32, device=_DEV) + kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) + outs = {} + for engine in ("auto", "reference"): + wide = torch.full((2, width), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, pre_idx, kv_lens, wide, compress_ratio=1, engine=engine) + torch.cuda.synchronize() + outs[engine] = wide.reshape(-1)[: 2 * top_k].view(2, top_k).clone() + packed_a, packed_r = outs["auto"], outs["reference"] + # row 0 (kernel row): same value multiset via the SAME packed convention + ga = torch.sort( + torch.gather(logits[0], 0, packed_a[0].to(torch.int64)) + 0.0, descending=True + ).values + gr = torch.sort( + torch.gather(logits[0], 0, packed_r[0].to(torch.int64)) + 0.0, descending=True + ).values + assert torch.equal(ga, gr), "wide-buffer packing convention diverged between engines" + # row 1 (short row): identity + -1 tail at the packed location, bit-equal + expect = torch.cat( + [ + torch.arange(300, dtype=torch.int32, device=_DEV), + torch.full((top_k - 300,), -1, dtype=torch.int32, device=_DEV), + ] + ) + assert torch.equal(torch.sort(packed_a[1, :300]).values, expect[:300]) + assert bool((packed_a[1, 300:] == -1).all()) + assert torch.equal(torch.sort(packed_r[1, :300]).values, expect[:300]) + assert bool((packed_r[1, 300:] == -1).all()) + + def test_selfsampling_topk_varlen_guards(): logits = torch.randn((2, 8192), dtype=torch.float32, device=_DEV) pre_idx = torch.zeros((2, 512), dtype=torch.int32, device=_DEV) @@ -334,6 +415,13 @@ def test_selfsampling_topk_varlen_guards(): ss_host.run_varlen(logits, pre_idx, kv, indices, compress_ratio=2) with pytest.raises(RuntimeError, match="CUDA tensor"): ss_host.run_varlen(logits, pre_idx, kv.cpu(), indices) + with pytest.raises(RuntimeError, match="num_rows"): + # request-level-shaped indices under MTP: MUST be rejected (the + # kernel grid comes from logits rows — silent OOB writes otherwise) + ss_host.run_varlen(logits, pre_idx[:1], kv[:1], indices[:1], next_n=2) + with pytest.raises(RuntimeError, match="contiguous"): + strided = torch.zeros((2, 2), dtype=torch.int32, device=_DEV)[:, 0] + ss_host.run_varlen(logits, pre_idx, strided, indices) def test_selfsampling_topk_varlen_cuda_graph(): From 06af8550c1c8a1d33e35e05266fbfe110af2cabb Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:56:54 +0000 Subject: [PATCH 13/40] [TRTLLM-15293][chore] self-sampling top-K: apply pre-commit formatting Release-Check (PR_Github #67203) failed on pre-commit: ruff-format reflow on 3 files, one ruff F841 (unused next_n unpack in the varlen differential test), one codespell hit (statics -> static fields). Formatting-only plus the two mechanical fixes; kernel exactness re-verified on GPU after the reflow: full unit file 65/65 passed (sm100, standalone overlay stack). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 4 +- .../top_k/gvr_topk_decode_self_sampling.py | 50 ++++++++++++++++--- .../gvr_topk_decode_self_sampling_host.py | 46 ++++++++++++++--- .../parallel/test_gvr_selfsampling_topk.py | 49 +++++++++++------- 4 files changed, 115 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index f3e8f219d673..905b171f0469 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -708,9 +708,7 @@ def __init__( ) self._selfsampling_run_varlen = None if self._use_self_sampling_topk: - from ....cute_dsl_kernels.blackwell.top_k import ( - selfsampling_topk_run_varlen, - ) + from ....cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen self._selfsampling_run_varlen = selfsampling_topk_run_varlen self.mtp_index_share = sparse_params.mtp_index_share diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 32d96906832e..55c6e1a4fd61 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1333,8 +1333,18 @@ class GvrMainKernel: """CuTeDSL port of gvr_main (kernel.cu L377).""" def __init__( - self, blk: int, u: int, minb: int, nbs: int, kpt: int, split: bool, tshg: bool = False, - varlen: bool = False, next_n: int = 1, cr_shift: int = 0, r_const: int = 1, + self, + blk: int, + u: int, + minb: int, + nbs: int, + kpt: int, + split: bool, + tshg: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + r_const: int = 1, ): assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177, measured)" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) @@ -2883,8 +2893,27 @@ def __call__( ): b = logits.shape[0] self.kern( - logits, pre_idx, out, ws, n, npad, k, scap_dead, cmp_dead, R, SMP, TGT, Q, SS2, TGT2, - kv_lens, aim_base, sfac, amin, sd_en, tsh_en, + logits, + pre_idx, + out, + ws, + n, + npad, + k, + scap_dead, + cmp_dead, + R, + SMP, + TGT, + Q, + SS2, + TGT2, + kv_lens, + aim_base, + sfac, + amin, + sd_en, + tsh_en, ).launch(grid=(R, b, 1), block=(self.blk, 1, 1), stream=stream, min_blocks_per_mp=self.minb) @@ -2910,8 +2939,17 @@ def get_compiled(tpl, options_extra: str = ""): else: blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl kern = GvrMainKernel( - blk, u, minb, nbs, kpt, bool(split), bool(tshg), - varlen=True, next_n=next_n, cr_shift=cr_shift, r_const=r_const, + blk, + u, + minb, + nbs, + kpt, + bool(split), + bool(tshg), + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + r_const=r_const, ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 88ebd084cdb6..6db792b5005a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -664,9 +664,23 @@ def route_streaming(b, n, npad, k, force_main=False): return { "kernel": "clus", "tpl": (1024, u_, 1, SNB, cs), - "rt": {"n": n, "npad": npad, "k": k, "SCAP": scap, "CMP": cmp_, - "SMP": smp, "TGT": tgt, "Q": q_, "SS2": ss2, "TGT2": tgt2}, - "grid": (cs, b), "cluster": cs, "block": 1024, "smem": smc, "ws": False, + "rt": { + "n": n, + "npad": npad, + "k": k, + "SCAP": scap, + "CMP": cmp_, + "SMP": smp, + "TGT": tgt, + "Q": q_, + "SS2": ss2, + "TGT2": tgt2, + }, + "grid": (cs, b), + "cluster": cs, + "block": 1024, + "smem": smc, + "ws": False, } smem_main = (scap + 4) * (8 if (R > 1 or b <= 296) else 4) + (cmp_ + 1) * 8 @@ -676,9 +690,24 @@ def _main(blk_, minb_, u_, split_): return { "kernel": "main", "tpl": (blk_, u_, minb_, SNB, kpt, split_, tshg), - "rt": {"n": n, "npad": npad, "k": k, "SCAP_": scap, "CMP_": cmp_, "R": R, - "SMP": smp, "TGT": tgt, "Q": q_, "SS2": ss2, "TGT2": tgt2}, - "grid": (R, b), "cluster": 1, "block": blk_, "smem": smem_main, "ws": True, + "rt": { + "n": n, + "npad": npad, + "k": k, + "SCAP_": scap, + "CMP_": cmp_, + "R": R, + "SMP": smp, + "TGT": tgt, + "Q": q_, + "SS2": ss2, + "TGT2": tgt2, + }, + "grid": (R, b), + "cluster": 1, + "block": blk_, + "smem": smem_main, + "ws": True, } if big: @@ -934,6 +963,7 @@ def _dummy_kv(dev_index, device): _DUMMY_KV[dev_index] = t return t + # hot-path local bindings (each torch. lookup costs ~0.1 us; the B1 # battery runs on EVERY call — mirror of main.cpp's "sub-100ns predicted # branches" intent within Python's reach; measured in notes/ct_op_NOTES.md) @@ -963,6 +993,7 @@ def _build_launcher(b, n, npad, k): if fam == "main": dev = _device() raw = dev.get_compiled(tpl) + # compiled ABI: (logits, pre_idx, out, ws, n, npad, k, SCAP_, CMP_, # R, SMP, TGT, Q, SS2, TGT2, # kv_lens, aim_base, sfac, amin, sd_en, tsh_en) @@ -1322,8 +1353,7 @@ def run_varlen( # ---- reference engine (differential oracle): b=1 host loop -------------- if _is_capturing(): raise RuntimeError( - "run_varlen reference engine reads kv_lens on host, " - "illegal under CUDA graph capture" + "run_varlen reference engine reads kv_lens on host, illegal under CUDA graph capture" ) # match the engine's flat-packed output convention for wider-than-k # buffers (pack ONCE from the tensor base, then slice per row) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 736407f00b5f..1690cec36414 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -128,9 +128,7 @@ def test_selfsampling_topk_exactness(batch_size, top_k, n_valid): @pytest.mark.parametrize("batch_size", [1, 4], ids=lambda b: f"bs{b}") -@pytest.mark.parametrize( - "top_k,n_valid", _SHORT_CASES, ids=[f"k{k}_n{n}" for k, n in _SHORT_CASES] -) +@pytest.mark.parametrize("top_k,n_valid", _SHORT_CASES, ids=[f"k{k}_n{n}" for k, n in _SHORT_CASES]) def test_selfsampling_topk_short_path(batch_size, top_k, n_valid): gen = torch.Generator(device=_DEV).manual_seed(top_k + n_valid) npad = (n_valid + 63) // 64 * 64 @@ -175,9 +173,7 @@ def test_selfsampling_topk_values_short_path(): top_k, n_valid, bs = 1024, 512, 4 gen = torch.Generator(device=_DEV).manual_seed(42) logits = torch.randn((bs, n_valid), generator=gen, dtype=torch.float32, device=_DEV) - pre_idx = torch.randint( - 0, n_valid, (bs, top_k), generator=gen, dtype=torch.int32, device=_DEV - ) + pre_idx = torch.randint(0, n_valid, (bs, top_k), generator=gen, dtype=torch.int32, device=_DEV) indices = torch.full((bs, top_k), -7, dtype=torch.int32, device=_DEV) values = torch.full((bs, top_k), 7.0, dtype=torch.float32, device=_DEV) ss_host.run(logits, pre_idx, n_valid, indices, values) @@ -235,8 +231,14 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut ) kv_lens = torch.tensor(kv, dtype=torch.int32, device=_DEV) ss_host.run_varlen( - logits, pre_idx, kv_lens, indices, - next_n=next_n, compress_ratio=cr, values=values, engine=engine, + logits, + pre_idx, + kv_lens, + indices, + next_n=next_n, + compress_ratio=cr, + values=values, + engine=engine, ) torch.cuda.synchronize() fmin = torch.finfo(torch.float32).min @@ -255,7 +257,9 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut assert int(torch.unique(idx).numel()) == top_k ref = torch.topk(logits[r, :n], top_k).values got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values - assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), f"row {r} inexact" + assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), ( + f"row {r} inexact" + ) if values is not None: assert torch.equal(values[r], torch.gather(logits[r], 0, idx)) @@ -283,7 +287,7 @@ def test_selfsampling_topk_varlen_engine_matches_reference(): equal the reference loop's on a mixed batch (deep SPLIT rows, tsh band, short rows, compressed space).""" kv = [524288, 131075, 32800, 2000, 65540, 8192, 262144, 900] - top_k, next_n, cr = 1024, 1, 4 + top_k, cr = 1024, 4 rows = len(kv) n_r = [(v - 1 + 1) // cr for v in kv] npad = (max(n_r) + 63) // 64 * 64 @@ -474,9 +478,7 @@ def refresh(step): assert int(idx.min()) >= 0 and int(idx.max()) < n assert int(torch.unique(idx).numel()) == top_k ref = torch.topk(logits[r, :n], top_k).values - got = torch.sort( - torch.gather(logits[r], 0, idx) + 0.0, descending=True - ).values + got = torch.sort(torch.gather(logits[r], 0, idx) + 0.0, descending=True).values assert torch.equal(got, torch.sort(ref + 0.0, descending=True).values), ( f"replay step {step} row {r} inexact (n={n})" ) @@ -513,8 +515,19 @@ def test_selfsampling_route_factorization(): for k in (512, 1024, 2048): ns = set() for c in ( - 2 * k, 3 * k, 4 * k + 64, 2560, 4096, 8192, 16384, - 4 * 1024, 4 * 4096, 4 * 32768, 65536, 131072, 262144, + 2 * k, + 3 * k, + 4 * k + 64, + 2560, + 4096, + 8192, + 16384, + 4 * 1024, + 4 * 4096, + 4 * 32768, + 65536, + 131072, + 262144, ): ns.update(v for v in range(c - 4, c + 5) if k < v <= npad) ns.update(range(k + 1, npad + 1, 4999)) @@ -524,14 +537,16 @@ def test_selfsampling_route_factorization(): ns.add(k + 1 + s % (npad - k - 1)) for n in sorted(ns): assert ss_host.route_split(b, n, npad, k) == ss_host.route(b, n, npad, k), ( - b, n, k, + b, + n, + k, ) checked += 1 assert checked > 10_000 def test_selfsampling_route_bands_contiguous(): - """route_bands must tile the envelope contiguously with n-free statics.""" + """route_bands must tile the envelope contiguously with n-free static fields.""" bands = ss_host.route_bands(8, 262144, 1024) assert bands[0][0] == 1025 and bands[-1][1] == 262144 for (_, h1, _), (l2, _, _) in zip(bands, bands[1:]): From fa12c67d6bd5ff9bf5c3b96bccd42821bf590aee Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 06:39:42 +0000 Subject: [PATCH 14/40] [TRTLLM-15293][chore] self-sampling top-K: demote run/run_ws to testing-only surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Public package exports now carry ONLY the production contract (selfsampling_topk_run_varlen: per-request device kv_lens, no batch-uniformity assumption — mirrors the single-op shape of the CUDA indexer_topk_decode integration). run/run_ws keep serving as the bench DPS contract and the reference-oracle plumbing (_run_impl is what engine="reference" walks row by row), but are no longer package-exported and carry TESTING/BENCH ONLY docstring warnings. Tests already import the host module directly — zero test churn. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../cute_dsl_kernels/blackwell/top_k/__init__.py | 6 ------ .../top_k/gvr_topk_decode_self_sampling_host.py | 10 ++++++++-- 2 files changed, 8 insertions(+), 8 deletions(-) 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 5e2cf2001e74..d5cab389489b 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 @@ -20,10 +20,7 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel -from .gvr_topk_decode_self_sampling_host import route as selfsampling_topk_route -from .gvr_topk_decode_self_sampling_host import run as selfsampling_topk_run from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen -from .gvr_topk_decode_self_sampling_host import run_ws as selfsampling_topk_run_ws from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -38,8 +35,5 @@ "DirectTopKKernel", "tiered_topk", "is_tiered_topk_supported", - "selfsampling_topk_route", - "selfsampling_topk_run", "selfsampling_topk_run_varlen", - "selfsampling_topk_run_ws", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 6db792b5005a..363123633e83 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1161,7 +1161,11 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): # exports (main.cpp:98-124) # --------------------------------------------------------------------------- def run(logits, pre_idx, n_valid, indices, values=None): - """Fast 4-arg form: signature-identical to the original candidate. + """TESTING/BENCH ONLY — production callers must use ``run_varlen`` (per-request + device kv_lens; this entry assumes one batch-uniform host ``n_valid``, + which real serving batches do not satisfy). + + Fast 4-arg form: signature-identical to the original candidate. ``values`` (optional DPS output, default None = OFF) mirrors the production values writeback; see _run_impl. Default per-device slab workspace resolved FIRST (main.cpp:99-102 -- @@ -1178,7 +1182,9 @@ def run(logits, pre_idx, n_valid, indices, values=None): def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): - """Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" + """TESTING/BENCH ONLY — production callers must use ``run_varlen(workspace=...)``. + + Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" validate_run_ws(workspace, logits) _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) From 18a287e4a3be1649b2b99e912cb69afa2b2f0f61 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:48:27 +0000 Subject: [PATCH 15/40] [TRTLLM-15293][fix] self-sampling top-K: dtype gate, arm observability, hint/domain tests Framework-integration audit hardening (all GPU-verified, 70/70 unit file): - dispatch gate now also requires fp32 logits (falls through loudly instead of feeding a non-fp32 tensor into the fp32-typed DSL engine; production DSA logits are always fp32 today, this is belt-and-braces for future paths) - logger.info_once on first engagement + logger.warning_once on first hardware-format fall-through: operators can tell which arm served without profiling - new tests: hints containing -1 (the production short-row pad tail that flows back through heuristic_prev_topk; engine's unsigned-compare guard verified exact) and route() domain up to 8192 rows (max_batch x next_n exceeds the bench grid's b<=1024 envelope) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 29 +++++++++++++++---- .../parallel/test_gvr_selfsampling_topk.py | 26 ++++++++++++++++- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 905b171f0469..195d20c0aba2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1833,19 +1833,38 @@ def sparse_attn_indexer( if not metadata.use_cute_dsl_topk: heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens] - if ( + _ss_ready = ( self._use_self_sampling_topk and self._enable_heuristic_topk and pre_idx is not None + ) + if _ss_ready and not ( # engine hardware-format gate (falls through otherwise): - # row-major logits with a float4-aligned row stride and a - # 16B-aligned base (the DSL paged-MQA arena view — column- - # sliced from a 256-aligned buffer — satisfies this; odd - # max_seq_len DeepGEMM layouts do not) + # fp32 row-major logits with a float4-aligned row stride + # and a 16B-aligned base (the DSL paged-MQA arena view — + # column-sliced from a 256-aligned buffer — satisfies + # this; odd max_seq_len DeepGEMM layouts do not) + logits_decode.dtype == torch.float32 and logits_decode.stride(1) == 1 and logits_decode.stride(0) % 4 == 0 and logits_decode.data_ptr() % 16 == 0 ): + logger.warning_once( + "TRTLLM_GVR_SELF_SAMPLING=1 but the decode logits do " + "not satisfy the engine's hardware-format gate " + f"(dtype={logits_decode.dtype}, " + f"strides={tuple(logits_decode.stride())}); falling " + "through to the in-tree top-K path.", + key="selfsampling_topk_fallthrough", + ) + _ss_ready = False + if _ss_ready: + logger.info_once( + "self-sampling GVR top-K engaged " + f"(K={self.index_topk}, cr={self.compress_ratio}, " + f"next_n={next_n}).", + key="selfsampling_topk_engaged", + ) # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): # one launch for the batch; per-row n from device kv_lens, # capture-stable tuning from the max-seq-len engine diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 1690cec36414..9aed6429ba7f 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -184,7 +184,9 @@ def test_selfsampling_topk_values_short_path(): assert bool((indices[:, n_valid:] == -1).all()) -@pytest.mark.parametrize("hint_kind", ["all_zero", "all_same", "all_max", "half_dup"]) +@pytest.mark.parametrize( + "hint_kind", ["all_zero", "all_same", "all_max", "half_dup", "minus_one_tail", "all_minus_one"] +) @pytest.mark.parametrize( "top_k,n_valid", [(512, 8192), (2048, 131075)], ids=["k512_n8192", "k2048_n131075"] ) @@ -200,6 +202,16 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): pre_idx = torch.full((2, top_k), 1234, dtype=torch.int32, device=_DEV) elif hint_kind == "all_max": pre_idx = torch.full((2, top_k), n_valid - 1, dtype=torch.int32, device=_DEV) + elif hint_kind == "minus_one_tail": + # production short-row pad convention writes -1 tails into prev_topk; + # the next step feeds them back as hints — must stay in-bounds + gen = torch.Generator(device=_DEV).manual_seed(2) + pre_idx = torch.randint( + 0, n_valid, (2, top_k), generator=gen, dtype=torch.int32, device=_DEV + ) + pre_idx[:, top_k // 3 :] = -1 + elif hint_kind == "all_minus_one": + pre_idx = torch.full((2, top_k), -1, dtype=torch.int32, device=_DEV) else: gen = torch.Generator(device=_DEV).manual_seed(1) pre_idx = torch.randint( @@ -504,6 +516,18 @@ def test_selfsampling_topk_guards(): ss_host.run(logits, pre_idx[:0], 8192, indices) +def test_selfsampling_route_large_batch_domain(): + """Production num_rows = max_batch_size * next_n can exceed the bench + grid's b<=1024 envelope (e.g. 1024 * 4 = 4096 rows). route() is a pure + function — assert the full domain stays well-formed up to 8192 rows.""" + for k in (512, 1024, 2048): + for n in (k + 1, 4096, 65536, 262144): + npad = (n + 63) // 64 * 64 + for b in (1536, 2048, 4096, 8192): + r = ss_host.route(b, n, npad, k) + assert r is not None and len(r) >= 2, (b, n, k, r) + + def test_selfsampling_route_factorization(): """Two-time-scale dispatch groundwork: route() must factor losslessly into route_static (constant on n-bands, freezable at capture time) and From 893eead779238b7046b416cc616a05218b232291 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:48:30 +0000 Subject: [PATCH 16/40] [TRTLLM-15293][perf] self-sampling top-K: warmup_varlen pre-compile helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors warmup_heuristic_topk_decode / warmup_cute_dsl_radix_topk: one tiny real launch per requested num_rows compiles the varlen engine's envelope tuples so no live request pays the first-touch DSL JIT (measured: 9.1 s cold compile for two tuples on a fresh cache; idempotent re-call 0 ms; post-warmup first real call 0.1 ms). Exposed as a module-level helper — CUDA-graph capture warmup already compiles the captured batch sizes, and wiring an automatic init hook needs max_seq_len plumbing that is not available at Indexer.__init__ time (follow-up). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 363123633e83..ddd10af84f08 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1406,3 +1406,54 @@ def run_varlen( "validate_run_ws", "kernel_view", ] + + +# -------------------------------------------------------------------------- +# warmup: pre-compile the varlen engine for an engine envelope so no live +# request pays the first-touch DSL JIT (mirrors warmup_heuristic_topk_decode +# and warmup_cute_dsl_radix_topk). Idempotent per (device, geometry) key. +# CUDA-graph capture warmup naturally compiles the captured batch sizes; +# this covers the eager/first-touch path (num_rows defaults to (1,)). +_VARLEN_WARMUP_DONE: set = set() +_VARLEN_WARMUP_LOCK = threading.Lock() + + +def warmup_varlen(top_k, max_seq_len, compress_ratio=1, next_n=1, num_rows_list=(1,)): + """TESTING/INIT ONLY — compile the varlen engine's envelope tuples. + + One tiny real launch per requested ``num_rows`` (compile keys do not + depend on tensor contents). Uses the current CUDA device. + """ + dev = torch.cuda.current_device() + key = ( + dev, + int(top_k), + int(max_seq_len), + int(compress_ratio), + int(next_n), + tuple(int(r) for r in num_rows_list), + ) + with _VARLEN_WARMUP_LOCK: + if key in _VARLEN_WARMUP_DONE: + return + _VARLEN_WARMUP_DONE.add(key) + n_env = max(1, int(max_seq_len) // int(compress_ratio)) + npad = (n_env + 63) // 64 * 64 + for rows in key[5]: + rows = max(int(next_n), rows - rows % int(next_n) or int(next_n)) + batch = rows // int(next_n) + logits = torch.zeros((rows, npad), dtype=torch.float32, device=dev) + kv_lens = torch.full((batch,), int(max_seq_len), dtype=torch.int32, device=dev) + pre_idx = torch.zeros((batch, int(top_k)), dtype=torch.int32, device=dev) + out = torch.empty((rows, int(top_k)), dtype=torch.int32, device=dev) + run_varlen( + logits, + pre_idx, + kv_lens, + out, + next_n=int(next_n), + compress_ratio=int(compress_ratio), + max_seq_len=int(max_seq_len), + ) + del logits, kv_lens, pre_idx, out + torch.cuda.synchronize() From 0d60f17e8f4b495d22e5b5692dcad44de05b6218 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:26:13 +0000 Subject: [PATCH 17/40] [TRTLLM-15293][perf] self-sampling top-K: wire warmup into ModelEngine.warmup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the warmup story: DSAtrtllmAttentionMetadata.warmup_selfsampling_topk mirrors warmup_cute_dsl_radix_topk (same ModelEngine hook, which is where max_seq_len is actually available — Indexer.__init__ is not). Gated on the env flag + the same init-contract conditions; compiles the eager first-touch (num_rows=next_n) tuple via warmup_varlen so no live request pays the DSL JIT. Also: seam init comment now states the two-guard reality (format gate falls through with a one-time warning; in-engine contract violations raise), and run_varlen documents the inherited NaN-ordering limitation (finite inputs incl. +/-inf are tie-aware exact). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 6 ++-- .../attention_backend/sparse/dsa/metadata.py | 35 +++++++++++++++++++ .../gvr_topk_decode_self_sampling_host.py | 6 ++++ .../_torch/pyexecutor/model_engine.py | 2 ++ 4 files changed, 47 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 195d20c0aba2..e3a6626ee177 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -697,8 +697,10 @@ def __init__( # kv_lens, request-level raw prev-top-K hints, per-row MTP window, # in-kernel n <= topK short path); tuning is frozen from # indexer_max_seq_len at capture time, so the launch is - # CUDA-graph-replay safe. Contract violations raise loudly rather - # than silently falling back — this flag is an explicit experiment. + # CUDA-graph-replay safe. Two guard mechanisms coexist: the + # dispatch-site hardware-format gate falls through to the in-tree + # path (with a one-time warning), while contract violations inside + # the engine raise loudly — this flag is an explicit experiment. self._use_self_sampling_topk = ( os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" and IS_CUTLASS_DSL_AVAILABLE diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index a41631cda0e1..e8587ec238d3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional @@ -277,6 +278,40 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: num_sms=self.num_sms, ) + def warmup_selfsampling_topk(self, next_n: int) -> None: + """Pre-compile the self-sampling GVR varlen engine during warmup. + + Mirrors ``warmup_cute_dsl_radix_topk``: captured geometries compile in + the warmup-step forwards; this covers the eager first-touch (bs=1) + tuple so a live request never pays the DSL JIT. No-op unless the + opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + """ + if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": + return + if not self.enable_heuristic_topk or self.kv_cache_manager is None: + return + top_k = getattr(self.sparse_metadata_params, "index_topk", None) + if not top_k or int(top_k) not in (512, 1024, 2048): + return + cr = int(self._indexer_compress_ratio) if self._indexer_compress_ratio else 1 + if cr not in (1, 4): + return + try: + from ....cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling_host as _ss_host, + ) + except ImportError: + return + # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len + # is compressed — same multiply-back as the dispatch seam) + _ss_host.warmup_varlen( + int(top_k), + int(self.get_indexer_max_seq_len()) * cr, + compress_ratio=cr, + next_n=int(next_n), + num_rows_list=(int(next_n),), + ) + def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. # Especially for the changes in the _preprocess_inputs() of model_engine.py. diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index ddd10af84f08..fb26aa58eae6 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1225,6 +1225,12 @@ def run_varlen( ``kv_lens.max()`` host read (documented sync, refused under capture). ``engine="reference"`` keeps the b=1 host-loop reference implementation — the differential oracle the in-kernel engine is validated against. + + KNOWN LIMITATION: on rows containing NaN logits the selected index SET + can differ from ``heuristicTopKDecode.cu`` (both kernels order NaNs + implementation-specifically; inherited from the translation campaign's + probe battery). Finite inputs — including +/-inf and denormals — are + tie-aware exact. """ if not (isinstance(kv_lens, _TENSOR) and kv_lens.is_cuda): raise RuntimeError("kv_lens must be a CUDA tensor") diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b99ccbeeea1c..2f5a9e05eb7a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1697,6 +1697,8 @@ def _warmup_cute_dsl_radix_topk(self) -> None: if isinstance(attn_meta, DSAtrtllmAttentionMetadata): next_n = 1 + self.original_max_draft_len attn_meta.warmup_cute_dsl_radix_topk(next_n) + if hasattr(attn_meta, "warmup_selfsampling_topk"): + attn_meta.warmup_selfsampling_topk(next_n) def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]): From cde20646e804b16ffe52f3d068980517bc3b5d13 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:09:15 +0000 Subject: [PATCH 18/40] [TRTLLM-15293][fix] self-sampling top-K: engine-level fp32 contract + next_n=3 coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit run_varlen now rejects non-fp32 logits with a clear contract error (the dispatch seam falls through before this; direct callers previously hit an opaque CuTe typing failure) — regression-tested. The varlen differential gains a next_n=3 (MTP2) case: no production config uses it today, but the per-row window formula is now verified to generalize (both engines exact), closing the gap noted in the PR's MTP-coverage section. GPU-verified: 3/3 targeted tests green. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../gvr_topk_decode_self_sampling_host.py | 5 +++++ .../parallel/test_gvr_selfsampling_topk.py | 18 +++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index fb26aa58eae6..4d403fd391fa 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1232,6 +1232,11 @@ def run_varlen( probe battery). Finite inputs — including +/-inf and denormals — are tie-aware exact. """ + if logits.dtype is not torch.float32: + raise RuntimeError( + f"logits must be float32 (got {logits.dtype}); bf16/fp16 paths " + "are a follow-up — see the PR roadmap" + ) if not (isinstance(kv_lens, _TENSOR) and kv_lens.is_cuda): raise RuntimeError("kv_lens must be a CUDA tensor") if kv_lens.dtype is not _I32: diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 9aed6429ba7f..1ea545bbfdde 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -284,8 +284,11 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut ([131075, 32800, 2000], 1, 4, 512), # v4-style compressed index space ([9000, 5001], 2, 1, 512), # MTP: n varies per row within a request ([65540], 4, 4, 1024), # MTP: compressed-boundary-crossing rows + ([40000, 7003], 3, 4, 512), # MTP2 (next_n=3): no production config + # uses it today, but the window formula must generalize (gap noted + # in the PR's MTP-coverage section) ], - ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4"], + ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], ) def test_selfsampling_topk_varlen(kv, next_n, cr, top_k, engine): """run_varlen production contract: per-row n from device kv_lens with the @@ -588,3 +591,16 @@ def test_selfsampling_dispatch_is_pure_and_total(): r = ss_host.route(4, n, npad, k) assert r["kernel"] in ("main", "reg", "clus", "reg_clus") assert r["block"] >= 128 and r["grid"][0] >= 1 + + +@pytest.mark.skipif(getSMVersion() < 100, reason="sm100+") +def test_selfsampling_topk_varlen_rejects_non_fp32(): + """Engine-level dtype contract: bf16/fp16 logits must raise a clear + error (the dispatch seam falls through before this; direct callers get + the loud contract message instead of a CuTe typing failure).""" + logits = torch.randn(1, 8192, device=_DEV, dtype=torch.bfloat16) + kv = torch.tensor([8000], dtype=torch.int32, device=_DEV) + pre = torch.zeros(1, 512, dtype=torch.int32, device=_DEV) + out = torch.empty(1, 512, dtype=torch.int32, device=_DEV) + with pytest.raises(RuntimeError, match="float32"): + ss_host.run_varlen(logits, pre, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) From cdf33967bf363732afc8645b17608d445b0fd3d5 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:43:57 +0000 Subject: [PATCH 19/40] [TRTLLM-15293][fix] self-sampling top-K: review round 1 (warmup coverage + robustness) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adopted from the first review round: - warmup_varlen records its done-key only after every launch succeeds, so a failed or interrupted warmup retries instead of short-circuiting to an uncompiled engine; one max-geometry allocation is reused across row counts (prefix views — compile keys depend on shapes only). - The ModelEngine warmup hook now passes the configured CUDA-graph batch sizes: the varlen launcher is keyed by the exact row count, so capture no longer depends on a prior warmup-forward having compiled its key, and warmed geometries never pay first-touch JIT. Best-effort under OOM (warn + lazy JIT) since dispatch works without warmup outside capture. - route/route_streaming reject b < 1 with a contract error instead of ZeroDivisionError. - run/_run_impl: the launch try block no longer swallows routing/compile errors as 'launch failed', and the values epilogue matches run_varlen (clamp + mask instead of a context-poisoning gather assert). - Stale module docstring refreshed (run_varlen IS the wired production entry); warmup_varlen exported in __all__; public entry points annotated. - Test SM gate aligned with the production dispatch gate (>= 100; the module previously skipped everywhere except exactly sm_100). GPU-verified: full battery 73/73; smoke 8/8 (b<1 guards, multi-rows warmup + idempotency, CUDA-graph capture at a warmed row count exact, unwarmed capture raises loudly, injected-failure warmup retries with a real recompile). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 42 +++-- .../gvr_topk_decode_self_sampling_host.py | 148 +++++++++++------- .../_torch/pyexecutor/model_engine.py | 3 +- .../parallel/test_gvr_selfsampling_topk.py | 9 +- 4 files changed, 131 insertions(+), 71 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index e8587ec238d3..332ef1d5cde8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata +from tensorrt_llm.logger import logger from .cache_manager import is_dsa_cache_manager from .indexer import ( @@ -278,13 +279,19 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: num_sms=self.num_sms, ) - def warmup_selfsampling_topk(self, next_n: int) -> None: + def warmup_selfsampling_topk( + self, next_n: int, batch_sizes: Optional[List[int]] = None + ) -> None: """Pre-compile the self-sampling GVR varlen engine during warmup. - Mirrors ``warmup_cute_dsl_radix_topk``: captured geometries compile in - the warmup-step forwards; this covers the eager first-touch (bs=1) - tuple so a live request never pays the DSL JIT. No-op unless the - opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + Mirrors ``warmup_cute_dsl_radix_topk``. The varlen launcher is keyed + by the exact row count, so this warms the eager first-touch tuple + (bs=1) plus every configured CUDA-graph batch size — capture then + never depends on a prior warmup-forward having compiled its key, and + a live request never pays the DSL JIT for those geometries. Eager + batches outside ``batch_sizes`` still compile lazily on first touch. + No-op unless the opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects + the engine. """ if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": return @@ -302,15 +309,26 @@ def warmup_selfsampling_topk(self, next_n: int) -> None: ) except ImportError: return + rows = {int(next_n)} + for bs in batch_sizes or (): + rows.add(int(bs) * int(next_n)) # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len # is compressed — same multiply-back as the dispatch seam) - _ss_host.warmup_varlen( - int(top_k), - int(self.get_indexer_max_seq_len()) * cr, - compress_ratio=cr, - next_n=int(next_n), - num_rows_list=(int(next_n),), - ) + try: + _ss_host.warmup_varlen( + int(top_k), + int(self.get_indexer_max_seq_len()) * cr, + compress_ratio=cr, + next_n=int(next_n), + num_rows_list=tuple(sorted(rows)), + ) + except torch.cuda.OutOfMemoryError: + # warmup is best-effort: the dispatch works without it (engines + # JIT lazily outside capture), so do not fail engine init + logger.warning( + "self-sampling GVR warmup ran out of memory; varlen engines " + "will JIT-compile lazily on first touch instead." + ) def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 4d403fd391fa..32226830ab9b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -38,18 +38,20 @@ on real V3.2 decode captures raw prev-step hints land on MORE of the current top-K than +1-shifted ones (mean overlap 0.773 vs 0.536 across 15 cells x 14 consecutive step-pairs, the gap widening with ISL), so one offset-free hint -convention serves all three models. The production decode engine -instead reads per-request ``seq_lens`` on-device with per-row MTP offsets -(sync-free, CUDA-graph-replay safe with growing KV); adopting that per-row -contract inside these kernels is tracked follow-up work. Until then this -module is exercised standalone (unit tests / benchmarking) and must not be -substituted for the tiered path under continuous batching, MTP -(``next_n > 1``), or CUDA-graph capture. +convention serves all three models. The production per-row contract +(per-request ``kv_lens`` read on-device, per-row MTP offsets — sync-free and +CUDA-graph-replay safe with growing KV) is implemented by ``run_varlen``, +which is the entry the opt-in DSA dispatch seam calls. The batch-uniform +``run``/``run_ws`` entries keep the original standalone contract (one +host-side ``n_valid`` for the whole batch), are exercised for unit tests and +benchmarking only, and must not be substituted for the tiered path under +continuous batching, MTP (``next_n > 1``), or CUDA-graph capture. """ import math import operator import threading +from collections.abc import Sequence import torch @@ -148,8 +150,10 @@ def _device(): BLKC = 1024 # L2374 CTA size of the clustered register path -def route(b, n, npad, k): +def route(b: int, n: int, npad: int, k: int) -> dict[str, object]: """Mirror of gvr_topk_launch (kernel.cu L2754-3197). Pure. See module doc.""" + if b < 1: + raise RuntimeError(f"route requires b >= 1, got {b}") wide = b <= 148 # L2757 # ================= register-resident block (L2758-2949) ================= @@ -490,7 +494,7 @@ def _main(BLK, MINB, U, SPLIT): _DYN_SMEM = ("reg", "regimg") # smem depends on CMP/IMGW -> recomputed per n -def route_static(b, n, npad, k): +def route_static(b: int, n: int, npad: int, k: int) -> dict[str, object]: """route() with the n-continuous fields redacted (see _DYN_RT/_DYN_SMEM). Constant on maximal n-intervals ("bands"); every redacted field is reconstructible from (static, n) by route_dynamic.""" @@ -503,7 +507,7 @@ def route_static(b, n, npad, k): return st -def route_dynamic(static, n): +def route_dynamic(static: dict[str, object], n: int) -> tuple[dict[str, object], int]: """Recompute the redacted n-continuous scalars from (static, n). Returns (rt_updates, smem). Transcribed independently from route() — the factorization fuzz is the equivalence proof, and the device-side @@ -581,7 +585,7 @@ def route_dynamic(static, n): ) -def route_split(b, n, npad, k): +def route_split(b: int, n: int, npad: int, k: int) -> dict[str, object]: """route_static + route_dynamic recombined — must equal route() exactly (the factorization fuzz in the unit tests asserts this).""" st = route_static(b, n, npad, k) @@ -592,7 +596,9 @@ def route_split(b, n, npad, k): return plan -def route_streaming(b, n, npad, k, force_main=False): +def route_streaming( + b: int, n: int, npad: int, k: int, force_main: bool = False +) -> dict[str, object]: """route() restricted to its STREAMING half (main / clus) — the varlen capture policy: per-row kernels must be picked from the families that are correct for ANY row length, so the register-resident specialists are @@ -601,6 +607,8 @@ def route_streaming(b, n, npad, k, force_main=False): 110,003/110,003 agreement). force_main additionally skips the clus rounding (v1 varlen engine ships the gvr_main port first; the raw min(r1, r2) R then matches the CUDA else-branch exactly).""" + if b < 1: + raise RuntimeError(f"route_streaming requires b >= 1, got {b}") R = 1 if b <= 32: r1 = max(148 // b, 1) @@ -759,7 +767,9 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): return lc -def route_bands(b, npad, k, n_lo=None, n_hi=None): +def route_bands( + b: int, npad: int, k: int, n_lo: int | None = None, n_hi: int | None = None +) -> list[tuple[int, int, dict[str, object]]]: """Enumerate maximal n-intervals on which route_static is constant. Dense O(n_hi - n_lo) scan of the pure host dispatch — an offline / engine-init tool (seconds for the 262144-token envelope), NOT a hot @@ -1142,11 +1152,11 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): key = (b, n, npad, k) lc = _LAUNCH_CACHE.get(key) + if lc is None: + lc = _build_launcher(b, n, npad, k) + _LAUNCH_CACHE[key] = lc + fn, args, needs_ws = lc try: - if lc is None: - lc = _build_launcher(b, n, npad, k) - _LAUNCH_CACHE[key] = lc - fn, args, needs_ws = lc if needs_ws: fn(logits, pre_idx, indices, ws, *args) else: @@ -1154,13 +1164,23 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): except Exception as e: raise RuntimeError(f"gvr_topk launch failed (b={b} n={n} npad={npad} k={k}): {e}") from e if values is not None: - values.copy_(logits.gather(1, indices.to(torch.int64))) + # same epilogue as run_varlen: a (never-expected) negative index + # degrades to -FLT_MAX instead of a context-poisoning device assert + idx64 = indices.to(torch.int64) + values.copy_(logits.gather(1, idx64.clamp_min(0))) + values.masked_fill_(indices < 0, torch.finfo(_F32).min) # --------------------------------------------------------------------------- # exports (main.cpp:98-124) # --------------------------------------------------------------------------- -def run(logits, pre_idx, n_valid, indices, values=None): +def run( + logits: torch.Tensor, + pre_idx: torch.Tensor, + n_valid: int, + indices: torch.Tensor, + values: torch.Tensor | None = None, +) -> None: """TESTING/BENCH ONLY — production callers must use ``run_varlen`` (per-request device kv_lens; this entry assumes one batch-uniform host ``n_valid``, which real serving batches do not satisfy). @@ -1181,7 +1201,14 @@ def run(logits, pre_idx, n_valid, indices, values=None): _run_impl(logits, pre_idx, n_valid, indices, ws, values) -def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): +def run_ws( + logits: torch.Tensor, + pre_idx: torch.Tensor, + n_valid: int, + indices: torch.Tensor, + workspace: torch.Tensor, + values: torch.Tensor | None = None, +) -> None: """TESTING/BENCH ONLY — production callers must use ``run_varlen(workspace=...)``. Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" @@ -1190,17 +1217,17 @@ def run_ws(logits, pre_idx, n_valid, indices, workspace, values=None): def run_varlen( - logits, - pre_idx, - kv_lens, - indices, - next_n=1, - compress_ratio=1, - values=None, - max_seq_len=None, - engine="auto", - workspace=None, -): + logits: torch.Tensor, + pre_idx: torch.Tensor, + kv_lens: torch.Tensor, + indices: torch.Tensor, + next_n: int = 1, + compress_ratio: int = 1, + values: torch.Tensor | None = None, + max_seq_len: int | None = None, + engine: str = "auto", + workspace: torch.Tensor | None = None, +) -> None: """Production-contract varlen entry (per-row device kv_lens). Row semantics (mirror of ``heuristicTopKDecode.cu`` and the in-tree @@ -1411,6 +1438,7 @@ def run_varlen( "run", "run_ws", "run_varlen", + "warmup_varlen", "workspace_bytes", "WS_BYTES", "default_workspace", @@ -1429,42 +1457,52 @@ def run_varlen( _VARLEN_WARMUP_LOCK = threading.Lock() -def warmup_varlen(top_k, max_seq_len, compress_ratio=1, next_n=1, num_rows_list=(1,)): +def warmup_varlen( + top_k: int, + max_seq_len: int, + compress_ratio: int = 1, + next_n: int = 1, + num_rows_list: Sequence[int] = (1,), +) -> None: """TESTING/INIT ONLY — compile the varlen engine's envelope tuples. One tiny real launch per requested ``num_rows`` (compile keys do not - depend on tensor contents). Uses the current CUDA device. + depend on tensor contents). Uses the current CUDA device. The done-key + is recorded only after every launch succeeds, so a failed or interrupted + warmup is retried on the next call instead of short-circuiting to an + uncompiled engine. """ dev = torch.cuda.current_device() - key = ( - dev, - int(top_k), - int(max_seq_len), - int(compress_ratio), - int(next_n), - tuple(int(r) for r in num_rows_list), - ) + nn = max(1, int(next_n)) + # round each request down to a next_n multiple (min next_n) and dedup + rows_list = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) + if not rows_list: + return + key = (dev, int(top_k), int(max_seq_len), int(compress_ratio), nn, tuple(rows_list)) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: return - _VARLEN_WARMUP_DONE.add(key) n_env = max(1, int(max_seq_len) // int(compress_ratio)) npad = (n_env + 63) // 64 * 64 - for rows in key[5]: - rows = max(int(next_n), rows - rows % int(next_n) or int(next_n)) - batch = rows // int(next_n) - logits = torch.zeros((rows, npad), dtype=torch.float32, device=dev) - kv_lens = torch.full((batch,), int(max_seq_len), dtype=torch.int32, device=dev) - pre_idx = torch.zeros((batch, int(top_k)), dtype=torch.int32, device=dev) - out = torch.empty((rows, int(top_k)), dtype=torch.int32, device=dev) + rows_max = rows_list[-1] + # one allocation at the largest geometry; smaller row counts run on + # contiguous prefix views (compile keys depend on shapes only) + logits = torch.zeros((rows_max, npad), dtype=torch.float32, device=dev) + kv_lens = torch.full((rows_max // nn,), int(max_seq_len), dtype=torch.int32, device=dev) + pre_idx = torch.zeros((rows_max // nn, int(top_k)), dtype=torch.int32, device=dev) + out = torch.empty((rows_max, int(top_k)), dtype=torch.int32, device=dev) + for rows in rows_list: + batch = rows // nn run_varlen( - logits, - pre_idx, - kv_lens, - out, - next_n=int(next_n), + logits[:rows], + pre_idx[:batch], + kv_lens[:batch], + out[:rows], + next_n=nn, compress_ratio=int(compress_ratio), max_seq_len=int(max_seq_len), ) - del logits, kv_lens, pre_idx, out + del logits, kv_lens, pre_idx, out torch.cuda.synchronize() + with _VARLEN_WARMUP_LOCK: + _VARLEN_WARMUP_DONE.add(key) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2f5a9e05eb7a..a5d24a40955f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1698,7 +1698,8 @@ def _warmup_cute_dsl_radix_topk(self) -> None: next_n = 1 + self.original_max_draft_len attn_meta.warmup_cute_dsl_radix_topk(next_n) if hasattr(attn_meta, "warmup_selfsampling_topk"): - attn_meta.warmup_selfsampling_topk(next_n) + attn_meta.warmup_selfsampling_topk( + next_n, batch_sizes=self._cuda_graph_batch_sizes) def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]): diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 1ea545bbfdde..4278da4b4564 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -40,8 +40,12 @@ if not IS_CUTLASS_DSL_AVAILABLE: pytest.skip("cutlass DSL is required for gvr_selfsampling_topk tests", allow_module_level=True) -if getSMVersion() != 100: - pytest.skip("self-sampling GVR kernels target Blackwell sm_100", allow_module_level=True) +if getSMVersion() < 100: + pytest.skip( + "self-sampling GVR kernels require Blackwell (SM100+) — same gate as " + "the production dispatch", + allow_module_level=True, + ) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( gvr_topk_decode_self_sampling_host as ss_host, @@ -593,7 +597,6 @@ def test_selfsampling_dispatch_is_pure_and_total(): assert r["block"] >= 128 and r["grid"][0] >= 1 -@pytest.mark.skipif(getSMVersion() < 100, reason="sm100+") def test_selfsampling_topk_varlen_rejects_non_fp32(): """Engine-level dtype contract: bf16/fp16 logits must raise a clear error (the dispatch seam falls through before this; direct callers get From 74af5174e88945c89b23a27678c0c7fa85cbbb29 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:28:16 +0000 Subject: [PATCH 20/40] [TRTLLM-15293][fix] self-sampling top-K: warm the launcher key dispatch actually looks up MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The varlen launcher key includes the logits row stride (npad), and the two sides derived it differently: the DSL paged-MQA arena rounds each row up to 256 elements (CuteDSLPagedMQALogitsRunner: compute_block_kv=128, SPLIT_KV= 256), while warmup_varlen synthesized its logits at 64-element rounding. Whenever the indexer max_seq_len is not a 256 multiple, warmup compiled a variant dispatch never looks up: with CUDA graphs the pre-capture warmup forwards masked this (init just wasted a few seconds on phantom keys), but in eager serving the first live decode paid the full first-touch DSL JIT — exactly the stall the helper exists to remove. Never a correctness issue: a key miss lazily JITs the true variant. - warmup_varlen gains row_stride; the caller passes the producer's actual row stride and the done-key includes it. - metadata.warmup_selfsampling_topk mirrors the active producer: DSL paged-MQA -> round up to 256; DeepGEMM -> exact width (non-float4 widths fall through at the dispatch format gate, so nothing to warm). A future drift only degrades warmup back to unused keys, never to wrong results. - warmup now also applies the same hardware gates as the dispatch flag (CUTLASS DSL available, SM100+), so setting the env var on an unsupported stack can no longer make engine warmup compile Blackwell kernels. - New tests: capture immediately after a row_stride warmup on a column-sliced arena view (fails loudly with the old 64-rounding), and kv_lens < next_n zero-window rows (graph-padding dummies) pad-emit -1 while cohabiting normal requests stay exact. GPU-verified: full battery 75/75 (73 prior + 2 new). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 37 +++++++++++--- .../gvr_topk_decode_self_sampling_host.py | 21 ++++++-- .../parallel/test_gvr_selfsampling_topk.py | 50 +++++++++++++++++++ 3 files changed, 97 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 332ef1d5cde8..8383a04deaf7 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -285,16 +285,22 @@ def warmup_selfsampling_topk( """Pre-compile the self-sampling GVR varlen engine during warmup. Mirrors ``warmup_cute_dsl_radix_topk``. The varlen launcher is keyed - by the exact row count, so this warms the eager first-touch tuple - (bs=1) plus every configured CUDA-graph batch size — capture then - never depends on a prior warmup-forward having compiled its key, and - a live request never pays the DSL JIT for those geometries. Eager - batches outside ``batch_sizes`` still compile lazily on first touch. - No-op unless the opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects - the engine. + by the exact row count AND the logits row stride: rows cover the + eager first-touch tuple (bs=1) plus every configured CUDA-graph + batch size, and the stride mirrors what the active paged-MQA + producer emits (the DSL arena rounds the row up to 256 elements; + DeepGEMM is exact-width) so the warmed keys are the ones dispatch + actually looks up. Captured geometries are also compiled by the + pre-capture warmup forwards; eager batches outside ``batch_sizes`` + still compile lazily on first touch. No-op unless the opt-in gate + (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. """ if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": return + # same hardware gates as the dispatch flag (indexer __init__): never + # compile Blackwell kernels on unsupported stacks during warmup + if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() < 100: + return if not self.enable_heuristic_topk or self.kv_cache_manager is None: return top_k = getattr(self.sparse_metadata_params, "index_topk", None) @@ -312,15 +318,30 @@ def warmup_selfsampling_topk( rows = {int(next_n)} for bs in batch_sizes or (): rows.add(int(bs) * int(next_n)) + msl_c = int(self.get_indexer_max_seq_len()) + if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits: + # mirror the DSL paged-MQA arena stride (cute_dsl_custom_ops + # CuteDSLPagedMQALogitsRunner: compute_block_kv=128, SPLIT_KV= + # 2*128 -> rows round up to 256 elements). A drift here only + # degrades warmup to unused keys — dispatch still lazy-JITs the + # true key outside capture, so it can never become incorrect. + row_stride = (msl_c + 255) // 256 * 256 + else: + # DeepGEMM emits exact-width rows; a non-float4 width falls + # through at the dispatch format gate, so there is nothing to warm + row_stride = msl_c + if row_stride % 4: + return # helper takes max_seq_len in kv-token space (get_indexer_max_seq_len # is compressed — same multiply-back as the dispatch seam) try: _ss_host.warmup_varlen( int(top_k), - int(self.get_indexer_max_seq_len()) * cr, + msl_c * cr, compress_ratio=cr, next_n=int(next_n), num_rows_list=tuple(sorted(rows)), + row_stride=row_stride, ) except torch.cuda.OutOfMemoryError: # warmup is best-effort: the dispatch works without it (engines diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 32226830ab9b..a462d7b6e6d2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -1463,6 +1463,7 @@ def warmup_varlen( compress_ratio: int = 1, next_n: int = 1, num_rows_list: Sequence[int] = (1,), + row_stride: int | None = None, ) -> None: """TESTING/INIT ONLY — compile the varlen engine's envelope tuples. @@ -1471,6 +1472,13 @@ def warmup_varlen( is recorded only after every launch succeeds, so a failed or interrupted warmup is retried on the next call instead of short-circuiting to an uncompiled engine. + + ``row_stride`` must be the logits row stride the serving producer will + emit: the launcher key includes it, so a warmup at a different stride + compiles a variant dispatch never looks up. Callers that know the + producer layout (e.g. the DSL paged-MQA arena's 256-element rounding) + must pass it; the 64-element default only matches producers that round + the same way. """ dev = torch.cuda.current_device() nn = max(1, int(next_n)) @@ -1478,12 +1486,19 @@ def warmup_varlen( rows_list = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) if not rows_list: return - key = (dev, int(top_k), int(max_seq_len), int(compress_ratio), nn, tuple(rows_list)) + n_env = max(1, int(max_seq_len) // int(compress_ratio)) + if row_stride is None: + npad = (n_env + 63) // 64 * 64 + else: + npad = int(row_stride) + if npad < n_env or npad % 4: + raise RuntimeError( + f"row_stride must be a float4-multiple >= n_env={n_env}, got {row_stride}" + ) + key = (dev, int(top_k), int(max_seq_len), int(compress_ratio), nn, tuple(rows_list), npad) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: return - n_env = max(1, int(max_seq_len) // int(compress_ratio)) - npad = (n_env + 63) // 64 * 64 rows_max = rows_list[-1] # one allocation at the largest geometry; smaller row counts run on # contiguous prefix views (compile keys depend on shapes only) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 4278da4b4564..f5fa763362d1 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -607,3 +607,53 @@ def test_selfsampling_topk_varlen_rejects_non_fp32(): out = torch.empty(1, 512, dtype=torch.int32, device=_DEV) with pytest.raises(RuntimeError, match="float32"): ss_host.run_varlen(logits, pre, kv, out, next_n=1, compress_ratio=4, max_seq_len=32768) + + +def test_selfsampling_topk_varlen_zero_window_rows(): + """CUDA-graph padding dummy rows: kv_lens < next_n makes every MTP-window + row's valid length n <= 0. The kernel must clamp those rows onto the + zero-work short path and pad-emit -1 for the whole row, while the normal + request in the same batch stays exact.""" + torch.manual_seed(0) + k, nn, cr, msl = 512, 4, 4, 40000 + kv = torch.tensor([msl, 1], dtype=torch.int32, device=_DEV) + rows = kv.numel() * nn + npad = (msl // cr + 63) // 64 * 64 + logits = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) + pre = torch.zeros(kv.numel(), k, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + torch.cuda.synchronize() + assert (out[nn:] == -1).all().item(), "n<=0 rows must be fully -1-padded" + for r in range(nn): + n_r = (msl - nn + r + 1) // cr + ref = torch.topk(logits[r, :n_r], k).values.sort().values + got = logits[r].gather(0, out[r].long().clamp_min(0)).sort().values + assert torch.equal(ref, got) + + +def test_selfsampling_warmup_row_stride_matches_arena(): + """warmup_varlen(row_stride=...) must compile the SAME launcher key the + dispatch derives from a column-sliced arena view (row stride wider than + the logical width, like the DSL paged-MQA arena's 256-element rounding). + msl is chosen so 64-rounding != 256-rounding: with a mismatched warmup + stride, capture below hits the loud not-compiled raise.""" + k, msl = 512, 8300 + stride = (msl + 255) // 256 * 256 + rows = 2 + ss_host.warmup_varlen( + k, msl, compress_ratio=1, next_n=1, num_rows_list=(rows,), row_stride=stride + ) + arena = torch.randn(rows, stride, dtype=torch.float32, device=_DEV) + logits = arena[:, :msl] # non-contiguous column slice, like serving + kv = torch.full((rows,), msl, dtype=torch.int32, device=_DEV) + pre = torch.zeros(rows, k, dtype=torch.int32, device=_DEV) + out = torch.empty(rows, k, dtype=torch.int32, device=_DEV) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + ss_host.run_varlen(logits, pre, kv, out, max_seq_len=msl) + g.replay() + torch.cuda.synchronize() + ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values + got = arena.gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(ref, got) From facf2a26203da2528fcc6e117d24cc21a0b4c8a3 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 20 Aug 2026 05:15:30 +0000 Subject: [PATCH 21/40] [TRTLLM-15293][fix] self-sampling top-K: restrict SM gate to validated datacenter Blackwell MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous review round aligned the test gate with the dispatch gate at >= 100, which silently admits consumer Blackwell (sm_120/121): those parts lack thread-block clusters, the kernels have never been validated there, and CI stages like RTX5080/5090 and GB10 would start collecting the suite. Tighten all three gates (dispatch flag, warmup hook, test module) to the validated set {sm_100, sm_103} — B200/B300-class datacenter parts — which was the reviewer's suggested alternative. Spot-checked on sm_100: module still collected, 4/4 targeted tests green, hooks clean. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa/indexer.py | 5 ++++- .../_torch/attention_backend/sparse/dsa/metadata.py | 4 ++-- .../_torch/thop/parallel/test_gvr_selfsampling_topk.py | 7 ++++--- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index e3a6626ee177..82a38ca41346 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -704,7 +704,10 @@ def __init__( self._use_self_sampling_topk = ( os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" and IS_CUTLASS_DSL_AVAILABLE - and get_sm_version() >= 100 + # validated datacenter Blackwell only (B200/B300); consumer + # Blackwell (sm_120/121) lacks thread-block clusters and is + # not a supported target for these kernels + and get_sm_version() in (100, 103) and sparse_params.index_topk in (512, 1024, 2048) and compress_ratio in (1, 4) ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 8383a04deaf7..2e5863fb4f29 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -298,8 +298,8 @@ def warmup_selfsampling_topk( if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": return # same hardware gates as the dispatch flag (indexer __init__): never - # compile Blackwell kernels on unsupported stacks during warmup - if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() < 100: + # compile these kernels on unsupported stacks during warmup + if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() not in (100, 103): return if not self.enable_heuristic_topk or self.kv_cache_manager is None: return diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index f5fa763362d1..9657ab3f3aa2 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -40,10 +40,11 @@ if not IS_CUTLASS_DSL_AVAILABLE: pytest.skip("cutlass DSL is required for gvr_selfsampling_topk tests", allow_module_level=True) -if getSMVersion() < 100: +if getSMVersion() not in (100, 103): pytest.skip( - "self-sampling GVR kernels require Blackwell (SM100+) — same gate as " - "the production dispatch", + "self-sampling GVR kernels target datacenter Blackwell (sm_100/103) " + "— same gate as the production dispatch; consumer Blackwell " + "(sm_120/121) lacks thread-block clusters", allow_module_level=True, ) From 834d29a3994f73c4b2849475d0304a2c8ff891c6 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 04:52:31 +0000 Subject: [PATCH 22/40] [TRTLLM-15293][fix] self-sampling top-K: batch-rows admission envelope (fall through above 32 rows) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At high in-flight row counts the varlen streaming engine stops splitting rows across CTAs (route_streaming keeps multi-CTA R only for b <= 32, R=2 up to b <= 74, then one CTA per row) and long rows lose ~5x per call to the in-tree per-row-split kernels — measured in DSv4 Pro 1M-ISL MTP7 EP16 serving: rows=304, 2.63 ms vs 0.49 ms per top-K call, ~4x TPOT end to end. Gate the DSA dispatch on a shared MAX_VARLEN_ROWS=32 admission envelope (info_once fall-through to the in-tree path), drop over-envelope row counts from warmup_varlen so engine init never compiles keys dispatch cannot admit, and document the envelope on run_varlen. Large batches keep taking the in-tree path until a throughput tier lands (see roadmap). Two new tests: the envelope constant matches route_streaming's multi-CTA region (rows=304 collapses to R=1), and warmup drops over-envelope row counts while still warming the admitted prefix. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 23 +++++++++++- .../attention_backend/sparse/dsa/metadata.py | 6 ++- .../blackwell/top_k/__init__.py | 2 + .../gvr_topk_decode_self_sampling_host.py | 19 ++++++++++ .../parallel/test_gvr_selfsampling_topk.py | 37 +++++++++++++++++++ 5 files changed, 84 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 82a38ca41346..634ecba6bfef 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -712,10 +712,15 @@ def __init__( and compress_ratio in (1, 4) ) self._selfsampling_run_varlen = None + self._selfsampling_max_rows = 0 if self._use_self_sampling_topk: - from ....cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen + from ....cute_dsl_kernels.blackwell.top_k import ( + SELFSAMPLING_TOPK_MAX_ROWS, + selfsampling_topk_run_varlen, + ) self._selfsampling_run_varlen = selfsampling_topk_run_varlen + self._selfsampling_max_rows = SELFSAMPLING_TOPK_MAX_ROWS self.mtp_index_share = sparse_params.mtp_index_share if self._enable_heuristic_topk and layer_idx == 0: @@ -1843,6 +1848,22 @@ def sparse_attn_indexer( and self._enable_heuristic_topk and pre_idx is not None ) + if _ss_ready and num_gen_tokens > self._selfsampling_max_rows: + # batch-rows admission envelope: the varlen engine splits + # a row across CTAs only for small batches; beyond + # MAX_VARLEN_ROWS it runs one CTA per row and loses ~5x + # per call to the in-tree per-row-split kernels at long + # rows (DSv4 Pro 1M-ISL MTP7, rows=304: 2.63 ms vs + # 0.49 ms). Large batches take the in-tree path until a + # throughput tier lands. + logger.info_once( + "self-sampling GVR top-K: decode batch rows exceed " + f"the validated envelope ({num_gen_tokens} > " + f"{self._selfsampling_max_rows}); large batches take " + "the in-tree top-K path.", + key="selfsampling_topk_rows_fallthrough", + ) + _ss_ready = False if _ss_ready and not ( # engine hardware-format gate (falls through otherwise): # fp32 row-major logits with a float4-aligned row stride diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 2e5863fb4f29..56acc517fc72 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -292,8 +292,10 @@ def warmup_selfsampling_topk( DeepGEMM is exact-width) so the warmed keys are the ones dispatch actually looks up. Captured geometries are also compiled by the pre-capture warmup forwards; eager batches outside ``batch_sizes`` - still compile lazily on first touch. No-op unless the opt-in gate - (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + still compile lazily on first touch. Row counts above the engine's + admission envelope (``MAX_VARLEN_ROWS``) are dropped by the helper — + those batches dispatch to the in-tree path instead. No-op unless the + opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. """ if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": 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 d5cab389489b..b5aef709096f 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 @@ -20,6 +20,7 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_self_sampling_host import MAX_VARLEN_ROWS as SELFSAMPLING_TOPK_MAX_ROWS from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -36,4 +37,5 @@ "tiered_topk", "is_tiered_topk_supported", "selfsampling_topk_run_varlen", + "SELFSAMPLING_TOPK_MAX_ROWS", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index a462d7b6e6d2..a8ed11234ed4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -149,6 +149,16 @@ def _device(): CMPC = 4096 # L2372 crossing-bin slots per CTA, clustered register path BLKC = 1024 # L2374 CTA size of the clustered register path +# Dispatch-policy admission envelope (not a kernel.cu constant): the +# streaming varlen engine multi-CTA-splits a row only for small batches +# (route_streaming: R = min(148//b, ...) for b <= 32); past that R collapses +# toward one CTA per row and long rows lose badly to the in-tree per-row +# split kernels (measured 5.4x per call at rows=304, n=262144, K=1024, +# B200: 2.63 ms vs 0.49 ms). Integrations must fall through to the in-tree +# path above this row count; warmup_varlen drops larger requests so engine +# init never compiles keys dispatch will not admit. +MAX_VARLEN_ROWS = 32 + def route(b: int, n: int, npad: int, k: int) -> dict[str, object]: """Mirror of gvr_topk_launch (kernel.cu L2754-3197). Pure. See module doc.""" @@ -1258,6 +1268,12 @@ def run_varlen( implementation-specifically; inherited from the translation campaign's probe battery). Finite inputs — including +/-inf and denormals — are tie-aware exact. + + PERFORMANCE ENVELOPE: correct for any ``num_rows``, but the streaming + engine splits a row across CTAs only for small batches — past + ``MAX_VARLEN_ROWS`` it runs one CTA per row and long rows lose ~5x to + the in-tree per-row-split kernels. Serving integrations must gate on + ``MAX_VARLEN_ROWS`` and fall through to the in-tree path above it. """ if logits.dtype is not torch.float32: raise RuntimeError( @@ -1484,6 +1500,9 @@ def warmup_varlen( nn = max(1, int(next_n)) # round each request down to a next_n multiple (min next_n) and dedup rows_list = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) + # admission envelope: dispatch falls through to the in-tree path above + # MAX_VARLEN_ROWS, so compiling larger keys would only burn init time + rows_list = [r for r in rows_list if r <= MAX_VARLEN_ROWS] if not rows_list: return n_env = max(1, int(max_seq_len) // int(compress_ratio)) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 9657ab3f3aa2..d0f4a15d33ee 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -658,3 +658,40 @@ def test_selfsampling_warmup_row_stride_matches_arena(): ref = torch.topk(arena[:, :msl], k, dim=1).values.sort(dim=1).values got = arena.gather(1, out.long().clamp_min(0)).sort(dim=1).values assert torch.equal(ref, got) + + +def test_selfsampling_rows_envelope_constant(): + """The rows admission envelope is one shared constant, and it matches + route_streaming's multi-CTA region: inside the envelope a long row is + split across CTAs (R > 1); far outside it the split collapses to one + CTA per row — the regression measured at rows=304, n=262144 (5.4x vs + the in-tree per-row-split kernel).""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import SELFSAMPLING_TOPK_MAX_ROWS + + assert SELFSAMPLING_TOPK_MAX_ROWS == ss_host.MAX_VARLEN_ROWS == 32 + n = 262144 + plan_in = ss_host.route_streaming(ss_host.MAX_VARLEN_ROWS, n, n, 1024, force_main=True) + assert int(plan_in["grid"][0]) > 1, "inside the envelope long rows must split" + plan_out = ss_host.route_streaming(304, n, n, 1024, force_main=True) + assert int(plan_out["grid"][0]) == 1, "rows=304 collapses to one CTA per row" + + +def test_selfsampling_warmup_drops_rows_beyond_envelope(): + """warmup_varlen must not compile engines for row counts the dispatch + envelope never admits (rows > MAX_VARLEN_ROWS): a CUDA-graph batch-size + list with large batches warms only the admitted prefix.""" + k, msl, nn = 512, 8192, 4 + stride = (msl + 255) // 256 * 256 + before = set(ss_host._VARLEN_CACHE.keys()) + ss_host.warmup_varlen( + k, + msl, + compress_ratio=1, + next_n=nn, + num_rows_list=(nn, 8 * nn, 76 * nn), # rows 4, 32, 304 + row_stride=stride, + ) + new_rows = {key[0] for key in ss_host._VARLEN_CACHE if key not in before} + assert not any(r > ss_host.MAX_VARLEN_ROWS for r in new_rows), new_rows + have_rows = {key[0] for key in ss_host._VARLEN_CACHE} + assert {nn, 8 * nn} <= have_rows, "in-envelope rows must still be warmed" From 70ece97249364215a616854687c61de1ac04d509 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:22:54 +0000 Subject: [PATCH 23/40] [TRTLLM-15293][fix] self-sampling top-K: warm the full rows envelope; single-row width gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two residuals from the production audit. (1) Warmup previously covered only {next_n} and the CUDA-graph batch sizes, so eager mixed prefill+decode batches at other in-flight generation counts paid a multi-second first-touch JIT on the serving path (4-13 reachable unwarmed r_const bands per model with padded graph lists). Warm every admissible row count instead — next_n multiples up to MAX_VARLEN_ROWS — which bounds init cost at <= 32/next_n tiny launches over the same few engine compiles and leaves nothing to compile at serving time inside the envelope. (2) The dispatch format gate checked only stride(0), but single-row batches derive their row window from shape[1] (arena last-row safety), so a non-float4 indexer_max_seq_len with next_n==1 raised inside run_varlen instead of falling through; gate that width explicitly. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa/indexer.py | 5 +++++ .../_torch/attention_backend/sparse/dsa/metadata.py | 12 ++++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 634ecba6bfef..61251479b32c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1874,6 +1874,11 @@ def sparse_attn_indexer( and logits_decode.stride(1) == 1 and logits_decode.stride(0) % 4 == 0 and logits_decode.data_ptr() % 16 == 0 + # single-row batches derive their row window from + # shape[1] (arena last-row safety), so that width must + # satisfy the same float4 rule — otherwise run_varlen + # raises instead of falling through + and (num_gen_tokens > 1 or logits_decode.shape[1] % 4 == 0) ): logger.warning_once( "TRTLLM_GVR_SELF_SAMPLING=1 but the decode logits do " diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 56acc517fc72..7204ec15ee55 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -317,9 +317,17 @@ def warmup_selfsampling_topk( ) except ImportError: return - rows = {int(next_n)} + # warm EVERY admissible row count (next_n multiples up to the + # engine's admission envelope): eager mixed prefill+decode batches + # arrive at arbitrary in-flight generation counts, and each unwarmed + # r_const band costs a multi-second first-touch JIT on the serving + # path. The envelope keeps this bounded (<= 32/next_n launches, + # ~4-13 distinct engine compiles); over-envelope rows dispatch + # in-tree and need no warmup (the helper drops them anyway). + nn = int(next_n) + rows = set(range(nn, int(_ss_host.MAX_VARLEN_ROWS) + 1, nn)) or {nn} for bs in batch_sizes or (): - rows.add(int(bs) * int(next_n)) + rows.add(int(bs) * nn) msl_c = int(self.get_indexer_max_seq_len()) if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits: # mirror the DSL paged-MQA arena stride (cute_dsl_custom_ops From 8f3f7ddae6b33cfacff40166bb4ad19124417006 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:18:56 +0000 Subject: [PATCH 24/40] [TRTLLM-15293][perf] self-sampling top-K: enable TSH-floor staging for small batches; dedup dead TSHG key slot The varlen launcher gated TSH-floor staging on num_rows > 15, stranding small batches in SPLIT-main without the staged floor. On real deep-layer captures this is a distribution-dependent ~6x tail (v4_pro_512k L46/L52, n4=32768, rows 1-8: 142-151 us -> 25 us with staging; tie-aware exact preserved, healthy layers and n4 > 32768 rows unchanged within noise). The kernel compiles the TSH machinery whenever SPLIT and gates it per row at runtime, so this is a runtime-scalar-only change: no new engines, no compile-key change, CUDA-graph safe. Also normalize the dead TSHG slot out of the varlen compile key (the ctor overrides it under SPLIT), so row counts differing only in that slot share one engine instead of compiling twice. Varlen/envelope test subset: 22/22. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling_host.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index a8ed11234ed4..a1f0164f0e75 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -755,7 +755,10 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): r_const = rt["R"] cr_shift = 0 if cr == 1 else 2 dev = _device() - fn = dev.get_compiled(tpl + (next_n, cr_shift, r_const)) + # TSHG (tpl[6]) is dead under varlen (the ctor compiles the TSH + # machinery in whenever SPLIT); normalize it out of the compile key so + # row counts differing only in that slot share one engine + fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const)) big = num_rows * r_const <= 148 aim_base = ( ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) @@ -769,7 +772,13 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): ) amin = 3 * k if r_const == 2 else (7 * k) // 2 sd_en = 1 if (k > 1024 and not big) else 0 - tsh_en = 1 if (tpl[5] and num_rows > 15 and k <= 1024) else 0 + # TSH-floor staging: gate on SPLIT and K only. The old num_rows > 15 + # condition stranded small batches (rows <= 15) in SPLIT-main without + # the staged floor — a distribution-dependent 6x tail on real deep-layer + # captures (v4_pro_512k L46/L52, n4 = 32768, rows 1-8: 142-151 us vs + # 25 us with staging; healthy layers and n4 > 32768 rows unaffected — + # the kernel gates TSH per row at runtime anyway). + tsh_en = 1 if (tpl[5] and k <= 1024) else 0 pre = (0, npad, k, rt["SCAP_"], rt["CMP_"], r_const, 0, 0, 0, 0, 0) tail = (aim_base, sfac, amin, sd_en, tsh_en) lc = (fn, pre, tail) From c8957cf261ff42a652fd1e2984dab58c3a493579 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:47:09 +0000 Subject: [PATCH 25/40] [TRTLLM-15293][perf] self-sampling top-K: per-row varlen port of the clustered register-resident family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lift the varlen production path's force-main restriction for its first specialist family: GvrRegClusKernel gains the same per-row varlen mode already shipped in GvrMainKernel — n is re-derived per row in-kernel from device kv_lens (production heuristicTopKDecode contract), hints map to the request row, and short rows (n <= k) emit identity + (-1) tail from rank 0 and skip the body entirely (a zero-work pass would reach the degenerate crossing-overflow emitter). All derived quantities are pure functions of the row, so the whole-body guard is cluster-uniform and the cluster barriers stay aligned. The varlen launcher admits the family exactly where the free route() picks it (route parity tier 1) — a pure function of the capture-stable cache key, so CUDA-graph replay safety is unchanged; its whole admission window (n4 <= 32768) fits capture-frozen envelopes. On real captures this recovers most of the streaming-only integration tax in the family's window: the deep-layer distribution tail (v4_pro_512k L46/L52, rows 8) goes 147 us (pre-TSH-fix) / 25 us (TSH fix) -> 20.7 us, against 16-18 us for the batch-uniform standalone. Full self-sampling UT suite green (79/79) incl. two new tests: heterogeneous kv_lens + MTP row windows vs the reference oracle, and cluster-family CUDA-graph capture/replay exactness. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 910 ++++++++++-------- .../gvr_topk_decode_self_sampling_host.py | 43 +- .../parallel/test_gvr_selfsampling_topk.py | 67 ++ 3 files changed, 604 insertions(+), 416 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 55c6e1a4fd61..897f4e6c7b33 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -5450,19 +5450,44 @@ def _val__regclus(frags, s: int): class GvrRegClusKernel: """gvr_reg_clus (kernel.cu L2376-2379).""" - def __init__(self, blk: int, vpt: int, cs: int, pdl: bool = False): + def __init__( + self, + blk: int, + vpt: int, + cs: int, + pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, + ): assert blk == BLKC, "all instantiations BLK=BLKC=1024 (spec §4b)" assert vpt in (1, 2, 4) and cs in (2, 4, 8) self.blk = blk self.vpt = vpt self.cs = cs self.pdl = bool(pdl) + # per-row varlen mode (production heuristicTopKDecode contract, same + # semantics as GvrMainKernel): n is re-derived PER ROW in-kernel from + # a device kv_lens tensor; the scalar n launch arg becomes the + # envelope clamp bound. next_n / cr_shift are compile-time. + self.varlen = bool(varlen) + self.next_n = int(next_n) + self.cr_shift = int(cr_shift) + if self.varlen: + assert self.next_n >= 1 and self.cr_shift in (0, 2) self.S = vpt * 4 # L2381 self.span = blk * vpt # L2382 (float4 per CTA) # ------------------------------------------------------------------ @cute.kernel - def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + def kern( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + kv_lens: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + ): BLK = cutlass.const_expr(self.blk) VPT = cutlass.const_expr(self.vpt) CS = cutlass.const_expr(self.cs) @@ -5476,6 +5501,46 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c rank, row, _ = cute.arch.block_idx() # L2396-2397 (P11: bx=rank) lane = tid & cutlass.Int32(31) + # ================= per-row varlen prologue (varlen mode only) ========= + # Same production contract as GvrMainKernel (see its prologue): row r + # serves request r // next_n with + # n = (kv_lens[req] - next_n + r % next_n + 1) >> cr_shift, clamped to + # the envelope launch arg n (the launcher admits this family only when + # the envelope fits its capacity window, so per-row n never exceeds + # capacity). Every quantity is a pure function of `row`, so all CS + # ranks of a row's cluster (and all threads) compute identical values + # -- the whole-body guard below is cluster-uniform and the cluster + # barriers inside remain aligned. Short rows (n <= k) emit + # identity + (-1) tail from rank 0 here and SKIP the body entirely: a + # zero-work pass would reach the degenerate crossing-overflow emitter + # and poison the output. + short = cutlass.Int32(0) + prow = row + if cutlass.const_expr(self.varlen): + kq = cutlass.Int32(pre_idx.shape[1]) + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + prow = req + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > n: + nv = n + if nv <= kq: + short = cutlass.Int32(1) + if short == cutlass.Int32(0): + n = nv + if short != cutlass.Int32(0): + if rank == cutlass.Int32(0): + if tid < kq: + ov = cutlass.Int32(-1) + if tid < nv: + ov = tid + out[row, tid] = ov + # ------------------------------------------------------------------ # Predeclarations (DSL AST rule: every scalar (re)assigned under a # dynamic if/while must pre-exist with a stable type; constant inits @@ -5536,278 +5601,419 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c LOQ = cutlass.Float32(0.0) qv = cutlass.Float32(0.0) - npad = cutlass.Int32(logits.shape[1]) # noqa: F841 - k = cutlass.Int32(pre_idx.shape[1]) - out_row = out[row, None] - x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base - p_addr = pre_idx[row, None].iterator.toint() - - # ---- shared-memory window (map in module docstring) ---- - sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) - sbase = sptr.toint() # Int32 shared addr - - s_res = _smem_view__regclus(cutlass.Int32, sbase, 0, 6) - s_cnt = _smem_view__regclus(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_o2 - s_kmm = _smem_view__regclus(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax - s_ws = _smem_view__regclus(cutlass.Int32, sbase, 16, 32) - s_wmn = _smem_view__regclus(cutlass.Uint32, sbase, 48, 32) - s_wmx = _smem_view__regclus(cutlass.Uint32, sbase, 80, 32) - s_hist = _smem_view__regclus(cutlass.Int32, sbase, W_HIST, NB__regclus) - s_mrg = _smem_view__regclus(cutlass.Int32, sbase, W_MRG, NB__regclus) - s_hoff = _smem_view__regclus(cutlass.Int32, sbase, W_HOFF, NB__regclus) - s_ck = _smem_view__regclus(cutlass.Uint32, sbase, W_CK, CMPC) - s_ci = _smem_view__regclus(cutlass.Int32, sbase, W_CI, CMPC, align=4) - # raw byte bases for DSMEM (mapa) addressing - hist_addr = sbase + cutlass.Int32(W_HIST * 4) - ck_addr = sbase + cutlass.Int32(W_CK * 4) - ci_addr = sbase + cutlass.Int32(W_CI * 4) - - n4 = n >> cutlass.Int32(2) # L2405 - ntail = n - (n4 << cutlass.Int32(2)) # L2406 - base4 = rank * cutlass.Int32(self.span) # L2407 - tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L2425 - - # ---- P0: redundant hint gather, EVERY CTA (L2410-2413; k<=BLK by - # dispatch gate L2897). One coalesced word per thread, NO cluster - # barrier — GMIN/GMAX identical everywhere by construction. - if tid < k: - pv0 = ld_g_i32(p_addr, tid) - - # ---- P1: row load — predicated flat float4[VPT] batch (L2415-2424; - # the CUDA has NO exact-fit peel here, guard is per-load). Issue all - # loads first (op43 L1), then -INFINITY-fill missed slots (op43 L2). - atom128 = g2r_atom_f32(128, invariant=True) - frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] - for u in cutlass.range_constexpr(VPT): - i = base4 + tid + cutlass.Int32(u * self.blk) - if i < n4: - ld_g_f32x4(atom128, x_addr, i, frags[u]) - for u in cutlass.range_constexpr(VPT): - i = base4 + tid + cutlass.Int32(u * self.blk) - if i >= n4: # -INFINITY fill L2421 - for z in cutlass.range_constexpr(4): - frags[u][z] = cutlass.Float32(_NEG_INF__regclus) - # tail element: rank 0 only (L2425-2426) - if rank == cutlass.Int32(0): - if tid < ntail: - tval = ldg_f32(x_addr, tix) - - # ---- P2: init (L2428-2429). NB__regclus == BLK -> single-pass hist clear. - if tid == cutlass.Int32(0): - s_cnt[0] = cutlass.Int32(0) - s_cnt[1] = cutlass.Int32(0) - for z in cutlass.range_constexpr(NB__regclus // self.blk): - s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - - # ---- P3: GMIN/GMAX from the hint (L2431-2445), ONE barrier fold. - lmin = cutlass.Uint32(0xFFFFFFFF) - lmax = cutlass.Uint32(0) - if cutlass.Uint32(pv0) < cutlass.Uint32(n): - uk = fkey(ldg_f32(x_addr, pv0)) # __ldg(X+pv0) L2433 - lmin = uk - lmax = uk - lmin = warp_min_u32(lmin) - lmax = warp_max_u32(lmax) - if lane == cutlass.Int32(0): - s_wmn[tid >> cutlass.Int32(5)] = lmin - s_wmx[tid >> cutlass.Int32(5)] = lmax - cute.arch.barrier() # L2438 - a = cutlass.Uint32(0xFFFFFFFF) - c = cutlass.Uint32(0) - if lane < cutlass.Int32(NW): - a = cutlass.Uint32(s_wmn[lane]) - c = cutlass.Uint32(s_wmx[lane]) - lmin = warp_min_u32(a) - lmax = warp_max_u32(c) - Tv = invkey(lmin) - GMAX = invkey(lmax) + if short == cutlass.Int32(0): + npad = cutlass.Int32(logits.shape[1]) # noqa: F841 + k = cutlass.Int32(pre_idx.shape[1]) + out_row = out[row, None] + x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base + p_addr = pre_idx[prow, None].iterator.toint() # request-level under varlen + + # ---- shared-memory window (map in module docstring) ---- + sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) + sbase = sptr.toint() # Int32 shared addr + + s_res = _smem_view__regclus(cutlass.Int32, sbase, 0, 6) + s_cnt = _smem_view__regclus(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_o2 + s_kmm = _smem_view__regclus(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax + s_ws = _smem_view__regclus(cutlass.Int32, sbase, 16, 32) + s_wmn = _smem_view__regclus(cutlass.Uint32, sbase, 48, 32) + s_wmx = _smem_view__regclus(cutlass.Uint32, sbase, 80, 32) + s_hist = _smem_view__regclus(cutlass.Int32, sbase, W_HIST, NB__regclus) + s_mrg = _smem_view__regclus(cutlass.Int32, sbase, W_MRG, NB__regclus) + s_hoff = _smem_view__regclus(cutlass.Int32, sbase, W_HOFF, NB__regclus) + s_ck = _smem_view__regclus(cutlass.Uint32, sbase, W_CK, CMPC) + s_ci = _smem_view__regclus(cutlass.Int32, sbase, W_CI, CMPC, align=4) + # raw byte bases for DSMEM (mapa) addressing + hist_addr = sbase + cutlass.Int32(W_HIST * 4) + ck_addr = sbase + cutlass.Int32(W_CK * 4) + ci_addr = sbase + cutlass.Int32(W_CI * 4) + + n4 = n >> cutlass.Int32(2) # L2405 + ntail = n - (n4 << cutlass.Int32(2)) # L2406 + base4 = rank * cutlass.Int32(self.span) # L2407 + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L2425 + + # ---- P0: redundant hint gather, EVERY CTA (L2410-2413; k<=BLK by + # dispatch gate L2897). One coalesced word per thread, NO cluster + # barrier — GMIN/GMAX identical everywhere by construction. + if tid < k: + pv0 = ld_g_i32(p_addr, tid) + + # ---- P1: row load — predicated flat float4[VPT] batch (L2415-2424; + # the CUDA has NO exact-fit peel here, guard is per-load). Issue all + # loads first (op43 L1), then -INFINITY-fill missed slots (op43 L2). + atom128 = g2r_atom_f32(128, invariant=True) + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + for u in cutlass.range_constexpr(VPT): + i = base4 + tid + cutlass.Int32(u * self.blk) + if i < n4: + ld_g_f32x4(atom128, x_addr, i, frags[u]) + for u in cutlass.range_constexpr(VPT): + i = base4 + tid + cutlass.Int32(u * self.blk) + if i >= n4: # -INFINITY fill L2421 + for z in cutlass.range_constexpr(4): + frags[u][z] = cutlass.Float32(_NEG_INF__regclus) + # tail element: rank 0 only (L2425-2426) + if rank == cutlass.Int32(0): + if tid < ntail: + tval = ldg_f32(x_addr, tix) - # ---- collapse guard, NaN-safe (L2446-2453) - okc = cutlass.Int32(0) - if Tv < GMAX: - if (GMAX - Tv) > cutlass.Float32(1e-30): - okc = cutlass.Int32(1) - if okc == cutlass.Int32(0): - Tv = cutlass.Float32(SENT_LO) - GMAX = cutlass.Float32(SENT_HI) + # ---- P2: init (L2428-2429). NB__regclus == BLK -> single-pass hist clear. + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + s_cnt[1] = cutlass.Int32(0) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + + # ---- P3: GMIN/GMAX from the hint (L2431-2445), ONE barrier fold. + lmin = cutlass.Uint32(0xFFFFFFFF) + lmax = cutlass.Uint32(0) + if cutlass.Uint32(pv0) < cutlass.Uint32(n): + uk = fkey(ldg_f32(x_addr, pv0)) # __ldg(X+pv0) L2433 + lmin = uk + lmax = uk + lmin = warp_min_u32(lmin) + lmax = warp_max_u32(lmax) + if lane == cutlass.Int32(0): + s_wmn[tid >> cutlass.Int32(5)] = lmin + s_wmx[tid >> cutlass.Int32(5)] = lmax + cute.arch.barrier() # L2438 + a = cutlass.Uint32(0xFFFFFFFF) + c = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + a = cutlass.Uint32(s_wmn[lane]) + c = cutlass.Uint32(s_wmx[lane]) + lmin = warp_min_u32(a) + lmax = warp_max_u32(c) + Tv = invkey(lmin) + GMAX = invkey(lmax) + + # ---- collapse guard, NaN-safe (L2446-2453) + okc = cutlass.Int32(0) + if Tv < GMAX: + if (GMAX - Tv) > cutlass.Float32(1e-30): + okc = cutlass.Int32(1) + if okc == cutlass.Int32(0): + Tv = cutlass.Float32(SENT_LO) + GMAX = cutlass.Float32(SENT_HI) + + # ---- bin transform constants (L2454-2467): branchless trash bin. + WD = (GMAX - Tv) * cutlass.Float32(1.0 / float(NB__regclus - 2)) + wsel = cutlass.Float32(1e-30) + if WD > cutlass.Float32(0.0): + wsel = WD + SC = cutlass.Float32(1.0) / wsel + CQ0 = cutlass.Float32(1.0) - Tv * SC + CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf__regclus(CQ0) + cutlass.Float32(1.0)) - # ---- bin transform constants (L2454-2467): branchless trash bin. - WD = (GMAX - Tv) * cutlass.Float32(1.0 / float(NB__regclus - 2)) - wsel = cutlass.Float32(1e-30) - if WD > cutlass.Float32(0.0): - wsel = WD - SC = cutlass.Float32(1.0) / wsel - CQ0 = cutlass.Float32(1.0) - Tv * SC - CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf__regclus(CQ0) + cutlass.Float32(1.0)) - - # ---- P4: histogram (L2469-2472); tval add UNCONDITIONAL (trash bin - # swallows -INFINITY via the saturating cvt). - for s in cutlass.range_constexpr(S): - qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) + # ---- P4: histogram (L2469-2472); tval add UNCONDITIONAL (trash bin + # swallows -INFINITY via the saturating cvt). + for s in cutlass.range_constexpr(S): + qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + qv = _fmaf__regclus(tval, SC, CQ) bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - qv = _fmaf__regclus(tval, SC, CQ) - bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) - atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - - # ---- P5: cluster merge (L2474-2484) - _cluster_sync_aligned() # L2474 - for z in cutlass.range_constexpr(NB__regclus // self.blk): - i = tid + cutlass.Int32(z * self.blk) - # CS-unrolled remote u32 loads: batch-issue, then fold (#pragma - # unroll L2477; one mapa per (i, r) exactly like map_shared_rank) - hvals = [] - for r in cutlass.range_constexpr(CS): - ma = _mapa_shared_cluster_addr( - hist_addr + (i << cutlass.Int32(2)), cutlass.Int32(r) - ) - hvals.append(_ld_shared_cluster_i32(ma)) - tot_a = cutlass.Int32(0) - pre_a = cutlass.Int32(0) - for r in cutlass.range_constexpr(CS): - if cutlass.Int32(r) < rank: - pre_a = pre_a + hvals[r] # rank-exclusive - tot_a = tot_a + hvals[r] - s_mrg[i] = tot_a - s_hoff[i] = pre_a - - # ---- P6: scan (L2485-2492) - cute.arch.barrier() # L2485 - scan_cross_w(s_mrg, s_ws, k, tid, s_res, blk=self.blk, nb=NB__regclus) - cute.arch.barrier() # L2487 - above = s_res[RES_ABOVE] - m = s_res[RES_M] - Bv = s_res[RES_B] - need = k - above - whole = cutlass.Int32(0) - if need >= m: - whole = cutlass.Int32(1) - degen = cutlass.Int32(0) - if m > cutlass.Int32(CS * CMPC): - degen = cutlass.Int32(1) - for z in cutlass.range_constexpr(NB__regclus // self.blk): - i = tid + cutlass.Int32(z * self.blk) - s_mrg[i] = s_mrg[i] + s_hoff[i] # L2491 global cursor - cute.arch.barrier() # L2492 - # ---- P7: register sweep emit (L2494-2527, !degen) - if degen == cutlass.Int32(0): - LOQ = cutlass.Float32(Bv) # L2495 - lim1 = above - if whole == cutlass.Int32(1): - lim1 = above + m # L2496 - for s in cutlass.range_constexpr(S): - qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) # bit-identical L2499 + # ---- P5: cluster merge (L2474-2484) + _cluster_sync_aligned() # L2474 + for z in cutlass.range_constexpr(NB__regclus // self.blk): + i = tid + cutlass.Int32(z * self.blk) + # CS-unrolled remote u32 loads: batch-issue, then fold (#pragma + # unroll L2477; one mapa per (i, r) exactly like map_shared_rank) + hvals = [] + for r in cutlass.range_constexpr(CS): + ma = _mapa_shared_cluster_addr( + hist_addr + (i << cutlass.Int32(2)), cutlass.Int32(r) + ) + hvals.append(_ld_shared_cluster_i32(ma)) + tot_a = cutlass.Int32(0) + pre_a = cutlass.Int32(0) + for r in cutlass.range_constexpr(CS): + if cutlass.Int32(r) < rank: + pre_a = pre_a + hvals[r] # rank-exclusive + tot_a = tot_a + hvals[r] + s_mrg[i] = tot_a + s_hoff[i] = pre_a + + # ---- P6: scan (L2485-2492) + cute.arch.barrier() # L2485 + scan_cross_w(s_mrg, s_ws, k, tid, s_res, blk=self.blk, nb=NB__regclus) + cute.arch.barrier() # L2487 + above = s_res[RES_ABOVE] + m = s_res[RES_M] + Bv = s_res[RES_B] + need = k - above + whole = cutlass.Int32(0) + if need >= m: + whole = cutlass.Int32(1) + degen = cutlass.Int32(0) + if m > cutlass.Int32(CS * CMPC): + degen = cutlass.Int32(1) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + i = tid + cutlass.Int32(z * self.blk) + s_mrg[i] = s_mrg[i] + s_hoff[i] # L2491 global cursor + cute.arch.barrier() # L2492 + + # ---- P7: register sweep emit (L2494-2527, !degen) + if degen == cutlass.Int32(0): + LOQ = cutlass.Float32(Bv) # L2495 + lim1 = above + if whole == cutlass.Int32(1): + lim1 = above + m # L2496 + for s in cutlass.range_constexpr(S): + qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) # bit-identical L2499 + if qv >= LOQ: + bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) + p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + idx = ( + (base4 + tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if p < lim1: + out_row[p] = idx + else: + if whole == cutlass.Int32(0): + # crossing overflow -> striped DSMEM slabs; TWO + # separate u32 remote stores (NOT packed, L2507-10) + q2i = p - above + rnk = q2i >> cutlass.Int32(LCMPC) + j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) + _st_shared_cluster_i32( + _mapa_shared_cluster_addr(ck_addr + j, rnk), + fkey(_val__regclus(frags, s)), + ) + _st_shared_cluster_i32( + _mapa_shared_cluster_addr(ci_addr + j, rnk), idx + ) + # tail element (L2514-2526): tval == -INF fails q>=LOQ elsewhere + qv = _fmaf__regclus(tval, SC, CQ) if qv >= LOQ: bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - idx = ( - (base4 + tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) - ) + cutlass.Int32(s % 4) if p < lim1: - out_row[p] = idx + out_row[p] = tix else: if whole == cutlass.Int32(0): - # crossing overflow -> striped DSMEM slabs; TWO - # separate u32 remote stores (NOT packed, L2507-10) q2i = p - above rnk = q2i >> cutlass.Int32(LCMPC) j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) _st_shared_cluster_i32( - _mapa_shared_cluster_addr(ck_addr + j, rnk), - fkey(_val__regclus(frags, s)), + _mapa_shared_cluster_addr(ck_addr + j, rnk), fkey(tval) ) - _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), idx) - # tail element (L2514-2526): tval == -INF fails q>=LOQ elsewhere - qv = _fmaf__regclus(tval, SC, CQ) - if qv >= LOQ: - bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) - p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - if p < lim1: - out_row[p] = tix - else: - if whole == cutlass.Int32(0): - q2i = p - above - rnk = q2i >> cutlass.Int32(LCMPC) - j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) - _st_shared_cluster_i32( - _mapa_shared_cluster_addr(ck_addr + j, rnk), fkey(tval) - ) - _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), tix) + _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), tix) - # ---- P8 (L2529-2530): release staging to rank 0 - cute.arch.barrier() # L2529 - _cluster_sync_aligned() # L2530 + # ---- P8 (L2529-2530): release staging to rank 0 + cute.arch.barrier() # L2529 + _cluster_sync_aligned() # L2530 - # ---- P9: rank-0 selection (L2532-2647) - if rank == cutlass.Int32(0): - if whole == cutlass.Int32(0): - mc = m - if degen == cutlass.Int32(1): - mc = cutlass.Int32(0) # L2533 - if degen == cutlass.Int32(0): - if mc <= cutlass.Int32(QUADC__regclus): - # (1) quad-96: all candidates LOCAL (96 < CMPC), - # O(mc^2) slot-order tie-broken rank (L2535-2543) - i = tid - while i < mc: - uq = cutlass.Uint32(s_ck[i]) - rnk = cutlass.Int32(0) - j = cutlass.Int32(0) - while j < mc: - vq = cutlass.Uint32(s_ck[j]) - tinc = cutlass.Int32(0) - if vq > uq: - tinc = cutlass.Int32(1) - if vq == uq: - if j < i: + # ---- P9: rank-0 selection (L2532-2647) + if rank == cutlass.Int32(0): + if whole == cutlass.Int32(0): + mc = m + if degen == cutlass.Int32(1): + mc = cutlass.Int32(0) # L2533 + if degen == cutlass.Int32(0): + if mc <= cutlass.Int32(QUADC__regclus): + # (1) quad-96: all candidates LOCAL (96 < CMPC), + # O(mc^2) slot-order tie-broken rank (L2535-2543) + i = tid + while i < mc: + uq = cutlass.Uint32(s_ck[i]) + rnk = cutlass.Int32(0) + j = cutlass.Int32(0) + while j < mc: + vq = cutlass.Uint32(s_ck[j]) + tinc = cutlass.Int32(0) + if vq > uq: tinc = cutlass.Int32(1) - rnk = rnk + tinc - j = j + cutlass.Int32(1) - if rnk < need: - out_row[above + rnk] = s_ci[i] - i = i + cutlass.Int32(BLK) - else: - # (2) key-space narrowing over striped DSMEM slabs - # (L2544-2596): slot = i & (CMPC-1), rank = i >> LCMPC - if tid == cutlass.Int32(0): - s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) - s_kmm[1] = cutlass.Uint32(0) - cute.arch.barrier() # L2546 - i = tid - while i < mc: - kv = cutlass.Uint32( - _ld_shared_cluster_i32( - _mapa_shared_cluster_addr( - ck_addr - + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), - i >> cutlass.Int32(LCMPC), + if vq == uq: + if j < i: + tinc = cutlass.Int32(1) + rnk = rnk + tinc + j = j + cutlass.Int32(1) + if rnk < need: + out_row[above + rnk] = s_ci[i] + i = i + cutlass.Int32(BLK) + else: + # (2) key-space narrowing over striped DSMEM slabs + # (L2544-2596): slot = i & (CMPC-1), rank = i >> LCMPC + if tid == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + cute.arch.barrier() # L2546 + i = tid + while i < mc: + kv = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), + i >> cutlass.Int32(LCMPC), + ) ) ) - ) - atomic_min_cta(s_kmm.iterator, kv) - atomic_max_cta(s_kmm.iterator + 1, kv) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2551 - rlo = cutlass.Uint32(s_kmm[0]) - rhi = cutlass.Uint32(s_kmm[1]) - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - aboveC = cutlass.Int32(0) - needC = need - mm = mc + atomic_min_cta(s_kmm.iterator, kv) + atomic_max_cta(s_kmm.iterator + 1, kv) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L2551 + rlo = cutlass.Uint32(s_kmm[0]) + rhi = cutlass.Uint32(s_kmm[1]) + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + while done == cutlass.Int32(0): # <=6 levels L2553 + if needC == mm: # L2554 + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(6): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2w = cutlass.Int32(32) - clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = cutlass.Int32(0) + if b2w > cutlass.Int32(LNB): + sh2 = b2w - cutlass.Int32(LNB) + for z in cutlass.range_constexpr(NB__regclus // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + cute.arch.barrier() # L2563 + i = tid + while i < mc: + unar = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ( + (i & cutlass.Int32(CMPC - 1)) + << cutlass.Int32(2) + ), + i >> cutlass.Int32(LCMPC), + ) + ) + ) + if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): + if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): + bnn = ( + cutlass.Uint32(unar) - cutlass.Uint32(rlo) + ) >> cutlass.Uint32(sh2) + bnn = _umin_u32__regclus( + bnn, cutlass.Uint32(NB__regclus - 1) + ) + atomic_add_cta( + s_hist.iterator + cutlass.Int32(bnn), + cutlass.Int32(1), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L2568 + find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) + cute.arch.barrier() # L2570 + aboveC = aboveC + s_res[RES_ABOVE] + needC = needC - s_res[RES_ABOVE] + mm = s_res[RES_M] + b_lv = s_res[RES_B] + nlo = cutlass.Uint32(rlo) + ( + cutlass.Uint32(b_lv) << cutlass.Uint32(sh2) + ) + if b_lv != cutlass.Int32(NB__regclus - 1): + rhi = nlo + ( + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) + - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + cute.arch.barrier() # L2576 + # two-predicate ballot emit over the striped slabs + lml = cutlass.Int32(cute.arch.lanemask_lt()) + it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it = cutlass.Int32(0) + while it < it2: + i = it * cutlass.Int32(BLK) + tid + uke = cutlass.Uint32(0) + idv = cutlass.Int32(0) + if i < mc: # predicated remote + uke = cutlass.Uint32( + _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ck_addr + + ( + (i & cutlass.Int32(CMPC - 1)) + << cutlass.Int32(2) + ), + i >> cutlass.Int32(LCMPC), + ) + ) + ) + idv = _ld_shared_cluster_i32( + _mapa_shared_cluster_addr( + ci_addr + + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), + i >> cutlass.Int32(LCMPC), + ) + ) + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + if i < mc: + if cutlass.Int64(cutlass.Uint32(uke)) > ethr: + q1f = cutlass.Int32(1) + if cutlass.Int64(cutlass.Uint32(uke)) == ethr: + q2f = cutlass.Int32(1) + n1 = ballot(q1f == cutlass.Int32(1)) + n2 = ballot(q2f == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_cnt.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_cnt.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1f == cutlass.Int32(1): + if p1e < aboveC: + out_row[above + p1e] = idv + if q2f == cutlass.Int32(1): + if p2e < needC: + out_row[above + aboveC + p2e] = idv + it = it + cutlass.Int32(1) + else: + # (3) degen safety net (L2597-2645): crossing bin larger + # than the whole cluster buffer -> exact whole-row + # key-space narrowing by rank 0 alone, <=8 levels. + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + aboveC = cutlass.Int32(0) # above2 + needC = k # need2 + mm = n # m2 + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) lev = cutlass.Int32(0) done = cutlass.Int32(0) - while done == cutlass.Int32(0): # <=6 levels L2553 - if needC == mm: # L2554 + while done == cutlass.Int32(0): + if needC == mm: # L2603 ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) aboveC = aboveC + mm needC = cutlass.Int32(0) + tie_m = cutlass.Int32(0) done = cutlass.Int32(1) if done == cutlass.Int32(0): if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): ethr = cutlass.Int64(cutlass.Uint32(rlo)) done = cutlass.Int32(1) - if lev >= cutlass.Int32(6): + if lev >= cutlass.Int32(8): # L2605 ethr = cutlass.Int64(cutlass.Uint32(rlo)) done = cutlass.Int32(1) if done == cutlass.Int32(0): @@ -5820,21 +6026,10 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c sh2 = b2w - cutlass.Int32(LNB) for z in cutlass.range_constexpr(NB__regclus // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L2563 + cute.arch.barrier() # L2612 i = tid - while i < mc: - unar = cutlass.Uint32( - _ld_shared_cluster_i32( - _mapa_shared_cluster_addr( - ck_addr - + ( - (i & cutlass.Int32(CMPC - 1)) - << cutlass.Int32(2) - ), - i >> cutlass.Int32(LCMPC), - ) - ) - ) + while i < n: # whole-row bin L2613 + unar = fkey(ldg_f32(x_addr, i)) if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): bnn = ( @@ -5848,9 +6043,9 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2568 + cute.arch.barrier() # L2618 find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) - cute.arch.barrier() # L2570 + cute.arch.barrier() # L2620 aboveC = aboveC + s_res[RES_ABOVE] needC = needC - s_res[RES_ABOVE] mm = s_res[RES_M] @@ -5865,39 +6060,29 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c ) rlo = nlo lev = lev + cutlass.Int32(1) - cute.arch.barrier() # L2576 - # two-predicate ballot emit over the striped slabs + cute.arch.barrier() # L2626 + nA = k # tie_m ? above2 : k + if tie_m == cutlass.Int32(1): + nA = aboveC + nT = cutlass.Int32(0) + if tie_m == cutlass.Int32(1): + nT = needC lml = cutlass.Int32(cute.arch.lanemask_lt()) - it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it2 = (n + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) it = cutlass.Int32(0) - while it < it2: + while it < it2: # L2628-2645 i = it * cutlass.Int32(BLK) + tid uke = cutlass.Uint32(0) - idv = cutlass.Int32(0) - if i < mc: # predicated remote - uke = cutlass.Uint32( - _ld_shared_cluster_i32( - _mapa_shared_cluster_addr( - ck_addr - + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), - i >> cutlass.Int32(LCMPC), - ) - ) - ) - idv = _ld_shared_cluster_i32( - _mapa_shared_cluster_addr( - ci_addr - + ((i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2)), - i >> cutlass.Int32(LCMPC), - ) - ) + if i < n: + uke = fkey(ldg_f32(x_addr, i)) q1f = cutlass.Int32(0) q2f = cutlass.Int32(0) - if i < mc: + if i < n: if cutlass.Int64(cutlass.Uint32(uke)) > ethr: q1f = cutlass.Int32(1) - if cutlass.Int64(cutlass.Uint32(uke)) == ethr: - q2f = cutlass.Int32(1) + if tie_m == cutlass.Int32(1): + if cutlass.Int64(cutlass.Uint32(uke)) == ethr: + q2f = cutlass.Int32(1) n1 = ballot(q1f == cutlass.Int32(1)) n2 = ballot(q2f == cutlass.Int32(1)) b1 = cutlass.Int32(0) @@ -5912,134 +6097,30 @@ def kern(self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: c p1e = b1 + popc(n1 & lml) p2e = b2 + popc(n2 & lml) if q1f == cutlass.Int32(1): - if p1e < aboveC: - out_row[above + p1e] = idv + if p1e < nA: + out_row[p1e] = i if q2f == cutlass.Int32(1): - if p2e < needC: - out_row[above + aboveC + p2e] = idv + if p2e < nT: + out_row[nA + p2e] = i it = it + cutlass.Int32(1) - else: - # (3) degen safety net (L2597-2645): crossing bin larger - # than the whole cluster buffer -> exact whole-row - # key-space narrowing by rank 0 alone, <=8 levels. - rlo = cutlass.Uint32(0) - rhi = cutlass.Uint32(0xFFFFFFFF) - aboveC = cutlass.Int32(0) # above2 - needC = k # need2 - mm = n # m2 - ethr = cutlass.Int64(0) - tie_m = cutlass.Int32(1) - lev = cutlass.Int32(0) - done = cutlass.Int32(0) - while done == cutlass.Int32(0): - if needC == mm: # L2603 - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) - aboveC = aboveC + mm - needC = cutlass.Int32(0) - tie_m = cutlass.Int32(0) - done = cutlass.Int32(1) - if done == cutlass.Int32(0): - if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - done = cutlass.Int32(1) - if lev >= cutlass.Int32(8): # L2605 - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - done = cutlass.Int32(1) - if done == cutlass.Int32(0): - d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) - b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) - sh2 = cutlass.Int32(0) - if b2w > cutlass.Int32(LNB): - sh2 = b2w - cutlass.Int32(LNB) - for z in cutlass.range_constexpr(NB__regclus // self.blk): - s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L2612 - i = tid - while i < n: # whole-row bin L2613 - unar = fkey(ldg_f32(x_addr, i)) - if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): - if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): - bnn = ( - cutlass.Uint32(unar) - cutlass.Uint32(rlo) - ) >> cutlass.Uint32(sh2) - bnn = _umin_u32__regclus( - bnn, cutlass.Uint32(NB__regclus - 1) - ) - atomic_add_cta( - s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) - ) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2618 - find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) - cute.arch.barrier() # L2620 - aboveC = aboveC + s_res[RES_ABOVE] - needC = needC - s_res[RES_ABOVE] - mm = s_res[RES_M] - b_lv = s_res[RES_B] - nlo = cutlass.Uint32(rlo) + ( - cutlass.Uint32(b_lv) << cutlass.Uint32(sh2) - ) - if b_lv != cutlass.Int32(NB__regclus - 1): - rhi = nlo + ( - (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) - ) - rlo = nlo - lev = lev + cutlass.Int32(1) - cute.arch.barrier() # L2626 - nA = k # tie_m ? above2 : k - if tie_m == cutlass.Int32(1): - nA = aboveC - nT = cutlass.Int32(0) - if tie_m == cutlass.Int32(1): - nT = needC - lml = cutlass.Int32(cute.arch.lanemask_lt()) - it2 = (n + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) - it = cutlass.Int32(0) - while it < it2: # L2628-2645 - i = it * cutlass.Int32(BLK) + tid - uke = cutlass.Uint32(0) - if i < n: - uke = fkey(ldg_f32(x_addr, i)) - q1f = cutlass.Int32(0) - q2f = cutlass.Int32(0) - if i < n: - if cutlass.Int64(cutlass.Uint32(uke)) > ethr: - q1f = cutlass.Int32(1) - if tie_m == cutlass.Int32(1): - if cutlass.Int64(cutlass.Uint32(uke)) == ethr: - q2f = cutlass.Int32(1) - n1 = ballot(q1f == cutlass.Int32(1)) - n2 = ballot(q2f == cutlass.Int32(1)) - b1 = cutlass.Int32(0) - b2 = cutlass.Int32(0) - if lane == cutlass.Int32(0): - if n1 != cutlass.Int32(0): - b1 = atomic_add_cta(s_cnt.iterator, popc(n1)) - if n2 != cutlass.Int32(0): - b2 = atomic_add_cta(s_cnt.iterator + 1, popc(n2)) - b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) - b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) - p1e = b1 + popc(n1 & lml) - p2e = b2 + popc(n2 & lml) - if q1f == cutlass.Int32(1): - if p1e < nA: - out_row[p1e] = i - if q2f == cutlass.Int32(1): - if p2e < nT: - out_row[nA + p2e] = i - it = it + cutlass.Int32(1) - # ---- P10: FINAL cluster rendezvous (L2648) — ALL ranks reach it; - # keeps peers resident until rank 0 has read their ck/ci. - _cluster_sync_aligned() + # ---- P10: FINAL cluster rendezvous (L2648) — ALL ranks reach it; + # keeps peers resident until rank 0 has read their ck/ci. + _cluster_sync_aligned() # ------------------------------------------------------------------ @cute.jit def __call__( - self, logits: cute.Tensor, pre_idx: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, stream + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + kv_lens: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + stream, ): - b = logits.shape[0] - self.kern(logits, pre_idx, out, n).launch( + b = out.shape[0] + self.kern(logits, pre_idx, kv_lens, out, n).launch( grid=(self.cs, b, 1), block=(self.blk, 1, 1), cluster=(self.cs, 1, 1), @@ -6056,15 +6137,17 @@ def __call__( _COMPILE_CACHE__regclus: dict = {} -def get_compiled__regclus(tpl, dump_dir=None, pdl=False): +def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" - key = (tuple(tpl), bool(pdl)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) compiled = _COMPILE_CACHE__regclus.get(key) if compiled is None: from cutlass.cute import runtime as _crt blk, vpt, cs = tpl - kernel = GvrRegClusKernel(blk, vpt, cs, pdl=pdl) + kernel = GvrRegClusKernel( + blk, vpt, cs, pdl=pdl, varlen=varlen, next_n=next_n, cr_shift=cr_shift + ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() nb3_, nc3_ = cute.sym_int(), cute.sym_int() @@ -6077,12 +6160,23 @@ def get_compiled__regclus(tpl, dump_dir=None, pdl=False): out_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (nb3_, nc3_), stride_order=(1, 0), assumed_align=16 ) + v0_ = cute.sym_int() + kv_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (v0_,), stride_order=(0,), assumed_align=4 + ) fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) opts = "--enable-tvm-ffi" if dump_dir: opts += f" --keep-ptx --keep-cubin --dump-dir {dump_dir}" compiled = cute.compile( - kernel, lg_fake, pi_fake, out_fake, cutlass.Int32(0), stream=fake_stream, options=opts + kernel, + lg_fake, + pi_fake, + kv_fake, + out_fake, + cutlass.Int32(0), + stream=fake_stream, + options=opts, ) _COMPILE_CACHE__regclus[key] = compiled return compiled diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index a1f0164f0e75..e0e40c7ccfab 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -749,12 +749,28 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): hit = _VARLEN_CACHE.get(key) if hit is not None: return hit - plan = route_streaming(num_rows, max(min(n_env, npad), k + 1), npad, k, force_main=True) + n_eff = max(min(n_env, npad), k + 1) + cr_shift = 0 if cr == 1 else 2 + dev = _device() + # ---- route() parity, family tier 1: clustered register-resident -------- + # Admit reg_clus exactly where the free route picks it (886-real-capture + # grid: this family recovers the deep-layer distribution tail and the + # large-N small-rows band; its whole admission window n4 <= 32768 fits + # capture-frozen envelopes). The choice is a pure function of this cache + # key, so CUDA-graph replay safety is unchanged; per-row n / short-row + # handling lives in-kernel (GvrMainKernel varlen discipline). + plan_free = route(num_rows, n_eff, npad, k) + if plan_free["kernel"] == "reg_clus": + fn = dev.get_compiled__regclus( + tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + ) + lc = ("reg_clus", fn, n_eff) + _VARLEN_CACHE[key] = lc + return lc + plan = route_streaming(num_rows, n_eff, npad, k, force_main=True) tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) rt = plan["rt"] r_const = rt["R"] - cr_shift = 0 if cr == 1 else 2 - dev = _device() # TSHG (tpl[6]) is dead under varlen (the ctor compiles the TSH # machinery in whenever SPLIT); normalize it out of the compile key so # row counts differing only in that slot share one engine @@ -781,7 +797,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): tsh_en = 1 if (tpl[5] and k <= 1024) else 0 pre = (0, npad, k, rt["SCAP_"], rt["CMP_"], r_const, 0, 0, 0, 0, 0) tail = (aim_base, sfac, amin, sd_en, tsh_en) - lc = (fn, pre, tail) + lc = ("main", fn, pre, tail) _VARLEN_CACHE[key] = lc return lc @@ -1066,9 +1082,16 @@ def fn(lg, pi, o, w, *a, _raw=raw): return (fn, args, False) if fam == "reg_clus": dev = _device() - # compiled ABI: (logits, pre_idx, out, n) -- smem/k derived in-module + # compiled ABI: (logits, pre_idx, kv_lens, out, n) -- kv_lens is the + # dead varlen slot in batch-uniform mode (dummy, gvr_main precedent); + # smem/k derived in-module fn = dev.get_compiled__regclus(tpl) - return (fn, (rt["n"],), False) + n_arg = rt["n"] + + def _call(lg, pi, idx, _fn=fn, _n=n_arg): + _fn(lg, pi, _dummy_kv(lg.get_device(), lg.device), idx, _n) + + return (_call, (), False) # unreachable: route() only emits the five families above raise RuntimeError(f"unknown dispatch family {fam!r}") @@ -1403,14 +1426,18 @@ def run_varlen( "before CUDA graph capture" ) lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) - fn, pre, tail = lc idx = indices if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) vals = values if vals is not None and vals.shape[1] != k: vals = vals.reshape(-1)[: num_rows * k].view(num_rows, k) - fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) + if lc[0] == "reg_clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) + lc[1](lg, pre_idx, kv_lens, idx, lc[2]) + else: + _, fn, pre, tail = lc + fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) if vals is not None: idx64 = idx.to(torch.int64) vals.copy_(lg.gather(1, idx64.clamp_min(0))) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index d0f4a15d33ee..f137ecdb6251 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -695,3 +695,70 @@ def test_selfsampling_warmup_drops_rows_beyond_envelope(): assert not any(r > ss_host.MAX_VARLEN_ROWS for r in new_rows), new_rows have_rows = {key[0] for key in ss_host._VARLEN_CACHE} assert {nn, 8 * nn} <= have_rows, "in-envelope rows must still be warmed" + + +def test_selfsampling_varlen_regclus_parity_and_oracle(): + """The varlen launcher must admit the clustered register-resident family + exactly where the free route picks it (route() parity tier 1), and the + per-row varlen port must match the reference oracle on a heterogeneous + batch: long rows, a short row (n <= k, in-kernel identity + -1 tail) and + a zero-window row, under MTP row windows (next_n=4).""" + k, msl_c, nn, cr = 1024, 131072, 4, 4 + npad = msl_c # 256-aligned already + assert ss_host.route(8, msl_c, npad, k)["kernel"] == "reg_clus" + batch = 3 + rows = batch * nn + torch.manual_seed(7) + lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) + kv = torch.tensor([msl_c * cr, 900, nn - 1], dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + key = (rows, npad, k, msl_c, nn, cr) + assert ss_host._VARLEN_CACHE[key][0] == "reg_clus", ss_host._VARLEN_CACHE[key][0] + ss_host.run_varlen( + lg, + pre, + kv, + ref, + next_n=nn, + compress_ratio=cr, + max_seq_len=msl_c * cr, + engine="reference", + ) + torch.cuda.synchronize() + for r in range(rows): + if (ref[r] >= 0).any(): + row = lg[r].float() + got = row[out[r].long().clamp_min(0)].sort().values + want = row[ref[r].long().clamp_min(0)].sort().values + assert torch.equal(got, want), f"row {r} value multiset mismatch" + assert torch.equal(out[r] < 0, ref[r] < 0), f"row {r} pad mask mismatch" + else: + assert torch.equal(out[r], ref[r]), f"row {r} expected all -1" + + +def test_selfsampling_varlen_regclus_cuda_graph(): + """Cluster-family varlen engine must be CUDA-graph capturable: warmed + engine, capture one launch, replay twice, tie-aware exact each time.""" + k, msl_c, cr = 1024, 131072, 4 + rows = 8 + torch.manual_seed(11) + lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) + kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + out.fill_(-7) + with torch.cuda.graph(g): + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values + for _ in range(2): + out.fill_(-7) + g.replay() + torch.cuda.synchronize() + got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(got, ref_v) From bcf6c472205e1e69a29f245e80cd6f36e3db67b8 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:52:10 +0000 Subject: [PATCH 26/40] =?UTF-8?q?[TRTLLM-15293][perf]=20self-sampling=20to?= =?UTF-8?q?p-K:=20full-range=20dispatch=20=E2=80=94=20every=20row=20count?= =?UTF-8?q?=20and=20envelope=20goes=20through=20the=20self-sampling=20engi?= =?UTF-8?q?nes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the rows admission envelope: with TRTLLM_GVR_SELF_SAMPLING=1 the self-sampling engines now serve the full production range (BS 1..1024+ x next_n, envelopes to 1M kv tokens) instead of falling through to the in-tree path above 32 rows — restoring the standalone design intent. Warmup switches from row-count filtering to band-aware enumeration: one representative row per distinct engine compile key (bounded time and memory for arbitrarily large CUDA-graph batch lists), then per-row-count launcher-cache population (pure host work) so capture at any requested geometry finds its key immediately — preserving the warm-the-key-dispatch- looks-up discipline the arena-stride capture test pins. Production-lessons checklist re-validated on this change: capture-stable pure-function dispatch (family from the launcher key, per-row n in-kernel from device kv_lens), loud-fail on uncompiled keys under capture, two-pass pre-capture warmup coverage, producer row-stride key mirroring, workspace slab grow-never/keep-alive, cluster-family graph capture, heterogeneous kv_lens + MTP row-window oracle tests. Full suite 79/79. Real-capture record (vs the in-tree production op, cold-L2, all exact): rows 64..1024 wins 2.9-4.7x at N=262144, 3.1-3.7x on the former deep-layer pathology band, 1.6-3.5x on DSv3.2 163k; short-N cells (N <= ~2k) trade 0.67-0.94x at microsecond scale — accepted by design: enabled means self-sampling, everywhere. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 23 +----- .../attention_backend/sparse/dsa/metadata.py | 9 ++- .../blackwell/top_k/__init__.py | 2 - .../gvr_topk_decode_self_sampling_host.py | 62 ++++++++++----- .../parallel/test_gvr_selfsampling_topk.py | 76 ++++++++++--------- 5 files changed, 88 insertions(+), 84 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 61251479b32c..2963a4808502 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -712,15 +712,10 @@ def __init__( and compress_ratio in (1, 4) ) self._selfsampling_run_varlen = None - self._selfsampling_max_rows = 0 if self._use_self_sampling_topk: - from ....cute_dsl_kernels.blackwell.top_k import ( - SELFSAMPLING_TOPK_MAX_ROWS, - selfsampling_topk_run_varlen, - ) + from ....cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen self._selfsampling_run_varlen = selfsampling_topk_run_varlen - self._selfsampling_max_rows = SELFSAMPLING_TOPK_MAX_ROWS self.mtp_index_share = sparse_params.mtp_index_share if self._enable_heuristic_topk and layer_idx == 0: @@ -1848,22 +1843,6 @@ def sparse_attn_indexer( and self._enable_heuristic_topk and pre_idx is not None ) - if _ss_ready and num_gen_tokens > self._selfsampling_max_rows: - # batch-rows admission envelope: the varlen engine splits - # a row across CTAs only for small batches; beyond - # MAX_VARLEN_ROWS it runs one CTA per row and loses ~5x - # per call to the in-tree per-row-split kernels at long - # rows (DSv4 Pro 1M-ISL MTP7, rows=304: 2.63 ms vs - # 0.49 ms). Large batches take the in-tree path until a - # throughput tier lands. - logger.info_once( - "self-sampling GVR top-K: decode batch rows exceed " - f"the validated envelope ({num_gen_tokens} > " - f"{self._selfsampling_max_rows}); large batches take " - "the in-tree top-K path.", - key="selfsampling_topk_rows_fallthrough", - ) - _ss_ready = False if _ss_ready and not ( # engine hardware-format gate (falls through otherwise): # fp32 row-major logits with a float4-aligned row stride diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 7204ec15ee55..c6f39dc3373e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -292,10 +292,11 @@ def warmup_selfsampling_topk( DeepGEMM is exact-width) so the warmed keys are the ones dispatch actually looks up. Captured geometries are also compiled by the pre-capture warmup forwards; eager batches outside ``batch_sizes`` - still compile lazily on first touch. Row counts above the engine's - admission envelope (``MAX_VARLEN_ROWS``) are dropped by the helper — - those batches dispatch to the in-tree path instead. No-op unless the - opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + still compile lazily on first touch. The helper enumerates one + representative row per distinct engine compile key (band-aware), so + arbitrarily large batch lists warm in bounded time and memory. + No-op unless the opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects + the engine. """ if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": 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 b5aef709096f..d5cab389489b 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 @@ -20,7 +20,6 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel -from .gvr_topk_decode_self_sampling_host import MAX_VARLEN_ROWS as SELFSAMPLING_TOPK_MAX_ROWS from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -37,5 +36,4 @@ "tiered_topk", "is_tiered_topk_supported", "selfsampling_topk_run_varlen", - "SELFSAMPLING_TOPK_MAX_ROWS", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index e0e40c7ccfab..47d79dea8b04 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -149,16 +149,6 @@ def _device(): CMPC = 4096 # L2372 crossing-bin slots per CTA, clustered register path BLKC = 1024 # L2374 CTA size of the clustered register path -# Dispatch-policy admission envelope (not a kernel.cu constant): the -# streaming varlen engine multi-CTA-splits a row only for small batches -# (route_streaming: R = min(148//b, ...) for b <= 32); past that R collapses -# toward one CTA per row and long rows lose badly to the in-tree per-row -# split kernels (measured 5.4x per call at rows=304, n=262144, K=1024, -# B200: 2.63 ms vs 0.49 ms). Integrations must fall through to the in-tree -# path above this row count; warmup_varlen drops larger requests so engine -# init never compiles keys dispatch will not admit. -MAX_VARLEN_ROWS = 32 - def route(b: int, n: int, npad: int, k: int) -> dict[str, object]: """Mirror of gvr_topk_launch (kernel.cu L2754-3197). Pure. See module doc.""" @@ -1301,11 +1291,10 @@ def run_varlen( probe battery). Finite inputs — including +/-inf and denormals — are tie-aware exact. - PERFORMANCE ENVELOPE: correct for any ``num_rows``, but the streaming - engine splits a row across CTAs only for small batches — past - ``MAX_VARLEN_ROWS`` it runs one CTA per row and long rows lose ~5x to - the in-tree per-row-split kernels. Serving integrations must gate on - ``MAX_VARLEN_ROWS`` and fall through to the in-tree path above it. + FULL-RANGE PRODUCTION CONTRACT: correct and dispatched for any + ``num_rows`` (BS 1..1024+ x next_n) and any envelope up to 1M kv tokens. + Family selection (streaming main / clustered register-resident) is a + pure function of the capture-stable launcher key. """ if logits.dtype is not torch.float32: raise RuntimeError( @@ -1535,10 +1524,38 @@ def warmup_varlen( dev = torch.cuda.current_device() nn = max(1, int(next_n)) # round each request down to a next_n multiple (min next_n) and dedup - rows_list = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) - # admission envelope: dispatch falls through to the in-tree path above - # MAX_VARLEN_ROWS, so compiling larger keys would only burn init time - rows_list = [r for r in rows_list if r <= MAX_VARLEN_ROWS] + req_rows = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) + if not req_rows: + return + # BAND-AWARE enumeration: the engine compile key depends on the plan's + # constexpr tuple (+ r_const family axis), NOT on the exact row count, so + # warming ONE representative row per distinct engine key covers every row + # count up to the largest request. Representatives are the first row of + # each band, which keeps the warmup allocation bounded (~a few hundred + # rows) even when CUDA-graph batch lists reach thousands of rows. + n_env_c = max(1, int(max_seq_len) // int(compress_ratio)) + npad_c = (n_env_c + 63) // 64 * 64 if row_stride is None else int(row_stride) + seen_keys = set() + rows_list = [] + r = nn + r_max = req_rows[-1] + while r <= r_max: + plan_free = route(r, max(min(n_env_c, npad_c), int(top_k) + 1), npad_c, int(top_k)) + if plan_free["kernel"] == "reg_clus": + ekey = ("reg_clus", tuple(plan_free["tpl"])) + else: + p = route_streaming( + r, + max(min(n_env_c, npad_c), int(top_k) + 1), + npad_c, + int(top_k), + force_main=True, + ) + ekey = ("main", tuple(p["tpl"][:6]), p["rt"]["R"]) + if ekey not in seen_keys: + seen_keys.add(ekey) + rows_list.append(r) + r += nn if not rows_list: return n_env = max(1, int(max_seq_len) // int(compress_ratio)) @@ -1574,5 +1591,12 @@ def warmup_varlen( ) del logits, kv_lens, pre_idx, out torch.cuda.synchronize() + # band launches compiled every ENGINE; now populate the per-row-count + # LAUNCHER cache entries for the exact requested row counts (pure host + # work, zero allocation/launch — engines hit the compile cache), so a + # CUDA-graph capture at any requested geometry finds its key immediately. + n_env_l = min(max(int(max_seq_len) >> (0 if int(compress_ratio) == 1 else 2), 1), npad) + for r in req_rows: + _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio)) with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index f137ecdb6251..4b52217718ae 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -660,43 +660,6 @@ def test_selfsampling_warmup_row_stride_matches_arena(): assert torch.equal(ref, got) -def test_selfsampling_rows_envelope_constant(): - """The rows admission envelope is one shared constant, and it matches - route_streaming's multi-CTA region: inside the envelope a long row is - split across CTAs (R > 1); far outside it the split collapses to one - CTA per row — the regression measured at rows=304, n=262144 (5.4x vs - the in-tree per-row-split kernel).""" - from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import SELFSAMPLING_TOPK_MAX_ROWS - - assert SELFSAMPLING_TOPK_MAX_ROWS == ss_host.MAX_VARLEN_ROWS == 32 - n = 262144 - plan_in = ss_host.route_streaming(ss_host.MAX_VARLEN_ROWS, n, n, 1024, force_main=True) - assert int(plan_in["grid"][0]) > 1, "inside the envelope long rows must split" - plan_out = ss_host.route_streaming(304, n, n, 1024, force_main=True) - assert int(plan_out["grid"][0]) == 1, "rows=304 collapses to one CTA per row" - - -def test_selfsampling_warmup_drops_rows_beyond_envelope(): - """warmup_varlen must not compile engines for row counts the dispatch - envelope never admits (rows > MAX_VARLEN_ROWS): a CUDA-graph batch-size - list with large batches warms only the admitted prefix.""" - k, msl, nn = 512, 8192, 4 - stride = (msl + 255) // 256 * 256 - before = set(ss_host._VARLEN_CACHE.keys()) - ss_host.warmup_varlen( - k, - msl, - compress_ratio=1, - next_n=nn, - num_rows_list=(nn, 8 * nn, 76 * nn), # rows 4, 32, 304 - row_stride=stride, - ) - new_rows = {key[0] for key in ss_host._VARLEN_CACHE if key not in before} - assert not any(r > ss_host.MAX_VARLEN_ROWS for r in new_rows), new_rows - have_rows = {key[0] for key in ss_host._VARLEN_CACHE} - assert {nn, 8 * nn} <= have_rows, "in-envelope rows must still be warmed" - - def test_selfsampling_varlen_regclus_parity_and_oracle(): """The varlen launcher must admit the clustered register-resident family exactly where the free route picks it (route() parity tier 1), and the @@ -762,3 +725,42 @@ def test_selfsampling_varlen_regclus_cuda_graph(): torch.cuda.synchronize() got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values assert torch.equal(got, ref_v) + + +def test_selfsampling_varlen_full_row_range(): + """Full-range production contract: with self-sampling enabled, EVERY row + count dispatches to the self-sampling engines (no rows-based fall-through) + — spot-check the throughput end (rows 304 and 1024) for tie-aware + exactness at a mid-size envelope.""" + k, msl_c, cr = 1024, 65536, 4 + torch.manual_seed(3) + for rows in (304, 1024): + lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) + kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + torch.cuda.synchronize() + key = (rows, msl_c, k, msl_c, 1, cr) + assert key in ss_host._VARLEN_CACHE, "row count must dispatch in-engine" + ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values + got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(got, ref_v), f"rows={rows} value multiset mismatch" + + +def test_selfsampling_warmup_band_enumeration_bounded(): + """warmup_varlen must cover arbitrarily large batch lists by warming one + representative row per distinct engine compile key — bounded time and + memory (the representative rows saturate a few hundred, never the + requested thousands).""" + k, msl = 512, 65536 # kv tokens, cr=4 -> n_env 16384: few bands, fast + before = len(ss_host._VARLEN_WARMUP_DONE) + ss_host.warmup_varlen( + k, + msl, + compress_ratio=4, + next_n=4, + num_rows_list=(4096,), + row_stride=(msl // 4 + 255) // 256 * 256, + ) + assert len(ss_host._VARLEN_WARMUP_DONE) == before + 1 From 303ac14ad772b35073429c278c7828e7d8e282e4 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:57:19 +0000 Subject: [PATCH 27/40] [TRTLLM-15293][test] self-sampling top-K: heterogeneous-lengths row-independence test at throughput scale MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pin the varlen row-independence contract on the streaming main family: one 304-row batch mixing random per-request kv_lens with full-length, short (n <= k), zero-window and boundary rows under MTP row windows (next_n=4) — every row must match its own per-prefix torch.topk value multiset, short rows the in-kernel identity + (-1) tail, zero-window rows all -1. No engine may assume batch-uniform indexer-logits lengths. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../parallel/test_gvr_selfsampling_topk.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 4b52217718ae..514e34b4205d 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -764,3 +764,41 @@ def test_selfsampling_warmup_band_enumeration_bounded(): row_stride=(msl // 4 + 255) // 256 * 256, ) assert len(ss_host._VARLEN_WARMUP_DONE) == before + 1 + + +def test_selfsampling_varlen_heterogeneous_lengths_main(): + """Row-independence contract at throughput scale on the streaming main + family: rows are naturally independent tasks — per-row data AND length + (random kv_lens spanning long / mid / short / n<=k / zero-window rows in + one 304-row batch, MTP row windows via next_n=4). Every row must match + its own per-prefix torch.topk value multiset; short rows must be + identity + (-1) tail; zero-window rows all -1.""" + k, msl_c, cr, nn = 1024, 65536, 4, 4 + batch = 76 + rows = batch * nn # 304 -> route_streaming main family + torch.manual_seed(5) + lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) + kv = torch.randint(1, msl_c * cr, (batch,), dtype=torch.int32, device=_DEV) + kv[0] = msl_c * cr # full length + kv[1] = 900 # short (n <= k) + kv[2] = nn - 1 # zero-window (every row of the request empty) + kv[3] = k * cr + nn # just above the short path + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + torch.cuda.synchronize() + kl = kv.tolist() + for r in range(rows): + n_r = max(kl[r // nn] - nn + (r % nn) + 1, 0) // cr + n_r = min(n_r, msl_c) + if n_r <= 0: + assert (out[r] == -1).all(), f"row {r}: zero-window must be all -1" + continue + if n_r <= k: + want = list(range(n_r)) + [-1] * (k - n_r) + assert out[r].tolist() == want, f"row {r}: short path identity+pad" + continue + row = lg[r].float()[:n_r] + ref_v = torch.topk(row, k).values.sort().values + got = row[out[r].long()].sort().values + assert torch.equal(got, ref_v), f"row {r}: value multiset mismatch (n={n_r})" From 4699e5d49f54156f4f6f92692ea443d89b399ed8 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:29:04 +0000 Subject: [PATCH 28/40] [TRTLLM-15293][perf] self-sampling top-K: per-row varlen port of the register-resident family (reg/regimg) Route-parity tier 2: the varlen launcher now admits the register-resident family (and its img-window flavor) exactly where the free route picks it, instead of forcing those bands through the streaming main kernel. This family owns the small/mid-N band across all row counts in the standalone dispatch, where main pays a 1.5x tax. Kernel side (GvrRegClusKernel discipline): per-row n re-derived in-kernel from device kv_lens (request windows via next_n/cr_shift, envelope clamp); short rows (n <= k) emit identity + (-1) tail in-kernel and skip the body -- k can exceed BLK on this family, so the emit is a strided loop, not the reg_clus single predicate. kv_lens joins the ABI as a dead dummy slot in batch-uniform mode (gvr_main precedent). CMP/QC/smem stay envelope-derived launch constants: in-kernel they are pure capacity clamps, a fast-path threshold and the launch smem size -- safe upper bounds for every per-row n <= envelope. Production lessons re-checked: capture-stable dispatch (family remains a pure function of the launcher cache key), warmup band enumeration gains the reg engine key, loud-fail on uncompiled keys under capture unchanged. Tests: route-parity + heterogeneous oracle (reg, k>BLK strided short-row emit, regimg window) and CUDA-graph capture/replay; suite 82/82. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 1334 +++++++++-------- .../gvr_topk_decode_self_sampling_host.py | 34 +- .../parallel/test_gvr_selfsampling_topk.py | 81 + 3 files changed, 827 insertions(+), 622 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 897f4e6c7b33..b5b5aaa21bb9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -3256,6 +3256,9 @@ def __init__( img: bool, nbh: int = NB__reg, pdl: bool = False, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, ): assert blk in (256, 512, 1024) and vpt in (1, 2, 4) assert nbh in (256, 512, 1024, 2048) @@ -3269,6 +3272,15 @@ def __init__( self.img = bool(img) self.nbh = nbh self.pdl = bool(pdl) + # per-row varlen mode (production heuristicTopKDecode contract, same + # semantics as GvrMainKernel/GvrRegClusKernel): n is re-derived PER + # ROW in-kernel from a device kv_lens tensor; the scalar n launch arg + # becomes the envelope clamp bound. next_n / cr_shift compile-time. + self.varlen = bool(varlen) + self.next_n = int(next_n) + self.cr_shift = int(cr_shift) + if self.varlen: + assert self.next_n >= 1 and self.cr_shift in (0, 2) # derived compile-time constants (kernel.cu L1292-1294, L1359, L1389, L1485) self.S = vpt * 4 self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) @@ -3282,6 +3294,7 @@ def kern( self, logits: cute.Tensor, pre_idx: cute.Tensor, + kv_lens: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, cmp_: cutlass.Int32, @@ -3303,6 +3316,45 @@ def kern( row, _, _ = cute.arch.block_idx() lane = tid & cutlass.Int32(31) + # ================= per-row varlen prologue (varlen mode only) ========= + # Production heuristicTopKDecode contract (GvrMainKernel / + # GvrRegClusKernel discipline): row r serves request r // next_n with + # n = (kv_lens[req] - next_n + r % next_n + 1) >> cr_shift, clamped to + # the envelope launch arg n (the launcher admits this family only when + # the envelope fits its capacity window, so per-row n never exceeds + # capacity). One CTA per row, so the whole-body guard below is + # trivially block-uniform. Short rows (n <= k) emit identity + (-1) + # tail here (k can exceed BLK on this family -> strided loop, unlike + # the reg_clus k <= BLK single predicate) and SKIP the body entirely: + # a zero-work pass would reach the degenerate emitter and poison out. + short = cutlass.Int32(0) + prow = row + if cutlass.const_expr(self.varlen): + kq = cutlass.Int32(pre_idx.shape[1]) + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + prow = req + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > n: + nv = n + if nv <= kq: + short = cutlass.Int32(1) + if short == cutlass.Int32(0): + n = nv + if short != cutlass.Int32(0): + i = tid + while i < kq: + ov = cutlass.Int32(-1) + if i < nv: + ov = i + out[row, i] = ov + i = i + cutlass.Int32(BLK) + # ------------------------------------------------------------------ # Predeclarations: the DSL AST transformer requires every scalar that # is (re)assigned under a dynamic if/while region to pre-exist with a @@ -3383,346 +3435,393 @@ def kern( qt2 = cutlass.Float32(0.0) qt3 = cutlass.Float32(0.0) - npad = cutlass.Int32(logits.shape[1]) # noqa: F841 - k = cutlass.Int32(pre_idx.shape[1]) - out_row = out[row, None] - x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base - p_addr = pre_idx[row, None].iterator.toint() - - # ---- shared-memory window (map in module docstring) ---- - sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) - sbase = sptr.toint() # Int32 shared addr - - s_res = _smem_view(cutlass.Int32, sbase, 0, 6) - s_cnt = _smem_view(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_oc - s_kmm = _smem_view(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax - s_e12 = _smem_view(cutlass.Int32, sbase, 10, 2) # [0]=s_e1 [1]=s_e2 - s_ws = _smem_view(cutlass.Int32, sbase, 16, 32) - s_wmn = _smem_view(cutlass.Uint32, sbase, 48, 32) - s_wmx = _smem_view(cutlass.Uint32, sbase, 80, 32) - s_hist = _smem_view(cutlass.Int32, sbase, STATIC_WORDS, self.nbh) - ck_base = sbase + cutlass.Int32((STATIC_WORDS + self.nbh) * 4) - ck = cute.make_tensor( - cute.make_ptr(cutlass.Uint32, ck_base, cute.AddressSpace.smem, assumed_align=16), - cute.make_layout((65536,)), - ) # typed view, no bound - ci = cute.make_tensor( - cute.make_ptr( - cutlass.Int32, - ck_base + cmp_ * cutlass.Int32(4), - cute.AddressSpace.smem, - assumed_align=4, - ), - cute.make_layout((65536,)), - ) - img_f = cute.make_tensor( # aliases ck/ci (L1299) - cute.make_ptr(cutlass.Float32, ck_base, cute.AddressSpace.smem, assumed_align=16), - cute.make_layout((65536,)), - ) - bm = cute.make_tensor( # aliases ck (L1409) - cute.make_ptr(cutlass.Int32, ck_base, cute.AddressSpace.smem, assumed_align=16), - cute.make_layout((65536,)), - ) - - n4 = n >> cutlass.Int32(2) - ntail = n - (n4 << cutlass.Int32(2)) - tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L1351 - - # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any dependent - # gather (L1314-1322); compiled out under DEG. - pvs = [] - if cutlass.const_expr(not self.deg): - for t in cutlass.range_constexpr(KPT): - pv = cutlass.Int32(-1) - j = tid + cutlass.Int32(t * self.blk) - if j < k: - pv = ld_g_i32(p_addr, j) - pvs.append(pv) - - # ---- row load: exact-fit peel + float4[VPT] register batch (L1327-1350) - atom128 = g2r_atom_f32(128, invariant=True) - frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] - if n4 >= cutlass.Int32(self.blk * self.vpt): # block-uniform peel - for u in cutlass.range_constexpr(VPT): - ld_g_f32x4(atom128, x_addr, tid + cutlass.Int32(u * self.blk), frags[u]) - else: # predicated flat batch - for u in cutlass.range_constexpr(VPT): - i = tid + cutlass.Int32(u * self.blk) - if i < n4: - ld_g_f32x4(atom128, x_addr, i, frags[u]) - for u in cutlass.range_constexpr(VPT): - i = tid + cutlass.Int32(u * self.blk) - if i >= n4: # -INFINITY fill L1346 - for q in cutlass.range_constexpr(4): - frags[u][q] = cutlass.Float32(_NEG_INF__reg) - - tval = cutlass.Float32(_NEG_INF__reg) - if tid < ntail: - tval = ldg_f32(x_addr, tix) # L1352 - - # ---- init (L1391-1392) - if tid == cutlass.Int32(0): - s_cnt[0] = cutlass.Int32(0) - s_cnt[1] = cutlass.Int32(0) - for z in cutlass.range_constexpr(self.nbh // self.blk): - s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - - # ---- bracket: 4 mutually exclusive compile-time arms (L1393-1451) - lmin = cutlass.Uint32(0xFFFFFFFF) - lmax = cutlass.Uint32(0) - if cutlass.const_expr(self.use_img): - fatom = _f32_smem_atom() - for u in cutlass.range_constexpr(VPT): # VPT == 1 here - i = tid + cutlass.Int32(u * self.blk) - if i < n4: - _sts128_f32(fatom, frags[u], ck_base, i * cutlass.Int32(16)) - if tid < ntail: - img_f[tix] = tval - cute.arch.barrier() # L1400 - for t in cutlass.range_constexpr(KPT): - p = pvs[t] - if cutlass.Uint32(p) < cutlass.Uint32(n): - uk = fkey(img_f[p]) - if uk < lmin: - lmin = uk - if uk > lmax: - lmax = uk - cute.arch.barrier() # L1406 (img dies) - elif cutlass.const_expr(self.use_bm): - nbw = (n + cutlass.Int32(31)) >> cutlass.Int32(5) - i = tid - while i < nbw: # bitmap clear (L1410) - bm[i] = cutlass.Int32(0) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1411 - for t in cutlass.range_constexpr(KPT): - p = pvs[t] - if cutlass.Uint32(p) < cutlass.Uint32(n): - atomic_or_cta( - bm.iterator + (p >> cutlass.Int32(5)), - cutlass.Int32(1) << (p & cutlass.Int32(31)), - ) - cute.arch.barrier() # L1417 - lmn = cutlass.Float32(_POS_INF) - lmx = cutlass.Float32(_NEG_INF__reg) - for u in cutlass.range_constexpr(VPT): - base = (tid + cutlass.Int32(u * self.blk)) << cutlass.Int32(2) - w = cutlass.Uint32(0) - if cutlass.Uint32(base) < cutlass.Uint32(n): - w = cutlass.Uint32(bm[base >> cutlass.Int32(5)]) >> cutlass.Uint32( - base & cutlass.Int32(31) - ) - for cbit in cutlass.range_constexpr(4): - if (w & cutlass.Uint32(1 << cbit)) != cutlass.Uint32(0): - lmn = fmin_f32(lmn, _val(frags, 4 * u + cbit)) - lmx = fmax_f32(lmx, _val(frags, 4 * u + cbit)) - if tid < ntail: - wt = cutlass.Uint32(bm[tix >> cutlass.Int32(5)]) >> cutlass.Uint32( - tix & cutlass.Int32(31) - ) - if (wt & cutlass.Uint32(1)) != cutlass.Uint32(0): - lmn = fmin_f32(lmn, tval) - lmx = fmax_f32(lmx, tval) - lmin = fkey(lmn) - lmax = fkey(lmx) # monotone (L1428) - cute.arch.barrier() # L1429 (bm dies) - elif cutlass.const_expr(self.deg): - lmn = cutlass.Float32(_POS_INF) - lmx = cutlass.Float32(_NEG_INF__reg) - for s in cutlass.range_constexpr(S): # L1436-1439 - v = _val(frags, s) - if v > cutlass.Float32(_NEG_INF__reg): - lmn = fmin_f32(lmn, v) - lmx = fmax_f32(lmx, v) - if tid < ntail: - lmn = fmin_f32(lmn, tval) - lmx = fmax_f32(lmx, tval) - lmin = fkey(lmn) - lmax = fkey(lmx) - else: - # default: KPT scattered fkey ldg gathers, batch-then-fold (L1443-1450) - xs = [] - for t in cutlass.range_constexpr(KPT): - xv = cutlass.Float32(0.0) - if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): - xv = ldg_f32(x_addr, pvs[t]) - xs.append(xv) - for t in cutlass.range_constexpr(KPT): - if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): - uk = fkey(xs[t]) - if uk < lmin: - lmin = uk - if uk > lmax: - lmax = uk - - # ---- block min/max in ONE barrier (L1452-1462); publishes hist clear - lmin = warp_min_u32(lmin) - lmax = warp_max_u32(lmax) - if lane == cutlass.Int32(0): - s_wmn[tid >> cutlass.Int32(5)] = lmin - s_wmx[tid >> cutlass.Int32(5)] = lmax - cute.arch.barrier() # L1456 - a = cutlass.Uint32(0xFFFFFFFF) + # wrap-scoping additions (varlen whole-body guard): names whose + # first assignment moves under the dynamic `if short == 0` region + # below and are reassigned deeper — same rule as the block above. + a = cutlass.Uint32(0) c = cutlass.Uint32(0) - if lane < cutlass.Int32(NW): - a = cutlass.Uint32(s_wmn[lane]) - c = cutlass.Uint32(s_wmx[lane]) - lmin = warp_min_u32(a) - lmax = warp_max_u32(c) - Tv = invkey(lmin) - GMAX = invkey(lmax) - - # ---- collapse guard, NaN-safe (L1464-1471) + lmin = cutlass.Uint32(0) + lmax = cutlass.Uint32(0) + esc = cutlass.Int32(0) okc = cutlass.Int32(0) - if Tv < GMAX: - if (GMAX - Tv) > cutlass.Float32(1e-30): - okc = cutlass.Int32(1) - if okc == cutlass.Int32(0): - Tv = cutlass.Float32(SENT_LO) - GMAX = cutlass.Float32(SENT_HI) - - # ---- bin transform constants (L1485-1511) - BRL = cutlass.const_expr(self.brl) # noqa: F841 - OFFf = cutlass.Float32(1.0 if self.brl else 0.0) - recip = 1.0 / float(self.nbh - (2 if self.brl else 0)) - WD = (GMAX - Tv) * cutlass.Float32(recip) - wsel = cutlass.Float32(1e-30) - if WD > cutlass.Float32(0.0): - wsel = WD - # rcp.approx (single MUFU.RCP) — the CUDA arm's exact lowering of - # `1.0f / wsel`; the previous `1.0 / wsel` spelling emitted the IEEE - # div.rn Newton triple + slowpath CALL on the barrier-bounded chain - # feeding all S classify FMULs. Output exactness is SC-invariant - # (any SC > 0 preserves the sign/monotonicity invariants, L1485-1511) - # and the WD > 0 arm is now bit-identical to CUDA's MUFU.RCP. - SC = cute.arch.rcp_approx(wsel) - QCAPf = cutlass.Float32(float(self.nbh - 1)) - CQ0 = OFFf - Tv * SC - CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf(CQ0) + cutlass.Float32(1.0)) - - # ---- histogram (L1513-1526) - if cutlass.const_expr(self.brl): - # fix-2 P4 (GATED, removable as one hunk): A1 ported to the BRL - # classify arm — hist base pinned ONCE via the same - # _smem_addr_reg__reg identity-mov used in the !BRL arm below, and the - # result-discarded classify atomics spelled as resultless - # red.shared (_red_shared_add1__reg). Value-identical: same +1 to the - # same byte address (hb + 4*bn == &s_hist[bn]), same .relaxed.cta - # ordering; the q/bn computations are untouched so classify/emit - # bit-identity (BRL requirement) is preserved. Emit-path hist - # atomics (results used, L1630+) are NOT touched. - hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) - for s in cutlass.range_constexpr(S): - q = _fmaf__reg(_val(frags, s), SC, CQ) - bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) - _red_shared_add1__reg(hb + (cutlass.Int32(bn) << cutlass.Int32(2))) - qt = _fmaf__reg(tval, SC, CQ) # unconditional (L1517) - bnt = _umin_u32(f2u_rz(qt), cutlass.Uint32(self.nbh - 1)) - _red_shared_add1__reg(hb + (cutlass.Int32(bnt) << cutlass.Int32(2))) - else: - # hist base pinned ONCE (byte addr, +STATIC_BYTES = word 128 map); - # each site below is then LEA + ATOMS exactly like the CUDA arm - # instead of re-deriving the shared window per divergent block. - hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) - for s in cutlass.range_constexpr(S): - q = _submul_asm(_val(frags, s), Tv, SC) # anti-CSE classify - if q >= cutlass.Float32(0.0): - _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(q, QCAPf)) << cutlass.Int32(2))) - qt = _submul_asm(tval, Tv, SC) - if qt >= cutlass.Float32(0.0): - _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(qt, QCAPf)) << cutlass.Int32(2))) - cute.arch.barrier() # L1527 - - # ---- crossing-bin find (L1528-1538) - if cutlass.const_expr(self.cur or self.nbh > 1024): - scan_cross_w(s_hist, s_ws, k, tid, s_res, blk=self.blk, nb=self.nbh) - else: - find_cross(s_hist, k, tid, s_res, nb=self.nbh) - cute.arch.barrier() # L1535 - above = s_res[RES_ABOVE] - m = s_res[RES_M] - Bv = s_res[RES_B] - need = k - above whole = cutlass.Int32(0) - if need >= m: - whole = cutlass.Int32(1) + tval = cutlass.Float32(0.0) + wsel = cutlass.Float32(0.0) + GMAX = cutlass.Float32(0.0) + Tv = cutlass.Float32(0.0) + lmn = cutlass.Float32(0.0) + lmx = cutlass.Float32(0.0) - # ---- prod-fix ESCAPE (L1540-1617): 32-step key-space bisection - esc = cutlass.Int32(0) - if whole == cutlass.Int32(0): - if m > cmp_: - esc = cutlass.Int32(1) - if esc == cutlass.Int32(1): + if short == cutlass.Int32(0): + npad = cutlass.Int32(logits.shape[1]) # noqa: F841 + k = cutlass.Int32(pre_idx.shape[1]) + out_row = out[row, None] + x_addr = logits[row, None].iterator.toint() # Int64 gmem byte base + p_addr = pre_idx[prow, None].iterator.toint() # request-level under varlen + + # ---- shared-memory window (map in module docstring) ---- + sptr = cute.arch.get_dyn_smem(cutlass.Int32, alignment=16) + sbase = sptr.toint() # Int32 shared addr + + s_res = _smem_view(cutlass.Int32, sbase, 0, 6) + s_cnt = _smem_view(cutlass.Int32, sbase, 6, 2) # [0]=s_o1 [1]=s_oc + s_kmm = _smem_view(cutlass.Uint32, sbase, 8, 2) # [0]=s_kmin [1]=s_kmax + s_e12 = _smem_view(cutlass.Int32, sbase, 10, 2) # [0]=s_e1 [1]=s_e2 + s_ws = _smem_view(cutlass.Int32, sbase, 16, 32) + s_wmn = _smem_view(cutlass.Uint32, sbase, 48, 32) + s_wmx = _smem_view(cutlass.Uint32, sbase, 80, 32) + s_hist = _smem_view(cutlass.Int32, sbase, STATIC_WORDS, self.nbh) + ck_base = sbase + cutlass.Int32((STATIC_WORDS + self.nbh) * 4) + ck = cute.make_tensor( + cute.make_ptr(cutlass.Uint32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) # typed view, no bound + ci = cute.make_tensor( + cute.make_ptr( + cutlass.Int32, + ck_base + cmp_ * cutlass.Int32(4), + cute.AddressSpace.smem, + assumed_align=4, + ), + cute.make_layout((65536,)), + ) + img_f = cute.make_tensor( # aliases ck/ci (L1299) + cute.make_ptr(cutlass.Float32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) + bm = cute.make_tensor( # aliases ck (L1409) + cute.make_ptr(cutlass.Int32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) + + n4 = n >> cutlass.Int32(2) + ntail = n - (n4 << cutlass.Int32(2)) + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L1351 + + # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any dependent + # gather (L1314-1322); compiled out under DEG. + pvs = [] + if cutlass.const_expr(not self.deg): + for t in cutlass.range_constexpr(KPT): + pv = cutlass.Int32(-1) + j = tid + cutlass.Int32(t * self.blk) + if j < k: + pv = ld_g_i32(p_addr, j) + pvs.append(pv) + + # ---- row load: exact-fit peel + float4[VPT] register batch (L1327-1350) + atom128 = g2r_atom_f32(128, invariant=True) + frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + if n4 >= cutlass.Int32(self.blk * self.vpt): # block-uniform peel + for u in cutlass.range_constexpr(VPT): + ld_g_f32x4(atom128, x_addr, tid + cutlass.Int32(u * self.blk), frags[u]) + else: # predicated flat batch + for u in cutlass.range_constexpr(VPT): + i = tid + cutlass.Int32(u * self.blk) + if i < n4: + ld_g_f32x4(atom128, x_addr, i, frags[u]) + for u in cutlass.range_constexpr(VPT): + i = tid + cutlass.Int32(u * self.blk) + if i >= n4: # -INFINITY fill L1346 + for q in cutlass.range_constexpr(4): + frags[u][q] = cutlass.Float32(_NEG_INF__reg) + + tval = cutlass.Float32(_NEG_INF__reg) + if tid < ntail: + tval = ldg_f32(x_addr, tix) # L1352 + + # ---- init (L1391-1392) if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) s_cnt[1] = cutlass.Int32(0) - # DEVIATION (race fix, see notes): the CUDA zeroes s_o1/s_oc - # again between the nA read (L1581) and the emit (L1584) with - # only ONE barrier pair around both — a read/write race that - # nvcc's schedule happens to win and ptxas' does not (observed - # ~8% row corruption under CTA co-residency). We instead emit - # through the path-exclusive s_e1/s_e2 slots, zeroed HERE under - # the existing L1550 barrier; the racy mid-emit rezero is - # dropped. Barrier count unchanged. - s_e12[0] = cutlass.Int32(0) - s_e12[1] = cutlass.Int32(0) - cute.arch.barrier() # L1550 - klo = cutlass.Uint32(0) - bit = cutlass.Int32(31) - while bit >= cutlass.Int32(0): - kt = klo | (cutlass.Uint32(1) << cutlass.Uint32(bit)) - cnt = cutlass.Int32(0) + for z in cutlass.range_constexpr(self.nbh // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + + # ---- bracket: 4 mutually exclusive compile-time arms (L1393-1451) + lmin = cutlass.Uint32(0xFFFFFFFF) + lmax = cutlass.Uint32(0) + if cutlass.const_expr(self.use_img): + fatom = _f32_smem_atom() + for u in cutlass.range_constexpr(VPT): # VPT == 1 here + i = tid + cutlass.Int32(u * self.blk) + if i < n4: + _sts128_f32(fatom, frags[u], ck_base, i * cutlass.Int32(16)) + if tid < ntail: + img_f[tix] = tval + cute.arch.barrier() # L1400 + for t in cutlass.range_constexpr(KPT): + p = pvs[t] + if cutlass.Uint32(p) < cutlass.Uint32(n): + uk = fkey(img_f[p]) + if uk < lmin: + lmin = uk + if uk > lmax: + lmax = uk + cute.arch.barrier() # L1406 (img dies) + elif cutlass.const_expr(self.use_bm): + nbw = (n + cutlass.Int32(31)) >> cutlass.Int32(5) + i = tid + while i < nbw: # bitmap clear (L1410) + bm[i] = cutlass.Int32(0) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1411 + for t in cutlass.range_constexpr(KPT): + p = pvs[t] + if cutlass.Uint32(p) < cutlass.Uint32(n): + atomic_or_cta( + bm.iterator + (p >> cutlass.Int32(5)), + cutlass.Int32(1) << (p & cutlass.Int32(31)), + ) + cute.arch.barrier() # L1417 + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for u in cutlass.range_constexpr(VPT): + base = (tid + cutlass.Int32(u * self.blk)) << cutlass.Int32(2) + w = cutlass.Uint32(0) + if cutlass.Uint32(base) < cutlass.Uint32(n): + w = cutlass.Uint32(bm[base >> cutlass.Int32(5)]) >> cutlass.Uint32( + base & cutlass.Int32(31) + ) + for cbit in cutlass.range_constexpr(4): + if (w & cutlass.Uint32(1 << cbit)) != cutlass.Uint32(0): + lmn = fmin_f32(lmn, _val(frags, 4 * u + cbit)) + lmx = fmax_f32(lmx, _val(frags, 4 * u + cbit)) + if tid < ntail: + wt = cutlass.Uint32(bm[tix >> cutlass.Int32(5)]) >> cutlass.Uint32( + tix & cutlass.Int32(31) + ) + if (wt & cutlass.Uint32(1)) != cutlass.Uint32(0): + lmn = fmin_f32(lmn, tval) + lmx = fmax_f32(lmx, tval) + lmin = fkey(lmn) + lmax = fkey(lmx) # monotone (L1428) + cute.arch.barrier() # L1429 (bm dies) + elif cutlass.const_expr(self.deg): + lmn = cutlass.Float32(_POS_INF) + lmx = cutlass.Float32(_NEG_INF__reg) + for s in cutlass.range_constexpr(S): # L1436-1439 + v = _val(frags, s) + if v > cutlass.Float32(_NEG_INF__reg): + lmn = fmin_f32(lmn, v) + lmx = fmax_f32(lmx, v) + if tid < ntail: + lmn = fmin_f32(lmn, tval) + lmx = fmax_f32(lmx, tval) + lmin = fkey(lmn) + lmax = fkey(lmx) + else: + # default: KPT scattered fkey ldg gathers, batch-then-fold (L1443-1450) + xs = [] + for t in cutlass.range_constexpr(KPT): + xv = cutlass.Float32(0.0) + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): + xv = ldg_f32(x_addr, pvs[t]) + xs.append(xv) + for t in cutlass.range_constexpr(KPT): + if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): + uk = fkey(xs[t]) + if uk < lmin: + lmin = uk + if uk > lmax: + lmax = uk + + # ---- block min/max in ONE barrier (L1452-1462); publishes hist clear + lmin = warp_min_u32(lmin) + lmax = warp_max_u32(lmax) + if lane == cutlass.Int32(0): + s_wmn[tid >> cutlass.Int32(5)] = lmin + s_wmx[tid >> cutlass.Int32(5)] = lmax + cute.arch.barrier() # L1456 + a = cutlass.Uint32(0xFFFFFFFF) + c = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + a = cutlass.Uint32(s_wmn[lane]) + c = cutlass.Uint32(s_wmx[lane]) + lmin = warp_min_u32(a) + lmax = warp_max_u32(c) + Tv = invkey(lmin) + GMAX = invkey(lmax) + + # ---- collapse guard, NaN-safe (L1464-1471) + okc = cutlass.Int32(0) + if Tv < GMAX: + if (GMAX - Tv) > cutlass.Float32(1e-30): + okc = cutlass.Int32(1) + if okc == cutlass.Int32(0): + Tv = cutlass.Float32(SENT_LO) + GMAX = cutlass.Float32(SENT_HI) + + # ---- bin transform constants (L1485-1511) + BRL = cutlass.const_expr(self.brl) # noqa: F841 + OFFf = cutlass.Float32(1.0 if self.brl else 0.0) + recip = 1.0 / float(self.nbh - (2 if self.brl else 0)) + WD = (GMAX - Tv) * cutlass.Float32(recip) + wsel = cutlass.Float32(1e-30) + if WD > cutlass.Float32(0.0): + wsel = WD + # rcp.approx (single MUFU.RCP) — the CUDA arm's exact lowering of + # `1.0f / wsel`; the previous `1.0 / wsel` spelling emitted the IEEE + # div.rn Newton triple + slowpath CALL on the barrier-bounded chain + # feeding all S classify FMULs. Output exactness is SC-invariant + # (any SC > 0 preserves the sign/monotonicity invariants, L1485-1511) + # and the WD > 0 arm is now bit-identical to CUDA's MUFU.RCP. + SC = cute.arch.rcp_approx(wsel) + QCAPf = cutlass.Float32(float(self.nbh - 1)) + CQ0 = OFFf - Tv * SC + CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf(CQ0) + cutlass.Float32(1.0)) + + # ---- histogram (L1513-1526) + if cutlass.const_expr(self.brl): + # fix-2 P4 (GATED, removable as one hunk): A1 ported to the BRL + # classify arm — hist base pinned ONCE via the same + # _smem_addr_reg__reg identity-mov used in the !BRL arm below, and the + # result-discarded classify atomics spelled as resultless + # red.shared (_red_shared_add1__reg). Value-identical: same +1 to the + # same byte address (hb + 4*bn == &s_hist[bn]), same .relaxed.cta + # ordering; the q/bn computations are untouched so classify/emit + # bit-identity (BRL requirement) is preserved. Emit-path hist + # atomics (results used, L1630+) are NOT touched. + hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) + for s in cutlass.range_constexpr(S): + q = _fmaf__reg(_val(frags, s), SC, CQ) + bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) + _red_shared_add1__reg(hb + (cutlass.Int32(bn) << cutlass.Int32(2))) + qt = _fmaf__reg(tval, SC, CQ) # unconditional (L1517) + bnt = _umin_u32(f2u_rz(qt), cutlass.Uint32(self.nbh - 1)) + _red_shared_add1__reg(hb + (cutlass.Int32(bnt) << cutlass.Int32(2))) + else: + # hist base pinned ONCE (byte addr, +STATIC_BYTES = word 128 map); + # each site below is then LEA + ATOMS exactly like the CUDA arm + # instead of re-deriving the shared window per divergent block. + hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) + for s in cutlass.range_constexpr(S): + q = _submul_asm(_val(frags, s), Tv, SC) # anti-CSE classify + if q >= cutlass.Float32(0.0): + _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(q, QCAPf)) << cutlass.Int32(2))) + qt = _submul_asm(tval, Tv, SC) + if qt >= cutlass.Float32(0.0): + _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(qt, QCAPf)) << cutlass.Int32(2))) + cute.arch.barrier() # L1527 + + # ---- crossing-bin find (L1528-1538) + if cutlass.const_expr(self.cur or self.nbh > 1024): + scan_cross_w(s_hist, s_ws, k, tid, s_res, blk=self.blk, nb=self.nbh) + else: + find_cross(s_hist, k, tid, s_res, nb=self.nbh) + cute.arch.barrier() # L1535 + above = s_res[RES_ABOVE] + m = s_res[RES_M] + Bv = s_res[RES_B] + need = k - above + whole = cutlass.Int32(0) + if need >= m: + whole = cutlass.Int32(1) + + # ---- prod-fix ESCAPE (L1540-1617): 32-step key-space bisection + esc = cutlass.Int32(0) + if whole == cutlass.Int32(0): + if m > cmp_: + esc = cutlass.Int32(1) + if esc == cutlass.Int32(1): + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + s_cnt[1] = cutlass.Int32(0) + # DEVIATION (race fix, see notes): the CUDA zeroes s_o1/s_oc + # again between the nA read (L1581) and the emit (L1584) with + # only ONE barrier pair around both — a read/write race that + # nvcc's schedule happens to win and ptxas' does not (observed + # ~8% row corruption under CTA co-residency). We instead emit + # through the path-exclusive s_e1/s_e2 slots, zeroed HERE under + # the existing L1550 barrier; the racy mid-emit rezero is + # dropped. Barrier count unchanged. + s_e12[0] = cutlass.Int32(0) + s_e12[1] = cutlass.Int32(0) + cute.arch.barrier() # L1550 + klo = cutlass.Uint32(0) + bit = cutlass.Int32(31) + while bit >= cutlass.Int32(0): + kt = klo | (cutlass.Uint32(1) << cutlass.Uint32(bit)) + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(S): + ix = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if ix < n: + if fkey(_val(frags, s)) >= kt: + cnt = cnt + cutlass.Int32(1) + if tid < ntail: + if fkey(tval) >= kt: + cnt = cnt + cutlass.Int32(1) + cnt = cutlass.Int32(warp_add_i32(cnt)) + if lane == cutlass.Int32(0): + if cnt != cutlass.Int32(0): + atomic_add_cta(s_cnt.iterator, cnt) + cute.arch.barrier() # L1563 + if s_cnt[0] >= k: + klo = kt + cute.arch.barrier() # L1565 + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + cute.arch.barrier() # L1567 + bit = bit - cutlass.Int32(1) + ethr = cutlass.Int64(klo) # k-th largest key + abv = cutlass.Int32(0) for s in cutlass.range_constexpr(S): ix = ( (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) ) + cutlass.Int32(s % 4) if ix < n: - if fkey(_val(frags, s)) >= kt: - cnt = cnt + cutlass.Int32(1) + if cutlass.Int64(fkey(_val(frags, s))) > ethr: + abv = abv + cutlass.Int32(1) if tid < ntail: - if fkey(tval) >= kt: - cnt = cnt + cutlass.Int32(1) - cnt = cutlass.Int32(warp_add_i32(cnt)) - if lane == cutlass.Int32(0): - if cnt != cutlass.Int32(0): - atomic_add_cta(s_cnt.iterator, cnt) - cute.arch.barrier() # L1563 - if s_cnt[0] >= k: - klo = kt - cute.arch.barrier() # L1565 - if tid == cutlass.Int32(0): - s_cnt[0] = cutlass.Int32(0) - cute.arch.barrier() # L1567 - bit = bit - cutlass.Int32(1) - ethr = cutlass.Int64(klo) # k-th largest key - abv = cutlass.Int32(0) - for s in cutlass.range_constexpr(S): - ix = ( - (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) - ) + cutlass.Int32(s % 4) - if ix < n: - if cutlass.Int64(fkey(_val(frags, s))) > ethr: + if cutlass.Int64(fkey(tval)) > ethr: abv = abv + cutlass.Int32(1) - if tid < ntail: - if cutlass.Int64(fkey(tval)) > ethr: - abv = abv + cutlass.Int32(1) - abv = cutlass.Int32(warp_add_i32(abv)) - if lane == cutlass.Int32(0): - if abv != cutlass.Int32(0): - atomic_add_cta(s_cnt.iterator + 1, abv) - cute.arch.barrier() # L1580 - nA = s_cnt[1] - nT = k - nA - # (rezero dropped — emit counters live in s_e12, see race-fix note) - cute.arch.barrier() # L1583 - lml = cutlass.Int32(cute.arch.lanemask_lt()) - for s in cutlass.range_constexpr(S): - ixv = ( - (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) - ) + cutlass.Int32(s % 4) + abv = cutlass.Int32(warp_add_i32(abv)) + if lane == cutlass.Int32(0): + if abv != cutlass.Int32(0): + atomic_add_cta(s_cnt.iterator + 1, abv) + cute.arch.barrier() # L1580 + nA = s_cnt[1] + nT = k - nA + # (rezero dropped — emit counters live in s_e12, see race-fix note) + cute.arch.barrier() # L1583 + lml = cutlass.Int32(cute.arch.lanemask_lt()) + for s in cutlass.range_constexpr(S): + ixv = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + u64 = cutlass.Int64(-1) + if ixv < n: + u64 = cutlass.Int64(fkey(_val(frags, s))) + q1e = cutlass.Int32(0) + q2e = cutlass.Int32(0) + if u64 > ethr: + q1e = cutlass.Int32(1) + if u64 == ethr: + q2e = cutlass.Int32(1) + n1 = ballot(q1e == cutlass.Int32(1)) + n2 = ballot(q2e == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_e12.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1e == cutlass.Int32(1): + if p1e < nA: + out_row[p1e] = ixv + if q2e == cutlass.Int32(1): + if p2e < nT: + out_row[nA + p2e] = ixv + # tail element (L1601-1615) u64 = cutlass.Int64(-1) - if ixv < n: - u64 = cutlass.Int64(fkey(_val(frags, s))) + if tid < ntail: + u64 = cutlass.Int64(fkey(tval)) q1e = cutlass.Int32(0) q2e = cutlass.Int32(0) if u64 > ethr: @@ -3744,313 +3843,284 @@ def kern( p2e = b2 + popc(n2 & lml) if q1e == cutlass.Int32(1): if p1e < nA: - out_row[p1e] = ixv + out_row[p1e] = tix if q2e == cutlass.Int32(1): if p2e < nT: - out_row[nA + p2e] = ixv - # tail element (L1601-1615) - u64 = cutlass.Int64(-1) - if tid < ntail: - u64 = cutlass.Int64(fkey(tval)) - q1e = cutlass.Int32(0) - q2e = cutlass.Int32(0) - if u64 > ethr: - q1e = cutlass.Int32(1) - if u64 == ethr: - q2e = cutlass.Int32(1) - n1 = ballot(q1e == cutlass.Int32(1)) - n2 = ballot(q2e == cutlass.Int32(1)) - b1 = cutlass.Int32(0) - b2 = cutlass.Int32(0) - if lane == cutlass.Int32(0): - if n1 != cutlass.Int32(0): - b1 = atomic_add_cta(s_e12.iterator, popc(n1)) - if n2 != cutlass.Int32(0): - b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) - b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) - b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) - p1e = b1 + popc(n1 & lml) - p2e = b2 + popc(n2 & lml) - if q1e == cutlass.Int32(1): - if p1e < nA: - out_row[p1e] = tix - if q2e == cutlass.Int32(1): - if p2e < nT: - out_row[nA + p2e] = tix - # (CUDA returns here — everything below is the else-arm) - else: - # ---- emit (L1619-1701) - if cutlass.const_expr(self.cur): - LOQ = cutlass.Float32(Bv) # int->float cvt (L1626) - lim1 = above - if whole == cutlass.Int32(1): - lim1 = above + m - for s in cutlass.range_constexpr(S): + out_row[nA + p2e] = tix + # (CUDA returns here — everything below is the else-arm) + else: + # ---- emit (L1619-1701) + if cutlass.const_expr(self.cur): + LOQ = cutlass.Float32(Bv) # int->float cvt (L1626) + lim1 = above + if whole == cutlass.Int32(1): + lim1 = above + m + for s in cutlass.range_constexpr(S): + if cutlass.const_expr(self.brl): + q = _fmaf__reg(_val(frags, s), SC, CQ) # bit-identical to classify + else: + q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) # L1630 emit spelling + idx = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + p = cutlass.Int32(0) + if q >= LOQ: + bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) + p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + if p < lim1: + out_row[p] = idx + else: + if whole == cutlass.Int32(0): + q2i = p - above + if q2i < cmp_: # escape-made-safe guard + ck[q2i] = fkey(_val(frags, s)) + ci[q2i] = idx + # tail (L1640-1647) if cutlass.const_expr(self.brl): - q = _fmaf__reg(_val(frags, s), SC, CQ) # bit-identical to classify + qt2 = _fmaf__reg(tval, SC, CQ) else: - q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) # L1630 emit spelling - idx = ( - (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) - ) + cutlass.Int32(s % 4) + qt2 = _fmaf__reg(tval - Tv, SC, OFFf) p = cutlass.Int32(0) - if q >= LOQ: - bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) + if qt2 >= LOQ: + bn = _umin_u32(f2u_rz(qt2), cutlass.Uint32(self.nbh - 1)) p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) if p < lim1: - out_row[p] = idx + out_row[p] = tix else: if whole == cutlass.Int32(0): q2i = p - above - if q2i < cmp_: # escape-made-safe guard - ck[q2i] = fkey(_val(frags, s)) - ci[q2i] = idx - # tail (L1640-1647) - if cutlass.const_expr(self.brl): - qt2 = _fmaf__reg(tval, SC, CQ) + if q2i < cmp_: + ck[q2i] = fkey(tval) + ci[q2i] = tix else: - qt2 = _fmaf__reg(tval - Tv, SC, OFFf) - p = cutlass.Int32(0) - if qt2 >= LOQ: - bn = _umin_u32(f2u_rz(qt2), cutlass.Uint32(self.nbh - 1)) - p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - if p < lim1: - out_row[p] = tix + # two-mask ballot emit (L1650-1701) + HIf = cutlass.Float32(_POS_INF) + LOf = cutlass.Float32(_POS_INF) + if whole == cutlass.Int32(1): + HIf = cutlass.Float32(Bv) else: - if whole == cutlass.Int32(0): - q2i = p - above - if q2i < cmp_: - ck[q2i] = fkey(tval) - ci[q2i] = tix - else: - # two-mask ballot emit (L1650-1701) - HIf = cutlass.Float32(_POS_INF) - LOf = cutlass.Float32(_POS_INF) - if whole == cutlass.Int32(1): - HIf = cutlass.Float32(Bv) - else: - if Bv < cutlass.Int32(self.nbh - 1): - HIf = cutlass.Float32(Bv + cutlass.Int32(1)) - LOf = cutlass.Float32(Bv) - m1 = cutlass.Int32(0) - m2 = cutlass.Int32(0) - for s in cutlass.range_constexpr(S): + if Bv < cutlass.Int32(self.nbh - 1): + HIf = cutlass.Float32(Bv + cutlass.Int32(1)) + LOf = cutlass.Float32(Bv) + m1 = cutlass.Int32(0) + m2 = cutlass.Int32(0) + for s in cutlass.range_constexpr(S): + if cutlass.const_expr(self.brl): + q = _fmaf__reg(_val(frags, s), SC, CQ) + else: + q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) + if q >= HIf: + m1 = m1 | cutlass.Int32(1 << s) + else: + if q >= LOf: + m2 = m2 | cutlass.Int32(1 << s) if cutlass.const_expr(self.brl): - q = _fmaf__reg(_val(frags, s), SC, CQ) + qt3 = _fmaf__reg(tval, SC, CQ) else: - q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) - if q >= HIf: - m1 = m1 | cutlass.Int32(1 << s) + qt3 = _fmaf__reg(tval - Tv, SC, OFFf) + t1 = cutlass.Int32(0) + t2 = cutlass.Int32(0) + if qt3 >= HIf: + t1 = cutlass.Int32(1) else: - if q >= LOf: - m2 = m2 | cutlass.Int32(1 << s) - if cutlass.const_expr(self.brl): - qt3 = _fmaf__reg(tval, SC, CQ) - else: - qt3 = _fmaf__reg(tval - Tv, SC, OFFf) - t1 = cutlass.Int32(0) - t2 = cutlass.Int32(0) - if qt3 >= HIf: - t1 = cutlass.Int32(1) - else: - if qt3 >= LOf: - t2 = cutlass.Int32(1) - c1 = popc(m1) + t1 - c2 = popc(m2) + t2 - s1, s2 = warp_incl_scan_add2(c1, c2, lane) # L1669-1673 - b1 = cutlass.Int32(0) - b2 = cutlass.Int32(0) - if lane == cutlass.Int32(31): - b1 = atomic_add_cta(s_cnt.iterator, s1) - b2 = atomic_add_cta(s_cnt.iterator + 1, s2) - b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(31)) - b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(31)) - p1 = b1 + (s1 - c1) - p2 = b2 + (s2 - c2) - lim1 = above - if whole == cutlass.Int32(1): - lim1 = k - wm = m1 # sparse set-bit walk - while wm != cutlass.Int32(0): - sdyn = ffs_m1(wm) - idx = ( - (tid + (sdyn >> cutlass.Int32(2)) * cutlass.Int32(self.blk)) - << cutlass.Int32(2) - ) + (sdyn & cutlass.Int32(3)) - if p1 < lim1: - out_row[p1] = idx - p1 = p1 + cutlass.Int32(1) - wm = wm & (wm - cutlass.Int32(1)) - if t1 == cutlass.Int32(1): - if p1 < lim1: - out_row[p1] = tix - p1 = p1 + cutlass.Int32(1) - if m2 != cutlass.Int32(0): # static-unrolled (L1689) - for s in cutlass.range_constexpr(S): - if (m2 & cutlass.Int32(1 << s)) != cutlass.Int32(0): - idx = ( - (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) - ) + cutlass.Int32(s % 4) - if p2 < cmp_: - ck[p2] = fkey(_val(frags, s)) - ci[p2] = idx - p2 = p2 + cutlass.Int32(1) - if t2 == cutlass.Int32(1): - if p2 < cmp_: - ck[p2] = fkey(tval) - ci[p2] = tix - p2 = p2 + cutlass.Int32(1) - - # ---- refine (skipped when whole — CUDA returned inside emit) - if whole == cutlass.Int32(0): - cute.arch.barrier() # L1703 - if cutlass.const_expr(self.cur): - mc = m - if mc > cmp_: - mc = cmp_ - else: - mc = s_cnt[1] - if mc > cmp_: - mc = cmp_ - quad = cutlass.Int32(0) - if mc >= m: - if mc <= qc: - quad = cutlass.Int32(1) - if quad == cutlass.Int32(1): - # O(mc^2) index-tie-broken rank (L1706-1718) - i = tid - while i < mc: - uq = cutlass.Uint32(ck[i]) - r = cutlass.Int32(0) - j = cutlass.Int32(0) - while j < mc: - vq = cutlass.Uint32(ck[j]) - tinc = cutlass.Int32(0) - if vq > uq: - tinc = cutlass.Int32(1) - if vq == uq: - if j < i: + if qt3 >= LOf: + t2 = cutlass.Int32(1) + c1 = popc(m1) + t1 + c2 = popc(m2) + t2 + s1, s2 = warp_incl_scan_add2(c1, c2, lane) # L1669-1673 + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(31): + b1 = atomic_add_cta(s_cnt.iterator, s1) + b2 = atomic_add_cta(s_cnt.iterator + 1, s2) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(31)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(31)) + p1 = b1 + (s1 - c1) + p2 = b2 + (s2 - c2) + lim1 = above + if whole == cutlass.Int32(1): + lim1 = k + wm = m1 # sparse set-bit walk + while wm != cutlass.Int32(0): + sdyn = ffs_m1(wm) + idx = ( + (tid + (sdyn >> cutlass.Int32(2)) * cutlass.Int32(self.blk)) + << cutlass.Int32(2) + ) + (sdyn & cutlass.Int32(3)) + if p1 < lim1: + out_row[p1] = idx + p1 = p1 + cutlass.Int32(1) + wm = wm & (wm - cutlass.Int32(1)) + if t1 == cutlass.Int32(1): + if p1 < lim1: + out_row[p1] = tix + p1 = p1 + cutlass.Int32(1) + if m2 != cutlass.Int32(0): # static-unrolled (L1689) + for s in cutlass.range_constexpr(S): + if (m2 & cutlass.Int32(1 << s)) != cutlass.Int32(0): + idx = ( + (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) + ) + cutlass.Int32(s % 4) + if p2 < cmp_: + ck[p2] = fkey(_val(frags, s)) + ci[p2] = idx + p2 = p2 + cutlass.Int32(1) + if t2 == cutlass.Int32(1): + if p2 < cmp_: + ck[p2] = fkey(tval) + ci[p2] = tix + p2 = p2 + cutlass.Int32(1) + + # ---- refine (skipped when whole — CUDA returned inside emit) + if whole == cutlass.Int32(0): + cute.arch.barrier() # L1703 + if cutlass.const_expr(self.cur): + mc = m + if mc > cmp_: + mc = cmp_ + else: + mc = s_cnt[1] + if mc > cmp_: + mc = cmp_ + quad = cutlass.Int32(0) + if mc >= m: + if mc <= qc: + quad = cutlass.Int32(1) + if quad == cutlass.Int32(1): + # O(mc^2) index-tie-broken rank (L1706-1718) + i = tid + while i < mc: + uq = cutlass.Uint32(ck[i]) + r = cutlass.Int32(0) + j = cutlass.Int32(0) + while j < mc: + vq = cutlass.Uint32(ck[j]) + tinc = cutlass.Int32(0) + if vq > uq: tinc = cutlass.Int32(1) - r = r + tinc - j = j + cutlass.Int32(1) - if r < need: - out_row[above + r] = ci[i] - i = i + cutlass.Int32(BLK) - else: - # ---- fallback: exact key-space narrowing (L1720-1775) - if tid == cutlass.Int32(0): - s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) - s_kmm[1] = cutlass.Uint32(0) - cute.arch.barrier() # L1724 - i = tid - while i < mc: - kv = cutlass.Uint32(ck[i]) - atomic_min_cta(s_kmm.iterator, kv) - atomic_max_cta(s_kmm.iterator + 1, kv) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1726 - rlo = cutlass.Uint32(s_kmm[0]) - rhi = cutlass.Uint32(s_kmm[1]) - ethr = cutlass.Int64(rlo) - aboveC = cutlass.Int32(0) - needC = need - mm = mc - lev = cutlass.Int32(0) - done = cutlass.Int32(0) - while done == cutlass.Int32(0): - if needC == mm: # L1730 - ethr = cutlass.Int64(rlo) - cutlass.Int64(1) - aboveC = aboveC + mm - needC = cutlass.Int32(0) - done = cutlass.Int32(1) - if done == cutlass.Int32(0): - if rlo >= rhi: # L1731 - ethr = cutlass.Int64(rlo) - done = cutlass.Int32(1) - if lev >= cutlass.Int32(6): # L1732 - ethr = cutlass.Int64(rlo) + if vq == uq: + if j < i: + tinc = cutlass.Int32(1) + r = r + tinc + j = j + cutlass.Int32(1) + if r < need: + out_row[above + r] = ci[i] + i = i + cutlass.Int32(BLK) + else: + # ---- fallback: exact key-space narrowing (L1720-1775) + if tid == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + cute.arch.barrier() # L1724 + i = tid + while i < mc: + kv = cutlass.Uint32(ck[i]) + atomic_min_cta(s_kmm.iterator, kv) + atomic_max_cta(s_kmm.iterator + 1, kv) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1726 + rlo = cutlass.Uint32(s_kmm[0]) + rhi = cutlass.Uint32(s_kmm[1]) + ethr = cutlass.Int64(rlo) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + lev = cutlass.Int32(0) + done = cutlass.Int32(0) + while done == cutlass.Int32(0): + if needC == mm: # L1730 + ethr = cutlass.Int64(rlo) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) done = cutlass.Int32(1) - if done == cutlass.Int32(0): - d2 = rhi - rlo - b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) - sh2 = cutlass.Int32(0) - if b2w > cutlass.Int32(LNBH): - sh2 = b2w - cutlass.Int32(LNBH) - for z in cutlass.range_constexpr(self.nbh // self.blk): - s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L1737 - i = tid - while i < mc: - unar = cutlass.Uint32(ck[i]) - if unar >= rlo: - if unar <= rhi: - bnn = (unar - rlo) >> cutlass.Uint32(sh2) - bnn = _umin_u32(bnn, cutlass.Uint32(self.nbh - 1)) - atomic_add_cta( - s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) - ) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1743 - if cutlass.const_expr(self.nbh > 1024): - scan_cross_w( - s_hist, s_ws, needC, tid, s_res, blk=self.blk, nb=self.nbh - ) - else: - find_cross(s_hist, needC, tid, s_res, nb=self.nbh) - cute.arch.barrier() # L1746 - aboveC = aboveC + s_res[RES_ABOVE] - needC = needC - s_res[RES_ABOVE] - mm = s_res[RES_M] - b_lv = s_res[RES_B] - nlo = rlo + (cutlass.Uint32(b_lv) << cutlass.Uint32(sh2)) - if b_lv != cutlass.Int32(self.nbh - 1): - rhi = nlo + ( - (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) - ) - rlo = nlo - lev = lev + cutlass.Int32(1) - # final two-predicate ballot emit (L1752-1775) - if tid == cutlass.Int32(0): - s_e12[0] = cutlass.Int32(0) - s_e12[1] = cutlass.Int32(0) - cute.arch.barrier() # L1753 - lml = cutlass.Int32(cute.arch.lanemask_lt()) - it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) - it = cutlass.Int32(0) - while it < it2: - i = it * cutlass.Int32(BLK) + tid - uke = cutlass.Uint32(0) - idv = cutlass.Int32(0) - if i < mc: - uke = cutlass.Uint32(ck[i]) - idv = ci[i] - q1f = cutlass.Int32(0) - q2f = cutlass.Int32(0) - if i < mc: - if cutlass.Int64(uke) > ethr: - q1f = cutlass.Int32(1) - if cutlass.Int64(uke) == ethr: - q2f = cutlass.Int32(1) - n1 = ballot(q1f == cutlass.Int32(1)) - n2 = ballot(q2f == cutlass.Int32(1)) - b1 = cutlass.Int32(0) - b2 = cutlass.Int32(0) - if lane == cutlass.Int32(0): - if n1 != cutlass.Int32(0): - b1 = atomic_add_cta(s_e12.iterator, popc(n1)) - if n2 != cutlass.Int32(0): - b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) - b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) - b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) - p1e = b1 + popc(n1 & lml) - p2e = b2 + popc(n2 & lml) - if q1f == cutlass.Int32(1): - if p1e < aboveC: - out_row[above + p1e] = idv - if q2f == cutlass.Int32(1): - if p2e < needC: - out_row[above + aboveC + p2e] = idv - it = it + cutlass.Int32(1) + if done == cutlass.Int32(0): + if rlo >= rhi: # L1731 + ethr = cutlass.Int64(rlo) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(6): # L1732 + ethr = cutlass.Int64(rlo) + done = cutlass.Int32(1) + if done == cutlass.Int32(0): + d2 = rhi - rlo + b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + sh2 = cutlass.Int32(0) + if b2w > cutlass.Int32(LNBH): + sh2 = b2w - cutlass.Int32(LNBH) + for z in cutlass.range_constexpr(self.nbh // self.blk): + s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) + cute.arch.barrier() # L1737 + i = tid + while i < mc: + unar = cutlass.Uint32(ck[i]) + if unar >= rlo: + if unar <= rhi: + bnn = (unar - rlo) >> cutlass.Uint32(sh2) + bnn = _umin_u32(bnn, cutlass.Uint32(self.nbh - 1)) + atomic_add_cta( + s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # L1743 + if cutlass.const_expr(self.nbh > 1024): + scan_cross_w( + s_hist, s_ws, needC, tid, s_res, blk=self.blk, nb=self.nbh + ) + else: + find_cross(s_hist, needC, tid, s_res, nb=self.nbh) + cute.arch.barrier() # L1746 + aboveC = aboveC + s_res[RES_ABOVE] + needC = needC - s_res[RES_ABOVE] + mm = s_res[RES_M] + b_lv = s_res[RES_B] + nlo = rlo + (cutlass.Uint32(b_lv) << cutlass.Uint32(sh2)) + if b_lv != cutlass.Int32(self.nbh - 1): + rhi = nlo + ( + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + # final two-predicate ballot emit (L1752-1775) + if tid == cutlass.Int32(0): + s_e12[0] = cutlass.Int32(0) + s_e12[1] = cutlass.Int32(0) + cute.arch.barrier() # L1753 + lml = cutlass.Int32(cute.arch.lanemask_lt()) + it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) + it = cutlass.Int32(0) + while it < it2: + i = it * cutlass.Int32(BLK) + tid + uke = cutlass.Uint32(0) + idv = cutlass.Int32(0) + if i < mc: + uke = cutlass.Uint32(ck[i]) + idv = ci[i] + q1f = cutlass.Int32(0) + q2f = cutlass.Int32(0) + if i < mc: + if cutlass.Int64(uke) > ethr: + q1f = cutlass.Int32(1) + if cutlass.Int64(uke) == ethr: + q2f = cutlass.Int32(1) + n1 = ballot(q1f == cutlass.Int32(1)) + n2 = ballot(q2f == cutlass.Int32(1)) + b1 = cutlass.Int32(0) + b2 = cutlass.Int32(0) + if lane == cutlass.Int32(0): + if n1 != cutlass.Int32(0): + b1 = atomic_add_cta(s_e12.iterator, popc(n1)) + if n2 != cutlass.Int32(0): + b2 = atomic_add_cta(s_e12.iterator + 1, popc(n2)) + b1 = cute.arch.shuffle_sync(b1, cutlass.Int32(0)) + b2 = cute.arch.shuffle_sync(b2, cutlass.Int32(0)) + p1e = b1 + popc(n1 & lml) + p2e = b2 + popc(n2 & lml) + if q1f == cutlass.Int32(1): + if p1e < aboveC: + out_row[above + p1e] = idv + if q2f == cutlass.Int32(1): + if p2e < needC: + out_row[above + aboveC + p2e] = idv + it = it + cutlass.Int32(1) # ------------------------------------------------------------------ @cute.jit @@ -4058,6 +4128,7 @@ def __call__( self, logits: cute.Tensor, pre_idx: cute.Tensor, + kv_lens: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, cmp_: cutlass.Int32, @@ -4066,7 +4137,7 @@ def __call__( stream, ): b = logits.shape[0] - self.kern(logits, pre_idx, out, n, cmp_, qc, smem_bytes).launch( + self.kern(logits, pre_idx, kv_lens, out, n, cmp_, qc, smem_bytes).launch( grid=(b, 1, 1), block=(self.blk, 1, 1), stream=stream, @@ -4082,16 +4153,29 @@ def __call__( _COMPILE_CACHE__reg: dict = {} -def get_compiled__reg(tpl, dump_dir=None, pdl=False): +def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0): """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, MINB, KPT, CUR, DEG, IMG, NBH).""" - key = (tuple(tpl), bool(pdl)) + key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift)) compiled = _COMPILE_CACHE__reg.get(key) if compiled is None: from cutlass.cute import runtime as _crt blk, vpt, minb, kpt, cur, deg, img, nbh = tpl - kernel = GvrTopkRegKernel(blk, vpt, minb, kpt, cur, deg, img, nbh, pdl=pdl) + kernel = GvrTopkRegKernel( + blk, + vpt, + minb, + kpt, + cur, + deg, + img, + nbh, + pdl=pdl, + varlen=varlen, + next_n=next_n, + cr_shift=cr_shift, + ) nb_, nc_ = cute.sym_int(), cute.sym_int() nb2_, nc2_ = cute.sym_int(), cute.sym_int() nb3_, nc3_ = cute.sym_int(), cute.sym_int() @@ -4104,6 +4188,10 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False): out_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (nb3_, nc3_), stride_order=(1, 0), assumed_align=16 ) + v0_ = cute.sym_int() + kv_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (v0_,), stride_order=(0,), assumed_align=4 + ) fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) opts = "--enable-tvm-ffi" if dump_dir: @@ -4113,6 +4201,7 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False): kernel, lg_fake, pi_fake, + kv_fake, out_fake, cutlass.Int32(0), cutlass.Int32(0), @@ -4143,7 +4232,12 @@ def reg_topk(logits, pre_idx, n, out, rd=None): assert rt["IMGOFF"] == tpl[7], (rt["IMGOFF"], tpl[7]) # IMGOFF == NBH compiled = get_compiled__reg(tpl) smem = STATIC_BYTES + rd["smem"] - compiled(logits, pre_idx, out, int(n), rt["CMP"], rt["QC"], smem) + try: + from .gvr_topk_decode_self_sampling_host import _dummy_kv + except ImportError: + from gvr_topk_decode_self_sampling_host import _dummy_kv + kv = _dummy_kv(logits.get_device(), logits.device) # dead varlen ABI slot + compiled(logits, pre_idx, kv, out, int(n), rt["CMP"], rt["QC"], smem) return out diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 47d79dea8b04..a02c015ac4f0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -757,6 +757,25 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): lc = ("reg_clus", fn, n_eff) _VARLEN_CACHE[key] = lc return lc + # ---- route() parity, family tier 2: register-resident (+img flavor) ---- + # Same admission rule as tier 1: exactly where the free route picks + # reg/regimg (the whole small/mid-N band across all row counts). CMP/QC/ + # smem are envelope-derived launch constants -- in-kernel they are pure + # capacity clamps (CMP), a fast-path threshold (QC) and the launch smem + # size, all safe upper bounds for every per-row n <= envelope; per-row n + # / short-row handling lives in-kernel (GvrRegClusKernel discipline). + if plan_free["kernel"] in ("reg", "regimg"): + fn = dev.get_compiled__reg( + tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift + ) + rt_f = plan_free["rt"] + lc = ( + "reg", + fn, + (rt_f["n"], rt_f["CMP"], rt_f["QC"], dev.STATIC_BYTES + plan_free["smem"]), + ) + _VARLEN_CACHE[key] = lc + return lc plan = route_streaming(num_rows, n_eff, npad, k, force_main=True) tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) rt = plan["rt"] @@ -1021,8 +1040,14 @@ def _build_launcher(b, n, npad, k): rt = rd["rt"] if fam in ("reg", "regimg"): dev = _device() - fn = dev.get_compiled__reg(tpl) - # compiled ABI: (logits, pre_idx, out, n, CMP, QC, smem_total) + raw = dev.get_compiled__reg(tpl) + + # compiled ABI: (logits, pre_idx, kv_lens, out, n, CMP, QC, + # smem_total) -- kv_lens is the dead varlen slot in batch-uniform + # mode (dummy, gvr_main/reg_clus precedent) + def fn(lg, pi, o, *a, _raw=raw): + _raw(lg, pi, _dummy_kv(lg.get_device(), lg.device), o, *a) + args = (rt["n"], rt["CMP"], rt["QC"], dev.STATIC_BYTES + rd["smem"]) return (fn, args, False) if fam == "main": @@ -1424,6 +1449,9 @@ def run_varlen( if lc[0] == "reg_clus": # compiled ABI: (logits, pre_idx, kv_lens, out, n_envelope) lc[1](lg, pre_idx, kv_lens, idx, lc[2]) + elif lc[0] == "reg": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) + lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) else: _, fn, pre, tail = lc fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) @@ -1543,6 +1571,8 @@ def warmup_varlen( plan_free = route(r, max(min(n_env_c, npad_c), int(top_k) + 1), npad_c, int(top_k)) if plan_free["kernel"] == "reg_clus": ekey = ("reg_clus", tuple(plan_free["tpl"])) + elif plan_free["kernel"] in ("reg", "regimg"): + ekey = ("reg", tuple(plan_free["tpl"])) else: p = route_streaming( r, diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 514e34b4205d..b380fac507ba 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -727,6 +727,87 @@ def test_selfsampling_varlen_regclus_cuda_graph(): assert torch.equal(got, ref_v) +def test_selfsampling_varlen_reg_parity_and_oracle(): + """The varlen launcher must admit the register-resident family (and its + img flavor) exactly where the free route picks it (route() parity tier 2), + and the per-row varlen port must match the reference oracle on a + heterogeneous batch: long rows, a short row (n <= k; k can exceed BLK on + this family, exercising the strided identity + -1 tail loop) and a + zero-window row, under MTP row windows (next_n=4).""" + torch.manual_seed(13) + # (k, msl_c, cr, expected free-route family) + cases = [ + (512, 6144, 4, "reg"), # v4-style small-N band + (2048, 8192, 1, "reg"), # k > BLK: strided short-row emit + (512, 3072, 4, "regimg"), # img window (n4 in (512, 1024]) + ] + nn = 4 + batch = 3 + rows = batch * nn + for k, msl_c, cr, want in cases: + npad = (msl_c + 63) // 64 * 64 + fam = ss_host.route(rows, msl_c, npad, k)["kernel"] + assert fam == want, (fam, want, k, msl_c) + msl = msl_c * cr + lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) + kv = torch.tensor( + [msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV + ) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) + key = (rows, npad, k, msl_c, nn, cr) + assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] + ss_host.run_varlen( + lg, + pre, + kv, + ref, + next_n=nn, + compress_ratio=cr, + max_seq_len=msl, + engine="reference", + ) + torch.cuda.synchronize() + for r in range(rows): + if (ref[r] >= 0).any(): + row = lg[r].float() + got = row[out[r].long().clamp_min(0)].sort().values + want_v = row[ref[r].long().clamp_min(0)].sort().values + assert torch.equal(got, want_v), f"k={k} row {r} value multiset mismatch" + assert torch.equal(out[r] < 0, ref[r] < 0), f"k={k} row {r} pad mask mismatch" + else: + assert torch.equal(out[r], ref[r]), f"k={k} row {r} expected all -1" + + +def test_selfsampling_varlen_reg_cuda_graph(): + """Register-family varlen engine must be CUDA-graph capturable: warmed + engine, capture one launch, replay twice, tie-aware exact each time.""" + k, msl_c, cr = 512, 4096, 4 + rows = 16 + torch.manual_seed(17) + lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) + kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + torch.cuda.synchronize() + key = (rows, msl_c, k, msl_c, 1, cr) + assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] + g = torch.cuda.CUDAGraph() + out.fill_(-7) + with torch.cuda.graph(g): + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values + for _ in range(2): + out.fill_(-7) + g.replay() + torch.cuda.synchronize() + got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(got, ref_v) + + def test_selfsampling_varlen_full_row_range(): """Full-range production contract: with self-sampling enabled, EVERY row count dispatches to the self-sampling engines (no rows-based fall-through) From 6d1354a13d1d65229f26e5fabdfefc09eb35e39d Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:57:56 +0000 Subject: [PATCH 29/40] [TRTLLM-15293][perf] self-sampling top-K: MUFU.RCP spelling for the remaining clus/reg_clus reciprocals Convert the last three IEEE-div reciprocal sites (clus sample-hist scale, clus window scale, reg_clus bin-transform scale) to cute.arch.rcp_approx, completing the fix-4 CUDA --use_fast_math parity (bare MUFU.RCP instead of the div.rn rcp+Newton+CALL chain). All three feed classify bucketing only, which is scale-invariant for any SC > 0; inputs are clamped positive finite by the existing guards. UT 82/82; paired same-GPU A/B on clus/reg_clus cells is noise-level (<=1.3%), no regression (the fix-1 spill hazard that kept main's blk!=512 arm on div.rn does not reproduce in these families). The cute.math.div(..., fastmath=True) spelling (per DSL-team guidance) is not available on cutlass 4.5.0 (this PR's validated stack); adopt it when the tree moves to DSL >= 4.6. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index b5b5aaa21bb9..b2bc70cc66cb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -4694,7 +4694,10 @@ def kern( sok = cutlass.Int32(1) if sok != cutlass.Int32(0): # L1932-1947 sample hist w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) - sc_s = cutlass.Float32(1.0) / w + # CUDA --use_fast_math lowers `1.0f / w` to a bare MUFU.RCP; the + # plain `/` spelling emits the IEEE div.rn rcp+Newton+CALL chain + # (fix-4). w > 0 by the sok guard; bucketing is SC-invariant. + sc_s = cute.arch.rcp_approx(w) if shas != cutlass.Int32(0): for t in cutlass.range_constexpr(4): bq = C.f2s_rz((fsa[t] - SMIN) * sc_s) @@ -4838,7 +4841,9 @@ def kern( wdok = cutlass.Int32(1) if wdok == cutlass.Int32(0): WD = cutlass.Float32(1e-30) - SC = cutlass.Float32(1.0) / WD + # MUFU.RCP spelling (fix-4): WD >= 1e-30 finite by the wdok + # clamp; classify bucketing is SC-invariant for any SC > 0. + SC = cute.arch.rcp_approx(WD) # ---- P3 row pass over OWNED CHUNKS (L2033-2121) ---- g = rank + cutlass.Int32(0) @@ -5796,7 +5801,9 @@ def kern( wsel = cutlass.Float32(1e-30) if WD > cutlass.Float32(0.0): wsel = WD - SC = cutlass.Float32(1.0) / wsel + # MUFU.RCP spelling (fix-4, mirrors the reg family's L3668 site): + # wsel >= 1e-30 finite; bucketing is SC-invariant for any SC > 0. + SC = cute.arch.rcp_approx(wsel) CQ0 = cutlass.Float32(1.0) - Tv * SC CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf__regclus(CQ0) + cutlass.Float32(1.0)) From 8244a62f17569a0d1e91cd3f3178611478bfd1c7 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:32:03 +0000 Subject: [PATCH 30/40] [TRTLLM-15293][perf] self-sampling top-K: per-row varlen port of the cluster-split family (clus) Route-parity tier 3, completing the family coverage: the varlen production entry now admits every family exactly where the free route picks it (main / reg_clus / reg / regimg / clus) -- no force-main fallback band remains. GvrClusKernel gains the per-row varlen mode (GvrMainKernel / GvrRegClusKernel discipline): n is re-derived per row from device kv_lens; the sampling-ladder scalars (SMP/TGT/Q/SS2/TGT2, dead launch slots in this mode) are re-derived per row by the route_dynamic clus mirror (exact-integer isqrt aim ladder, Int64 target products, QUAD 16-multiplier geometry). SCAP/CMP are pure functions of (rows, CS, k) -- never of n -- so the envelope launch values are the per-row values; smem extents unchanged. The 'big' occupancy flag is launch-computed in the jit wrapper (rows * CS <= 148). One deviation from route_dynamic, documented in-kernel: the QUAD schedule is computed for every non-short row instead of only n > SCAP -- the host only ever launches this family with n > SCAP, so the SMP == 0 no-sample path is untested; per-row n <= SCAP rows get a small valid schedule (sampling only steers the rung; exactness is schedule-invariant, and sample positions stay under the row's own n). Short rows (n <= k) emit identity + (-1) tail from cluster rank 0 and skip the whole body -- 'short' is a pure function of 'row', identical across all CS ranks, so every cluster barrier stays aligned. Verification: UT 84/84 (2 new: heterogeneous parity+oracle over CS=2 and CS=4 clusters with mid rows below the standalone admission floor, and a CUDA-graph capture/replay test). Real-capture clus zone (153 cells / 255 cases from the 886-grid census): all exact, integration tax (same-session CUDA-events pairing, cmp886_r3 protocol) 1.319 -> 1.137 gm (worst 1.216), now at the same intrinsic per-row level as the other families. Graph-replay protocol on clus shapes: tax 1.096-1.124. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 1552 +++++++++-------- .../gvr_topk_decode_self_sampling_host.py | 41 +- .../parallel/test_gvr_selfsampling_topk.py | 79 + 3 files changed, 974 insertions(+), 698 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index b2bc70cc66cb..f334f192fdcc 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -4321,12 +4321,33 @@ class GvrClusKernel: """CuTeDSL port of gvr_clus (kernel.cu L1795).""" def __init__( - self, blk: int, u: int, minb: int, nbs: int, cs: int, scap: int = 8192, cmp_: int = 2048 + self, + blk: int, + u: int, + minb: int, + nbs: int, + cs: int, + scap: int = 8192, + cmp_: int = 2048, + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, ): assert blk == 1024, "gvr_clus is always BLK=1024 (dispatch L3132)" assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1) (L1796)" assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177)" assert u in (1, 2, 4, 8) and cs in (2, 4, 8) + # per-row varlen mode (production heuristicTopKDecode contract, same + # semantics as GvrMainKernel / GvrRegClusKernel): n and the sampling- + # ladder scalars are re-derived PER ROW in-kernel from a device + # kv_lens tensor; the scalar launch args become the envelope clamp + # bound (n) and dead slots (SMP/TGT/Q/SS2/TGT2). + self.varlen = bool(varlen) + self.next_n = int(next_n) + self.cr_shift = int(cr_shift) + if self.varlen: + assert self.next_n >= 1 and self.cr_shift in (0, 2) + self.lcs = cs.bit_length() - 1 # log2(CS) for the per-row Q shift self.blk = blk self.u = u self.minb = minb @@ -4502,6 +4523,7 @@ def kern( self, logits: cute.Tensor, pre_idx: cute.Tensor, + kv_lens: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, npad: cutlass.Int32, @@ -4513,6 +4535,7 @@ def kern( Q: cutlass.Int32, SS2: cutlass.Int32, TGT2: cutlass.Int32, + bigf: cutlass.Int32, ): BLK = self.blk U = self.u @@ -4528,6 +4551,114 @@ def kern( row = by lane = tidx & cutlass.Int32(31) + # ================= per-row varlen prologue (varlen mode only) ========= + # Production contract (GvrMainKernel / GvrRegClusKernel discipline): + # row r serves request r // next_n with + # n = (kv_lens[req] - next_n + r % next_n + 1) >> cr_shift, clamped to + # the envelope launch arg n. The sampling-ladder scalars + # (SMP/TGT/Q/SS2/TGT2 — dead launch args in this mode) are re-derived + # from this row's n by the route_dynamic() clus formulas below, with + # ONE deviation: the QUAD sample geometry is computed for every + # non-short row instead of only n > SCAP — the host only ever launches + # this family with n > SCAP, so the SMP == 0 no-sample path is + # untested; short-of-envelope rows get a small valid schedule instead + # (sampling only steers the rung; exactness is schedule-invariant). + # Every quantity is a pure function of `row`, so all CS ranks of a + # row's cluster (and all threads) compute identical values — the + # whole-body guard below is cluster-uniform and the cluster barriers + # inside remain aligned. Short rows (n <= k) emit identity + (-1) + # tail from rank 0 here and SKIP the body entirely. + short = cutlass.Int32(0) + prow = row + if cutlass.const_expr(self.varlen): + kq = cutlass.Int32(pre_idx.shape[1]) + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + prow = req + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > n: + nv = n + if nv <= kq: + short = cutlass.Int32(1) + if short == cutlass.Int32(0): + n = nv + # ---- aim ladder (route_dynamic mirror; isqrt discipline and + # Int64 target products exactly as the GvrMainKernel prologue). + x6 = cutlass.Int32(6) * nv + ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) + while ri * ri > x6: + ri = ri - cutlass.Int32(1) + while (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: + ri = ri + cutlass.Int32(1) + r6 = ri + if x6 - ri * ri > ri: + r6 = ri + cutlass.Int32(1) + # aim_base: R = CS > 1 always for this family; bigf is the + # launch-computed occupancy flag (num_rows * CS <= 148). + aim = k << cutlass.Int32(1) + if bigf == cutlass.Int32(0): + if k >= cutlass.Int32(1024): + aim = (cutlass.Int32(11) * k) >> cutlass.Int32(3) + else: + aim = (cutlass.Int32(3) * k) >> cutlass.Int32(1) + if r6 > aim: + aim = r6 + amin = cutlass.Int32(3) * k + if cutlass.const_expr(self.cs != 2): + amin = (cutlass.Int32(7) * k) >> cutlass.Int32(1) + if aim < amin: + aim = amin + if aim > (SCAP >> cutlass.Int32(1)): + aim = SCAP >> cutlass.Int32(1) + if aim < k: + aim = k + # ---- QUAD sample geometry (route_dynamic clus override, + # always-on per the deviation note above; k <= 1024 for this + # family so sfac has no k > 1024 arm). + n4v = nv >> cutlass.Int32(2) + sel = cutlass.Int32(32 if self.cs == 2 else 16) * nv // aim + if sel < cutlass.Int32(256): + sel = cutlass.Int32(256) + nh = nv >> cutlass.Int32(1) + if sel > nh: + sel = nh + quads = sel >> cutlass.Int32(4) + if quads < cutlass.Int32(1): + quads = cutlass.Int32(1) + quarter = n4v >> cutlass.Int32(2) + if quarter < cutlass.Int32(1): + quarter = cutlass.Int32(1) + if quads > quarter: + quads = quarter + SS2 = quarter // quads + if SS2 < cutlass.Int32(1): + SS2 = cutlass.Int32(1) + SMP = quarter // SS2 + if SMP < cutlass.Int32(1): + SMP = cutlass.Int32(1) + smp16 = cutlass.Int64(SMP) * cutlass.Int64(16) + tgt64 = cutlass.Int64(aim) * smp16 // cutlass.Int64(nv) + TGT = cutlass.Int32(tgt64) + if TGT < cutlass.Int32(1): + TGT = cutlass.Int32(1) + tgt264 = cutlass.Int64(k) * smp16 // cutlass.Int64(nv) + TGT2 = cutlass.Int32(tgt264) + if TGT2 < cutlass.Int32(1): + TGT2 = cutlass.Int32(1) + Q = (n4v + cutlass.Int32(CS - 1)) >> cutlass.Int32(self.lcs) + if short != cutlass.Int32(0): + if rank == cutlass.Int32(0): + if tidx < kq: + ov = cutlass.Int32(-1) + if tidx < nv: + ov = tidx + out[row, tidx] = ov + # ---- shared memory (CUDA dynamic map order L1802-1814, then static allocs) ---- smem = SmemAllocator() s_hist = smem.allocate_tensor( # hist[NBS] @ blob start @@ -4577,139 +4708,63 @@ def kern( cute.make_layout((self.cmp,)), ) - # ---- row bases (L1830-1833) ---- - row64 = cutlass.Int64(row) - x_addr = logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) - p_addr = pre_idx.iterator.toint() + row64 * cutlass.Int64(k) * cutlass.Int64(4) - out_row = out[row, None] - - # ---- interleaved chunk ownership (L1835-1846) ---- - n4 = n >> cutlass.Int32(2) - nCh = (n4 + cutlass.Int32(STEPC - 1)) // cutlass.Int32(STEPC) - nFullG = n4 // cutlass.Int32(STEPC) - tail0 = n4 << cutlass.Int32(2) - tailn = cutlass.Int32(0) - if rank == cutlass.Int32(0): - tailn = n - tail0 - - if tidx == cutlass.Int32(0): # L1848 - s_res[C.RES_B2] = cutlass.Int32(-1) - s_res[C.RES_B3] = cutlass.Int32(-1) - s_scal[0] = cutlass.Int32(0) # s_bufn - if tidx < cutlass.Int32(self.hb): # L1849 (HB<=BLK) - s_hist[tidx] = cutlass.Int32(0) - - # ============ P1: QUAD sample (hint gather LAZY, L1851-1878) ======== - # one 64B line = 4 float4 per location, TWO threads: tid takes the - # lower pair at p4, tid+SMP the upper pair at p4+2 (L1868-1869). - atom128 = C.g2r_atom_f32(128, invariant=True) - fsa = cute.make_fragment((4,), cutlass.Float32) - fsb = cute.make_fragment((4,), cutlass.Float32) - smp2 = SMP * cutlass.Int32(2) - shas = cutlass.Int32(0) - if tidx < smp2: - shas = cutlass.Int32(1) - if shas != cutlass.Int32(0): - p4 = tidx * SS2 * cutlass.Int32(4) - if tidx >= SMP: - p4 = (tidx - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) - C.ld_g_f32x4(atom128, x_addr, p4, fsa) - C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) - - # ============ P2: quantile rung, redundant per CTA (L1883-1997) ===== - smn = cutlass.Float32(float("inf")) - smx = cutlass.Float32(float("-inf")) - if shas != cutlass.Int32(0): # L1928-1932 - for t in cutlass.range_constexpr(4): - smn = C.fmin_f32(smn, fsa[t]) - smx = C.fmax_f32(smx, fsa[t]) - for t in cutlass.range_constexpr(4): - smn = C.fmin_f32(smn, fsb[t]) - smx = C.fmax_f32(smx, fsb[t]) - fma_ = cute.make_fragment((4,), cutlass.Float32) # mop-up pair bufs - fmb_ = cute.make_fragment((4,), cutlass.Float32) - j = tidx + cutlass.Int32(BLK) # mop-up L1891-1897 - while j < smp2: - p4 = j * SS2 * cutlass.Int32(4) - if j >= SMP: - p4 = (j - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) - C.ld_g_f32x4(atom128, x_addr, p4, fma_) - C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) - for t in cutlass.range_constexpr(4): - smn = C.fmin_f32(smn, fma_[t]) - smx = C.fmax_f32(smx, fma_[t]) - for t in cutlass.range_constexpr(4): - smn = C.fmin_f32(smn, fmb_[t]) - smx = C.fmax_f32(smx, fmb_[t]) - j = j + cutlass.Int32(BLK) - a0 = C.warp_min_u32(C.fkey(smn)) # L1898-1901 - c0m = C.warp_max_u32(C.fkey(smx)) - if lane == cutlass.Int32(0): - s_wmn[tidx >> cutlass.Int32(5)] = a0 - s_wmx[tidx >> cutlass.Int32(5)] = c0m - cute.arch.barrier() # ---- barrier L1902 ---- - - # PRIME-LATE (L1903-1916): every rank's sample has landed; prime NOW. - lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L1904 - pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(PFD)] - for uu in cutlass.range_constexpr(PFD): # clamped prime L1906 - i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) - ic = i_ - if ic >= n4: - ic = lim4 - C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - # asm prefetch gate (L1912-1916): DEEP rows only; empty for U<=PFD - if cutlass.const_expr(U > PFD): - gpp = cutlass.Int32(0) - if n4 >= cutlass.Int32(32768): - if (rank + cutlass.Int32(1)) * cutlass.Int32(STEPC) <= n4: - gpp = cutlass.Int32(1) - if gpp != cutlass.Int32(0): - for uu in cutlass.range_constexpr(PFD, U): - C._prefetch_l2( - x_addr - + cutlass.Int64( - rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) - ) - * cutlass.Int64(16) - ) - - # cross-warp sample reduce (L1917-1923) - av = cutlass.Uint32(0xFFFFFFFF) - cv = cutlass.Uint32(0) - if lane < cutlass.Int32(NW): - av = s_wmn[lane] - cv = s_wmx[lane] - SMIN = C.invkey(C.warp_min_u32(av)) - SMAX = C.invkey(C.warp_max_u32(cv)) + # ---- whole-body short-row guard (cluster-uniform: `short` is a pure + # function of `row`, identical across all CS ranks and threads, so + # every cluster barrier below stays aligned; short rows already + # emitted identity + -1 tail in the prologue) ---- + if short == cutlass.Int32(0): + # ---- row bases (L1830-1833; pre_idx is request-level under varlen) ---- + row64 = cutlass.Int64(row) + x_addr = logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) + p_addr = pre_idx.iterator.toint() + cutlass.Int64(prow) * cutlass.Int64(k) * cutlass.Int64(4) + out_row = out[row, None] - GMIN = cutlass.Float32(C.SENT_LO) # L1925-1926 - GMAX = cutlass.Float32(C.SENT_HI) - T = cutlass.Float32(_NEG_INF__clus) - HIC = cutlass.Float32(_NEG_INF__clus) - w = cutlass.Float32(0.0) - sok = cutlass.Int32(0) # L1930 - if SMP > cutlass.Int32(0): - if SMAX > SMIN: - sok = cutlass.Int32(1) - if sok != cutlass.Int32(0): # L1932-1947 sample hist - w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) - # CUDA --use_fast_math lowers `1.0f / w` to a bare MUFU.RCP; the - # plain `/` spelling emits the IEEE div.rn rcp+Newton+CALL chain - # (fix-4). w > 0 by the sok guard; bucketing is SC-invariant. - sc_s = cute.arch.rcp_approx(w) + # ---- interleaved chunk ownership (L1835-1846) ---- + n4 = n >> cutlass.Int32(2) + nCh = (n4 + cutlass.Int32(STEPC - 1)) // cutlass.Int32(STEPC) + nFullG = n4 // cutlass.Int32(STEPC) + tail0 = n4 << cutlass.Int32(2) + tailn = cutlass.Int32(0) + if rank == cutlass.Int32(0): + tailn = n - tail0 + + if tidx == cutlass.Int32(0): # L1848 + s_res[C.RES_B2] = cutlass.Int32(-1) + s_res[C.RES_B3] = cutlass.Int32(-1) + s_scal[0] = cutlass.Int32(0) # s_bufn + if tidx < cutlass.Int32(self.hb): # L1849 (HB<=BLK) + s_hist[tidx] = cutlass.Int32(0) + + # ============ P1: QUAD sample (hint gather LAZY, L1851-1878) ======== + # one 64B line = 4 float4 per location, TWO threads: tid takes the + # lower pair at p4, tid+SMP the upper pair at p4+2 (L1868-1869). + atom128 = C.g2r_atom_f32(128, invariant=True) + fsa = cute.make_fragment((4,), cutlass.Float32) + fsb = cute.make_fragment((4,), cutlass.Float32) + smp2 = SMP * cutlass.Int32(2) + shas = cutlass.Int32(0) + if tidx < smp2: + shas = cutlass.Int32(1) if shas != cutlass.Int32(0): + p4 = tidx * SS2 * cutlass.Int32(4) + if tidx >= SMP: + p4 = (tidx - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fsa) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) + + # ============ P2: quantile rung, redundant per CTA (L1883-1997) ===== + smn = cutlass.Float32(float("inf")) + smx = cutlass.Float32(float("-inf")) + if shas != cutlass.Int32(0): # L1928-1932 for t in cutlass.range_constexpr(4): - bq = C.f2s_rz((fsa[t] - SMIN) * sc_s) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + smn = C.fmin_f32(smn, fsa[t]) + smx = C.fmax_f32(smx, fsa[t]) for t in cutlass.range_constexpr(4): - bq = C.f2s_rz((fsb[t] - SMIN) * sc_s) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) - j = tidx + cutlass.Int32(BLK) # mop-up reloads + smn = C.fmin_f32(smn, fsb[t]) + smx = C.fmax_f32(smx, fsb[t]) + fma_ = cute.make_fragment((4,), cutlass.Float32) # mop-up pair bufs + fmb_ = cute.make_fragment((4,), cutlass.Float32) + j = tidx + cutlass.Int32(BLK) # mop-up L1891-1897 while j < smp2: p4 = j * SS2 * cutlass.Int32(4) if j >= SMP: @@ -4717,590 +4772,671 @@ def kern( C.ld_g_f32x4(atom128, x_addr, p4, fma_) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) for t in cutlass.range_constexpr(4): - bq = C.f2s_rz((fma_[t] - SMIN) * sc_s) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + smn = C.fmin_f32(smn, fma_[t]) + smx = C.fmax_f32(smx, fma_[t]) for t in cutlass.range_constexpr(4): - bq = C.f2s_rz((fmb_[t] - SMIN) * sc_s) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + smn = C.fmin_f32(smn, fmb_[t]) + smx = C.fmax_f32(smx, fmb_[t]) j = j + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1948 ---- - # triple-target ZERO scan (L1952-1955): TGT / TGT2 / 2*TGT - C.scan_cross0( - s_hist, - TGT, - tidx, - s_res, - TGT2, - TGT * cutlass.Int32(2), - s_hist, - nb=NBS, - zero=True, - two=True, - three=True, - ) - cute.arch.barrier() # ---- barrier L1956 ---- + a0 = C.warp_min_u32(C.fkey(smn)) # L1898-1901 + c0m = C.warp_max_u32(C.fkey(smx)) + if lane == cutlass.Int32(0): + s_wmn[tidx >> cutlass.Int32(5)] = a0 + s_wmx[tidx >> cutlass.Int32(5)] = c0m + cute.arch.barrier() # ---- barrier L1902 ---- + + # PRIME-LATE (L1903-1916): every rank's sample has landed; prime NOW. + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L1904 + pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(PFD)] + for uu in cutlass.range_constexpr(PFD): # clamped prime L1906 + i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + # asm prefetch gate (L1912-1916): DEEP rows only; empty for U<=PFD + if cutlass.const_expr(U > PFD): + gpp = cutlass.Int32(0) + if n4 >= cutlass.Int32(32768): + if (rank + cutlass.Int32(1)) * cutlass.Int32(STEPC) <= n4: + gpp = cutlass.Int32(1) + if gpp != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD, U): + C._prefetch_l2( + x_addr + + cutlass.Int64( + rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ) + * cutlass.Int64(16) + ) - tot0 = s_res[C.RES_TOT] - b1v = s_res[C.RES_B] - if sok != cutlass.Int32(0): # L1957 - if tot0 >= TGT: - T = _fmaf__clus(cutlass.Float32(b1v), w, SMIN) - needg = cutlass.Int32(1) # L1958-1963 - if T > cutlass.Float32(_NEG_INF__clus): - needg = cutlass.Int32(0) - if needg != cutlass.Int32(0): - # degenerate sample: identical on every rank of the cluster - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers - T = GMIN - if sok != cutlass.Int32(0): # L1964-1982 HIC - if tot0 >= TGT: - b2v = s_res[C.RES_B2] - if b2v >= cutlass.Int32(0): - Tk = _fmaf__clus(cutlass.Float32(b2v), w, SMIN) - up = C.fmax_f32(Tk - T, cutlass.Float32(0.0)) - # heavy-tail cap by T - T3 (rank-TGT..rank-2TGT distance) + # cross-warp sample reduce (L1917-1923) + av = cutlass.Uint32(0xFFFFFFFF) + cv = cutlass.Uint32(0) + if lane < cutlass.Int32(NW): + av = s_wmn[lane] + cv = s_wmx[lane] + SMIN = C.invkey(C.warp_min_u32(av)) + SMAX = C.invkey(C.warp_max_u32(cv)) + + GMIN = cutlass.Float32(C.SENT_LO) # L1925-1926 + GMAX = cutlass.Float32(C.SENT_HI) + T = cutlass.Float32(_NEG_INF__clus) + HIC = cutlass.Float32(_NEG_INF__clus) + w = cutlass.Float32(0.0) + sok = cutlass.Int32(0) # L1930 + if SMP > cutlass.Int32(0): + if SMAX > SMIN: + sok = cutlass.Int32(1) + if sok != cutlass.Int32(0): # L1932-1947 sample hist + w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) + # CUDA --use_fast_math lowers `1.0f / w` to a bare MUFU.RCP; the + # plain `/` spelling emits the IEEE div.rn rcp+Newton+CALL chain + # (fix-4). w > 0 by the sok guard; bucketing is SC-invariant. + sc_s = cute.arch.rcp_approx(w) + if shas != cutlass.Int32(0): + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsa[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fsb[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = tidx + cutlass.Int32(BLK) # mop-up reloads + while j < smp2: + p4 = j * SS2 * cutlass.Int32(4) + if j >= SMP: + p4 = (j - SMP) * SS2 * cutlass.Int32(4) + cutlass.Int32(2) + C.ld_g_f32x4(atom128, x_addr, p4, fma_) + C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fmb_) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fma_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + for t in cutlass.range_constexpr(4): + bq = C.f2s_rz((fmb_[t] - SMIN) * sc_s) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + j = j + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L1948 ---- + # triple-target ZERO scan (L1952-1955): TGT / TGT2 / 2*TGT + C.scan_cross0( + s_hist, + TGT, + tidx, + s_res, + TGT2, + TGT * cutlass.Int32(2), + s_hist, + nb=NBS, + zero=True, + two=True, + three=True, + ) + cute.arch.barrier() # ---- barrier L1956 ---- + + tot0 = s_res[C.RES_TOT] + b1v = s_res[C.RES_B] + if sok != cutlass.Int32(0): # L1957 + if tot0 >= TGT: + T = _fmaf__clus(cutlass.Float32(b1v), w, SMIN) + needg = cutlass.Int32(1) # L1958-1963 + if T > cutlass.Float32(_NEG_INF__clus): + needg = cutlass.Int32(0) + if needg != cutlass.Int32(0): + # degenerate sample: identical on every rank of the cluster + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers + T = GMIN + if sok != cutlass.Int32(0): # L1964-1982 HIC + if tot0 >= TGT: + b2v = s_res[C.RES_B2] + if b2v >= cutlass.Int32(0): + Tk = _fmaf__clus(cutlass.Float32(b2v), w, SMIN) + up = C.fmax_f32(Tk - T, cutlass.Float32(0.0)) + # heavy-tail cap by T - T3 (rank-TGT..rank-2TGT distance) + if tot0 >= TGT * cutlass.Int32(2): + b3v = s_res[C.RES_B3] + if b3v >= cutlass.Int32(0): + T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) + if T > T3: + up = C.fmin_f32(up, cutlass.Float32(2.0) * (T - T3)) + HIC = C.fmax_f32( + _fmaf__clus(cutlass.Float32(4.0), up, T), + _fmaf__clus(cutlass.Float32(8.0), w, T), + ) + # r4 (a000) ladder floor in SHARED (64-reg wall, L1983-1996) + if tidx == cutlass.Int32(0): + t5 = cutlass.Float32(_NEG_INF__clus) + if sok != cutlass.Int32(0): if tot0 >= TGT * cutlass.Int32(2): b3v = s_res[C.RES_B3] if b3v >= cutlass.Int32(0): - T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) - if T > T3: - up = C.fmin_f32(up, cutlass.Float32(2.0) * (T - T3)) - HIC = C.fmax_f32( - _fmaf__clus(cutlass.Float32(4.0), up, T), - _fmaf__clus(cutlass.Float32(8.0), w, T), - ) - # r4 (a000) ladder floor in SHARED (64-reg wall, L1983-1996) - if tidx == cutlass.Int32(0): - t5 = cutlass.Float32(_NEG_INF__clus) - if sok != cutlass.Int32(0): - if tot0 >= TGT * cutlass.Int32(2): - b3v = s_res[C.RES_B3] - if b3v >= cutlass.Int32(0): - if T > GMIN: - T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) - if T3 < T: - t5 = T3 - s_tsh[0] = t5 - - # ============ attempt loop (L1999-2162) — MUST NOT unroll =========== - listN = cutlass.Int32(0) - above = cutlass.Int32(0) - m = cutlass.Int32(0) - need = cutlass.Int32(0) - B = cutlass.Int32(0) - SC = cutlass.Float32(1.0) - TF = T - complete = cutlass.Int32(0) - valid = cutlass.Int32(0) - - fr = [ - cute.make_fragment((4,), cutlass.Float32) for _ in range(U - PFD) - ] # explicit batch (op43 L1) - # (empty for U<=PFD — every row-pass float4 then comes from pf[]) - att = cutlass.Int32(0) - running = cutlass.Int32(1) - while running != cutlass.Int32(0): - if att > cutlass.Int32(0): # retry preamble L2005-2024 - # EXACTNESS: re-prime pf[] (stale roll data, L2006-2015) - if rank < nFullG: - for uu in cutlass.range_constexpr(PFD): - C.ld_g_f32x4( - atom128, - x_addr, - rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK), - pf[uu], - ) - else: - for uu in cutlass.range_constexpr(PFD): - i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) - ic = i_ - if ic >= n4: - ic = lim4 - C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - C._cluster_sync_aligned() # ==== clus.sync L2016 ==== - if tidx < cutlass.Int32(NBS): # L2021-2022 - s_hist[tidx] = cutlass.Int32(0) - if tidx == cutlass.Int32(0): - s_scal[0] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2024 ---- + if T > GMIN: + T3 = _fmaf__clus(cutlass.Float32(b3v), w, SMIN) + if T3 < T: + t5 = T3 + s_tsh[0] = t5 - TF = T # window L2026-2031 - hi = C.fmax_f32(GMAX, T) - if HIC > T: - if HIC < hi: - hi = HIC - WD = (hi - T) * cutlass.Float32(1.0 / 256.0) - wdok = cutlass.Int32(0) - if WD > cutlass.Float32(0.0): - wdok = cutlass.Int32(1) - if wdok == cutlass.Int32(0): - WD = cutlass.Float32(1e-30) - # MUFU.RCP spelling (fix-4): WD >= 1e-30 finite by the wdok - # clamp; classify bucketing is SC-invariant for any SC > 0. - SC = cute.arch.rcp_approx(WD) - - # ---- P3 row pass over OWNED CHUNKS (L2033-2121) ---- - g = rank + cutlass.Int32(0) - while g < nCh: - i0 = g * cutlass.Int32(STEPC) + tidx - M = cutlass.Int32(0) - isfull = cutlass.Int32(0) - if g < nFullG: - isfull = cutlass.Int32(1) - if isfull != cutlass.Int32(0): # full body L2042-2049 - for uu in cutlass.range_constexpr(PFD, U): - C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) - for uu in cutlass.range_constexpr(U): - if cutlass.const_expr(uu < PFD): - vv = pf[uu] - else: - vv = fr[uu - PFD] - for q in cutlass.range_constexpr(4): - M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) - else: # partial body L2051-2064 - for uu in cutlass.range_constexpr(PFD, U): - i_ = i0 + cutlass.Int32(uu * BLK) - ic = i_ - if ic >= n4: - ic = lim4 # clamp in [n, npad) - C.ld_g_f32x4(atom128, x_addr, ic, fr[uu - PFD]) - for uu in cutlass.range_constexpr(U): - if cutlass.const_expr(uu < PFD): - vv = pf[uu] - else: - vv = fr[uu - PFD] - i_ = i0 + cutlass.Int32(uu * BLK) - okq = cutlass.Int32(0) - if i_ < n4: - okq = cutlass.Int32(1) - if okq != cutlass.Int32(0): # +inf-pad escape, ok-gated - for q in cutlass.range_constexpr(4): - M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) - # ROLL THE PREFETCH FORWARD (L2066-2081): next OWNED chunk, - # issued before the reservation and the survivor walk. - g2 = g + cutlass.Int32(CS) - if g2 < nCh: - j0 = g2 * cutlass.Int32(STEPC) + tidx - infull = cutlass.Int32(0) - if g2 < nFullG: - infull = cutlass.Int32(1) - if infull != cutlass.Int32(0): + # ============ attempt loop (L1999-2162) — MUST NOT unroll =========== + listN = cutlass.Int32(0) + above = cutlass.Int32(0) + m = cutlass.Int32(0) + need = cutlass.Int32(0) + B = cutlass.Int32(0) + SC = cutlass.Float32(1.0) + TF = T + complete = cutlass.Int32(0) + valid = cutlass.Int32(0) + + fr = [ + cute.make_fragment((4,), cutlass.Float32) for _ in range(U - PFD) + ] # explicit batch (op43 L1) + # (empty for U<=PFD — every row-pass float4 then comes from pf[]) + att = cutlass.Int32(0) + running = cutlass.Int32(1) + while running != cutlass.Int32(0): + if att > cutlass.Int32(0): # retry preamble L2005-2024 + # EXACTNESS: re-prime pf[] (stale roll data, L2006-2015) + if rank < nFullG: for uu in cutlass.range_constexpr(PFD): - C.ld_g_f32x4(atom128, x_addr, j0 + cutlass.Int32(uu * BLK), pf[uu]) + C.ld_g_f32x4( + atom128, + x_addr, + rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK), + pf[uu], + ) else: for uu in cutlass.range_constexpr(PFD): - j_ = j0 + cutlass.Int32(uu * BLK) - jc = j_ - if jc >= n4: - jc = lim4 - C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) - # warp-aggregated slot reservation (L2082-2095) - cnt = cutlass.Int32(C.popc(M)) - inc = C.warp_incl_scan_add(cnt, lane) - bpos = cutlass.Int32(0) - if lane == cutlass.Int32(31): - if inc != cutlass.Int32(0): - bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) - pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) - # survivor bit-walk, software-pipelined ONE deep (L2105-2119); - # reload X[idx] — never hold the U float4s across the walk - if M != cutlass.Int32(0): - bp = C.ffs_m1(M) - M = M & (M - cutlass.Int32(1)) - idx = ( - (i0 + (bp >> cutlass.Int32(2)) * cutlass.Int32(BLK)) << cutlass.Int32(2) - ) + (bp & cutlass.Int32(3)) - xv = C.ldg_f32(x_addr, idx) - while M != cutlass.Int32(0): - bp2 = C.ffs_m1(M) - M = M & (M - cutlass.Int32(1)) - idx2 = ( - (i0 + (bp2 >> cutlass.Int32(2)) * cutlass.Int32(BLK)) - << cutlass.Int32(2) - ) + (bp2 & cutlass.Int32(3)) - xv2 = C.ldg_f32(x_addr, idx2) - pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) - idx = idx2 - xv = xv2 - pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) - g = g + cutlass.Int32(CS) - # rank-0 scalar tail (L2122-2130): per-thread atomics, bound-check - i = tidx - while i < tailn: - x = C.ldg_f32(x_addr, tail0 + i) - if x >= TF: - bq = C.f2s_rz((x - TF) * SC) # signed form L2125 - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) - post = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) - if post < SCAP: - s_cbuf2[post] = ( - cutlass.Uint64(cutlass.Uint32(tail0 + i)) << cutlass.Uint64(32) - ) | cutlass.Uint64(C.u32_of_f32(x)) - i = i + cutlass.Int32(BLK) - - # ---- cluster merge (L2132-2148) ---- - C._cluster_sync_aligned() # ==== clus.sync L2135 ==== - myn = s_scal[0] # L2140 - self._merge_scan0_local(s_hist, s_mrg, rank, k, tidx, s_res) - cute.arch.barrier() # ---- barrier L2143 ---- - tot = s_res[C.RES_TOT] - acc = cutlass.Int32(0) - if tot >= k: - acc = cutlass.Int32(1) - if acc != cutlass.Int32(0): # L2145-2148 accept - valid = cutlass.Int32(1) - complete = cutlass.Int32(0) - if myn <= SCAP: - complete = cutlass.Int32(1) - listN = myn - above = s_res[C.RES_ABOVE] - m = s_res[C.RES_M] - need = k - s_res[C.RES_ABOVE] - B = s_res[C.RES_B] - running = cutlass.Int32(0) - else: - if att == cutlass.Int32(2): # L2149 - running = cutlass.Int32(0) - else: - # rung ladder (L2150-2161) — cluster-uniform on every arm - tshtaken = cutlass.Int32(0) - if att == cutlass.Int32(0): - T5 = s_tsh[0] - if T5 > cutlass.Float32(_NEG_INF__clus): - if T5 < TF: - T = T5 - tshtaken = cutlass.Int32(1) - if tshtaken == cutlass.Int32(0): - # LAZY GATHER — every rank computes identical GMIN - if GMIN == cutlass.Float32(C.SENT_LO): - GMIN, GMAX = C.gather_hint( - x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 - ) # 2 barriers inside - floorhit = cutlass.Int32(1) # L2160 - if T > GMIN: - floorhit = cutlass.Int32(0) - if floorhit != cutlass.Int32(0): - running = cutlass.Int32(0) - else: - T = GMIN # L2161 - att = att + cutlass.Int32(1) + i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 + C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) + C._cluster_sync_aligned() # ==== clus.sync L2016 ==== + if tidx < cutlass.Int32(NBS): # L2021-2022 + s_hist[tidx] = cutlass.Int32(0) + if tidx == cutlass.Int32(0): + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2024 ---- + + TF = T # window L2026-2031 + hi = C.fmax_f32(GMAX, T) + if HIC > T: + if HIC < hi: + hi = HIC + WD = (hi - T) * cutlass.Float32(1.0 / 256.0) + wdok = cutlass.Int32(0) + if WD > cutlass.Float32(0.0): + wdok = cutlass.Int32(1) + if wdok == cutlass.Int32(0): + WD = cutlass.Float32(1e-30) + # MUFU.RCP spelling (fix-4): WD >= 1e-30 finite by the wdok + # clamp; classify bucketing is SC-invariant for any SC > 0. + SC = cute.arch.rcp_approx(WD) - # ============ classification (L2165-2173) ============ - whole = cutlass.Int32(0) - if valid != cutlass.Int32(0): - if need >= m: - whole = cutlass.Int32(1) - lim1 = above - if whole != cutlass.Int32(0): - lim1 = above + m - degen = cutlass.Int32(0) - if valid == cutlass.Int32(0): - degen = cutlass.Int32(1) - if m > CMP: - degen = cutlass.Int32(1) - mc = cutlass.Int32(0) - if degen == cutlass.Int32(0): - mc = m - # crossing candidates land in RANK 0's ck64c via DSMEM (L2173) - rk64 = C._mapa_shared_cluster_addr(ck_addr, cutlass.Int32(0)) - - if degen == cutlass.Int32(0): - if complete != cutlass.Int32(0): - # ---- P5 emit from staged cbuf (L2176-2187) ---- - i = tidx - while i < listN: - pk64 = s_cbuf2[i] - vx = cutlass.Int32(cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF))) - idv = cutlass.Int32(pk64 >> cutlass.Uint64(32)) - xv = C.f32_of_i32(vx) - self._p5_emit(xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64) - i = i + cutlass.Int32(BLK) - else: - # ---- EXACTNESS re-sweep: OWNED CHUNKS + rank-0 true tail - # (L2188-2222) ---- + # ---- P3 row pass over OWNED CHUNKS (L2033-2121) ---- g = rank + cutlass.Int32(0) while g < nCh: - lo2 = (g * cutlass.Int32(STEPC)) << cutlass.Int32(2) - e4 = (g + cutlass.Int32(1)) * cutlass.Int32(STEPC) - if e4 > n4: - e4 = n4 - hi2 = e4 << cutlass.Int32(2) - i = lo2 + tidx - while i < hi2: - x = C.ldg_f32(x_addr, i) - if x >= TF: - self._p5_emit( - x, i, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 - ) - i = i + cutlass.Int32(BLK) + i0 = g * cutlass.Int32(STEPC) + tidx + M = cutlass.Int32(0) + isfull = cutlass.Int32(0) + if g < nFullG: + isfull = cutlass.Int32(1) + if isfull != cutlass.Int32(0): # full body L2042-2049 + for uu in cutlass.range_constexpr(PFD, U): + C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + else: # partial body L2051-2064 + for uu in cutlass.range_constexpr(PFD, U): + i_ = i0 + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= n4: + ic = lim4 # clamp in [n, npad) + C.ld_g_f32x4(atom128, x_addr, ic, fr[uu - PFD]) + for uu in cutlass.range_constexpr(U): + if cutlass.const_expr(uu < PFD): + vv = pf[uu] + else: + vv = fr[uu - PFD] + i_ = i0 + cutlass.Int32(uu * BLK) + okq = cutlass.Int32(0) + if i_ < n4: + okq = cutlass.Int32(1) + if okq != cutlass.Int32(0): # +inf-pad escape, ok-gated + for q in cutlass.range_constexpr(4): + M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + # ROLL THE PREFETCH FORWARD (L2066-2081): next OWNED chunk, + # issued before the reservation and the survivor walk. + g2 = g + cutlass.Int32(CS) + if g2 < nCh: + j0 = g2 * cutlass.Int32(STEPC) + tidx + infull = cutlass.Int32(0) + if g2 < nFullG: + infull = cutlass.Int32(1) + if infull != cutlass.Int32(0): + for uu in cutlass.range_constexpr(PFD): + C.ld_g_f32x4(atom128, x_addr, j0 + cutlass.Int32(uu * BLK), pf[uu]) + else: + for uu in cutlass.range_constexpr(PFD): + j_ = j0 + cutlass.Int32(uu * BLK) + jc = j_ + if jc >= n4: + jc = lim4 + C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) + # warp-aggregated slot reservation (L2082-2095) + cnt = cutlass.Int32(C.popc(M)) + inc = C.warp_incl_scan_add(cnt, lane) + bpos = cutlass.Int32(0) + if lane == cutlass.Int32(31): + if inc != cutlass.Int32(0): + bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) + pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) + # survivor bit-walk, software-pipelined ONE deep (L2105-2119); + # reload X[idx] — never hold the U float4s across the walk + if M != cutlass.Int32(0): + bp = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx = ( + (i0 + (bp >> cutlass.Int32(2)) * cutlass.Int32(BLK)) << cutlass.Int32(2) + ) + (bp & cutlass.Int32(3)) + xv = C.ldg_f32(x_addr, idx) + while M != cutlass.Int32(0): + bp2 = C.ffs_m1(M) + M = M & (M - cutlass.Int32(1)) + idx2 = ( + (i0 + (bp2 >> cutlass.Int32(2)) * cutlass.Int32(BLK)) + << cutlass.Int32(2) + ) + (bp2 & cutlass.Int32(3)) + xv2 = C.ldg_f32(x_addr, idx2) + pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) + idx = idx2 + xv = xv2 + pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) g = g + cutlass.Int32(CS) - t2 = tidx - while t2 < tailn: - ii = tail0 + t2 - x = C.ldg_f32(x_addr, ii) + # rank-0 scalar tail (L2122-2130): per-thread atomics, bound-check + i = tidx + while i < tailn: + x = C.ldg_f32(x_addr, tail0 + i) if x >= TF: - self._p5_emit( - x, ii, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 - ) - t2 = t2 + cutlass.Int32(BLK) + bq = C.f2s_rz((x - TF) * SC) # signed form L2125 + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + post = C.atomic_add_cta(s_scal.iterator + 0, cutlass.Int32(1)) + if post < SCAP: + s_cbuf2[post] = ( + cutlass.Uint64(cutlass.Uint32(tail0 + i)) << cutlass.Uint64(32) + ) | cutlass.Uint64(C.u32_of_f32(x)) + i = i + cutlass.Int32(BLK) - # ============ EXIT RENDEZVOUS (L2226) ============ - # all DSMEM traffic retired; the ONLY exit rendezvous. rank!=0 falls - # through to the kernel end (post-barrier asymmetric exit L2227); - # NO later cluster barrier. - C._cluster_sync_aligned() # ==== clus.sync L2226 ==== + # ---- cluster merge (L2132-2148) ---- + C._cluster_sync_aligned() # ==== clus.sync L2135 ==== + myn = s_scal[0] # L2140 + self._merge_scan0_local(s_hist, s_mrg, rank, k, tidx, s_res) + cute.arch.barrier() # ---- barrier L2143 ---- + tot = s_res[C.RES_TOT] + acc = cutlass.Int32(0) + if tot >= k: + acc = cutlass.Int32(1) + if acc != cutlass.Int32(0): # L2145-2148 accept + valid = cutlass.Int32(1) + complete = cutlass.Int32(0) + if myn <= SCAP: + complete = cutlass.Int32(1) + listN = myn + above = s_res[C.RES_ABOVE] + m = s_res[C.RES_M] + need = k - s_res[C.RES_ABOVE] + B = s_res[C.RES_B] + running = cutlass.Int32(0) + else: + if att == cutlass.Int32(2): # L2149 + running = cutlass.Int32(0) + else: + # rung ladder (L2150-2161) — cluster-uniform on every arm + tshtaken = cutlass.Int32(0) + if att == cutlass.Int32(0): + T5 = s_tsh[0] + if T5 > cutlass.Float32(_NEG_INF__clus): + if T5 < TF: + T = T5 + tshtaken = cutlass.Int32(1) + if tshtaken == cutlass.Int32(0): + # LAZY GATHER — every rank computes identical GMIN + if GMIN == cutlass.Float32(C.SENT_LO): + GMIN, GMAX = C.gather_hint( + x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 + ) # 2 barriers inside + floorhit = cutlass.Int32(1) # L2160 + if T > GMIN: + floorhit = cutlass.Int32(0) + if floorhit != cutlass.Int32(0): + running = cutlass.Int32(0) + else: + T = GMIN # L2161 + att = att + cutlass.Int32(1) - if rank == cutlass.Int32(0): + # ============ classification (L2165-2173) ============ + whole = cutlass.Int32(0) + if valid != cutlass.Int32(0): + if need >= m: + whole = cutlass.Int32(1) + lim1 = above + if whole != cutlass.Int32(0): + lim1 = above + m + degen = cutlass.Int32(0) + if valid == cutlass.Int32(0): + degen = cutlass.Int32(1) + if m > CMP: + degen = cutlass.Int32(1) + mc = cutlass.Int32(0) if degen == cutlass.Int32(0): - if whole == cutlass.Int32(0): - # ---- P6 rank-0 refine (L2229-2303) ---- - if mc <= cutlass.Int32(QUADC_CLUS__clus): # L2232-2247 O(mc^2) - mc2 = mc & cutlass.Int32(~1) - i = tidx - while i < mc: - # G1: re-assert Uint64 at every unsigned compare - # in/after dynamic loops (ct_main notes). - u64v = s_ck64[i] - r_ = cutlass.Int32(0) - jq = cutlass.Int32(0) - while jq < mc2: # ulonglong2 16B reads - vlo, vhi = C._lds_v2_u64(ck_addr + jq * cutlass.Int32(8)) - r_ = ( - r_ - + cutlass.Int32(vlo > cutlass.Uint64(u64v)) - + cutlass.Int32(vhi > cutlass.Uint64(u64v)) - ) - jq = jq + cutlass.Int32(2) - if mc2 < mc: # odd tail L2244 - r_ = r_ + cutlass.Int32( - cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) - ) - if r_ < need: - out_row[above + r_] = cutlass.Int32( - cutlass.Uint32( - cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) - ) + mc = m + # crossing candidates land in RANK 0's ck64c via DSMEM (L2173) + rk64 = C._mapa_shared_cluster_addr(ck_addr, cutlass.Int32(0)) + + if degen == cutlass.Int32(0): + if complete != cutlass.Int32(0): + # ---- P5 emit from staged cbuf (L2176-2187) ---- + i = tidx + while i < listN: + pk64 = s_cbuf2[i] + vx = cutlass.Int32(cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF))) + idv = cutlass.Int32(pk64 >> cutlass.Uint64(32)) + xv = C.f32_of_i32(vx) + self._p5_emit(xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64) + i = i + cutlass.Int32(BLK) + else: + # ---- EXACTNESS re-sweep: OWNED CHUNKS + rank-0 true tail + # (L2188-2222) ---- + g = rank + cutlass.Int32(0) + while g < nCh: + lo2 = (g * cutlass.Int32(STEPC)) << cutlass.Int32(2) + e4 = (g + cutlass.Int32(1)) * cutlass.Int32(STEPC) + if e4 > n4: + e4 = n4 + hi2 = e4 << cutlass.Int32(2) + i = lo2 + tidx + while i < hi2: + x = C.ldg_f32(x_addr, i) + if x >= TF: + self._p5_emit( + x, i, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 ) i = i + cutlass.Int32(BLK) - else: - # key-space narrowing over ck64c (L2249-2278) - if tidx == cutlass.Int32(0): - s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) - s_kmm[1] = cutlass.Uint32(0) - if tidx < cutlass.Int32(NBS): # cleared ONCE L2251 - s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2255 ---- - i = tidx - while i < mc: - kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) - C.atomic_min_cta(s_kmm.iterator + 0, kk) - C.atomic_max_cta(s_kmm.iterator + 1, kk) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2257 ---- - rlo = s_kmm[0] - rhi = s_kmm[1] - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - aboveC = cutlass.Int32(0) - needC = need - mm = mc - brk = cutlass.Int32(0) - lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L2258-2278 (<=6 lvl) - if needC == mm: - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) - aboveC = aboveC + mm - needC = cutlass.Int32(0) - brk = cutlass.Int32(1) - elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - brk = cutlass.Int32(1) - elif lev >= cutlass.Int32(6): - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - brk = cutlass.Int32(1) - else: - d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) - b2_ = cutlass.Int32(32) - C.clz_i32( - cutlass.Int32(d2 | cutlass.Uint32(1)) - ) - sh2 = b2_ - cutlass.Int32(self.lb) - if sh2 < cutlass.Int32(0): - sh2 = cutlass.Int32(0) - sh2u = cutlass.Uint32(sh2) - i = tidx - while i < mc: # L2266-2269 re-bin - uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) - if uq >= cutlass.Uint32(rlo): - if uq <= cutlass.Uint32(rhi): - du = (uq - cutlass.Uint32(rlo)) >> sh2u - if du > cutlass.Uint32(NBS - 1): - du = cutlass.Uint32(NBS - 1) - C.atomic_add_cta( - s_hist.iterator + cutlass.Int32(du), - cutlass.Int32(1), - ) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2270 ---- - C.scan_cross0( - s_hist, + g = g + cutlass.Int32(CS) + t2 = tidx + while t2 < tailn: + ii = tail0 + t2 + x = C.ldg_f32(x_addr, ii) + if x >= TF: + self._p5_emit( + x, ii, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 + ) + t2 = t2 + cutlass.Int32(BLK) + + # ============ EXIT RENDEZVOUS (L2226) ============ + # all DSMEM traffic retired; the ONLY exit rendezvous. rank!=0 falls + # through to the kernel end (post-barrier asymmetric exit L2227); + # NO later cluster barrier. + C._cluster_sync_aligned() # ==== clus.sync L2226 ==== + + if rank == cutlass.Int32(0): + if degen == cutlass.Int32(0): + if whole == cutlass.Int32(0): + # ---- P6 rank-0 refine (L2229-2303) ---- + if mc <= cutlass.Int32(QUADC_CLUS__clus): # L2232-2247 O(mc^2) + mc2 = mc & cutlass.Int32(~1) + i = tidx + while i < mc: + # G1: re-assert Uint64 at every unsigned compare + # in/after dynamic loops (ct_main notes). + u64v = s_ck64[i] + r_ = cutlass.Int32(0) + jq = cutlass.Int32(0) + while jq < mc2: # ulonglong2 16B reads + vlo, vhi = C._lds_v2_u64(ck_addr + jq * cutlass.Int32(8)) + r_ = ( + r_ + + cutlass.Int32(vlo > cutlass.Uint64(u64v)) + + cutlass.Int32(vhi > cutlass.Uint64(u64v)) + ) + jq = jq + cutlass.Int32(2) + if mc2 < mc: # odd tail L2244 + r_ = r_ + cutlass.Int32( + cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) + ) + if r_ < need: + out_row[above + r_] = cutlass.Int32( + cutlass.Uint32( + cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) + ) + ) + i = i + cutlass.Int32(BLK) + else: + # key-space narrowing over ck64c (L2249-2278) + if tidx == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + if tidx < cutlass.Int32(NBS): # cleared ONCE L2251 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2255 ---- + i = tidx + while i < mc: + kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + C.atomic_min_cta(s_kmm.iterator + 0, kk) + C.atomic_max_cta(s_kmm.iterator + 1, kk) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2257 ---- + rlo = s_kmm[0] + rhi = s_kmm[1] + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + aboveC = cutlass.Int32(0) + needC = need + mm = mc + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L2258-2278 (<=6 lvl) + if needC == mm: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + aboveC = aboveC + mm + needC = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(6): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + i = tidx + while i < mc: # L2266-2269 re-bin + uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), + cutlass.Int32(1), + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2270 ---- + C.scan_cross0( + s_hist, + needC, + tidx, + s_res, + cutlass.Int32(0), + cutlass.Int32(0), + s_hist, + nb=NBS, + zero=True, + ) + cute.arch.barrier() # ---- barrier L2273 ---- + aboveC = aboveC + s_res[C.RES_ABOVE] + needC = needC - s_res[C.RES_ABOVE] + mm = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): # L2276 + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L2279 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2281 ---- + it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L2282-2301 ballot emit + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + idv = cutlass.Int32(0) + if i < mc: + w64 = s_ck64[i] + iu = cutlass.Int64(cutlass.Uint32(w64 >> cutlass.Uint64(32))) + idv = cutlass.Int32( + cutlass.Uint32(w64 & cutlass.Uint64(0xFFFFFFFF)) + ) + if iu > ethr: + p1 = cutlass.Int32(1) + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, + p2, + idv, + above, + aboveC, + above + aboveC, needC, - tidx, - s_res, - cutlass.Int32(0), - cutlass.Int32(0), - s_hist, - nb=NBS, - zero=True, - ) - cute.arch.barrier() # ---- barrier L2273 ---- - aboveC = aboveC + s_res[C.RES_ABOVE] - needC = needC - s_res[C.RES_ABOVE] - mm = s_res[C.RES_M] - sB = s_res[C.RES_B] - nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) - if sB != cutlass.Int32(NBS - 1): # L2276 - rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) - rlo = nlo - lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L2279 - s_scal[1] = cutlass.Int32(0) - s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2281 ---- - it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) - it = cutlass.Int32(0) - while it < it2: # L2282-2301 ballot emit - i = it * cutlass.Int32(BLK) + tidx - p1 = cutlass.Int32(0) - p2 = cutlass.Int32(0) - idv = cutlass.Int32(0) - if i < mc: - w64 = s_ck64[i] - iu = cutlass.Int64(cutlass.Uint32(w64 >> cutlass.Uint64(32))) - idv = cutlass.Int32( - cutlass.Uint32(w64 & cutlass.Uint64(0xFFFFFFFF)) + out_row, + s_scal, + lane, ) - if iu > ethr: - p1 = cutlass.Int32(1) - if iu == ethr: - p2 = cutlass.Int32(1) - self._ballot_pair_emit( - p1, - p2, - idv, - above, - aboveC, - above + aboveC, - needC, - out_row, - s_scal, - lane, + it = it + cutlass.Int32(1) + else: + # ---- degen fallback: whole-row key-space narrowing + # (L2305-2355; per-level clear + scan_cross w/ ws) ---- + rlo = cutlass.Uint32(0) + rhi = cutlass.Uint32(0xFFFFFFFF) + above2 = cutlass.Int32(0) + need2 = k + m2 = n + ethr = cutlass.Int64(0) + tie_m = cutlass.Int32(1) + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # L2312-2331 (<=8 lvl) + if need2 == m2: + ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) + above2 = above2 + m2 + need2 = cutlass.Int32(0) + tie_m = cutlass.Int32(0) + brk = cutlass.Int32(1) + elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + elif lev >= cutlass.Int32(8): + ethr = cutlass.Int64(cutlass.Uint32(rlo)) + brk = cutlass.Int32(1) + else: + d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) + b2_ = cutlass.Int32(32) - C.clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + sh2 = b2_ - cutlass.Int32(self.lb) + if sh2 < cutlass.Int32(0): + sh2 = cutlass.Int32(0) + sh2u = cutlass.Uint32(sh2) + if tidx < cutlass.Int32(NBS): # per-level clear L2320 + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2321 ---- + i = tidx + while i < n: # L2322-2324 whole row + uq = C.fkey(C.ldg_f32(x_addr, i)) + if uq >= cutlass.Uint32(rlo): + if uq <= cutlass.Uint32(rhi): + du = (uq - cutlass.Uint32(rlo)) >> sh2u + if du > cutlass.Uint32(NBS - 1): + du = cutlass.Uint32(NBS - 1) + C.atomic_add_cta( + s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) + ) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # ---- barrier L2325 ---- + # block-parallel scan (ONE internal barrier; only use + # of ws in this kernel, L2326) + C.scan_cross( + s_hist, + s_ws, + need2, + tidx, + s_res, + cutlass.Int32(0), + blk=BLK, + nb=NBS, + two=False, ) - it = it + cutlass.Int32(1) - else: - # ---- degen fallback: whole-row key-space narrowing - # (L2305-2355; per-level clear + scan_cross w/ ws) ---- - rlo = cutlass.Uint32(0) - rhi = cutlass.Uint32(0xFFFFFFFF) - above2 = cutlass.Int32(0) - need2 = k - m2 = n - ethr = cutlass.Int64(0) - tie_m = cutlass.Int32(1) - brk = cutlass.Int32(0) - lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L2312-2331 (<=8 lvl) - if need2 == m2: - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) - above2 = above2 + m2 - need2 = cutlass.Int32(0) - tie_m = cutlass.Int32(0) - brk = cutlass.Int32(1) - elif cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - brk = cutlass.Int32(1) - elif lev >= cutlass.Int32(8): - ethr = cutlass.Int64(cutlass.Uint32(rlo)) - brk = cutlass.Int32(1) - else: - d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) - b2_ = cutlass.Int32(32) - C.clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) - sh2 = b2_ - cutlass.Int32(self.lb) - if sh2 < cutlass.Int32(0): - sh2 = cutlass.Int32(0) - sh2u = cutlass.Uint32(sh2) - if tidx < cutlass.Int32(NBS): # per-level clear L2320 - s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2321 ---- - i = tidx - while i < n: # L2322-2324 whole row + cute.arch.barrier() # ---- barrier L2327 ---- + above2 = above2 + s_res[C.RES_ABOVE] + need2 = need2 - s_res[C.RES_ABOVE] + m2 = s_res[C.RES_M] + sB = s_res[C.RES_B] + nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) + if sB != cutlass.Int32(NBS - 1): + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): # L2336 + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier L2337 ---- + nA = k # L2338 + nT = cutlass.Int32(0) + if tie_m != cutlass.Int32(0): + nA = above2 + nT = need2 + it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # L2340-2354 + i = it * cutlass.Int32(BLK) + tidx + p1 = cutlass.Int32(0) + p2 = cutlass.Int32(0) + if i < n: uq = C.fkey(C.ldg_f32(x_addr, i)) - if uq >= cutlass.Uint32(rlo): - if uq <= cutlass.Uint32(rhi): - du = (uq - cutlass.Uint32(rlo)) >> sh2u - if du > cutlass.Uint32(NBS - 1): - du = cutlass.Uint32(NBS - 1) - C.atomic_add_cta( - s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) - ) - i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2325 ---- - # block-parallel scan (ONE internal barrier; only use - # of ws in this kernel, L2326) - C.scan_cross( - s_hist, - s_ws, - need2, - tidx, - s_res, - cutlass.Int32(0), - blk=BLK, - nb=NBS, - two=False, + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) + self._ballot_pair_emit( + p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane ) - cute.arch.barrier() # ---- barrier L2327 ---- - above2 = above2 + s_res[C.RES_ABOVE] - need2 = need2 - s_res[C.RES_ABOVE] - m2 = s_res[C.RES_M] - sB = s_res[C.RES_B] - nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) - if sB != cutlass.Int32(NBS - 1): - rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) - rlo = nlo - lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L2336 - s_scal[1] = cutlass.Int32(0) - s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2337 ---- - nA = k # L2338 - nT = cutlass.Int32(0) - if tie_m != cutlass.Int32(0): - nA = above2 - nT = need2 - it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) - it = cutlass.Int32(0) - while it < it2: # L2340-2354 - i = it * cutlass.Int32(BLK) + tidx - p1 = cutlass.Int32(0) - p2 = cutlass.Int32(0) - if i < n: - uq = C.fkey(C.ldg_f32(x_addr, i)) - iu = cutlass.Int64(uq) - if iu > ethr: - p1 = cutlass.Int32(1) - if tie_m != cutlass.Int32(0): - if iu == ethr: - p2 = cutlass.Int32(1) - self._ballot_pair_emit( - p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane - ) - it = it + cutlass.Int32(1) + it = it + cutlass.Int32(1) # ------------------------------------------------------------------ # host launcher: grid dim3(CS, b) + cluster (CS,1,1) (probe P11); @@ -5311,6 +5447,7 @@ def __call__( self, logits: cute.Tensor, pre_idx: cute.Tensor, + kv_lens: cute.Tensor, out: cute.Tensor, n: cutlass.Int32, npad: cutlass.Int32, @@ -5324,8 +5461,16 @@ def __call__( TGT2: cutlass.Int32, stream, ): - b = logits.shape[0] - self.kern(logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2).launch( + # varlen grid rows = out.shape[0] (logits row count == out row count in + # both modes); bigf = the route() `big` occupancy flag, a pure function + # of (rows, CS) so it is launch-computed, not an ABI scalar. + b = out.shape[0] + bigf = cutlass.Int32(0) + if b * cutlass.Int32(self.cs) <= cutlass.Int32(148): + bigf = cutlass.Int32(1) + self.kern( + logits, pre_idx, kv_lens, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2, bigf + ).launch( grid=(self.cs, b, 1), block=(self.blk, 1, 1), cluster=(self.cs, 1, 1), @@ -5340,16 +5485,26 @@ def __call__( _COMPILE_CACHE__clus = {} -def get_compiled__clus(tpl, scap: int = 8192, cmp_: int = 2048, options_extra: str = ""): +def get_compiled__clus( + tpl, + scap: int = 8192, + cmp_: int = 2048, + options_extra: str = "", + varlen: bool = False, + next_n: int = 1, + cr_shift: int = 0, +): """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every reachable route has 8192/2048 — asserted by run__clus()).""" - key = (tuple(tpl), scap, cmp_, options_extra) + key = (tuple(tpl), scap, cmp_, options_extra, bool(varlen), int(next_n), int(cr_shift)) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: return hit blk, u, minb, nbs, cs = tpl - kern = GvrClusKernel(blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_) + kern = GvrClusKernel( + blk, u, minb, nbs, cs, scap=scap, cmp_=cmp_, varlen=varlen, next_n=next_n, cr_shift=cr_shift + ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() r2, c2 = cute.sym_int(), cute.sym_int() @@ -5362,11 +5517,16 @@ def get_compiled__clus(tpl, scap: int = 8192, cmp_: int = 2048, options_extra: s out_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (r2, c2), stride_order=(1, 0), assumed_align=16 ) + v0 = cute.sym_int() + kv_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (v0,), stride_order=(0,), assumed_align=4 + ) fake_stream = _crt.make_fake_stream(use_tvm_ffi_env_stream=True) compiled = cute.compile( kern, logits_fake, pre_fake, + kv_fake, out_fake, *([cutlass.Int32(0)] * 10), stream=fake_stream, @@ -5392,9 +5552,11 @@ def run__clus(logits, pre_idx, n: int, out): kobj = GvrClusKernel(*r["tpl"], scap=rt["SCAP"], cmp_=rt["CMP"]) assert r["smem"] == kobj.dyn_bytes, (r["smem"], kobj.dyn_bytes) fn = get_compiled__clus(tuple(r["tpl"]), scap=rt["SCAP"], cmp_=rt["CMP"]) + dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot fn( logits, pre_idx, + dkv, out, rt["n"], rt["npad"], @@ -5415,9 +5577,11 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): tpl with caller-supplied runtime scalars (must be route()-consistent for the same CS; U only changes the chunk geometry).""" fn = get_compiled__clus(tuple(tpl), scap=rt["SCAP"], cmp_=rt["CMP"]) + dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot fn( logits, pre_idx, + dkv, out, rt["n"], rt["npad"], diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index a02c015ac4f0..c761aca49f78 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -732,7 +732,7 @@ def _main(blk_, minb_, u_, split_): def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): """Capture-time varlen plan + compiled launcher (v1 = the gvr_main port, - universally correct across the envelope; the clus port follows). Every + universally correct across the envelope; specialist tiers below). Every choice here is a function of capture-stable quantities only — mirroring the in-tree runner's pick_tuning(graph_capture=...) discipline.""" key = (num_rows, npad, k, n_env, next_n, cr) @@ -776,6 +776,30 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): ) _VARLEN_CACHE[key] = lc return lc + # ---- route() parity, family tier 3: cluster split (clus) --------------- + # Same admission rule: exactly where the free route picks clus (the + # large-N mid-rows band). SCAP/CMP are launch-stable (pure functions of + # rows/CS/k — never of n) so the envelope values are the per-row values; + # the sampling-ladder scalars (SMP/TGT/Q/SS2/TGT2) are dead launch slots, + # re-derived per row in-kernel by the route_dynamic clus mirror + # (GvrMainKernel discipline). Per-row n / short-row handling in-kernel. + if plan_free["kernel"] == "clus": + rt_f = plan_free["rt"] + fn = dev.get_compiled__clus( + tuple(plan_free["tpl"]), + scap=rt_f["SCAP"], + cmp_=rt_f["CMP"], + varlen=True, + next_n=next_n, + cr_shift=cr_shift, + ) + lc = ( + "clus", + fn, + (n_eff, npad, k, rt_f["SCAP"], rt_f["CMP"], 0, 0, 0, 0, 0), + ) + _VARLEN_CACHE[key] = lc + return lc plan = route_streaming(num_rows, n_eff, npad, k, force_main=True) tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) rt = plan["rt"] @@ -1079,8 +1103,9 @@ def fn(lg, pi, o, w, *a, _raw=raw): if fam == "clus": dev = _device() # compile key carries the smem-extent scalars (scap/cmp_); compiled - # ABI: (logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, - # SS2, TGT2) -- NO workspace (spec §4c) + # ABI: (logits, pre_idx, kv_lens, out, n, npad, k, SCAP, CMP, SMP, + # TGT, Q, SS2, TGT2) -- NO workspace (spec §4c); kv_lens is the + # dead varlen slot in batch-uniform mode (dummy, reg_clus precedent) fn = dev.get_compiled__clus(tpl, scap=rt["SCAP"], cmp_=rt["CMP"]) args = ( rt["n"], @@ -1094,7 +1119,11 @@ def fn(lg, pi, o, w, *a, _raw=raw): rt["SS2"], rt["TGT2"], ) - return (fn, args, False) + + def _call(lg, pi, idx, _fn=fn, _args=args): + _fn(lg, pi, _dummy_kv(lg.get_device(), lg.device), idx, *_args) + + return (_call, (), False) if fam == "reg_clus": dev = _device() # compiled ABI: (logits, pre_idx, kv_lens, out, n) -- kv_lens is the @@ -1452,6 +1481,10 @@ def run_varlen( elif lc[0] == "reg": # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, CMP, QC, smem) lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) + elif lc[0] == "clus": + # compiled ABI: (logits, pre_idx, kv_lens, out, n_env, npad, k, + # SCAP, CMP, dead DYN x5) + lc[1](lg, pre_idx, kv_lens, idx, *lc[2]) else: _, fn, pre, tail = lc fn(lg, pre_idx, idx, ws, *pre, kv_lens, *tail) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index b380fac507ba..402c9bb29522 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -781,6 +781,85 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): assert torch.equal(out[r], ref[r]), f"k={k} row {r} expected all -1" +def test_selfsampling_varlen_clus_parity_and_oracle(): + """The varlen launcher must admit the cluster-split family exactly where + the free route picks it (route() parity tier 3), and the per-row varlen + port must match the reference oracle on a heterogeneous batch: full rows, + mid rows BELOW the family's standalone admission floor (n <= SCAP, + exercising the always-on per-row QUAD schedule), a short row (n <= k, + in-kernel identity + -1 tail from cluster rank 0) and a zero-window row, + under MTP row windows (next_n=4). Covers CS=2 and CS=4 clusters.""" + torch.manual_seed(23) + nn = 4 + # (rows, msl_c, k, expected cluster size) + cases = [(64, 131072, 1024, 2), (32, 131072, 1024, 4)] + cr = 4 + for rows, msl_c, k, want_cs in cases: + npad = msl_c + plan = ss_host.route(rows, msl_c, npad, k) + assert plan["kernel"] == "clus" and plan["cluster"] == want_cs, plan + batch = rows // nn + lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) + lens = [msl_c * cr, 900, nn - 1, 20000, 40000, 300000, msl_c * cr // 2, 5000] + kv = torch.tensor( + [lens[i % len(lens)] for i in range(batch)], dtype=torch.int32, device=_DEV + ) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) + key = (rows, npad, k, msl_c, nn, cr) + assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] + ss_host.run_varlen( + lg, + pre, + kv, + ref, + next_n=nn, + compress_ratio=cr, + max_seq_len=msl_c * cr, + engine="reference", + ) + torch.cuda.synchronize() + for r in range(rows): + if (ref[r] >= 0).any(): + row = lg[r].float() + got = row[out[r].long().clamp_min(0)].sort().values + want = row[ref[r].long().clamp_min(0)].sort().values + assert torch.equal(got, want), f"cs={want_cs} row {r} value multiset mismatch" + assert torch.equal(out[r] < 0, ref[r] < 0), f"cs={want_cs} row {r} pad mask" + else: + assert torch.equal(out[r], ref[r]), f"cs={want_cs} row {r} expected all -1" + + +def test_selfsampling_varlen_clus_cuda_graph(): + """Cluster-split-family varlen engine must be CUDA-graph capturable: + warmed engine, capture one launch, replay twice, tie-aware exact each + time.""" + k, msl_c, cr = 1024, 131072, 4 + rows = 32 + torch.manual_seed(29) + lg = torch.randn(rows, msl_c, dtype=torch.float32, device=_DEV) + pre = torch.randint(0, msl_c, (rows, k), dtype=torch.int32, device=_DEV) + kv = torch.full((rows,), msl_c * cr, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + torch.cuda.synchronize() + key = (rows, msl_c, k, msl_c, 1, cr) + assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] + g = torch.cuda.CUDAGraph() + out.fill_(-7) + with torch.cuda.graph(g): + ss_host.run_varlen(lg, pre, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) + ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values + for _ in range(2): + out.fill_(-7) + g.replay() + torch.cuda.synchronize() + got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values + assert torch.equal(got, ref_v) + + def test_selfsampling_varlen_reg_cuda_graph(): """Register-family varlen engine must be CUDA-graph capturable: warmed engine, capture one launch, replay twice, tie-aware exact each time.""" From 80484e8053b713c3814b4175c4587e89d9792f41 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 10:29:46 +0000 Subject: [PATCH 31/40] [TRTLLM-15293][perf] self-sampling top-K: cheapen the per-row varlen mirror chain The per-row re-derivation of the sampling-ladder scalars was the dominant varlen-vs-standalone kernel tax (l1p1 ncu: inst x1.58-1.73 on main shapes, x1.26-1.28 on clus -- all of it the DSL's runtime Int32/Int64 divides and isqrt fixup loops, redundantly issued by every thread). Two fixes, chosen per family by measured mirror share: - main: warp0 alone walks the ladder and publishes SMP/SS2/TGT/TGT2 through a 4-word smem block; the other warps spend the wait issuing register-free L2 hints for this CTA's own P3 slice, then everyone reads the scalars back after one CTA barrier. The chain itself swaps runtime divides for MUFU.RCP multiplies and collapses the isqrt fixup loops to single steps (f32 sqrt of an exactly-representable int is within 1). - clus: same cheap-chain spelling, kept all-thread in place. This family's mirror share is small, so a warp0+barrier hoist EXPOSES the chain's serial latency at ~1 CTA/SM (measured 1.14->1.31 tax regression) while the redundant form hides it across warps. Schedule quantities may drift +-1 vs the host double form; exactness is schedule-invariant (same argument as the clus port's always-sample deviation). A one-step window guard keeps SMP*SS2 inside the sampled row span. Q (chunk ownership) keeps its exact form in both families. Legacy (batch-uniform) codegen is untouched: every change sits under const_expr(varlen), including the smem slots. nsys verdict (interleaved eager sa/vl, cold-L2, 200 paired reps/cell, umbriel-b200-027): v32_128k r1 1.165->1.093, pro_1024k r8 1.184->1.103, pro_1024k r32 1.135->1.083, pro_512k r64 1.131->1.071; reg_clus control cells unchanged. Treated-cell tax gm 1.154->1.087. UT 84/84; exactness smoke (main x3 + clus x2 + short/long heterogeneous kv_lens) all exact under the tie-aware value-multiset oracle. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 254 ++++++++++++------ 1 file changed, 173 insertions(+), 81 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index f334f192fdcc..2d6a19521bf1 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1548,80 +1548,15 @@ def kern( Q = cutlass.Int32(0) if short == cutlass.Int32(0): n = nv - # ---- aim ladder (route_dynamic mirror) ---- - # r6 = int(0.5 + sqrt(6n)) computed EXACTLY in integers: f32 - # sqrt seed, fixup to isqrt, then round-half-up via the - # (x - r*r > r) test (bit-parity with the host double form, - # fuzz-proven over the whole n domain). - x6 = cutlass.Int32(6) * n - ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) - while ri * ri > x6: - ri = ri - cutlass.Int32(1) - while (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: - ri = ri + cutlass.Int32(1) - r6 = ri - if x6 - ri * ri > ri: - r6 = ri + cutlass.Int32(1) - aim = aim_base - if r6 > aim: - aim = r6 - if cutlass.const_expr(self.r_const > 1): - if aim < amin: - aim = amin - scap_c = cutlass.Int32(SCPB) # SCAP == SCPB for gvr_main (proven identity) - if aim > (scap_c >> cutlass.Int32(1)): - aim = scap_c >> cutlass.Int32(1) - if aim < k: - aim = k n4v = n >> cutlass.Int32(2) + # Ladder-scalar baselines only: the real SMP/SS2/TGT/TGT2 are + # derived by warp0 alone in the P2a block below (bit-identical + # formulas) and published through s_lad — every thread's local + # copies here are overwritten by the post-barrier smem read. SMP = cutlass.Int32(0) SS2 = cutlass.Int32(1) TGT = cutlass.Int32(0) TGT2 = cutlass.Int32(0) - # pair-sample gate: (n > SCAP or small_dense) and n4 >= 4; - # small_dense = k > 1024 and not big and n <= SCAP and n > 2k - # (k/big folded into the launch-constant sd_en flag). - gate = cutlass.Int32(0) - if n > scap_c: - gate = cutlass.Int32(1) - if sd_en != cutlass.Int32(0): - if n <= scap_c: - if n > (k << cutlass.Int32(1)): - gate = cutlass.Int32(1) - if n4v < cutlass.Int32(4): - gate = cutlass.Int32(0) - if gate != cutlass.Int32(0): - sel = sfac * n // aim - if sel < cutlass.Int32(256): - sel = cutlass.Int32(256) - nh = n >> cutlass.Int32(1) - if sel > nh: - sel = nh - pairs = sel >> cutlass.Int32(3) - if pairs < cutlass.Int32(1): - pairs = cutlass.Int32(1) - half = n4v >> cutlass.Int32(1) - if half < cutlass.Int32(1): - half = cutlass.Int32(1) - if pairs > half: - pairs = half - SS2 = half // pairs - if SS2 < cutlass.Int32(1): - SS2 = cutlass.Int32(1) - SMP = half // SS2 - if SMP < cutlass.Int32(1): - SMP = cutlass.Int32(1) - # TGT/TGT2 are 64-bit products in the CUDA host (aim*SMP*8 - # overflows i32 at large n) — mirror with Int64. - smp8 = cutlass.Int64(SMP) * cutlass.Int64(8) - tgt64 = cutlass.Int64(aim) * smp8 // cutlass.Int64(n) - TGT = cutlass.Int32(tgt64) - if TGT < cutlass.Int32(1): - TGT = cutlass.Int32(1) - tgt264 = cutlass.Int64(k) * smp8 // cutlass.Int64(n) - TGT2 = cutlass.Int32(tgt264) - if TGT2 < cutlass.Int32(1): - TGT2 = cutlass.Int32(1) if cutlass.const_expr(self.split): Q = (n4v + cutlass.Int32(self.r_const - 1)) // cutlass.Int32(self.r_const) else: @@ -1671,6 +1606,12 @@ def kern( s_kmm = smem.allocate_tensor( # L463 [0]=kmin [1]=kmax cutlass.Uint32, cute.make_ordered_layout((2,), order=(0,)), byte_alignment=8 ) + if cutlass.const_expr(self.varlen): + # P2a ladder broadcast slots: [0]=SMP [1]=SS2 [2]=TGT [3]=TGT2 + # (static like s_x4, so dyn_bytes keeps CUDA dispatch parity) + s_lad = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 + ) blob = smem.allocate_tensor( # dynamic-equivalent L447-456 cutlass.Int8, cute.make_ordered_layout((self.dyn_bytes,), order=(0,)), byte_alignment=16 ) @@ -1762,6 +1703,137 @@ def kern( if tidx < cutlass.Int32(self.hb): # L484-487 (HB<=BLK always) s_hist[tidx] = cutlass.Int32(0) + # ===== varlen P2a: warp0-only ladder mirror + register-free L2 hints = + # The sampling-ladder scalars are a pure function of the row, and the + # per-thread mirror chain cost more instructions than the rest of the + # kernel on 1-row launches (l1p1 ncu: inst x1.58-1.73, all of it the + # DSL's runtime Int32/Int64 divides + isqrt fixups, redundantly issued + # by every thread). warp0 alone walks the chain and publishes the four + # derived scalars through s_lad; the other warps spend the wait issuing + # L2 prefetch hints for this CTA's own P3 slice (register-free, so zero + # pressure on the 64-reg arms — the PRIME-LATE register loads below are + # untouched and simply hit L2). Values are bit-identical to the + # per-thread derivation this replaces. + if cutlass.const_expr(self.varlen): + if tidx < cutlass.Int32(32): + if short == cutlass.Int32(0): + # ---- aim ladder (P2b cheap mirror) ---- + # The ladder scalars steer the sampling rung only — + # exactness is schedule-invariant (retry/degen close every + # miss), so schedule-quantity drift of +-1 vs the host + # double form is sanctioned (same argument as the clus + # port's always-sample deviation). Serial latency is what + # matters here (this chain sits in front of a barrier): + # runtime divides become MUFU.RCP multiplies and the isqrt + # fixup loops collapse to single steps (the f32 sqrt of an + # exactly-representable int (6n <= 2^23) is within 1 of + # isqrt, so one correction per side suffices). Q (chunk + # ownership) stays exact — compile-time divisor, all-thread. + x6 = cutlass.Int32(6) * n + ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) + if ri * ri > x6: + ri = ri - cutlass.Int32(1) + if (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: + ri = ri + cutlass.Int32(1) + r6 = ri + if x6 - ri * ri > ri: + r6 = ri + cutlass.Int32(1) + aim = aim_base + if r6 > aim: + aim = r6 + if cutlass.const_expr(self.r_const > 1): + if aim < amin: + aim = amin + scap_c = cutlass.Int32(SCPB) # SCAP == SCPB for gvr_main (proven identity) + if aim > (scap_c >> cutlass.Int32(1)): + aim = scap_c >> cutlass.Int32(1) + if aim < k: + aim = k + n4w = n >> cutlass.Int32(2) + # pair-sample gate: (n > SCAP or small_dense) and n4 >= 4; + # small_dense = k > 1024 and not big and n <= SCAP and n > 2k + # (k/big folded into the launch-constant sd_en flag). + gate = cutlass.Int32(0) + if n > scap_c: + gate = cutlass.Int32(1) + if sd_en != cutlass.Int32(0): + if n <= scap_c: + if n > (k << cutlass.Int32(1)): + gate = cutlass.Int32(1) + if n4w < cutlass.Int32(4): + gate = cutlass.Int32(0) + if gate != cutlass.Int32(0): + # sel = sfac*n // aim via rcp (sfac*n <= 2^24: f32-exact + # to the last unit; quotient error < 1 => +-1 drift) + sel = cutlass.Int32( + cutlass.Float32(sfac * n) * cute.arch.rcp_approx(cutlass.Float32(aim)) + ) + if sel < cutlass.Int32(256): + sel = cutlass.Int32(256) + nh = n >> cutlass.Int32(1) + if sel > nh: + sel = nh + pairs = sel >> cutlass.Int32(3) + if pairs < cutlass.Int32(1): + pairs = cutlass.Int32(1) + half = n4w >> cutlass.Int32(1) + if half < cutlass.Int32(1): + half = cutlass.Int32(1) + if pairs > half: + pairs = half + SS2 = cutlass.Int32( + cutlass.Float32(half) * cute.arch.rcp_approx(cutlass.Float32(pairs)) + ) + if SS2 < cutlass.Int32(1): + SS2 = cutlass.Int32(1) + SMP = cutlass.Int32( + cutlass.Float32(half) * cute.arch.rcp_approx(cutlass.Float32(SS2)) + ) + # sample-window guard: the P1 gather indexes up to + # ~SMP*SS2*2 f32x4 lines; keep SMP*SS2 <= half so the + # window never walks past the row (approx error is + # bounded by +1, one decrement closes it) + if SMP * SS2 > half: + SMP = SMP - cutlass.Int32(1) + if SMP < cutlass.Int32(1): + SMP = cutlass.Int32(1) + # TGT/TGT2: i64 products // n -> f32 mul + one rcp(n). + # aim/SMP/k/n are all f32-exact here (<= 2^20); the + # quotients are <= 8*aim ~ 2^16, so the approx error + # stays far below 1 unit — +-1 at worst on the floor. + rn_ = cute.arch.rcp_approx(cutlass.Float32(n)) + smp8f = cutlass.Float32(SMP) * cutlass.Float32(8.0) + TGT = cutlass.Int32(cutlass.Float32(aim) * smp8f * rn_) + if TGT < cutlass.Int32(1): + TGT = cutlass.Int32(1) + TGT2 = cutlass.Int32(cutlass.Float32(k) * smp8f * rn_) + if TGT2 < cutlass.Int32(1): + TGT2 = cutlass.Int32(1) + if tidx == cutlass.Int32(0): + s_lad[0] = SMP + s_lad[1] = SS2 + s_lad[2] = TGT + s_lad[3] = TGT2 + # Register-free L2 hints for the first U-batch of this CTA's own + # P3 slice (clamped in-row, prefetch site #3 spelling): the data + # P3 touches first starts flowing while warp0 walks the chain. + # Short rows clamp every hint to the row's last line — harmless. + plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) + for uu in cutlass.range_constexpr(U): + # NOTE: names must not collide with the PRIME-LATE block's + # i_/ic — the DSL kills inner-scope names at region exit and + # a later same-name assignment inside a dynamic `if` trips + # "is None prior to this if". + pic = c0 + tidx + cutlass.Int32(uu * BLK) + if pic >= c1: + pic = plim4 + C._prefetch_l2(x_addr + cutlass.Int64(pic) * cutlass.Int64(16)) + cute.arch.barrier() # publish s_lad (also covers the smem inits) + SMP = s_lad[0] + SS2 = s_lad[1] + TGT = s_lad[2] + TGT2 = s_lad[3] + # ============ P1: sample prefetch (hint gather LAZY, L489-529) ======= atom128 = C.g2r_atom_f32(128, invariant=True) fsa = cute.make_fragment((4,), cutlass.Float32) @@ -4587,13 +4659,22 @@ def kern( short = cutlass.Int32(1) if short == cutlass.Int32(0): n = nv - # ---- aim ladder (route_dynamic mirror; isqrt discipline and - # Int64 target products exactly as the GvrMainKernel prologue). + # ---- aim ladder (P2b cheap mirror, all-thread) ---- + # Schedule quantities only (exactness is schedule-invariant, + # same argument as the always-sample deviation above): divides + # become MUFU.RCP multiplies and the isqrt fixup loops collapse + # to single steps (f32 sqrt of an exactly-representable int + # (6n <= 2^23) is within 1 of isqrt). All-thread on purpose: + # this family's mirror redundancy is small (inst x1.26-1.28) + # and a warp0+barrier hoist EXPOSES the chain's serial latency + # at ~1 CTA/SM — measured 1.14->1.31 tax regression on + # pro_1024k r32 — while the redundant form hides it across + # warps. Q (chunk ownership) keeps its exact shift form. x6 = cutlass.Int32(6) * nv ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) - while ri * ri > x6: + if ri * ri > x6: ri = ri - cutlass.Int32(1) - while (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: + if (ri + cutlass.Int32(1)) * (ri + cutlass.Int32(1)) <= x6: ri = ri + cutlass.Int32(1) r6 = ri if x6 - ri * ri > ri: @@ -4621,7 +4702,11 @@ def kern( # always-on per the deviation note above; k <= 1024 for this # family so sfac has no k > 1024 arm). n4v = nv >> cutlass.Int32(2) - sel = cutlass.Int32(32 if self.cs == 2 else 16) * nv // aim + sel = cutlass.Int32( + cutlass.Float32(nv) + * cutlass.Float32(32.0 if self.cs == 2 else 16.0) + * cute.arch.rcp_approx(cutlass.Float32(aim)) + ) if sel < cutlass.Int32(256): sel = cutlass.Int32(256) nh = nv >> cutlass.Int32(1) @@ -4635,19 +4720,26 @@ def kern( quarter = cutlass.Int32(1) if quads > quarter: quads = quarter - SS2 = quarter // quads + SS2 = cutlass.Int32( + cutlass.Float32(quarter) * cute.arch.rcp_approx(cutlass.Float32(quads)) + ) if SS2 < cutlass.Int32(1): SS2 = cutlass.Int32(1) - SMP = quarter // SS2 + SMP = cutlass.Int32( + cutlass.Float32(quarter) * cute.arch.rcp_approx(cutlass.Float32(SS2)) + ) + # sample-window guard: P1 indexes up to ~SMP*SS2*4 lines; keep + # SMP*SS2 <= quarter (approx error <= +1, one step closes it) + if SMP * SS2 > quarter: + SMP = SMP - cutlass.Int32(1) if SMP < cutlass.Int32(1): SMP = cutlass.Int32(1) - smp16 = cutlass.Int64(SMP) * cutlass.Int64(16) - tgt64 = cutlass.Int64(aim) * smp16 // cutlass.Int64(nv) - TGT = cutlass.Int32(tgt64) + rn_ = cute.arch.rcp_approx(cutlass.Float32(nv)) + smp16f = cutlass.Float32(SMP) * cutlass.Float32(16.0) + TGT = cutlass.Int32(cutlass.Float32(aim) * smp16f * rn_) if TGT < cutlass.Int32(1): TGT = cutlass.Int32(1) - tgt264 = cutlass.Int64(k) * smp16 // cutlass.Int64(nv) - TGT2 = cutlass.Int32(tgt264) + TGT2 = cutlass.Int32(cutlass.Float32(k) * smp16f * rn_) if TGT2 < cutlass.Int32(1): TGT2 = cutlass.Int32(1) Q = (n4v + cutlass.Int32(CS - 1)) >> cutlass.Int32(self.lcs) From 7fe0d51c72968435682fa78bc27627b114096af5 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:27:58 +0000 Subject: [PATCH 32/40] [TRTLLM-15293][chore] self-sampling top-K: pre-commit formatting for the TopK module edit Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/top_k.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 6e0320e44a3f..950b02c30f93 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -279,9 +279,7 @@ def _forward_decode_gvr( and scores.data_ptr() % 16 == 0 and (scores.shape[0] > 1 or scores.shape[1] % 4 == 0) ): - from ..cute_dsl_kernels.blackwell.top_k import ( - selfsampling_topk_run_varlen, - ) + from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_varlen logger.info_once( "self-sampling GVR top-K engaged " From e126c4cd7f6d1f55efc1f2f52b5c12486e7cfec2 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:42:27 +0000 Subject: [PATCH 33/40] [TRTLLM-15293][chore] self-sampling top-K: lint fixes for the kernel module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the dead per-row Q re-derivation in the clus varlen prologue (the Q launch slot has no in-kernel consumer in this family; chunk ownership derives from n4/STEPC — ruff F841), give the two torch-facing debug entries local torch imports (the DSL module is deliberately torch-free at import time — ruff F821), and take ruff-format's layout on the unit test. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../blackwell/top_k/gvr_topk_decode_self_sampling.py | 7 ++++++- .../_torch/thop/parallel/test_gvr_selfsampling_topk.py | 4 +--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 2d6a19521bf1..d38232d3d8e9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -4742,7 +4742,8 @@ def kern( TGT2 = cutlass.Int32(cutlass.Float32(k) * smp16f * rn_) if TGT2 < cutlass.Int32(1): TGT2 = cutlass.Int32(1) - Q = (n4v + cutlass.Int32(CS - 1)) >> cutlass.Int32(self.lcs) + # NB: the Q launch slot is dead in this family (no in-kernel + # consumer); chunk ownership derives from n4/STEPC below. if short != cutlass.Int32(0): if rank == cutlass.Int32(0): if tidx < kq: @@ -5632,6 +5633,8 @@ def run__clus(logits, pre_idx, n: int, out): """torch-facing single-call entry: routes (b, n, k) through ct_dispatch, asserts the shape lands on gvr_clus, launches the matching variant. gvr_clus takes NO workspace (spec §4c).""" + import torch # debug-entry only: module stays torch-free at import + try: from . import gvr_topk_decode_self_sampling_host as ct_dispatch except ImportError: @@ -5668,6 +5671,8 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): """Manual-lattice entry for route()-unreachable (U, CS) members: launches tpl with caller-supplied runtime scalars (must be route()-consistent for the same CS; U only changes the chunk geometry).""" + import torch # debug-entry only: module stays torch-free at import + fn = get_compiled__clus(tuple(tpl), scap=rt["SCAP"], cmp_=rt["CMP"]) dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot fn( diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 402c9bb29522..a9d0a7987596 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -751,9 +751,7 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): msl = msl_c * cr lg = torch.randn(rows, npad, dtype=torch.float32, device=_DEV) pre = torch.randint(0, msl_c, (batch, k), dtype=torch.int32, device=_DEV) - kv = torch.tensor( - [msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV - ) + kv = torch.tensor([msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ref = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, pre, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) From 3584c42f81e44ea540ca046085615029490c7aeb Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:16:07 +0000 Subject: [PATCH 34/40] [TRTLLM-15293][chore] self-sampling top-K: ruff-format line wrapping in the register family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure layout (long-line wrapping in GvrTopkRegKernel, no code change); validated against the exact pushed tree state this time — the previous lint pass ran on a local tree whose pending changes masked these lines. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 38 ++++++++++++++----- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index d38232d3d8e9..d49357340755 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -3938,7 +3938,9 @@ def kern( p = cutlass.Int32(0) if q >= LOQ: bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) - p = atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) + p = atomic_add_cta( + s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1) + ) if p < lim1: out_row[p] = idx else: @@ -4114,7 +4116,9 @@ def kern( done = cutlass.Int32(1) if done == cutlass.Int32(0): d2 = rhi - rlo - b2w = cutlass.Int32(32) - clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + b2w = cutlass.Int32(32) - clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) sh2 = cutlass.Int32(0) if b2w > cutlass.Int32(LNBH): sh2 = b2w - cutlass.Int32(LNBH) @@ -4129,7 +4133,8 @@ def kern( bnn = (unar - rlo) >> cutlass.Uint32(sh2) bnn = _umin_u32(bnn, cutlass.Uint32(self.nbh - 1)) atomic_add_cta( - s_hist.iterator + cutlass.Int32(bnn), cutlass.Int32(1) + s_hist.iterator + cutlass.Int32(bnn), + cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) cute.arch.barrier() # L1743 @@ -4147,7 +4152,8 @@ def kern( nlo = rlo + (cutlass.Uint32(b_lv) << cutlass.Uint32(sh2)) if b_lv != cutlass.Int32(self.nbh - 1): rhi = nlo + ( - (cutlass.Uint32(1) << cutlass.Uint32(sh2)) - cutlass.Uint32(1) + (cutlass.Uint32(1) << cutlass.Uint32(sh2)) + - cutlass.Uint32(1) ) rlo = nlo lev = lev + cutlass.Int32(1) @@ -4809,7 +4815,9 @@ def kern( # ---- row bases (L1830-1833; pre_idx is request-level under varlen) ---- row64 = cutlass.Int64(row) x_addr = logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) - p_addr = pre_idx.iterator.toint() + cutlass.Int64(prow) * cutlass.Int64(k) * cutlass.Int64(4) + p_addr = pre_idx.iterator.toint() + cutlass.Int64(prow) * cutlass.Int64( + k + ) * cutlass.Int64(4) out_row = out[row, None] # ---- interleaved chunk ownership (L1835-1846) ---- @@ -5084,7 +5092,9 @@ def kern( isfull = cutlass.Int32(1) if isfull != cutlass.Int32(0): # full body L2042-2049 for uu in cutlass.range_constexpr(PFD, U): - C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) + C.ld_g_f32x4( + atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD] + ) for uu in cutlass.range_constexpr(U): if cutlass.const_expr(uu < PFD): vv = pf[uu] @@ -5110,7 +5120,9 @@ def kern( okq = cutlass.Int32(1) if okq != cutlass.Int32(0): # +inf-pad escape, ok-gated for q in cutlass.range_constexpr(4): - M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + M = M | ( + cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q) + ) # ROLL THE PREFETCH FORWARD (L2066-2081): next OWNED chunk, # issued before the reservation and the survivor walk. g2 = g + cutlass.Int32(CS) @@ -5250,7 +5262,9 @@ def kern( vx = cutlass.Int32(cutlass.Uint32(pk64 & cutlass.Uint64(0xFFFFFFFF))) idv = cutlass.Int32(pk64 >> cutlass.Uint64(32)) xv = C.f32_of_i32(vx) - self._p5_emit(xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64) + self._p5_emit( + xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64 + ) i = i + cutlass.Int32(BLK) else: # ---- EXACTNESS re-sweep: OWNED CHUNKS + rank-0 true tail @@ -5395,7 +5409,9 @@ def kern( sB = s_res[C.RES_B] nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) if sB != cutlass.Int32(NBS - 1): # L2276 - rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rhi = nlo + ( + (cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1) + ) rlo = nlo lev = lev + cutlass.Int32(1) if tidx == cutlass.Int32(0): # L2279 @@ -5459,7 +5475,9 @@ def kern( brk = cutlass.Int32(1) else: d2 = cutlass.Uint32(rhi) - cutlass.Uint32(rlo) - b2_ = cutlass.Int32(32) - C.clz_i32(cutlass.Int32(d2 | cutlass.Uint32(1))) + b2_ = cutlass.Int32(32) - C.clz_i32( + cutlass.Int32(d2 | cutlass.Uint32(1)) + ) sh2 = b2_ - cutlass.Int32(self.lb) if sh2 < cutlass.Int32(0): sh2 = cutlass.Int32(0) From 6576326948713262f6938fe9af881b1b754fdafa Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:33:43 +0000 Subject: [PATCH 35/40] [TRTLLM-15293][fix] self-sampling top-K: port fragment creation to CUTLASS DSL 4.6.1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cute.make_fragment was removed in nvidia-cutlass-dsl 4.6.1 — the version this repo pins — so the self-sampling engines raised AttributeError at first JIT under the shipping dependency set (undetected by CI: the unit tests only collect on SM100/103 stages, which do not run this suite). cute.make_rmem_tensor is the same-signature replacement, already the convention in the in-tree DSL kernels, and exists in 4.5.0 as well, so every environment keeps working. Mechanical rename, 23 sites. Validated under BOTH pins: UT 84/84 on 4.6.1 (isolated overlay) and on 4.5.0; exactness smoke (main x3 + clus x2 + heterogeneous kv_lens) all exact on both; nsys spot-check in the normal range. A full-grid perf re-run under 4.6.1 is queued — all prior grid evidence compiled under 4.5.0, and codegen shifts measurably between compiler minors. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index d49357340755..8ad91bbaf4e2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -426,7 +426,7 @@ def ldg_f32(base_addr, idx): cute.AddressSpace.gmem, assumed_align=4, ) - frag = cute.make_fragment((1,), cutlass.Float32) + frag = cute.make_rmem_tensor((1,), cutlass.Float32) cute.copy(atom, cute.make_tensor(p, cute.make_layout((1,))), frag) return frag[0] @@ -788,7 +788,7 @@ def scan_cross0( atom = smem_atom_i32_128() hbase = s_hist.iterator.toint() # pass 1: span sum via NV uint4 LDS.128 (L243-251) - frags = [cute.make_fragment((4,), cutlass.Int32) for _ in range(NV)] + frags = [cute.make_rmem_tensor((4,), cutlass.Int32) for _ in range(NV)] sm = cutlass.Int32(0) for q in cutlass.range_constexpr(NV): boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) @@ -806,10 +806,10 @@ def scan_cross0( if cutlass.const_expr(HOLD): vv = frags[q] else: - vv = cute.make_fragment((4,), cutlass.Int32) # re-read span + vv = cute.make_rmem_tensor((4,), cutlass.Int32) # re-read span boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) lds128_i32(atom, hbase, boff, vv) - o4 = cute.make_fragment((4,), cutlass.Int32) + o4 = cute.make_rmem_tensor((4,), cutlass.Int32) for j in cutlass.range_constexpr(3, -1, -1): cq = vv[j] if cutlass.const_expr(zero): @@ -847,7 +847,7 @@ def scan_cross0( s_res[RES_B3] = gb after = after + cq if cutlass.const_expr(addf): # fold per-rank bin offset (L279-282) - av = cute.make_fragment((4,), cutlass.Int32) + av = cute.make_rmem_tensor((4,), cutlass.Int32) aoff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) lds128_i32(atom, s_addv.iterator.toint(), aoff, av) for j in cutlass.range_constexpr(4): @@ -933,7 +933,7 @@ def scan_cross_w(s_hist, s_ws, target, tidx, s_res, blk: cutlass.Constexpr, nb: NW = blk // 32 lane = tidx & cutlass.Int32(31) wid = tidx >> cutlass.Int32(5) - loc = cute.make_fragment((BPT,), cutlass.Int32) + loc = cute.make_rmem_tensor((BPT,), cutlass.Int32) base = tidx * cutlass.Int32(BPT) sm = cutlass.Int32(0) for i in cutlass.range_constexpr(BPT): # L297-300 (#pragma unroll) @@ -1023,7 +1023,7 @@ def merge_scan0( base = lane * cutlass.Int32(BPT) # descending walk: crossing pin + prefix-biased cursors into mrg for q in cutlass.range_constexpr(NV - 1, -1, -1): # L151-165 - o4 = cute.make_fragment((4,), cutlass.Int32) + o4 = cute.make_rmem_tensor((4,), cutlass.Int32) for j in cutlass.range_constexpr(3, -1, -1): cq = tot_r[q][j] o4[j] = after + pre_r[q][j] @@ -1278,7 +1278,7 @@ def _ldg_f32_rs(base_addr, idx, sc4): cute.AddressSpace.gmem, assumed_align=4, ) - frag = cute.make_fragment((1,), cutlass.Float32) + frag = cute.make_rmem_tensor((1,), cutlass.Float32) cute.copy(atom, cute.make_tensor(p, cute.make_layout((1,))), frag) return frag[0] @@ -1836,8 +1836,8 @@ def kern( # ============ P1: sample prefetch (hint gather LAZY, L489-529) ======= atom128 = C.g2r_atom_f32(128, invariant=True) - fsa = cute.make_fragment((4,), cutlass.Float32) - fsb = cute.make_fragment((4,), cutlass.Float32) + fsa = cute.make_rmem_tensor((4,), cutlass.Float32) + fsb = cute.make_rmem_tensor((4,), cutlass.Float32) shas = cutlass.Int32(0) if tidx < SMP: shas = cutlass.Int32(1) @@ -1856,8 +1856,8 @@ def kern( for t in cutlass.range_constexpr(4): smn = C.fmin_f32(smn, fsb[t]) smx = C.fmax_f32(smx, fsb[t]) - fma_ = cute.make_fragment((4,), cutlass.Float32) # strided-tail pair bufs - fmb_ = cute.make_fragment((4,), cutlass.Float32) + fma_ = cute.make_rmem_tensor((4,), cutlass.Float32) # strided-tail pair bufs + fmb_ = cute.make_rmem_tensor((4,), cutlass.Float32) j = tidx + cutlass.Int32(BLK) # L544-550 strided tail while j < SMP: p4 = j * SS2 * cutlass.Int32(2) @@ -1879,7 +1879,7 @@ def kern( # PRIME-LATE prefetch block (L556-616): strictly after the barrier. lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L524 - pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(max(PFD, 1))] + pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(PFD, 1))] if cutlass.const_expr(self.pf): fullsl = cutlass.Int32(0) if (c1 - c0) >= cutlass.Int32(BLK * U): @@ -2061,7 +2061,7 @@ def kern( alive = cutlass.Int32(1) fr = [ - cute.make_fragment((4,), cutlass.Float32) for _ in range(max(U - PFD, 1)) + cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(U - PFD, 1)) ] # explicit batch (op43 L1) att = cutlass.Int32(0) running = cutlass.Int32(1) @@ -3583,7 +3583,7 @@ def kern( # ---- row load: exact-fit peel + float4[VPT] register batch (L1327-1350) atom128 = g2r_atom_f32(128, invariant=True) - frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(VPT)] if n4 >= cutlass.Int32(self.blk * self.vpt): # block-uniform peel for u in cutlass.range_constexpr(VPT): ld_g_f32x4(atom128, x_addr, tid + cutlass.Int32(u * self.blk), frags[u]) @@ -4547,7 +4547,7 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): base = lane * cutlass.Int32(BPT) # descending walk: crossing pin + prefix-biased cursors into mrg for q in cutlass.range_constexpr(NV - 1, -1, -1): # L151-165 - o4 = cute.make_fragment((4,), cutlass.Int32) + o4 = cute.make_rmem_tensor((4,), cutlass.Int32) for j in cutlass.range_constexpr(3, -1, -1): cq = tot_r[q][j] o4[j] = after + pre_r[q][j] @@ -4840,8 +4840,8 @@ def kern( # one 64B line = 4 float4 per location, TWO threads: tid takes the # lower pair at p4, tid+SMP the upper pair at p4+2 (L1868-1869). atom128 = C.g2r_atom_f32(128, invariant=True) - fsa = cute.make_fragment((4,), cutlass.Float32) - fsb = cute.make_fragment((4,), cutlass.Float32) + fsa = cute.make_rmem_tensor((4,), cutlass.Float32) + fsb = cute.make_rmem_tensor((4,), cutlass.Float32) smp2 = SMP * cutlass.Int32(2) shas = cutlass.Int32(0) if tidx < smp2: @@ -4863,8 +4863,8 @@ def kern( for t in cutlass.range_constexpr(4): smn = C.fmin_f32(smn, fsb[t]) smx = C.fmax_f32(smx, fsb[t]) - fma_ = cute.make_fragment((4,), cutlass.Float32) # mop-up pair bufs - fmb_ = cute.make_fragment((4,), cutlass.Float32) + fma_ = cute.make_rmem_tensor((4,), cutlass.Float32) # mop-up pair bufs + fmb_ = cute.make_rmem_tensor((4,), cutlass.Float32) j = tidx + cutlass.Int32(BLK) # mop-up L1891-1897 while j < smp2: p4 = j * SS2 * cutlass.Int32(4) @@ -4888,7 +4888,7 @@ def kern( # PRIME-LATE (L1903-1916): every rank's sample has landed; prime NOW. lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L1904 - pf = [cute.make_fragment((4,), cutlass.Float32) for _ in range(PFD)] + pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(PFD)] for uu in cutlass.range_constexpr(PFD): # clamped prime L1906 i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) ic = i_ @@ -5037,7 +5037,7 @@ def kern( valid = cutlass.Int32(0) fr = [ - cute.make_fragment((4,), cutlass.Float32) for _ in range(U - PFD) + cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(U - PFD) ] # explicit batch (op43 L1) # (empty for U<=PFD — every row-pass float4 then comes from pf[]) att = cutlass.Int32(0) @@ -6021,7 +6021,7 @@ def kern( # the CUDA has NO exact-fit peel here, guard is per-load). Issue all # loads first (op43 L1), then -INFINITY-fill missed slots (op43 L2). atom128 = g2r_atom_f32(128, invariant=True) - frags = [cute.make_fragment((4,), cutlass.Float32) for _ in range(VPT)] + frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(VPT)] for u in cutlass.range_constexpr(VPT): i = base4 + tid + cutlass.Int32(u * self.blk) if i < n4: From 8374c6a52390094fd5da876024e7e08d8a427426 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 23 Aug 2026 09:21:54 +0000 Subject: [PATCH 36/40] [TRTLLM-15293][perf] self-sampling top-K: pin the 128-bit row loads to ld.global.nc.v4.f32 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DSL 4.6.1's NVVM rewrites adjacent 128-bit f32 copy-atom loads into v2.b64 register-PAIR loads plus mov.b64 unpacks (PTX: v4.b32 30->12, v2.b64 0->18 on gvr_clus). The even-aligned pair constraint fragments register allocation at the 64-register wall: the SAME ptxas 13.2 gives 64 regs / 4B spill from the 4.5.0 PTX but 63 regs / 80B spill from the 4.6.1 PTX — measured +20% on clus and +8.5% on the register family at runtime (the dominant share of the 4.6.1 grid regression, clus zone tax 1.068 -> 1.223). An inline-asm boundary pins the four-scalar-f32 shape on every DSL version; the isolated copy atom still lowers correctly, so this guards against the context-dependent rewrite only. nsys verdict under 4.6.1 (200 paired cold-L2 reps, replicated): clus r32 tax 1.163 -> 1.074 (vl 16.83 -> 13.87us), reg r12 1.207 -> 1.127, main and reg_clus controls neutral. UT 84/84 on 4.6.1 and 4.5.0; exactness smoke all exact on both pins. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 44 ++++++++++++++++--- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 8ad91bbaf4e2..20b3a1f6c9d6 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -402,19 +402,49 @@ def g2r_atom_i32(bits: int, invariant: bool = False): ) +@dsl_user_op +def _ld_g_nc_v4_f32(gaddr, *, loc=None, ip=None): + """Pinned `ld.global.nc.v4.f32` (CUDA `__ldg(const float4*)`). + + DSL 4.6.1's NVVM rewrites adjacent 128-bit f32 copy-atom loads into + v2.b64 register-PAIR loads plus mov.b64 unpacks (PTX diff: v4.b32 + 30->12, v2.b64 0->+18 on gvr_clus). The even-aligned pair constraint + fragments allocation at the 64-register wall: same source, same ptxas + 13.2 gives 64 regs / 4B spill from the 4.5.0 PTX but 63 regs / 80B + spill from the 4.6.1 PTX (clus +20% time, reg +8.5%). The asm boundary + pins the four-scalar-f32 shape on every DSL version.""" + from cutlass._mlir import ir as _ir + + st = _ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>") + r = mlir_llvm.inline_asm( + st, + [gaddr.ir_value(loc=loc, ip=ip)], + "ld.global.nc.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return tuple( + cutlass.Float32(mlir_llvm.extractvalue(T.f32(), r, [i], loc=loc, ip=ip)) for i in range(4) + ) + + def ld_g_f32x4(copy_atom, base_addr, v_idx, frag): """Load float4 #v_idx (16B units) from gmem byte base into frag[0..3]. op43 ct_tp.py:236-245 idiom. base_addr: Int64 byte address; frag: (4,) f32 fragment. Issue ALL batch members before consuming any (op43 lesson L1). + Pinned-asm form (see _ld_g_nc_v4_f32); copy_atom kept for call-site + compatibility. """ - p = cute.make_ptr( - cutlass.Float32, - base_addr + cutlass.Int64(v_idx) * cutlass.Int64(16), - cute.AddressSpace.gmem, - assumed_align=16, - ) - cute.copy(copy_atom, cute.make_tensor(p, cute.make_layout((4,))), frag) + v0, v1, v2, v3 = _ld_g_nc_v4_f32(base_addr + cutlass.Int64(v_idx) * cutlass.Int64(16)) + frag[0] = v0 + frag[1] = v1 + frag[2] = v2 + frag[3] = v3 def ldg_f32(base_addr, idx): From 01bb3107e1c3b7b20205b8a518f5e90ddf44bebb Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 23 Aug 2026 10:52:26 +0000 Subject: [PATCH 37/40] [TRTLLM-15293][perf] self-sampling top-K: drop the DSL 4.6.1 compile-time smem carveout in the register family DSL 4.6.1 moved the shared-memory carveout derivation into a compile-time MLIR attribute (smem.max_smem_per_mp, emitted when min_blocks_per_mp > 1) that only accounts for static smem. With the register family's dynamic launch smem this selects a 16 KiB shared-memory config, pinning each SM to a single resident CTA (achieved occupancy 67.6% -> 11.6%) and slowing saturating small-N grids up to 3.25x (N=1027 @ 1024 rows: 23.2us -> 7.1us varlen after the fix). Extend the existing _no_carveout() compile scope to also drop that attribute; clus/main/reg_clus launch with min_blocks_per_mp == 1 and are unaffected (verified neutral). Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 20b3a1f6c9d6..6a366dcb1fb0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -3338,10 +3338,29 @@ def _no_carveout(): orig = _cdsl._build_kernel_attrs _cdsl._build_kernel_attrs = lambda config: {} + # DSL >= 4.6.1 also derives a carveout inside the compiled artifact from + # the smem.max_smem_per_mp MLIR attribute (emitted when + # min_blocks_per_mp > 1). That calculation only sees the static smem, so + # with dynamic launch smem it selects a 16 KiB shared-memory config and + # pins the SM to one resident CTA. Drop the attribute as well so the + # driver default carveout stays in effect. + base = getattr(_cdsl, "CutlassBaseDSL", None) + orig_gen = getattr(base, "_generate_kernel_attrs", None) + + if orig_gen is not None: + + def _gen(self, config, _orig=orig_gen): + ret = _orig(self, config) + ret.pop("smem.max_smem_per_mp", None) + return ret + + base._generate_kernel_attrs = _gen try: yield finally: _cdsl._build_kernel_attrs = orig + if orig_gen is not None: + base._generate_kernel_attrs = orig_gen class GvrTopkRegKernel: From 41d61c93ae4cb95e8bac4907f9c4fd65f3d9c3c9 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:26:47 +0000 Subject: [PATCH 38/40] [TRTLLM-15293][chore] self-sampling top-K: slim code comments Condense narrative comment blocks to the technical invariant they encode and rewrite module/class docstrings as plain functional descriptions. Comment-and-docstring-only change: AST-verified equivalent (bare-string sections and assert messages normalized), tie-aware exactness smoke green. Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 27 +- .../attention_backend/sparse/dsa/metadata.py | 34 +- .../top_k/gvr_topk_decode_self_sampling.py | 1605 ++++++++--------- .../gvr_topk_decode_self_sampling_host.py | 644 +++---- .../parallel/test_gvr_selfsampling_topk.py | 15 +- 5 files changed, 1075 insertions(+), 1250 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index ce0095278d6d..177ba04fc1f2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -635,22 +635,19 @@ def __init__( self._enable_heuristic_topk = ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ) - # Opt-in self-sampling GVR top-K decode (standalone CuTeDSL modules, - # env-gated experimental path: TRTLLM_GVR_SELF_SAMPLING=1). Same - # operator contract as the tiered heuristic path (per-request device - # kv_lens, request-level raw prev-top-K hints, per-row MTP window, - # in-kernel n <= topK short path); tuning is frozen from - # indexer_max_seq_len at capture time, so the launch is - # CUDA-graph-replay safe. Two guard mechanisms coexist: the TopK - # module's hardware-format gate falls through to the CUDA GVR path - # (with a one-time warning), while contract violations inside the - # engine raise loudly — this flag is an explicit experiment. + # Opt-in self-sampling GVR top-K decode (CuTeDSL, env-gated: + # TRTLLM_GVR_SELF_SAMPLING=1). Same operator contract as the tiered + # heuristic path (per-request device kv_lens, raw prev-top-K hints, + # per-row MTP window, in-kernel n <= topK short path); tuning is + # frozen from indexer_max_seq_len at capture time, so the launch is + # CUDA-graph-replay safe. The TopK module's hardware-format gate + # falls through to the CUDA GVR path with a one-time warning; + # contract violations inside the engine raise. self._use_self_sampling_topk = ( os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" and IS_CUTLASS_DSL_AVAILABLE - # validated datacenter Blackwell only (B200/B300); consumer - # Blackwell (sm_120/121) lacks thread-block clusters and is - # not a supported target for these kernels + # datacenter Blackwell only; consumer Blackwell (sm_120/121) + # lacks thread-block clusters and get_sm_version() in (100, 103) and sparse_params.index_topk in (512, 1024, 2048) and compress_ratio in (1, 4) @@ -668,8 +665,8 @@ def __init__( else: decode_top_k_implementation = TopKImplementation.CUDA_RADIX if self._use_self_sampling_topk and self._enable_heuristic_topk: - # env opt-in promotes the self-sampling engine to the head of the - # decode chain (the GVR prior contract below is identical) + # env opt-in overrides the decode implementation; the GVR prior + # contract is identical decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR_V2 self.top_k = TopK( self.index_topk, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 783aabe4e64f..c647beaab653 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -345,17 +345,13 @@ def warmup_selfsampling_topk( Mirrors ``warmup_cute_dsl_radix_topk``. The varlen launcher is keyed by the exact row count AND the logits row stride: rows cover the - eager first-touch tuple (bs=1) plus every configured CUDA-graph - batch size, and the stride mirrors what the active paged-MQA - producer emits (the DSL arena rounds the row up to 256 elements; - DeepGEMM is exact-width) so the warmed keys are the ones dispatch - actually looks up. Captured geometries are also compiled by the - pre-capture warmup forwards; eager batches outside ``batch_sizes`` - still compile lazily on first touch. The helper enumerates one - representative row per distinct engine compile key (band-aware), so - arbitrarily large batch lists warm in bounded time and memory. - No-op unless the opt-in gate (TRTLLM_GVR_SELF_SAMPLING=1) selects - the engine. + small eager batches plus every configured CUDA-graph batch size, and + the stride mirrors what the active paged-MQA producer emits, so the + warmed keys are the ones dispatch actually looks up. Batches outside + this set still compile lazily on first touch. The helper enumerates + one representative row per distinct engine compile key, so large + batch lists warm in bounded time and memory. No-op unless the opt-in + gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. """ if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": return @@ -378,24 +374,18 @@ def warmup_selfsampling_topk( except ImportError: return nn = int(next_n) - # eager warm span: the launcher is keyed by exact row count, so warm - # the small-row counts eager mixed batches commonly produce (the - # engine-band compiles behind them are deduplicated by the helper); + # warm the small row counts eager mixed batches commonly produce; # larger eager row counts lazy-JIT on first touch, and CUDA-graph - # geometries are covered through ``batch_sizes`` below. 32 was the - # old admission-envelope bound; kept as the warm span after the - # full-range dispatch removed the envelope itself. + # geometries are covered through ``batch_sizes`` below eager_warm_rows = 32 rows = set(range(nn, eager_warm_rows + 1, nn)) or {nn} for bs in batch_sizes or (): rows.add(int(bs) * nn) msl_c = int(self.get_indexer_max_seq_len()) if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits: - # mirror the DSL paged-MQA arena stride (cute_dsl_custom_ops - # CuteDSLPagedMQALogitsRunner: compute_block_kv=128, SPLIT_KV= - # 2*128 -> rows round up to 256 elements). A drift here only - # degrades warmup to unused keys — dispatch still lazy-JITs the - # true key outside capture, so it can never become incorrect. + # mirror the DSL paged-MQA arena stride (rows round up to 256 + # elements). A drift here only degrades warmup to unused keys — + # dispatch still lazy-JITs the true key outside capture. row_stride = (msl_c + 255) // 256 * 256 else: # DeepGEMM emits exact-width rows; a non-float4 width falls diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 6a366dcb1fb0..543227e495fc 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -21,25 +21,10 @@ by count-crossing invariants, never by the estimate; the temporal hint (pre_idx) survives only as a degenerate-case anchor. -High-fidelity translation of the CUDA ``GVR-selfsampling-CUDA`` branch -(layer-10 head d76f4387c1): four kernel families -- sampling-ladder -slab/streaming (main), register-resident (reg), cluster streaming (clus), -clustered register-resident (regclus) -- merged into this single device -module for in-tree convention parity. The development/audit form (one -module per family, per-era commit history, SASS audit trail) lives on the -contributor fork branch ``GVR-selfsampling-CuTeDSL``. - -This file is GENERATED by a rename-only merge (collision symbols get a -``__`` suffix; no semantic edits) and is verified against the -per-family source by SASS bit-equality over every reachable compile tuple. -Do not hand-edit kernel bodies here without re-running that gate. - -Validation record (2026-08-17): 886-cell x 11-BS real-capture grid -(9,746 cases) vs the CUDA source arm -- 0 INEXACT, face-value gm 0.9743. -Known family-level exception (filed 2026-08-15): v32_8k + flash_32k at -BS >= 256, isolated-retest upper bound ~1.065 vs the 1.02 per-case bar -(compiler barrier-arrival raggedness; tracked for re-check on every -nvidia-cutlass-dsl version bump). +Four kernel families -- sampling-ladder slab/streaming (main), +register-resident (reg), cluster streaming (clus), clustered +register-resident (regclus) -- merged into a single device module; +collision symbols carry a ``__`` suffix. """ import contextlib @@ -56,56 +41,45 @@ from cutlass.cutlass_dsl import T, dsl_user_op from cutlass.utils.smem_allocator import SmemAllocator -# `C.` references throughout resolve against this module itself -# (the per-family sources imported the shared-units sibling as C). +# `C.` references throughout resolve against this module itself. C = sys.modules[__name__] # =========================================================================== -# ==== ct_common.py — shared device units ====================== +# ==== shared device units ===================================== # =========================================================================== -"""ct_common.py — op46 shared device-helper library (CuTeDSL port). +"""Device-helper library shared by the main / reg / clus / regclus families. -Translated ONCE from the frozen CUDA source -`op46_selfsampling_cutedsl/src_cuda/kernel.cu` L15-376 (helpers shared by -gvr_main / gvr_topk_reg / gvr_clus / gvr_reg_clus) per TRANSLATION_SPEC.md §5 -head, with every spelling pinned by probes/PROBE_RESULTS.md (P1..P15, all -BINDING) and op43 prior-art idioms (ct_tp/ct_gvr/ct_reg/ct_direct). - -Conventions for kernel translators ----------------------------------- +Conventions +----------- * Crossing-scan helpers write their scalar outputs into an Int32 smem tensor `s_res` using the slot map RES_B=0, RES_M=1, RES_ABOVE=2, RES_TOT=3, - RES_B2=4, RES_B3=5 (mirror of the CUDA `int *s_B,*s_m,*s_above,*s_tot` - out-params + the TWO/THREE extra pins). Slots are written ONLY on a pin, - exactly like the CUDA. -* Histograms are Int32 smem tensors (CUDA uint32_t): adds/scans are - bit-identical mod 2^32 and every compare the CUDA does against `target` - is already `(int)`-cast there. Totals < 2^31 by dispatch domain. + RES_B2=4, RES_B3=5. Slots are written ONLY on a pin. +* Histograms are Int32 smem tensors: adds/scans are bit-identical mod 2^32, + and totals stay below 2^31 by dispatch domain. * Warp-0-only helpers (find_cross / scan_cross0 / merge_scan0) contain NO - barrier (probe P14); scan_cross and scan_cross_w contain EXACTLY ONE - internal barrier (kernel.cu L199 / L306); gather_hint contains EXACTLY TWO - (L349 / L357). Do not add or drop any (op43 lesson L5). -* All warp collectives use the full mask FULLM = 0xffffffff (kernel.cu L15). + barrier; scan_cross and scan_cross_w contain EXACTLY ONE internal barrier; + gather_hint contains EXACTLY TWO. Do not add or drop any. +* All warp collectives use the full mask FULLM = 0xffffffff. """ # --------------------------------------------------------------------------- -# constants (kernel.cu L15-30) +# constants # --------------------------------------------------------------------------- FULLM = 0xFFFFFFFF -NB = 1024 # register-family base bin count (L16) -SNB = 256 # streaming-path bin count (L170-177) — MUST stay 256 -MAXC = 160 # multi-CTA SPLIT row cap (L17) -GCAP = 16384 # per-row slab capacity in int2 (L18) -QUADC = 96 # O(mc^2) rank gate, streaming/reg (L21) -QUADC_CLUS = 288 # clus + gvr_main gate (L28, spec §2 conflict resolution) -IDXB = 22 # packed candidate index bits (L29) +NB = 1024 # register-family base bin count +SNB = 256 # streaming-path bin count — MUST stay 256 +MAXC = 160 # multi-CTA SPLIT row cap +GCAP = 16384 # per-row slab capacity in int2 +QUADC = 96 # O(mc^2) rank gate, streaming/reg +QUADC_CLUS = 288 # clus + gvr_main gate +IDXB = 22 # packed candidate index bits IDXM = (1 << IDXB) - 1 -GVR_WS_OFF_OFF = MAXC * 8 # workspace g_off byte offset (L42) -GVR_WS_BUF_OFF = 2048 # workspace g_buf byte offset (L43) +GVR_WS_OFF_OFF = MAXC * 8 # workspace g_off byte offset +GVR_WS_BUF_OFF = 2048 # workspace g_buf byte offset -# degenerate-hint sentinels (kernel.cu L629, L356, L1014 exact-equality flag) +# degenerate-hint sentinels (exact-equality flag values) SENT_LO = -3.0e38 SENT_HI = 3.0e38 @@ -119,7 +93,7 @@ # --------------------------------------------------------------------------- -# float <-> u32 bitcasts (op43 ct_tp.py:145-151) +# float <-> u32 bitcasts # --------------------------------------------------------------------------- def u32_of_f32(v): """Raw fp32 bits as Uint32 (bit-cast, no conversion).""" @@ -140,9 +114,9 @@ def i32_of_f32(v): # --------------------------------------------------------------------------- -# fkey / invkey (kernel.cu L64-71) — order-preserving float->u32 radix key. +# fkey / invkey — order-preserving float->u32 radix key. # fkey: u ^ (((int32)u >> 31) | 0x80000000) [arithmetic-shift sign trick, -# spelled 0 - (u >> 31) on Uint32 per op43 ct_tp.py:171-175] +# spelled 0 - (u >> 31) on Uint32] # invkey: (K & 0x80000000) ? K ^ 0x80000000 : ~K [exact inverse] # Monotone over all finite floats and +-inf; min identity 0xffffffff, max 0. # --------------------------------------------------------------------------- @@ -153,12 +127,12 @@ def fkey_bits(u): def fkey(x): - """CUDA fkey(float) (L68-71). x: dynamic Float32 -> Uint32 key.""" + """CUDA fkey(float). x: dynamic Float32 -> Uint32 key.""" return fkey_bits(u32_of_f32(x)) def invkey_bits(K): - """CUDA invkey (L64-67) without the final bitcast: key -> fp32 bits.""" + """CUDA invkey without the final bitcast: key -> fp32 bits.""" s = K >> cutlass.Uint32(31) # 1 iff key top bit set m = (s - cutlass.Uint32(1)) | cutlass.Uint32(0x80000000) # s==1 -> m=0x80000000 (K^0x80000000); s==0 -> m=0xFFFFFFFF (~K) @@ -166,14 +140,13 @@ def invkey_bits(K): def invkey(K): - """CUDA invkey(uint32) (L64-67). K: dynamic Uint32 key -> Float32.""" + """CUDA invkey(uint32). K: dynamic Uint32 key -> Float32.""" return f32_of_u32(invkey_bits(K)) # --------------------------------------------------------------------------- -# warp redux wrappers (kernel.cu L59-62; scan_cross_w L307-315; probe P1/P13). -# Values passed to the u32 forms MUST be genuine cutlass.Uint32 — an Int32 -# silently lowers to redux.sync.{min,max}.s32 (probe P1 note). +# warp redux wrappers. Values passed to the u32 forms MUST be genuine +# cutlass.Uint32 — an Int32 silently lowers to redux.sync.{min,max}.s32. # --------------------------------------------------------------------------- def warp_min_u32(v): """__reduce_min_sync(FULLM, v) -> redux.sync.min.u32 (single inst).""" @@ -191,12 +164,12 @@ def warp_add_u32(v): def warp_add_i32(v): - """__reduce_add_sync on Int32 (scan_cross_w two-redux stage L314-315).""" + """__reduce_add_sync on Int32 (scan_cross_w two-redux stage).""" return cute.arch.warp_redux_sync(v, "add") def fmin_f32(a, b): - """fminf -> native min.f32 (probe P13; op43 inline-PTX no longer needed).""" + """fminf -> native min.f32.""" return cute.arch.fmin(a, b) @@ -225,15 +198,15 @@ def clz_i32(x): def ffs_m1(x): """__ffs(x) - 1 for x != 0 (bit index of lowest set bit). - Spelled popc((x & -x) - 1) per op43 ct_direct.py:210-221. Caller must - guarantee x != 0 (every kernel.cu use is inside a mask-walk loop). + Spelled popc((x & -x) - 1). Caller must guarantee x != 0 (every use is + inside a mask-walk loop). """ return cutlass.Int32(cute.arch.popc((x & (cutlass.Int32(0) - x)) - cutlass.Int32(1))) @cute.jit def hi_bit_or_zero(msk): - """CUDA `msk ? (31 - __clz(msk)) : 0` (find_cross L92/L102).""" + """CUDA `msk ? (31 - __clz(msk)) : 0`.""" r = cutlass.Int32(0) if msk != cutlass.Int32(0): r = cutlass.Int32(31) - clz_i32(msk) @@ -241,8 +214,7 @@ def hi_bit_or_zero(msk): # --------------------------------------------------------------------------- -# warp shfl scans (op43 ct_tp.py:186-206 + the TWO-interleaved variant that -# gvr_topk_reg L1669-1673 needs) +# warp shfl scans (plus the TWO-interleaved variant gvr_topk_reg needs) # --------------------------------------------------------------------------- @cute.jit def _shfl_up_add(val, lane, offset: cutlass.Constexpr): @@ -274,7 +246,7 @@ def _shfl_down_add(val, lane, offset: cutlass.Constexpr): @cute.jit def warp_incl_scan_add(val, lane): - """5-step inclusive __shfl_up_sync add scan (e.g. L144-146, L843-849).""" + """5-step inclusive __shfl_up_sync add scan.""" for o in [1, 2, 4, 8, 16]: val = _shfl_up_add(val, lane, o) return val @@ -282,7 +254,7 @@ def warp_incl_scan_add(val, lane): @cute.jit def warp_incl_scan_add2(v1, v2, lane): - """TWO interleaved inclusive shfl_up scans (gvr_topk_reg L1669-1673). + """TWO interleaved inclusive shfl_up scans. Per step o: shfl(v1); gated add; shfl(v2); gated add — the two dependency chains interleave so the second scan hides under the first's shfl latency @@ -300,15 +272,14 @@ def warp_incl_scan_add2(v1, v2, lane): @cute.jit def warp_suffix_scan_add(val, lane): - """5-step __shfl_down_sync suffix add scan (find_cross L86-89, L97-100).""" + """5-step __shfl_down_sync suffix add scan.""" for o in [1, 2, 4, 8, 16]: val = _shfl_down_add(val, lane, o) return val # --------------------------------------------------------------------------- -# CTA-scope shared-memory atomics (probe P7: returns OLD value, ATOMS not RED, -# warp-aggregation preserved; never sys-scope utils.distributed.atomicAdd) +# CTA-scope shared-memory atomics (return the OLD value; never sys scope) # --------------------------------------------------------------------------- def atomic_add_cta(ptr, val): """shared atomicAdd returning old value. ptr: cute Pointer @@ -319,7 +290,7 @@ def atomic_add_cta(ptr, val): def atomic_min_cta(ptr, val): - """shared atomicMin (s_kmin seeds, L1094-1101). Unsigned iff val is Uint32.""" + """shared atomicMin (s_kmin seeds). Unsigned iff val is Uint32.""" return cute.arch.atomic_min(ptr, val, sem="relaxed", scope="cta") @@ -334,15 +305,16 @@ def atomic_or_cta(ptr, val): # --------------------------------------------------------------------------- -# gpu-scope fences + global u64 atomicAdd (SPLIT slab protocol, probe P8) +# gpu-scope fences + global u64 atomicAdd (SPLIT slab protocol) # --------------------------------------------------------------------------- def threadfence_gpu(): - """__threadfence() == fence.acq_rel.gpu — use at BOTH L959 and L968.""" + """__threadfence() == fence.acq_rel.gpu — required on BOTH the release and + acquire sides of the slab hand-off.""" cute.arch.fence_acq_rel_gpu() def atomic_add_u64_gpu(ptr, val): - """atom.global.add.u64 returning the OLD value (L960 arrival RMW). + """atom.global.add.u64 returning the OLD value (arrival RMW). ptr: cute Pointer to an Int64 gmem word; val: cutlass.Int64. Packed arrival word: `cutlass.Int64(1 << 32) + cutlass.Int64(myn)`. @@ -351,13 +323,12 @@ def atomic_add_u64_gpu(ptr, val): # --------------------------------------------------------------------------- -# saturating converts (probe P4: native ctors emit cvt.rzi.{u32,s32}.f32) +# saturating converts (native ctors emit cvt.rzi.{u32,s32}.f32) # --------------------------------------------------------------------------- def f2u_rz(v): """__float2uint_rz: saturating (neg/-inf -> 0, huge -> 0xffffffff, NaN -> 0). - Native ctor spelling verified on this exact toolchain (P4). Dynamic values - only — host constants raise OverflowError on inf. + Dynamic values only — host constants raise OverflowError on inf. """ return cutlass.Uint32(v) @@ -368,7 +339,7 @@ def f2s_rz(v): # --------------------------------------------------------------------------- -# L2 prefetch escape hatch (op43 ct_gvr.py:42-54; kernel.cu sites L576/592/612) +# L2 prefetch escape hatch # --------------------------------------------------------------------------- @dsl_user_op def _prefetch_l2(gaddr, *, loc=None, ip=None): @@ -387,7 +358,7 @@ def _prefetch_l2(gaddr, *, loc=None, ip=None): # --------------------------------------------------------------------------- # global loads: 128-bit ldg (read-only) / plain, scalar forms, and __ldcg -# (L2-direct) vector forms for the slab consume (L929/945/978/1032/1172/1194) +# (L2-direct) vector forms for the slab consume # --------------------------------------------------------------------------- def g2r_atom_f32(bits: int, invariant: bool = True): """CopyG2ROp atom: bits=128 -> LDG.E.128[.CONSTANT], bits=32 -> scalar.""" @@ -406,13 +377,10 @@ def g2r_atom_i32(bits: int, invariant: bool = False): def _ld_g_nc_v4_f32(gaddr, *, loc=None, ip=None): """Pinned `ld.global.nc.v4.f32` (CUDA `__ldg(const float4*)`). - DSL 4.6.1's NVVM rewrites adjacent 128-bit f32 copy-atom loads into - v2.b64 register-PAIR loads plus mov.b64 unpacks (PTX diff: v4.b32 - 30->12, v2.b64 0->+18 on gvr_clus). The even-aligned pair constraint - fragments allocation at the 64-register wall: same source, same ptxas - 13.2 gives 64 regs / 4B spill from the 4.5.0 PTX but 63 regs / 80B - spill from the 4.6.1 PTX (clus +20% time, reg +8.5%). The asm boundary - pins the four-scalar-f32 shape on every DSL version.""" + The asm boundary pins the four-scalar-f32 shape: NVVM otherwise rewrites + adjacent 128-bit f32 copy-atom loads into v2.b64 register-pair loads, + whose even-aligned pair constraint fragments allocation at the + 64-register wall and induces spills.""" from cutlass._mlir import ir as _ir st = _ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>") @@ -435,10 +403,9 @@ def _ld_g_nc_v4_f32(gaddr, *, loc=None, ip=None): def ld_g_f32x4(copy_atom, base_addr, v_idx, frag): """Load float4 #v_idx (16B units) from gmem byte base into frag[0..3]. - op43 ct_tp.py:236-245 idiom. base_addr: Int64 byte address; frag: (4,) f32 - fragment. Issue ALL batch members before consuming any (op43 lesson L1). - Pinned-asm form (see _ld_g_nc_v4_f32); copy_atom kept for call-site - compatibility. + base_addr: Int64 byte address; frag: (4,) f32 fragment. Issue ALL batch + members before consuming any. Pinned-asm form (see _ld_g_nc_v4_f32); + copy_atom kept for call-site compatibility. """ v0, v1, v2, v3 = _ld_g_nc_v4_f32(base_addr + cutlass.Int64(v_idx) * cutlass.Int64(16)) frag[0] = v0 @@ -448,7 +415,7 @@ def ld_g_f32x4(copy_atom, base_addr, v_idx, frag): def ldg_f32(base_addr, idx): - """__ldg(X + idx): scalar read-only 4B gather (gather_hint L343).""" + """__ldg(X + idx): scalar read-only 4B gather.""" atom = g2r_atom_f32(32, invariant=True) p = cute.make_ptr( cutlass.Float32, @@ -462,7 +429,7 @@ def ldg_f32(base_addr, idx): def ld_g_i32(base_addr, idx): - """plain P[idx] scalar int32 load (gather_hint L340).""" + """plain P[idx] scalar int32 load.""" p = cute.make_ptr( cutlass.Int32, base_addr + cutlass.Int64(idx) * cutlass.Int64(4), @@ -476,7 +443,7 @@ def ld_g_i32(base_addr, idx): def _ldcg_v2_i32(gaddr, *, loc=None, ip=None): """__ldcg on an int2 (8B slab word): ld.global.cg.v2.u32 -> (x, y). - x = value bits, y = index (workspace g_buf layout, kernel.cu L38-41). + x = value bits, y = index (workspace g_buf layout). gaddr: Int64 byte address, 8B-aligned. """ ret = llvm.inline_asm( @@ -514,7 +481,7 @@ def _ldcg_v4_i32(gaddr, *, loc=None, ip=None): # --------------------------------------------------------------------------- -# 128-bit shared-memory ld/st (probe P5a copy-atom spelling) + ulonglong2 read +# 128-bit shared-memory ld/st (copy-atom spelling) + ulonglong2 read # --------------------------------------------------------------------------- def smem_atom_i32_128(): """CopyUniversalOp atom for ld/st.shared.v4.b32 on Int32 smem.""" @@ -539,7 +506,7 @@ def sts128_i32(copy_atom, frag, base_addr, byte_off): @dsl_user_op def _lds_v2_u64(saddr, *, loc=None, ip=None): - """ulonglong2 16B smem read (quad-rank L1080-1089, L2241): (lo, hi).""" + """ulonglong2 16B smem read (quad-rank path): (lo, hi).""" ret = llvm.inline_asm( llvm.StructType.get_literal([T.i64(), T.i64()]), [saddr.ir_value(loc=loc, ip=ip)], @@ -558,9 +525,8 @@ def _lds_v2_u64(saddr, *, loc=None, ip=None): # --------------------------------------------------------------------------- -# DSMEM op set (op43 ct_tp.py:44-125 + probe P5b 128-bit remote load). -# mapa returns a byte-addressed Int32 in the PEER's shared window; offset -# arithmetic after one mapa-per-rank is the proven op43/P5b form. +# DSMEM op set. mapa returns a byte-addressed Int32 in the PEER's shared +# window; offset arithmetic is applied after one mapa per rank. # --------------------------------------------------------------------------- @dsl_user_op def _mapa_shared_cluster(smem_ptr, peer_rank, *, loc=None, ip=None): @@ -635,7 +601,7 @@ def _ld_shared_cluster_f32(mapped_addr, *, loc=None, ip=None): @dsl_user_op def _ld_shared_cluster_v4_u32(mapped_addr, *, loc=None, ip=None): - """Single-shot remote 16B DSMEM load (probe P5b; merge_scan0 L136-137).""" + """Single-shot remote 16B DSMEM load.""" ret = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), [mapped_addr.ir_value(loc=loc, ip=ip)], @@ -685,7 +651,7 @@ def _st_shared_cluster_f32(mapped_addr, val, *, loc=None, ip=None): @dsl_user_op def _st_shared_cluster_u64(mapped_addr, val, *, loc=None, ip=None): - """ONE packed 8B DSMEM candidate push (op43 lesson L3; kernel.cu L2185). + """ONE packed 8B DSMEM candidate push (never split into two 4B stores). val = (Uint64(key) << 32) | Uint64(idx_bits). """ @@ -720,10 +686,10 @@ def _atom_shared_cluster_add_i32(mapped_addr, val, *, loc=None, ip=None): # --------------------------------------------------------------------------- -# aligned cluster barrier (op43 ct_reg.py:100-113; cg::cluster.sync() == +# aligned cluster barrier (cg::cluster.sync() == # barrier.cluster.{arrive,wait}.aligned). Writers use the FULL (releasing) -# arrive — cluster_arrive_relaxed has NO release and races DSMEM (known -# lesson). Never substitute the non-aligned cute.arch forms. +# arrive — cluster_arrive_relaxed has NO release and races DSMEM. Never +# substitute the non-aligned cute.arch forms. # --------------------------------------------------------------------------- @dsl_user_op def _cluster_arrive_aligned(*, loc=None, ip=None): @@ -737,16 +703,16 @@ def _cluster_wait_aligned(*, loc=None, ip=None): @cute.jit def _cluster_sync_aligned(): - """cg::cluster_group::sync() (kernel.cu L2016/2135/2226/2474/2530/2648).""" + """cg::cluster_group::sync().""" _cluster_arrive_aligned() _cluster_wait_aligned() # =========================================================================== -# find_cross (kernel.cu L73-109) +# find_cross # highest bin B with sum_{j>=B} hist[j] >= target; also total, m = hist[B], # above = sum_{j>B}. Warp-parallel (warp 0 only), bank-conflict free via the -# rotated indexing hist[lane*BPL + ((j+lane) & (BPL-1))] (L83 — DO NOT drop). +# rotated indexing hist[lane*BPL + ((j+lane) & (BPL-1))] — DO NOT drop it. # Non-destructive. NO barrier inside. # Writes s_res[RES_B/RES_M/RES_ABOVE] from the single pinning lane and # s_res[RES_TOT] from lane 0. @@ -783,17 +749,16 @@ def find_cross(s_hist, target, tidx, s_res, nb: cutlass.Constexpr): # =========================================================================== -# scan_cross0 (kernel.cu L218-286) -# Warp-0-only single-barrier vectorized suffix scan (streaming workhorse, -# NB_=256 at every production call site). Contains NO barrier — the caller -# pays exactly one after it. Leaves hist[j] = per-bin OUTPUT CURSOR -# (count strictly above bin j), or ZEROS when zero=True (folds the next -# phase's histogram clear). two/three pin extra crossing bins for -# target2/target3 into RES_B2/RES_B3. addf folds the per-rank bin-offset -# vector s_addv into the cursors (L279-282). -# HOLD register guard (L236-240): NV<=2 holds the span in regs across the -# scan; wider instantiations re-READ their span (no barrier needed — each -# lane only touches its own span). +# scan_cross0 +# Warp-0-only vectorized suffix scan (streaming workhorse, NB_=256 at every +# production call site). Contains NO barrier — the caller pays exactly one +# after it. Leaves hist[j] = per-bin OUTPUT CURSOR (count strictly above +# bin j), or ZEROS when zero=True (folds the next phase's histogram clear). +# two/three pin extra crossing bins for target2/target3 into RES_B2/RES_B3. +# addf folds the per-rank bin-offset vector s_addv into the cursors. +# HOLD register guard: NV<=2 holds the span in regs across the scan; wider +# instantiations re-READ their span (no barrier needed — each lane only +# touches its own span). # =========================================================================== @cute.jit def scan_cross0( @@ -812,26 +777,26 @@ def scan_cross0( ): BPT = nb // 32 # bins per lane (trace-time int) NV = BPT // 4 # 16B vectors per lane - HOLD = NV <= 2 # register-pressure guard (L240) + HOLD = NV <= 2 # register-pressure guard if tidx < cutlass.Int32(32): lane = tidx atom = smem_atom_i32_128() hbase = s_hist.iterator.toint() - # pass 1: span sum via NV uint4 LDS.128 (L243-251) + # pass 1: span sum via NV uint4 LDS.128 frags = [cute.make_rmem_tensor((4,), cutlass.Int32) for _ in range(NV)] sm = cutlass.Int32(0) for q in cutlass.range_constexpr(NV): boff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) lds128_i32(atom, hbase, boff, frags[q]) sm = sm + frags[q][0] + frags[q][1] + frags[q][2] + frags[q][3] - # 5-step inclusive shfl_up scan (L252-254) + # 5-step inclusive shfl_up scan w = warp_incl_scan_add(sm, lane) tot = cute.arch.shuffle_sync(w, cutlass.Int32(31)) after = tot - w # bins strictly above my span if lane == cutlass.Int32(0): s_res[RES_TOT] = tot base = lane * cutlass.Int32(BPT) - # pass 2: descending vector walk (L258-284) + # pass 2: descending vector walk for q in cutlass.range_constexpr(NV - 1, -1, -1): if cutlass.const_expr(HOLD): vv = frags[q] @@ -876,7 +841,7 @@ def scan_cross0( if cross3 != cutlass.Int32(0): s_res[RES_B3] = gb after = after + cq - if cutlass.const_expr(addf): # fold per-rank bin offset (L279-282) + if cutlass.const_expr(addf): # fold per-rank bin offset av = cute.make_rmem_tensor((4,), cutlass.Int32) aoff = (lane * cutlass.Int32(NV) + cutlass.Int32(q)) * cutlass.Int32(16) lds128_i32(atom, s_addv.iterator.toint(), aoff, av) @@ -887,11 +852,11 @@ def scan_cross0( # =========================================================================== -# scan_cross (kernel.cu L179-216) +# scan_cross # Block-parallel suffix scan over NB_ (<= BLK) bins. Leaves hist[j] = OUTPUT # CURSOR (count in bins > j) and pins the crossing bin. Warps that hold no -# bin skip the body. EXACTLY ONE internal barrier (L199); the caller pays its -# usual publish barrier after. Used by gvr_clus whole-row degen (L2326). +# bin skip the body. EXACTLY ONE internal barrier; the caller pays its usual +# publish barrier after. Used by the gvr_clus whole-row degenerate path. # =========================================================================== @cute.jit def scan_cross( @@ -910,17 +875,17 @@ def scan_cross( wid = tidx >> cutlass.Int32(5) c = cutlass.Int32(0) w = cutlass.Int32(0) - if tidx < cutlass.Int32(nb): # L189 + if tidx < cutlass.Int32(nb): c = s_hist[tidx] - w = warp_incl_scan_add(c, lane) # L191-193 + w = warp_incl_scan_add(c, lane) if lane == cutlass.Int32(31): - s_ws[wid] = w # L194 - cute.arch.barrier() # L199 — the ONE barrier - if tidx < cutlass.Int32(nb): # L200 + s_ws[wid] = w + cute.arch.barrier() # the ONE internal barrier + if tidx < cutlass.Int32(nb): v2 = cutlass.Int32(0) if lane < cutlass.Int32(NWU): v2 = s_ws[lane] - pre = warp_incl_scan_add(v2, lane) # L202-204 + pre = warp_incl_scan_add(v2, lane) tot = cute.arch.shuffle_sync(pre, cutlass.Int32(31)) off = cute.arch.shuffle_sync(pre - v2, wid) after = tot - (off + w) @@ -937,7 +902,7 @@ def scan_cross( s_res[RES_B] = tidx s_res[RES_ABOVE] = after s_res[RES_M] = c - if cutlass.const_expr(two): # L212-214 + if cutlass.const_expr(two): cross2 = cutlass.Int32(0) if after < target2: if (after + c) >= target2: @@ -949,13 +914,12 @@ def scan_cross( # =========================================================================== -# scan_cross_w (kernel.cu L288-327) +# scan_cross_w # Register-path block-parallel suffix scan for NB_ >= BLK: every thread owns # a private contiguous BPT = NB_/BLK span, so its read->write needs no -# barrier. EXACTLY ONE internal barrier (L306). The second stage is TWO -# REDUCTIONS, not a scan (L307-315): tot = redux_add(vv), off = redux_add -# ((lane < wid) ? vv : 0) — wid is warp-uniform so the masked operand stays -# convergent. +# barrier. EXACTLY ONE internal barrier. The second stage is TWO REDUCTIONS, +# not a scan: tot = redux_add(vv), off = redux_add((lane < wid) ? vv : 0) — +# wid is warp-uniform so the masked operand stays convergent. # =========================================================================== @cute.jit def scan_cross_w(s_hist, s_ws, target, tidx, s_res, blk: cutlass.Constexpr, nb: cutlass.Constexpr): @@ -966,25 +930,25 @@ def scan_cross_w(s_hist, s_ws, target, tidx, s_res, blk: cutlass.Constexpr, nb: loc = cute.make_rmem_tensor((BPT,), cutlass.Int32) base = tidx * cutlass.Int32(BPT) sm = cutlass.Int32(0) - for i in cutlass.range_constexpr(BPT): # L297-300 (#pragma unroll) + for i in cutlass.range_constexpr(BPT): loc[i] = s_hist[base + cutlass.Int32(i)] sm = sm + loc[i] - w = warp_incl_scan_add(sm, lane) # L301-304 + w = warp_incl_scan_add(sm, lane) if lane == cutlass.Int32(31): - s_ws[wid] = w # L305 - cute.arch.barrier() # L306 — the ONE barrier + s_ws[wid] = w + cute.arch.barrier() # the ONE internal barrier vv = cutlass.Int32(0) if lane < cutlass.Int32(NW): - vv = s_ws[lane] # L313 - tot = cutlass.Int32(warp_add_i32(vv)) # L314 + vv = s_ws[lane] + tot = cutlass.Int32(warp_add_i32(vv)) sel = cutlass.Int32(0) if lane < wid: sel = vv - off = cutlass.Int32(warp_add_i32(sel)) # L315 - after = tot - (off + w) # L316 + off = cutlass.Int32(warp_add_i32(sel)) + after = tot - (off + w) if tidx == cutlass.Int32(0): s_res[RES_TOT] = tot - for i in cutlass.range_constexpr(BPT - 1, -1, -1): # L318-326 + for i in cutlass.range_constexpr(BPT - 1, -1, -1): cq = loc[i] s_hist[base + cutlass.Int32(i)] = after # per-bin OUTPUT CURSOR gb = base + cutlass.Int32(i) @@ -1002,13 +966,13 @@ def scan_cross_w(s_hist, s_ws, target, tidx, s_res, blk: cutlass.Constexpr, nb: # =========================================================================== -# merge_scan0 (kernel.cu L111-168) +# merge_scan0 # Warp-0-fused cluster merge + suffix scan: each lane reads its BPT-bin span -# from EVERY rank's hist via 16B DSMEM loads (probe P5b), sums the cluster -# totals (and the r= blk//32 slots. Returns (gm, gx) Float32. -# op43 lessons L1/L2: both round trips are issued as predicated flat batches. +# Both round trips are issued as predicated flat batches. # =========================================================================== @cute.jit def gather_hint( @@ -1145,43 +1109,33 @@ def gather_hint( # =========================================================================== -# ==== ct_main.py — family: main =============================== +# ==== family: main ============================================ # =========================================================================== -"""ct_main.py — op46 gvr_main (streaming self-sampling GVR) CuTeDSL port. - -Ground truth: src_cuda/kernel.cu L377-1265 (frozen); phase contract, smem map -and barrier inventory per TRANSLATION_SPEC.md §5.1; every DSL spelling pinned -by probes/PROBE_RESULTS.md (BINDING) and op43 idioms (L1..L5 lessons). +"""gvr_main — streaming self-sampling GVR top-K. -Ctor knobs (compile-time, mirror of the CUDA template params, spec §4d): +Ctor knobs (compile-time, mirror of the CUDA template params): BLK ∈ {1024, 512, 256}, U ∈ {1,2,4,8}, MINB ∈ {1,2,4}, NBS = 256, KPT ∈ {1,2,4,8}, SPLIT ∈ {True, False} -Derived constexprs (kernel.cu L394-523, bit-identical): +Derived constexprs (bit-identical to the CUDA): HB=NBS; KBIG=(KPT>=2 && KPT*BLK>=2048); SCPB=(BLK>=1024)?(SPLIT?8192:16384) :(KBIG?8192:4096); CMPB=(BLK>=1024)?(KBIG?4096:2048):1024; SHD=!SPLIT; VSTG=SPLIT||BLK>=512; PFD=(MINB<=2)?min(U,4):0; PF=PFD>0; NATT=SPLIT?1:3. -Signature (ABI parity with kernel.cu L379-382 incl. dead SCAP_/CMP_): +Signature (ABI parity with the CUDA form incl. dead SCAP_/CMP_): run(logits[b,npad] f32, pre_idx[b,k] i32, out[b,k] i32, n, npad, k, SCAP_, CMP_, R, SMP, TGT, Q, SS2, TGT2, ws) -Grid dim3(R, b) native 2-D (probe P11); block BLK; min_blocks_per_mp=MINB is -the 64-register wall (probes P2/P15); smem via one SmemAllocator blob (all -extents compile-time), dynamic-equivalent region byte-identical to the host -formula L3149: (SCPB+4)*(VSTG?8:4) + (CMPB+1)*8. +Grid dim3(R, b) native 2-D; block BLK; min_blocks_per_mp=MINB is the +64-register wall; smem via one SmemAllocator blob (all extents compile-time), +dynamic-equivalent region byte-identical to the host dispatch formula: +(SCPB+4)*(VSTG?8:4) + (CMPB+1)*8. int2 staging convention: an int2 (x=value bits, y=index) is ONE little-endian Uint64 = (idx << 32) | value_bits, so cbuf2 / g_buf traffic is single u64 -ld/st (mirrors the CUDA ST.64/LD.64; __ldcg = ct_common._ldcg_v2_i32). - -Barrier inventory implemented (kernel.cu line cites, checklist per op43 L5): - L555, L655, L664, [retry: L753], L909, - SPLIT: L924, [overflow: L935], L956, L961, [last: L983, L985] - non-split: L992, [ladder: L1008 | L1017] - P6: L1075, L1097, L1100, per-level L1116/L1118, L1125 - degen A: L1162, per-level L1177/L1179, L1186 - degen B: L1221, per-level L1234/L1236, L1243 - + exactly 2 inside each gather_hint expansion (L349/L357, ct_common). -scan_cross0 contains NO barrier (probe P14 protocol). +ld/st (__ldcg = _ldcg_v2_i32). + +Barrier placement mirrors the CUDA source one-for-one (plus exactly 2 inside +each gather_hint expansion); scan_cross0 contains NO internal barrier. Do not +add or drop barriers. """ @@ -1190,16 +1144,15 @@ def gather_hint( IDXB__main = C.IDXB IDXM__main = C.IDXM QUADC_CLUS__main = C.QUADC_CLUS -WS_BYTES = C.GVR_WS_BUF_OFF + MAXC__main * GCAP__main * 8 # 20,973,568 (kernel.cu L44-46) +WS_BYTES = C.GVR_WS_BUF_OFF + MAXC__main * GCAP__main * 8 # 20,973,568 _NEG_INF = float("-inf") # --------------------------------------------------------------------------- -# single-rounding fma.rn.f32 (probe P6 emit spelling; ct_common has no fma). -# Used at every CUDA fmaf() site: T/Tk/T3 rung math (L665/L689/L722), HIC -# (L708), window terms. (x-T)*SC classify shapes stay plain sub+mul (P6: -# structurally uncontractible). +# single-rounding fma.rn.f32, used at every CUDA fmaf() site: T/Tk/T3 rung +# math, HIC, window terms. (x-T)*SC classify shapes stay plain sub+mul +# (structurally uncontractible). # --------------------------------------------------------------------------- @dsl_user_op def _fmaf(a, b, c, *, loc=None, ip=None): @@ -1216,14 +1169,14 @@ def _fmaf(a, b, c, *, loc=None, ip=None): def _st_g_u64(addr_i64, val_u64): - """plain st.global.u64 (slab publish L929/945/952, g_don restore L969).""" + """plain st.global.u64 (slab publish, g_don restore).""" p = cute.make_ptr(cutlass.Uint64, addr_i64, cute.AddressSpace.gmem, assumed_align=8) t = cute.make_tensor(p, cute.make_layout((1,))) t[0] = val_u64 def _st_g_u32(addr_i64, val_i32): - """plain st.global.u32 (g_off restore L969).""" + """plain st.global.u32 (g_off restore).""" p = cute.make_ptr(cutlass.Int32, addr_i64, cute.AddressSpace.gmem, assumed_align=4) t = cute.make_tensor(p, cute.make_layout((1,))) t[0] = val_i32 @@ -1231,11 +1184,11 @@ def _st_g_u32(addr_i64, val_i32): @dsl_user_op def _st_s_v2_u32(saddr_i32, lo_u32, hi_u32, *, loc=None, ip=None): - """st.shared.v2.u32 [saddr], {lo, hi} — the CUDA make_int2 STS.64 spelling - (kernel.cu L1018-1019). Byte-identical to the little-endian u64 pack - ((hi << 32) | lo) but keeps the two words as independent 32-bit registers, - so ptxas can coalesce the emission bit-walk's loop-carried (xv, idx) pair - straight into the store pair (drops 2 IMAD.MOV/iter; op46 SASS diff).""" + """st.shared.v2.u32 [saddr], {lo, hi} — the CUDA make_int2 STS.64 spelling. + Byte-identical to the little-endian u64 pack ((hi << 32) | lo) but keeps + the two words as independent 32-bit registers, so ptxas can coalesce the + emission bit-walk's loop-carried (xv, idx) pair straight into the store + pair.""" mlir_llvm.inline_asm( res=None, operands_=[ @@ -1256,7 +1209,7 @@ def _st_s_v2_u32(saddr_i32, lo_u32, hi_u32, *, loc=None, ip=None): def _pin_i64(v, *, loc=None, ip=None): """Opaque identity mov.b64: pins a loop-invariant Int64 so NVVM cannot rematerialize its defining chain (param ld.const + %ctaid reads + mul/add) - into every scf region body (PTX $L__BB0_123 evidence, op46 SASS diff).""" + into every scf region body.""" return cutlass.Int64( mlir_llvm.inline_asm( T.i64(), @@ -1291,16 +1244,15 @@ def _pin_i32(v, *, loc=None, ip=None): def _ldg_f32_rs(base_addr, idx, sc4): - """__ldg(X + idx) with the byte stride riding a register (fix-3 P1d). - - Identical to ct_common.ldg_f32 except `* 4` multiplies a caller-held - Int32: the row base is uniformized into URx by ptxas, IMAD.WIDE cannot - encode an immediate stride next to a UR addend, and a constant stride - register gets re-materialized INSIDE the survivor walk (its 26th - instruction; CUDA parity = 25). The caller loads the 4 from smem - (LDS results are opaque to ptxas value-tracking — asm movs and shfl - are NOT, see FIX2_CPU_VERIFY.md P1c falsification), so the register - stays live and the remat disappears.""" + """__ldg(X + idx) with the byte stride riding a register. + + Identical to ldg_f32 except `* 4` multiplies a caller-held Int32: the + row base is uniformized into URx by ptxas, IMAD.WIDE cannot encode an + immediate stride next to a UR addend, and a constant stride register + would otherwise be re-materialized inside the survivor walk. The caller + loads the 4 from smem (LDS results are opaque to ptxas value-tracking; + asm movs and shfl are not), so the register stays live and the remat + disappears.""" atom = C.g2r_atom_f32(32, invariant=True) p = cute.make_ptr( cutlass.Float32, @@ -1315,14 +1267,13 @@ def _ldg_f32_rs(base_addr, idx, sc4): @dsl_user_op def _smem_addr_reg(addr, *, loc=None, ip=None): - """Pin a CTA-shared 32-bit byte address in ONE register (ct_reg A1 donor). + """Pin a CTA-shared 32-bit byte address in ONE register. Identity `mov` behind an asm boundary: without it LLVM re-folds the `mov.b32 %r, __dynamic_shmem__0` symbol materialisation into EVERY use - site inside the divergent emission bit-walk — the residual +1 IMAD.MOV - per survivor at BS256 (VIOLATION_TRIAGE_20260814 §3). The asm result is - not duplicable, so the shared window is materialised exactly once. - Value-identical: a plain register copy.""" + site inside the divergent emission bit-walk (one extra IMAD.MOV per + survivor). The asm result is not duplicable, so the shared window is + materialised exactly once. Value-identical: a plain register copy.""" return cutlass.Int32( mlir_llvm.inline_asm( T.i32(), @@ -1340,12 +1291,12 @@ def _smem_addr_reg(addr, *, loc=None, ip=None): @dsl_user_op def _red_shared_add1(addr, *, loc=None, ip=None): - """CUDA `atomicAdd(&hist[bin], 1u)` with the result unused (ct_reg A1). + """CUDA `atomicAdd(&hist[bin], 1u)` with the result unused. `red` (not `atom`) is the result-less spelling — ptxas lowers it to the same ATOMS.POPC.INC.32 RZ the CUDA arm emits. Same ordering contract as - ct_common.atomic_add_cta (.relaxed scope .cta). Takes the final shared - byte address as a plain Int32 so ptxas fuses the shl+add into one LEA + atomic_add_cta (.relaxed scope .cta). Takes the final shared byte + address as a plain Int32 so ptxas fuses the shl+add into one LEA against the pinned `_smem_addr_reg` base.""" mlir_llvm.inline_asm( res=None, @@ -1360,7 +1311,7 @@ def _red_shared_add1(addr, *, loc=None, ip=None): class GvrMainKernel: - """CuTeDSL port of gvr_main (kernel.cu L377).""" + """gvr_main — streaming self-sampling GVR.""" def __init__( self, @@ -1376,7 +1327,7 @@ def __init__( cr_shift: int = 0, r_const: int = 1, ): - assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177, measured)" + assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) assert kpt in (1, 2, 4, 8) and minb in (1, 2, 4) self.blk = blk @@ -1397,28 +1348,26 @@ def __init__( self.r_const = int(r_const) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) and self.r_const >= 1 - # knife5 (layer 7): TSH-floor staging arm. SPLIT-only compile-time - # key; the CUDA form is a grid-uniform runtime gate over the same - # predicate (b > 15 && k <= 1024 && n4 <= 32768). varlen mode - # compiles the machinery in whenever SPLIT and gates it per row at - # runtime (tsh_en && n4 <= 32768) — mirroring the CUDA runtime gate. + # TSH-floor staging arm. SPLIT-only compile-time key; the CUDA form + # is a grid-uniform runtime gate over the same predicate + # (b > 15 && k <= 1024 && n4 <= 32768). varlen mode compiles the + # machinery in whenever SPLIT and gates it per row at runtime + # (tsh_en && n4 <= 32768) — mirroring the CUDA runtime gate. if self.varlen: self.tshg = bool(split) else: self.tshg = bool(tshg) and bool(split) - # derived constexprs (kernel.cu L394-523) - self.hb = nbs # L394 - self.kbig = (kpt >= 2) and (kpt * blk >= 2048) # L413 - self.scpb = ( - (8192 if split else 16384) if blk >= 1024 else (8192 if self.kbig else 4096) - ) # L422-423 - self.cmpb = (4096 if self.kbig else 2048) if blk >= 1024 else 1024 # L424 - self.shd = not split # L438 - self.vstg = split or blk >= 512 # L445 - self.pfd = (u if u < 4 else 4) if minb <= 2 else 0 # L522 - self.pf = self.pfd > 0 # L523 - self.natt = 1 if split else 3 # L733 - # smem blob byte map (kernel.cu L447-456): cbuf/cbuf2 alias @0, + # derived constexprs (bit-identical to the CUDA) + self.hb = nbs + self.kbig = (kpt >= 2) and (kpt * blk >= 2048) + self.scpb = (8192 if split else 16384) if blk >= 1024 else (8192 if self.kbig else 4096) + self.cmpb = (4096 if self.kbig else 2048) if blk >= 1024 else 1024 + self.shd = not split + self.vstg = split or blk >= 512 + self.pfd = (u if u < 4 else 4) if minb <= 2 else 0 + self.pf = self.pfd > 0 + self.natt = 1 if split else 3 + # smem blob byte map: cbuf/cbuf2 alias @0, # ck64 @ 4*(VSTG ? 2*(SCPB+4) : SCPB+4), size (CMPB+1)*8 self.ck_off = 4 * ((2 * (self.scpb + 4)) if self.vstg else (self.scpb + 4)) assert self.ck_off % 16 == 0, "ck64 must stay 16B aligned (ulonglong2)" @@ -1426,25 +1375,25 @@ def __init__( self.lb = self.nbs.bit_length() - 1 # log2(NBS)=8 # ------------------------------------------------------------------ - # GVR_EMITC (kernel.cu L869-883): classify+stage one survivor. - # Returns pos+1. Branchless trash slot min(pos, SCPB) (L866-868). + # GVR_EMITC: classify+stage one survivor. + # Returns pos+1. Branchless trash slot min(pos, SCPB). # ------------------------------------------------------------------ @cute.jit def _emitc(self, xv, idx, pos, TF, SC, hb, cb2, s_hist, s_cbuf, s_cbuf2): SCPB = self.scpb NBS = self.nbs if cutlass.const_expr(not self.split): - bn_u = C.f2u_rz((xv - TF) * SC) # saturating cvt.rzi (P4) + bn_u = C.f2u_rz((xv - TF) * SC) # saturating cvt.rzi if bn_u > cutlass.Uint32(NBS - 1): bn_u = cutlass.Uint32(NBS - 1) bn = cutlass.Int32(bn_u) if cutlass.const_expr(self.vstg): - # fix-2 P1: result unused -> resultless red off the pinned - # hist base (ct_reg A1 idiom) — no per-site smem-base refold + # result unused -> resultless red off the pinned hist base + # (no per-site smem-base refold) _red_shared_add1(hb + (bn << cutlass.Int32(2))) else: - # VSTG=False tuples sit at the 64-reg wall (fix-1 spill-5 - # lesson): keep the original spelling, no pinned base here + # VSTG=False tuples sit at the 64-register wall: keep the + # original spelling, no pinned base here C.atomic_add_cta(s_hist.iterator + bn, cutlass.Int32(1)) if cutlass.const_expr(not self.vstg): ps = pos @@ -1458,16 +1407,16 @@ def _emitc(self, xv, idx, pos, TF, SC, hb, cb2, s_hist, s_cbuf, s_cbuf2): if ps > cutlass.Int32(SCPB): ps = cutlass.Int32(SCPB) # int2 {value bits, idx} via st.shared.v2.u32 — same bytes as the - # former (idx << 32) | bits u64 pack (+0=bits, +4=idx), but no i64 - # materialization inside the bit-walk (kernel.cu L1018-1019 parity) - # fix-2 P1: address = one LEA off the pinned cb2 base + # (idx << 32) | bits u64 pack (+0=bits, +4=idx), but no i64 + # materialization inside the bit-walk; address = one LEA off the + # pinned cb2 base _st_s_v2_u32(cb2 + ps * cutlass.Int32(8), C.u32_of_f32(xv), cutlass.Uint32(idx)) return pos + cutlass.Int32(1) # ------------------------------------------------------------------ - # two-predicate warp-ballot emit step (shared by P6 L1124-1146 and both - # degen emits L1185-1210 / L1242-1263): q1 winners to out[base1+p] p> cutlass.Int32(2)) <= cutlass.Int32(32768): tsh_run = cutlass.Int32(1) - # ---- shared memory (one blob, compile-time offsets; spec §5.1 map) ---- + # ---- shared memory (one blob, compile-time offsets) ---- smem = SmemAllocator() s_hist = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((self.hb,), order=(0,)), byte_alignment=128 - ) # L446 __align__(16) - s_ws = smem.allocate_tensor( # L458 (unused; byte parity) # noqa: F841 + ) + s_ws = smem.allocate_tensor( # unused; byte parity # noqa: F841 cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 ) s_wmn = smem.allocate_tensor( @@ -1612,37 +1561,37 @@ def kern( s_wmx = smem.allocate_tensor( cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 ) - # ct_common crossing-scan result slots (RES_B/M/ABOVE/TOT/B2/B3) + # crossing-scan result slots (RES_B/M/ABOVE/TOT/B2/B3) s_res = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((8,), order=(0,)), byte_alignment=16 ) - # scalar block: [0]=s_bufn [1]=s_o1 [2]=s_o2 [3]=s_base (L459) + # scalar block: [0]=s_bufn [1]=s_o1 [2]=s_o2 [3]=s_base s_scal = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 ) - s_pk = smem.allocate_tensor( # L460 + s_pk = smem.allocate_tensor( cutlass.Int64, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=8 ) - s_tsh = smem.allocate_tensor( # L462 + s_tsh = smem.allocate_tensor( cutlass.Float32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 ) - # fix-3 P1d: STATIC smem word for the walk's byte stride — kept out - # of the blob so dyn_bytes stays equal to the CUDA dispatch's smem - # (test_main_smoke dispatch-parity assert). blk==512 VSTG only. + # STATIC smem word for the walk's byte stride — kept out of the blob + # so dyn_bytes stays equal to the CUDA dispatch's smem. blk==512 VSTG + # only. if cutlass.const_expr(self.vstg and self.blk == 512): s_x4 = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 ) - s_kmm = smem.allocate_tensor( # L463 [0]=kmin [1]=kmax + s_kmm = smem.allocate_tensor( # [0]=kmin [1]=kmax cutlass.Uint32, cute.make_ordered_layout((2,), order=(0,)), byte_alignment=8 ) if cutlass.const_expr(self.varlen): - # P2a ladder broadcast slots: [0]=SMP [1]=SS2 [2]=TGT [3]=TGT2 + # ladder broadcast slots: [0]=SMP [1]=SS2 [2]=TGT [3]=TGT2 # (static like s_x4, so dyn_bytes keeps CUDA dispatch parity) s_lad = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 ) - blob = smem.allocate_tensor( # dynamic-equivalent L447-456 + blob = smem.allocate_tensor( # dynamic-equivalent region cutlass.Int8, cute.make_ordered_layout((self.dyn_bytes,), order=(0,)), byte_alignment=16 ) sbase = blob.iterator.toint() @@ -1660,24 +1609,22 @@ def kern( cute.make_layout((CMPB + 1,)), ) - # fix-2 P1: emission smem bases pinned ONCE, outside the attempt/tile - # loops (asm identity mov, ct_reg A1 donor) — LLVM otherwise refolds - # the shared-window materialisation into every _emitc site inside the - # divergent bit-walk (residual +1 IMAD.MOV, VIOLATION_TRIAGE §3). - # VSTG-only: the (256,8,4,·) VSTG=False tuples keep their original - # spellings untouched (64-reg wall, fix-1 spill-5 lesson). + # emission smem bases pinned ONCE, outside the attempt/tile loops + # (asm identity mov) — LLVM otherwise refolds the shared-window + # materialisation into every _emitc site inside the divergent + # bit-walk. VSTG-only: the VSTG=False tuples keep their original + # spellings untouched (64-register wall). hb_pin = cutlass.Int32(0) cb2_pin = cutlass.Int32(0) if cutlass.const_expr(self.vstg): hb_pin = _smem_addr_reg(s_hist.iterator.toint()) cb2_pin = _smem_addr_reg(s_cbuf2.iterator.toint()) - # fix-3 P1d: park 4 in the dedicated smem word and load it back — + # park the stride 4 in the dedicated smem word and load it back — # the LDS result is opaque to ptxas, so the walk's stride register - # cannot be re-materialized in-loop (P1c's asm-mov and shfl forms - # were both folded by ptxas value-tracking). blk==512 family ONLY: - # the (256,8,4,·) arm is bit-frozen and 1024/SPLIT sit at the - # 64-reg wall. Threads are converged here (kernel prologue), so - # the one extra barrier is safe and costs ~nothing once per launch. + # cannot be re-materialized in-loop (asm-mov and shfl forms are + # folded by ptxas value-tracking). blk==512 family ONLY: the other + # arms sit at the 64-register wall. Threads are converged here + # (kernel prologue), so the one extra barrier is safe. x4_pin = cutlass.Int32(4) if cutlass.const_expr(self.vstg and self.blk == 512): if tidx == cutlass.Int32(0): @@ -1685,7 +1632,7 @@ def kern( cute.arch.barrier() x4_pin = s_x4[0] - # ---- row bases (L472-475) ---- + # ---- row bases ---- row64 = cutlass.Int64(row) # _pin_i64: keep the row base a REGISTER across the attempt/tile scf # regions (NVVM otherwise re-derives ld.param+%ctaid.y+mul per region) @@ -1699,66 +1646,60 @@ def kern( p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) out_row = out[row, None] ws_addr = ws.iterator.toint() - gdon_addr = ws_addr # L386-388 slab views + gdon_addr = ws_addr # slab views goff_addr = ws_addr + cutlass.Int64(C.GVR_WS_OFF_OFF) gbuf_addr = ws_addr + cutlass.Int64(C.GVR_WS_BUF_OFF) - # fix-2 P2 (SPLIT only): row-slab base pinned like x_addr above; the + # SPLIT only: row-slab base pinned like x_addr above; the # publish/gather/P5/degen consumers spell gbuf_row + i*8 instead of - # re-deriving gbuf_addr + (row64*GCAP__main + i)*8 per candidate (value- - # identical by i64 distributivity). CUDA parity: the (1024,4,1,256,2, - # SPLIT) SASS keeps this base in UR8/UR9 with ONE IMAD.WIDE.U32 per - # candidate (cand_K5a_l10 cuobjdump); DSL-side remat signature is - # hypothesis-driven — SASS verify pending drain window. + # re-deriving gbuf_addr + (row64*GCAP__main + i)*8 per candidate + # (value-identical by i64 distributivity). gbuf_row = cutlass.Int64(0) if cutlass.const_expr(self.split): gbuf_row = _pin_i64(gbuf_addr + row64 * cutlass.Int64(GCAP__main) * cutlass.Int64(8)) - n4 = n >> cutlass.Int32(2) # L477 + n4 = n >> cutlass.Int32(2) c0 = cutlass.Int32(0) c1 = n4 - if cutlass.const_expr(self.split): # L478-479 + if cutlass.const_expr(self.split): c0 = part * Q c1 = c0 + Q if c1 > n4: c1 = n4 - tail0 = n4 << cutlass.Int32(2) # L480 + tail0 = n4 << cutlass.Int32(2) tailn = cutlass.Int32(0) - if part == cutlass.Int32(0): # L481 + if part == cutlass.Int32(0): tailn = n - tail0 - if tidx == cutlass.Int32(0): # L483 + if tidx == cutlass.Int32(0): s_scal[0] = cutlass.Int32(0) # s_bufn s_res[C.RES_B2] = cutlass.Int32(-1) s_res[C.RES_B3] = cutlass.Int32(-1) - if tidx < cutlass.Int32(self.hb): # L484-487 (HB<=BLK always) + if tidx < cutlass.Int32(self.hb): # HB<=BLK always s_hist[tidx] = cutlass.Int32(0) - # ===== varlen P2a: warp0-only ladder mirror + register-free L2 hints = - # The sampling-ladder scalars are a pure function of the row, and the - # per-thread mirror chain cost more instructions than the rest of the - # kernel on 1-row launches (l1p1 ncu: inst x1.58-1.73, all of it the - # DSL's runtime Int32/Int64 divides + isqrt fixups, redundantly issued - # by every thread). warp0 alone walks the chain and publishes the four - # derived scalars through s_lad; the other warps spend the wait issuing - # L2 prefetch hints for this CTA's own P3 slice (register-free, so zero - # pressure on the 64-reg arms — the PRIME-LATE register loads below are + # ===== varlen: warp0-only ladder mirror + register-free L2 hints ===== + # The sampling-ladder scalars are a pure function of the row; issuing + # the mirror chain (runtime divides + isqrt fixups) per thread costs + # more instructions than the rest of the kernel on 1-row launches. + # warp0 alone walks the chain and publishes the four derived scalars + # through s_lad; the other warps spend the wait issuing L2 prefetch + # hints for this CTA's own P3 slice (register-free, so zero pressure + # on the 64-register arms — the PRIME-LATE register loads below are # untouched and simply hit L2). Values are bit-identical to the # per-thread derivation this replaces. if cutlass.const_expr(self.varlen): if tidx < cutlass.Int32(32): if short == cutlass.Int32(0): - # ---- aim ladder (P2b cheap mirror) ---- + # ---- aim ladder (cheap mirror) ---- # The ladder scalars steer the sampling rung only — # exactness is schedule-invariant (retry/degen close every - # miss), so schedule-quantity drift of +-1 vs the host - # double form is sanctioned (same argument as the clus - # port's always-sample deviation). Serial latency is what - # matters here (this chain sits in front of a barrier): - # runtime divides become MUFU.RCP multiplies and the isqrt - # fixup loops collapse to single steps (the f32 sqrt of an - # exactly-representable int (6n <= 2^23) is within 1 of - # isqrt, so one correction per side suffices). Q (chunk - # ownership) stays exact — compile-time divisor, all-thread. + # miss), so +-1 drift vs the host double form is allowed. + # Serial latency dominates (this chain sits in front of a + # barrier): runtime divides become MUFU.RCP multiplies and + # the isqrt fixup loops collapse to single steps (the f32 + # sqrt of an exactly-representable int (6n <= 2^23) is + # within 1 of isqrt, so one correction per side suffices). + # Q (chunk ownership) stays exact — compile-time divisor. x6 = cutlass.Int32(6) * n ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) if ri * ri > x6: @@ -1845,9 +1786,9 @@ def kern( s_lad[2] = TGT s_lad[3] = TGT2 # Register-free L2 hints for the first U-batch of this CTA's own - # P3 slice (clamped in-row, prefetch site #3 spelling): the data - # P3 touches first starts flowing while warp0 walks the chain. - # Short rows clamp every hint to the row's last line — harmless. + # P3 slice (clamped in-row): the data P3 touches first starts + # flowing while warp0 walks the chain. Short rows clamp every + # hint to the row's last line — harmless. plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) for uu in cutlass.range_constexpr(U): # NOTE: names must not collide with the PRIME-LATE block's @@ -1864,22 +1805,22 @@ def kern( TGT = s_lad[2] TGT2 = s_lad[3] - # ============ P1: sample prefetch (hint gather LAZY, L489-529) ======= + # ============ P1: sample prefetch (hint gather LAZY) ================= atom128 = C.g2r_atom_f32(128, invariant=True) fsa = cute.make_rmem_tensor((4,), cutlass.Float32) fsb = cute.make_rmem_tensor((4,), cutlass.Float32) shas = cutlass.Int32(0) if tidx < SMP: shas = cutlass.Int32(1) - if shas != cutlass.Int32(0): # L502-504 + if shas != cutlass.Int32(0): p4 = tidx * SS2 * cutlass.Int32(2) C.ld_g_f32x4(atom128, x_addr, p4, fsa) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) - # ============ P2: quantile rung from the sample (L531-727) =========== - smn = cutlass.Float32(float("inf")) # L538 + # ============ P2: quantile rung from the sample ====================== + smn = cutlass.Float32(float("inf")) smx = cutlass.Float32(float("-inf")) - if shas != cutlass.Int32(0): # L539-543 + if shas != cutlass.Int32(0): for t in cutlass.range_constexpr(4): smn = C.fmin_f32(smn, fsa[t]) smx = C.fmax_f32(smx, fsa[t]) @@ -1888,7 +1829,7 @@ def kern( smx = C.fmax_f32(smx, fsb[t]) fma_ = cute.make_rmem_tensor((4,), cutlass.Float32) # strided-tail pair bufs fmb_ = cute.make_rmem_tensor((4,), cutlass.Float32) - j = tidx + cutlass.Int32(BLK) # L544-550 strided tail + j = tidx + cutlass.Int32(BLK) # strided tail while j < SMP: p4 = j * SS2 * cutlass.Int32(2) C.ld_g_f32x4(atom128, x_addr, p4, fma_) @@ -1900,31 +1841,31 @@ def kern( smn = C.fmin_f32(smn, fmb_[t]) smx = C.fmax_f32(smx, fmb_[t]) j = j + cutlass.Int32(BLK) - a0 = C.warp_min_u32(C.fkey(smn)) # L551-554 + a0 = C.warp_min_u32(C.fkey(smn)) c0m = C.warp_max_u32(C.fkey(smx)) if lane == cutlass.Int32(0): s_wmn[tidx >> cutlass.Int32(5)] = a0 s_wmx[tidx >> cutlass.Int32(5)] = c0m - cute.arch.barrier() # ---- barrier L555 ---- + cute.arch.barrier() # ---- barrier (sample redux publish) ---- - # PRIME-LATE prefetch block (L556-616): strictly after the barrier. - lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L524 + # PRIME-LATE prefetch block: strictly after the barrier. + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(PFD, 1))] if cutlass.const_expr(self.pf): fullsl = cutlass.Int32(0) if (c1 - c0) >= cutlass.Int32(BLK * U): fullsl = cutlass.Int32(1) - if fullsl != cutlass.Int32(0): # L557-559 prime, full slice + if fullsl != cutlass.Int32(0): # prime, full slice for uu in cutlass.range_constexpr(PFD): C.ld_g_f32x4(atom128, x_addr, c0 + tidx + cutlass.Int32(uu * BLK), pf[uu]) - else: # L561-562 clamped prime + else: # clamped prime for uu in cutlass.range_constexpr(PFD): i_ = c0 + tidx + cutlass.Int32(uu * BLK) ic = i_ if ic >= c1: ic = lim4 C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - # asm prefetch site #1 (L573-577): gate (c1-c0)>=2*BLK*U && SMP>=160 + # asm prefetch site #1: gate (c1-c0)>=2*BLK*U && SMP>=160 g1 = cutlass.Int32(0) if (c1 - c0) >= cutlass.Int32(2 * BLK * U): if SMP >= cutlass.Int32(160): @@ -1939,14 +1880,14 @@ def kern( fullsl = cutlass.Int32(0) if (c1 - c0) >= cutlass.Int32(BLK * U): fullsl = cutlass.Int32(1) - if fullsl != cutlass.Int32(0): # site #2 (L589-592) + if fullsl != cutlass.Int32(0): # prefetch site #2 for uu in cutlass.range_constexpr(U): C._prefetch_l2( x_addr + cutlass.Int64(c0 + tidx + cutlass.Int32(uu * BLK)) * cutlass.Int64(16) ) else: - if SMP > cutlass.Int32(0): # site #3 knife4-L1 (L593-615) + if SMP > cutlass.Int32(0): # prefetch site #3 for uu in cutlass.range_constexpr(U): i_ = c0 + tidx + cutlass.Int32(uu * BLK) ic = i_ @@ -1954,7 +1895,7 @@ def kern( ic = lim4 C._prefetch_l2(x_addr + cutlass.Int64(ic) * cutlass.Int64(16)) - # cross-warp sample reduce (L617-623) + # cross-warp sample reduce av = cutlass.Uint32(0xFFFFFFFF) cv = cutlass.Uint32(0) if lane < cutlass.Int32(NW): @@ -1963,16 +1904,16 @@ def kern( SMIN = C.invkey(C.warp_min_u32(av)) SMAX = C.invkey(C.warp_max_u32(cv)) - GMIN = cutlass.Float32(C.SENT_LO) # L629 sentinels + GMIN = cutlass.Float32(C.SENT_LO) # sentinels GMAX = cutlass.Float32(C.SENT_HI) T = cutlass.Float32(_NEG_INF) HIC = cutlass.Float32(_NEG_INF) w = cutlass.Float32(0.0) - sok = cutlass.Int32(0) # L633 + sok = cutlass.Int32(0) if SMP > cutlass.Int32(0): if SMAX > SMIN: sok = cutlass.Int32(1) - if sok != cutlass.Int32(0): # L635-654 sample histogram + if sok != cutlass.Int32(0): # sample histogram w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) # rcp.approx.ftz.f32 = the CUDA arm's --use_fast_math 1.0f/w # (bare MUFU.RCP, no Newton refinement) — bitwise-aligned scale @@ -1988,7 +1929,7 @@ def kern( if bq > cutlass.Int32(NBS - 1): bq = cutlass.Int32(NBS - 1) C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) - j = tidx + cutlass.Int32(BLK) # L646-653 tail re-loads + j = tidx + cutlass.Int32(BLK) # tail re-loads while j < SMP: p4 = j * SS2 * cutlass.Int32(2) C.ld_g_f32x4(atom128, x_addr, p4, fma_) @@ -2004,9 +1945,9 @@ def kern( bq = cutlass.Int32(NBS - 1) C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) j = j + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L655 ---- - # triple-target ZERO scan (L659-660): TGT / TGT2 / 2*TGT - # (THREE = SHD || gated-SPLIT, knife5 layer 7) + cute.arch.barrier() # ---- barrier (sample histogram) ---- + # triple-target ZERO scan: TGT / TGT2 / 2*TGT + # (THREE = SHD || gated-SPLIT) C.scan_cross0( s_hist, TGT, @@ -2020,15 +1961,15 @@ def kern( two=True, three=(self.shd or self.tshg), ) - cute.arch.barrier() # ---- barrier L664 ---- + cute.arch.barrier() # ---- barrier (scan publish) ---- tot0 = s_res[C.RES_TOT] b1v = s_res[C.RES_B] - if sok != cutlass.Int32(0): # L665 + if sok != cutlass.Int32(0): if tot0 >= TGT: T = _fmaf(cutlass.Float32(b1v), w, SMIN) - Trung = T # L666 snapshot - needg = cutlass.Int32(1) # L667-675 degenerate sample + Trung = T # snapshot + needg = cutlass.Int32(1) # degenerate sample if T > cutlass.Float32(_NEG_INF): needg = cutlass.Int32(0) if needg != cutlass.Int32(0): @@ -2036,19 +1977,19 @@ def kern( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT ) # 2 barriers inside T = GMIN - if sok != cutlass.Int32(0): # L688-709 HIC tighten + if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: b2v = s_res[C.RES_B2] if b2v >= cutlass.Int32(0): Tk = _fmaf(cutlass.Float32(b2v), w, SMIN) anch = T if cutlass.const_expr(not self.split): - anch = C.fmin_f32(T, Trung) # L703 + anch = C.fmin_f32(T, Trung) d_ = C.fmax_f32(Tk - anch, cutlass.Float32(0.0)) HIC = C.fmax_f32( _fmaf(cutlass.Float32(4.0), d_, T), _fmaf(cutlass.Float32(8.0), w, T) - ) # L708 - if cutlass.const_expr(self.shd or self.tshg): # TSH floor (knife5: +gated SPLIT) + ) + if cutlass.const_expr(self.shd or self.tshg): # TSH floor (+gated SPLIT) if tidx == cutlass.Int32(0): t5 = cutlass.Float32(_NEG_INF) if sok != cutlass.Int32(0): @@ -2062,9 +2003,9 @@ def kern( s_tsh[0] = t5 if cutlass.const_expr(self.tshg): - # knife5 (layer 7) TSH-FLOOR STAGING: SPLIT has no retry ladder, - # so a rung overshoot (count(>=T) < k) used to hand the LAST CTA - # a single-CTA whole-row narrowing. Stage at the sample's + # TSH-FLOOR STAGING: SPLIT has no retry ladder, so a rung + # overshoot (count(>=T) < k) used to hand the LAST CTA a + # single-CTA whole-row narrowing. Stage at the sample's # rank-(2*TGT) floor instead: staged population ~aim -> ~2*aim, # and the merged histogram contains the k-crossing whenever # count(>=TSH) >= k. TSH miss falls to GMIN/degen unchanged. @@ -2077,7 +2018,7 @@ def kern( if t5s < T: T = t5s - # ============ attempt loop (L729-1019) — MUST NOT unroll ============ + # ============ attempt loop — MUST NOT unroll ============ listN = cutlass.Int32(0) above = cutlass.Int32(0) m = cutlass.Int32(0) @@ -2092,14 +2033,14 @@ def kern( fr = [ cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(U - PFD, 1)) - ] # explicit batch (op43 L1) + ] # explicit batch att = cutlass.Int32(0) running = cutlass.Int32(1) while running != cutlass.Int32(0): if cutlass.const_expr(not self.split): # SPLIT never retries (NATT=1) - if att > cutlass.Int32(0): # retry reset (L737-754) + if att > cutlass.Int32(0): # retry reset if cutlass.const_expr(self.pf): - # EXACTNESS: re-prime pf[] (stale roll data, L738-749) + # exactness: re-prime pf[] (holds stale roll data) fullsl = cutlass.Int32(0) if (c1 - c0) >= cutlass.Int32(BLK * U): fullsl = cutlass.Int32(1) @@ -2115,13 +2056,13 @@ def kern( if ic >= c1: ic = lim4 C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - if tidx < cutlass.Int32(NBS): # L750-751 + if tidx < cutlass.Int32(NBS): s_hist[tidx] = cutlass.Int32(0) - if tidx == cutlass.Int32(0): # L752 + if tidx == cutlass.Int32(0): s_scal[0] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L753 ---- + cute.arch.barrier() # ---- barrier (retry reset) ---- - TF = T # window (L756-761) + TF = T # window hi = C.fmax_f32(GMAX, T) if HIC > T: if HIC < hi: @@ -2132,30 +2073,26 @@ def kern( wdok = cutlass.Int32(1) if wdok == cutlass.Int32(0): WD = cutlass.Float32(1e-30) - # fix-4: CUDA compiles its own `1.0f / WD` here to a bare - # MUFU.RCP (approximate) — div.rn's dependent rcp+Newton+CALL - # chain serializes the attempt prologue and is the last - # CALL-class divergence (era-3 ncu: DSL executes ~4% FEWER - # instructions yet runs ~5% slower = latency-shaped, so the - # cost is the chain, not the instruction count). blk==512 - # ONLY: the fix-1 spill-5 revert was bisected on (256,8,4,·,4) - # — that family keeps the original div.rn spelling below. + # CUDA compiles its own `1.0f / WD` here to a bare MUFU.RCP + # (approximate); div.rn's dependent rcp+Newton+CALL chain + # serializes the attempt prologue. blk==512 ONLY: the + # (256,8,4,·) family keeps the original div.rn spelling below. if cutlass.const_expr(self.blk == 512): SC = cute.arch.rcp_approx(WD) else: SC = cutlass.Float32(1.0) / WD - # ---- P3 row pass (L763-908) ---- + # ---- P3 row pass ---- span = c1 - c0 step = cutlass.Int32(BLK * U) nFull = cutlass.Int32(0) rem = cutlass.Int32(0) - if span > cutlass.Int32(0): # L776-779 peel + if span > cutlass.Int32(0): # peel nFull = span // step rem = span - nFull * step # _pin_i32: the isfull peel predicate reads nFull every tile iter; # unpinned, NVVM re-derives the whole ld.param+shr/sel div chain - # at the loop head (v3 SASS evidence) + # at the loop head nFull = _pin_i32(nFull) nIt = nFull if rem > cutlass.Int32(0): @@ -2171,7 +2108,7 @@ def kern( isfull = cutlass.Int32(0) if it < nFull: isfull = cutlass.Int32(1) - if isfull != cutlass.Int32(0): # full body (L783-795) + if isfull != cutlass.Int32(0): # full body for uu in cutlass.range_constexpr(PFD, U): C.ld_g_f32x4(atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD]) for uu in cutlass.range_constexpr(U): @@ -2181,12 +2118,12 @@ def kern( vv = fr[uu - PFD] for q in cutlass.range_constexpr(4): M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) - else: # partial body (L796-815) + else: # partial body for uu in cutlass.range_constexpr(PFD, U): i_ = i0 + cutlass.Int32(uu * BLK) ic = i_ if ic >= c1: - ic = lim4 # clamped address (P10) + ic = lim4 # clamped address C.ld_g_f32x4(atom128, x_addr, ic, fr[uu - PFD]) for uu in cutlass.range_constexpr(U): if cutlass.const_expr(uu < PFD): @@ -2197,17 +2134,17 @@ def kern( okq = cutlass.Int32(0) if i_ < c1: okq = cutlass.Int32(1) - if okq != cutlass.Int32(0): # ok-gated (+inf pad fix L804-813) + if okq != cutlass.Int32(0): # ok-gated (+inf-pad escape) for q in cutlass.range_constexpr(4): M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) - # prefetch roll-forward BEFORE reservation/walk (L816-842) + # prefetch roll-forward BEFORE reservation/walk if cutlass.const_expr(self.pf): hasnext = cutlass.Int32(0) if it + cutlass.Int32(1) < nIt: hasnext = cutlass.Int32(1) if hasnext != cutlass.Int32(0): j0 = i0 + step - infull = cutlass.Int32(0) # warp-uniform peel L835 + infull = cutlass.Int32(0) # warp-uniform peel if it + cutlass.Int32(1) < nFull: infull = cutlass.Int32(1) if infull != cutlass.Int32(0): @@ -2220,7 +2157,7 @@ def kern( if jc >= c1: jc = lim4 C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) - # warp-aggregated reservation (L843-854) + # warp-aggregated reservation cnt = cutlass.Int32(C.popc(M)) inc = C.warp_incl_scan_add(cnt, lane) bpos = cutlass.Int32(0) @@ -2228,8 +2165,8 @@ def kern( if inc != cutlass.Int32(0): bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) - # survivor bit-walk, software-pipelined ONE deep (L884-898); - # reload X[idx] — do NOT hold the U float4s (+18% spill L855-859) + # survivor bit-walk, software-pipelined ONE deep; + # reload X[idx] — do NOT hold the U float4s (spills) if M != cutlass.Int32(0): bp = C.ffs_m1(M) M = M & (M - cutlass.Int32(1)) @@ -2260,7 +2197,7 @@ def kern( xv, idx, pos, TF, SC, hb_pin, cb2_pin, s_hist, s_cbuf, s_cbuf2 ) it = it + cutlass.Int32(1) - # scalar tail, part 0 only (L900-906) + # scalar tail, part 0 only i = tidx while i < tailn: x = C.ldg_f32(x_addr, tail0 + i) @@ -2270,12 +2207,12 @@ def kern( x, tail0 + i, post, TF, SC, hb_pin, cb2_pin, s_hist, s_cbuf, s_cbuf2 ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L909 ---- - myn = s_scal[0] # L911 + cute.arch.barrier() # ---- barrier (row pass) ---- + myn = s_scal[0] if cutlass.const_expr(self.split): - # ---- SLAB HAND-OFF (L913-989); exactly ONE attempt ---- - if tidx == cutlass.Int32(0): # L923 + # ---- SLAB HAND-OFF; exactly ONE attempt ---- + if tidx == cutlass.Int32(0): pgo = cute.make_ptr( cutlass.Int32, goff_addr + row64 * cutlass.Int64(4), @@ -2283,19 +2220,19 @@ def kern( assumed_align=4, ) s_scal[3] = cutlass.Int32(cute.arch.atomic_add(pgo, myn)) - cute.arch.barrier() # ---- barrier L924 ---- + cute.arch.barrier() # ---- barrier (slab offset) ---- base = s_scal[3] - if myn <= cutlass.Int32(SCPB): # L926-930 coalesced publish + if myn <= cutlass.Int32(SCPB): # coalesced publish i = tidx while i < myn: p = base + i if p < cutlass.Int32(GCAP__main): _st_g_u64(gbuf_row + cutlass.Int64(p) * cutlass.Int64(8), s_cbuf2[i]) i = i + cutlass.Int32(BLK) - else: # L931-955 overflow re-sweep + else: # overflow re-sweep if tidx == cutlass.Int32(0): s_scal[0] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L935 ---- + cute.arch.barrier() # ---- barrier (overflow reset) ---- lo2 = c0 << cutlass.Int32(2) hi2 = c1 << cutlass.Int32(2) i = lo2 + tidx @@ -2311,7 +2248,7 @@ def kern( | cutlass.Uint64(C.u32_of_f32(x)), ) i = i + cutlass.Int32(BLK) - i = tidx # true tail (L948-954) + i = tidx # true tail while i < tailn: x = C.ldg_f32(x_addr, tail0 + i) if x >= TF: @@ -2327,8 +2264,8 @@ def kern( | cutlass.Uint64(C.u32_of_f32(x)), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L956 ---- - if tidx == cutlass.Int32(0): # L959-960 release + RMW + cute.arch.barrier() # ---- barrier (slab publish) ---- + if tidx == cutlass.Int32(0): # release + RMW C.threadfence_gpu() pdon = cute.make_ptr( cutlass.Int64, @@ -2339,18 +2276,18 @@ def kern( s_pk[0] = C.atomic_add_u64_gpu( pdon, cutlass.Int64(1 << 32) + cutlass.Int64(myn) ) - cute.arch.barrier() # ---- barrier L961 ---- + cute.arch.barrier() # ---- barrier (arrival word) ---- pk = s_pk[0] - alive = cutlass.Int32(0) # L963 last-CTA test + alive = cutlass.Int32(0) # last-CTA test if cutlass.Int32(pk >> cutlass.Int64(32)) == R - cutlass.Int32(1): alive = cutlass.Int32(1) if alive != cutlass.Int32(0): - C.threadfence_gpu() # L968 acquire - if tidx == cutlass.Int32(0): # L969 ZERO-RESTORE + C.threadfence_gpu() # acquire + if tidx == cutlass.Int32(0): # ZERO-RESTORE _st_g_u32(goff_addr + row64 * cutlass.Int64(4), cutlass.Int32(0)) _st_g_u64(gdon_addr + row64 * cutlass.Int64(8), cutlass.Uint64(0)) total = cutlass.Int32(pk & cutlass.Int64(0xFFFFFFFF)) + myn - if total <= cutlass.Int32(GCAP__main): # L971-988 one-pass consume + if total <= cutlass.Int32(GCAP__main): # one-pass consume listN = total if total > cutlass.Int32(SCPB): fromg = cutlass.Int32(1) @@ -2366,10 +2303,10 @@ def kern( bq = C.f2s_rz((C.f32_of_i32(gvx) - TF) * SC) if bq > cutlass.Int32(NBS - 1): bq = cutlass.Int32(NBS - 1) - # fix-2 P2: resultless red off the P1 hist pin + # resultless red off the pinned hist base _red_shared_add1(hb_pin + (bq << cutlass.Int32(2))) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L983 ---- + cute.arch.barrier() # ---- barrier (slab histogram) ---- C.scan_cross0( s_hist, k, @@ -2381,17 +2318,17 @@ def kern( nb=NBS, zero=False, ) - cute.arch.barrier() # ---- barrier L985 ---- - if s_res[C.RES_TOT] >= k: # L986-987 + cute.arch.barrier() # ---- barrier (scan publish) ---- + if s_res[C.RES_TOT] >= k: valid = cutlass.Int32(1) complete = cutlass.Int32(1) above = s_res[C.RES_ABOVE] m = s_res[C.RES_M] need = k - above B = s_res[C.RES_B] - running = cutlass.Int32(0) # L989 break (NATT==1) + running = cutlass.Int32(0) # break (NATT==1) else: - # ---- non-split verify + rung ladder (L990-1018) ---- + # ---- non-split verify + rung ladder ---- C.scan_cross0( s_hist, k, @@ -2403,12 +2340,12 @@ def kern( nb=NBS, zero=False, ) - cute.arch.barrier() # ---- barrier L992 ---- + cute.arch.barrier() # ---- barrier (verify scan) ---- tot = s_res[C.RES_TOT] acc = cutlass.Int32(0) if tot >= k: acc = cutlass.Int32(1) - if acc != cutlass.Int32(0): # L994-998 accept + if acc != cutlass.Int32(0): # accept valid = cutlass.Int32(1) complete = cutlass.Int32(0) if myn <= cutlass.Int32(SCPB): @@ -2420,10 +2357,10 @@ def kern( B = s_res[C.RES_B] running = cutlass.Int32(0) else: - if att == cutlass.Int32(NATT - 1): # L999 ladder exhausted + if att == cutlass.Int32(NATT - 1): # ladder exhausted running = cutlass.Int32(0) else: - tshtaken = cutlass.Int32(0) # L1005-1010 TSH retry + tshtaken = cutlass.Int32(0) # TSH retry if cutlass.const_expr(self.shd): if att == cutlass.Int32(0): T5 = s_tsh[0] @@ -2432,24 +2369,24 @@ def kern( T = T5 tshtaken = cutlass.Int32(1) if tshtaken != cutlass.Int32(0): - cute.arch.barrier() # ---- barrier L1008 ---- + cute.arch.barrier() # ---- barrier (TSH retry) ---- else: - # LAZY GATHER (L1014, sentinel equality flag) + # LAZY GATHER (sentinel equality flag) if GMIN == cutlass.Float32(C.SENT_LO): GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=KPT ) - floorhit = cutlass.Int32(1) # L1015 + floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) if floorhit != cutlass.Int32(0): running = cutlass.Int32(0) else: - T = GMIN # L1016 - cute.arch.barrier() # ---- barrier L1017 ---- + T = GMIN + cute.arch.barrier() # ---- barrier (floor retry) ---- att = att + cutlass.Int32(1) - # ============ classification (L1021-1024) ============ + # ============ classification ============ if alive != cutlass.Int32(0): whole = cutlass.Int32(0) if valid != cutlass.Int32(0): @@ -2468,8 +2405,8 @@ def kern( mc = m if degen == cutlass.Int32(0): - # ---- P5 cursor emit (L1026-1071) ---- - if complete != cutlass.Int32(0): # L1028-1048 + # ---- P5 cursor emit ---- + if complete != cutlass.Int32(0): i = tidx while i < listN: idv = cutlass.Int32(0) @@ -2497,14 +2434,14 @@ def kern( vy = cutlass.Int32(pk64 >> cutlass.Uint64(32)) xv = C.f32_of_i32(vx) idv = vy - bq = C.f2s_rz((xv - TF) * SC) # L1034 + bq = C.f2s_rz((xv - TF) * SC) if bq > cutlass.Int32(NBS - 1): bq = cutlass.Int32(NBS - 1) else: - wpk = cutlass.Uint32(s_cbuf[i]) # L1036-1037 + wpk = cutlass.Uint32(s_cbuf[i]) idv = cutlass.Int32(wpk & cutlass.Uint32(IDXM__main)) bq = cutlass.Int32(wpk >> cutlass.Uint32(IDXB__main)) - if bq >= B: # L1039-1047 + if bq >= B: p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) if p < lim1: out_row[p] = idv @@ -2521,15 +2458,15 @@ def kern( ) | cutlass.Uint64(cutlass.Uint32(idv)) i = i + cutlass.Int32(BLK) else: - # collect overflow: scalar re-sweep, exact tail remap - # (L1049-1070) — zero extra live registers by design + # collect overflow: scalar re-sweep, exact tail remap — + # zero extra live registers by design lo2 = c0 << cutlass.Int32(2) hi2 = c1 << cutlass.Int32(2) i0_ = lo2 + tidx while i0_ < hi2 + tailn: i_ = i0_ if i0_ >= hi2: - i_ = tail0 + (i0_ - hi2) # L1057-1058 + i_ = tail0 + (i0_ - hi2) x = C.ldg_f32(x_addr, i_) if x >= TF: bq = C.f2s_rz((x - TF) * SC) @@ -2548,17 +2485,16 @@ def kern( ) | cutlass.Uint64(cutlass.Uint32(i_)) i0_ = i0_ + cutlass.Int32(BLK) - # ---- P6 refine (L1073-1147) ---- + # ---- P6 refine ---- if whole == cutlass.Int32(0): - cute.arch.barrier() # ---- barrier L1075 ---- - if mc <= cutlass.Int32(QUADC_CLUS__main): # L1077-1092 O(mc^2) rank + cute.arch.barrier() # ---- barrier (emit done) ---- + if mc <= cutlass.Int32(QUADC_CLUS__main): # O(mc^2) rank mc2 = mc & cutlass.Int32(~1) i = tidx while i < mc: # NOTE: values crossing a dynamic-while region are # re-wrapped SIGNED by the DSL — every u64 compare - # must re-assert Uint64 at the USE site (found via - # odd-tail rank corruption; see notes G1). + # must re-assert Uint64 at the USE site. u64v = s_ck64[i] r_ = cutlass.Int32(0) jq = cutlass.Int32(0) @@ -2570,7 +2506,7 @@ def kern( + cutlass.Int32(vhi > cutlass.Uint64(u64v)) ) jq = jq + cutlass.Int32(2) - if mc2 < mc: # odd tail L1089 + if mc2 < mc: # odd tail r_ = r_ + cutlass.Int32( cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) ) @@ -2582,20 +2518,20 @@ def kern( ) i = i + cutlass.Int32(BLK) else: - # key-space narrowing over ck64 (L1094-1123) + # key-space narrowing over ck64 if tidx == cutlass.Int32(0): s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) s_kmm[1] = cutlass.Uint32(0) - if tidx < cutlass.Int32(NBS): # cleared ONCE (L1096) + if tidx < cutlass.Int32(NBS): # cleared ONCE s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1097 ---- + cute.arch.barrier() # ---- barrier (narrowing init) ---- i = tidx while i < mc: kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) C.atomic_min_cta(s_kmm.iterator + 0, kk) C.atomic_max_cta(s_kmm.iterator + 1, kk) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1100 ---- + cute.arch.barrier() # ---- barrier (key range) ---- rlo = s_kmm[0] rhi = s_kmm[1] ethr = cutlass.Int64(cutlass.Uint32(rlo)) @@ -2604,7 +2540,7 @@ def kern( mm = mc brk = cutlass.Int32(0) lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L1103-1123 (<=6 levels) + while brk == cutlass.Int32(0): # <=6 levels if needC == mm: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) aboveC = aboveC + mm @@ -2626,7 +2562,7 @@ def kern( sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) i = tidx - while i < mc: # L1111-1115 re-bin + while i < mc: # re-bin uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -2638,7 +2574,7 @@ def kern( cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1116 ---- + cute.arch.barrier() # ---- barrier (level hist) ---- C.scan_cross0( s_hist, needC, @@ -2650,23 +2586,23 @@ def kern( nb=NBS, zero=True, ) - cute.arch.barrier() # ---- barrier L1118 ---- + cute.arch.barrier() # ---- barrier (level scan) ---- aboveC = aboveC + s_res[C.RES_ABOVE] needC = needC - s_res[C.RES_ABOVE] mm = s_res[C.RES_M] sB = s_res[C.RES_B] nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) - if sB != cutlass.Int32(NBS - 1): # L1121 + if sB != cutlass.Int32(NBS - 1): rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) rlo = nlo lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L1124 + if tidx == cutlass.Int32(0): s_scal[1] = cutlass.Int32(0) s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1125 ---- + cute.arch.barrier() # ---- barrier (emit counters) ---- it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) it = cutlass.Int32(0) - while it < it2: # L1127-1146 ballot emit + while it < it2: # ballot emit i = it * cutlass.Int32(BLK) + tidx p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) @@ -2695,13 +2631,12 @@ def kern( ) it = it + cutlass.Int32(1) else: - dga = cutlass.Int32(0) # L1156 gate: valid && complete + dga = cutlass.Int32(0) # gate: valid && complete if valid != cutlass.Int32(0): if complete != cutlass.Int32(0): dga = cutlass.Int32(1) if dga != cutlass.Int32(0): - # ---- degen A: narrowing over STAGED candidates - # (L1150-1212) ---- + # ---- degen A: narrowing over STAGED candidates ---- rlo = cutlass.Uint32(0) rhi = cutlass.Uint32(0xFFFFFFFF) above2 = cutlass.Int32(0) @@ -2709,12 +2644,12 @@ def kern( m2 = listN ethr = cutlass.Int64(0) tie_m = cutlass.Int32(1) - if tidx < cutlass.Int32(NBS): # L1161 + if tidx < cutlass.Int32(NBS): s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1162 ---- + cute.arch.barrier() # ---- barrier (degen A init) ---- brk = cutlass.Int32(0) lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L1163-1184 (<=8 levels) + while brk == cutlass.Int32(0): # <=8 levels if need2 == m2: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) above2 = above2 + m2 @@ -2737,7 +2672,7 @@ def kern( sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) i = tidx - while i < listN: # L1170-1176 + while i < listN: uq = cutlass.Uint32(0) if cutlass.const_expr(self.vstg): vx = cutlass.Int32(0) @@ -2772,7 +2707,7 @@ def kern( s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1177 ---- + cute.arch.barrier() # ---- barrier (level hist) ---- C.scan_cross0( s_hist, need2, @@ -2784,7 +2719,7 @@ def kern( nb=NBS, zero=True, ) - cute.arch.barrier() # ---- barrier L1179 ---- + cute.arch.barrier() # ---- barrier (level scan) ---- above2 = above2 + s_res[C.RES_ABOVE] need2 = need2 - s_res[C.RES_ABOVE] m2 = s_res[C.RES_M] @@ -2794,18 +2729,18 @@ def kern( rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) rlo = nlo lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L1185 + if tidx == cutlass.Int32(0): s_scal[1] = cutlass.Int32(0) s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1186 ---- - nA = k # L1187 + cute.arch.barrier() # ---- barrier (emit counters) ---- + nA = k nT = cutlass.Int32(0) if tie_m != cutlass.Int32(0): nA = above2 nT = need2 it2 = (listN + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) it = cutlass.Int32(0) - while it < it2: # L1189-1210 + while it < it2: i = it * cutlass.Int32(BLK) + tidx p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) @@ -2850,7 +2785,7 @@ def kern( ) it = it + cutlass.Int32(1) else: - # ---- degen B: whole-row narrowing (L1214-1264) ---- + # ---- degen B: whole-row narrowing ---- rlo = cutlass.Uint32(0) rhi = cutlass.Uint32(0xFFFFFFFF) above2 = cutlass.Int32(0) @@ -2858,12 +2793,12 @@ def kern( m2 = n ethr = cutlass.Int64(0) tie_m = cutlass.Int32(1) - if tidx < cutlass.Int32(NBS): # L1220 + if tidx < cutlass.Int32(NBS): s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1221 ---- + cute.arch.barrier() # ---- barrier (degen B init) ---- brk = cutlass.Int32(0) lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L1222-1241 (<=8 levels) + while brk == cutlass.Int32(0): # <=8 levels if need2 == m2: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) above2 = above2 + m2 @@ -2886,7 +2821,7 @@ def kern( sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) i = tidx - while i < n: # L1229-1233 whole row + while i < n: # whole row uq = C.fkey(C.ldg_f32(x_addr, i)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -2897,7 +2832,7 @@ def kern( s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1234 ---- + cute.arch.barrier() # ---- barrier (level hist) ---- C.scan_cross0( s_hist, need2, @@ -2909,7 +2844,7 @@ def kern( nb=NBS, zero=True, ) - cute.arch.barrier() # ---- barrier L1236 ---- + cute.arch.barrier() # ---- barrier (level scan) ---- above2 = above2 + s_res[C.RES_ABOVE] need2 = need2 - s_res[C.RES_ABOVE] m2 = s_res[C.RES_M] @@ -2919,10 +2854,10 @@ def kern( rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) rlo = nlo lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L1242 + if tidx == cutlass.Int32(0): s_scal[1] = cutlass.Int32(0) s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L1243 ---- + cute.arch.barrier() # ---- barrier (emit counters) ---- nA = k nT = cutlass.Int32(0) if tie_m != cutlass.Int32(0): @@ -2930,7 +2865,7 @@ def kern( nT = need2 it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) it = cutlass.Int32(0) - while it < it2: # L1246-1263 + while it < it2: i = it * cutlass.Int32(BLK) + tidx p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) @@ -2948,9 +2883,9 @@ def kern( it = it + cutlass.Int32(1) # ---- varlen short-row epilogue (production heuristicTopKDecode - # L72-84 convention): every valid position is in the top-K — emit - # identity indices and pad the tail with -1. The body above ran as - # a zero-work pass for these rows (n = 0, TGT = INT_MAX) so nothing + # convention): every valid position is in the top-K — emit identity + # indices and pad the tail with -1. The body above ran as a + # zero-work pass for these rows (n = 0, TGT = INT_MAX) so nothing # was written; only part 0 of a SPLIT row emits. if cutlass.const_expr(self.varlen): if short != cutlass.Int32(0): @@ -2965,7 +2900,7 @@ def kern( j = j + cutlass.Int32(BLK) # ------------------------------------------------------------------ - # host launcher (grid dim3(R, b) L2750; MINB wall via min_blocks_per_mp) + # host launcher (grid dim3(R, b); MINB wall via min_blocks_per_mp) # ------------------------------------------------------------------ @cute.jit def __call__( @@ -3151,14 +3086,10 @@ def _legacy_dummy_kv(ref): # =========================================================================== -# ==== ct_reg.py — family: reg ================================= +# ==== family: reg ============================================= # =========================================================================== -"""ct_reg.py — op46 gvr_topk_reg CuTeDSL translation. - -CUDA ground truth: src_cuda/kernel.cu L1269-1776 (register-resident exact -top-K, one CTA per row, histogram bins in FLOAT space). Contract: -TRANSLATION_SPEC.md §5.2. All probe verdicts (probes/PROBE_RESULTS.md) are -binding; op43 lessons L1-L5 applied throughout. +"""gvr_topk_reg — register-resident exact top-K, one CTA per row, histogram +bins in FLOAT space. Template knobs (CUDA `gvr_topk_reg`): ctor args of :class:`GvrTopkRegKernel`, read as `cutlass.const_expr(self.x)` @@ -3170,30 +3101,30 @@ def _legacy_dummy_kv(ref): __shared__ block is folded into the first 512 B so occupancy accounting matches nvcc's static+dynamic sum): - [0..5] s_res (ct_common slot map RES_B/M/ABOVE/TOT/B2/B3) + [0..5] s_res (shared slot map RES_B/M/ABOVE/TOT/B2/B3) [6..7] s_cnt (s_o1, s_oc) [8..9] s_kmm (s_kmin, s_kmax — Uint32) [10..11] s_e12 (s_e1, s_e2) [16..16+NW) ws (scan_cross_w workspace) [48..48+NW) wmn (Uint32 warp min partials) [80..80+NW) wmx (Uint32 warp max partials) - [128..128+NBH) hist (kernel.cu L1296) + [128..128+NBH) hist [128+NBH..128+NBH+CMP) ck (Uint32 crossing keys; CMP dynamic) [128+NBH+CMP..+2CMP) ci (Int32 crossing indices) - img/bm alias ck at word 128+NBH (kernel.cu L1299/L1409; IMGOFF==NBH) + img/bm alias ck at word 128+NBH (IMGOFF==NBH) Launch smem = 512 + dispatch_smem_bytes (dynamic Int32). -TOOLCHAIN GOTCHA (documented in notes/ct_reg_NOTES.md): dynamic launch smem -with min_blocks_per_mp>1 crashes cutlass_dsl._build_kernel_attrs (host ceil() -on a dynamic value while computing the PREFERRED_SHARED_MEMORY_CARVEOUT -hint). `_no_carveout()` scopes a monkeypatch around cute.compile dropping -ONLY that hint (CUDA __launch_bounds__ sets no carveout either); -`.reqntid`/`.minnctapersm` (the register wall) are unaffected — verified. +TOOLCHAIN GOTCHA: dynamic launch smem with min_blocks_per_mp>1 crashes +cutlass_dsl._build_kernel_attrs (host ceil() on a dynamic value while +computing the PREFERRED_SHARED_MEMORY_CARVEOUT hint). `_no_carveout()` +scopes a monkeypatch around cute.compile dropping ONLY that hint (CUDA +__launch_bounds__ sets no carveout either); `.reqntid`/`.minnctapersm` +(the register wall) are unaffected. """ -NB__reg = 1024 # kernel.cu L16 (NBH default) +NB__reg = 1024 # NBH default STATIC_WORDS = 128 # DSL smem prelude (static-__shared__ mirror) STATIC_BYTES = STATIC_WORDS * 4 _NEG_INF__reg = float("-inf") @@ -3201,11 +3132,11 @@ def _legacy_dummy_kv(ref): # --------------------------------------------------------------------------- -# module-local FP spellings (probe P6) +# module-local FP spellings # --------------------------------------------------------------------------- @dsl_user_op def _fmaf__reg(a, b, c, *, loc=None, ip=None): - """CUDA fmaf: single fma.rn.f32 (P6 emit spelling).""" + """CUDA fmaf: single fma.rn.f32.""" return cutlass.Float32( mlir_math.fma( a.ir_value(loc=loc, ip=ip), @@ -3220,12 +3151,12 @@ def _fmaf__reg(a, b, c, *, loc=None, ip=None): @dsl_user_op def _submul_asm(v, t, sc, *, loc=None, ip=None): - """(v - t) * sc with two roundings, opaque to CSE/contraction (P6). + """(v - t) * sc with two roundings, opaque to CSE/contraction. - Used at the !BRL classify site (kernel.cu L1521) so no sub-expression is - shared with the emit's `fmaf(v - T, SC, OFF)` — the CUDA deliberately - spells the two sites differently to stop nvcc holding all S q's live - across the barrier (L1506-1509). + Used at the !BRL classify site so no sub-expression is shared with the + emit's `fmaf(v - T, SC, OFF)` — the CUDA deliberately spells the two + sites differently to stop the compiler holding all S q's live across + the barrier. """ return cutlass.Float32( llvm.inline_asm( @@ -3244,15 +3175,14 @@ def _submul_asm(v, t, sc, *, loc=None, ip=None): @dsl_user_op def _smem_addr_reg__reg(addr, *, loc=None, ip=None): - """Pin a CTA-shared 32-bit byte address in ONE register (SASS audit fix). + """Pin a CTA-shared 32-bit byte address in ONE register. Identity `mov` behind an asm boundary: without it LLVM re-folds the `mov.b32 %r, __dynamic_shmem__0` symbol materialisation into EVERY use site, and ptxas then re-derives the CGA shared window (S2UR SR_CgaCtaId - + UMOV + ULEA, 3 instructions) inside each divergent classify block — - measured +24 warp-instructions/warp vs the CUDA arm, which keeps the - base in one UR. The asm result is not duplicable, so the window is - materialised exactly once. Value-identical: a plain register copy. + + UMOV + ULEA, 3 instructions) inside each divergent classify block. + The asm result is not duplicable, so the window is materialised exactly + once. Value-identical: a plain register copy. """ return cutlass.Int32( llvm.inline_asm( @@ -3274,11 +3204,10 @@ def _red_shared_add1__reg(addr, *, loc=None, ip=None): """CUDA classify `atomicAdd(&hist[bin], 1u)` with the result unused. `red` (not `atom`) is the result-less spelling — ptxas lowers it to the - same ATOMS.POPC.INC.32 RZ the CUDA arm emits (kernel.cu L1521-1526). - Same ordering contract as atomic_add_cta: .relaxed scope .cta. Takes the - final shared byte address as a plain Int32 so the address datapath stays - ordinary IR (ptxas fuses the shl+add into one LEA against the pinned - `_smem_addr_reg__reg` base). + same ATOMS.POPC.INC.32 RZ the CUDA arm emits. Same ordering contract as + atomic_add_cta: .relaxed scope .cta. Takes the final shared byte address + as a plain Int32 so the address datapath stays ordinary IR (ptxas fuses + the shl+add into one LEA against the pinned `_smem_addr_reg__reg` base). """ llvm.inline_asm( res=None, @@ -3364,7 +3293,7 @@ def _gen(self, config, _orig=orig_gen): class GvrTopkRegKernel: - """gvr_topk_reg (kernel.cu L1287).""" + """gvr_topk_reg.""" def __init__( self, @@ -3402,7 +3331,7 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) - # derived compile-time constants (kernel.cu L1292-1294, L1359, L1389, L1485) + # derived compile-time constants self.S = vpt * 4 self.lnbh = {256: 8, 512: 9, 2048: 11}.get(nbh, 10) self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 @@ -3431,7 +3360,7 @@ def kern( NW = cutlass.const_expr(self.blk // 32) if cutlass.const_expr(self.pdl): - cute.arch.griddepcontrol_wait() # L1291 (knob default off) + cute.arch.griddepcontrol_wait() # knob default off tid, _, _ = cute.arch.thread_idx() row, _, _ = cute.arch.block_idx() @@ -3606,21 +3535,21 @@ def kern( ), cute.make_layout((65536,)), ) - img_f = cute.make_tensor( # aliases ck/ci (L1299) + img_f = cute.make_tensor( # aliases ck/ci cute.make_ptr(cutlass.Float32, ck_base, cute.AddressSpace.smem, assumed_align=16), cute.make_layout((65536,)), ) - bm = cute.make_tensor( # aliases ck (L1409) + bm = cute.make_tensor( # aliases ck cute.make_ptr(cutlass.Int32, ck_base, cute.AddressSpace.smem, assumed_align=16), cute.make_layout((65536,)), ) n4 = n >> cutlass.Int32(2) ntail = n - (n4 << cutlass.Int32(2)) - tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L1351 + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` - # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any dependent - # gather (L1314-1322); compiled out under DEG. + # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any + # dependent gather; compiled out under DEG. pvs = [] if cutlass.const_expr(not self.deg): for t in cutlass.range_constexpr(KPT): @@ -3630,7 +3559,7 @@ def kern( pv = ld_g_i32(p_addr, j) pvs.append(pv) - # ---- row load: exact-fit peel + float4[VPT] register batch (L1327-1350) + # ---- row load: exact-fit peel + float4[VPT] register batch atom128 = g2r_atom_f32(128, invariant=True) frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(VPT)] if n4 >= cutlass.Int32(self.blk * self.vpt): # block-uniform peel @@ -3643,22 +3572,22 @@ def kern( ld_g_f32x4(atom128, x_addr, i, frags[u]) for u in cutlass.range_constexpr(VPT): i = tid + cutlass.Int32(u * self.blk) - if i >= n4: # -INFINITY fill L1346 + if i >= n4: # -INFINITY fill for q in cutlass.range_constexpr(4): frags[u][q] = cutlass.Float32(_NEG_INF__reg) tval = cutlass.Float32(_NEG_INF__reg) if tid < ntail: - tval = ldg_f32(x_addr, tix) # L1352 + tval = ldg_f32(x_addr, tix) - # ---- init (L1391-1392) + # ---- init if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) s_cnt[1] = cutlass.Int32(0) for z in cutlass.range_constexpr(self.nbh // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - # ---- bracket: 4 mutually exclusive compile-time arms (L1393-1451) + # ---- bracket: 4 mutually exclusive compile-time arms lmin = cutlass.Uint32(0xFFFFFFFF) lmax = cutlass.Uint32(0) if cutlass.const_expr(self.use_img): @@ -3669,7 +3598,7 @@ def kern( _sts128_f32(fatom, frags[u], ck_base, i * cutlass.Int32(16)) if tid < ntail: img_f[tix] = tval - cute.arch.barrier() # L1400 + cute.arch.barrier() # image staged for t in cutlass.range_constexpr(KPT): p = pvs[t] if cutlass.Uint32(p) < cutlass.Uint32(n): @@ -3678,14 +3607,14 @@ def kern( lmin = uk if uk > lmax: lmax = uk - cute.arch.barrier() # L1406 (img dies) + cute.arch.barrier() # img dies elif cutlass.const_expr(self.use_bm): nbw = (n + cutlass.Int32(31)) >> cutlass.Int32(5) i = tid - while i < nbw: # bitmap clear (L1410) + while i < nbw: # bitmap clear bm[i] = cutlass.Int32(0) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1411 + cute.arch.barrier() # bitmap cleared for t in cutlass.range_constexpr(KPT): p = pvs[t] if cutlass.Uint32(p) < cutlass.Uint32(n): @@ -3693,7 +3622,7 @@ def kern( bm.iterator + (p >> cutlass.Int32(5)), cutlass.Int32(1) << (p & cutlass.Int32(31)), ) - cute.arch.barrier() # L1417 + cute.arch.barrier() # bitmap set lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) for u in cutlass.range_constexpr(VPT): @@ -3715,12 +3644,12 @@ def kern( lmn = fmin_f32(lmn, tval) lmx = fmax_f32(lmx, tval) lmin = fkey(lmn) - lmax = fkey(lmx) # monotone (L1428) - cute.arch.barrier() # L1429 (bm dies) + lmax = fkey(lmx) # monotone + cute.arch.barrier() # bm dies elif cutlass.const_expr(self.deg): lmn = cutlass.Float32(_POS_INF) lmx = cutlass.Float32(_NEG_INF__reg) - for s in cutlass.range_constexpr(S): # L1436-1439 + for s in cutlass.range_constexpr(S): v = _val(frags, s) if v > cutlass.Float32(_NEG_INF__reg): lmn = fmin_f32(lmn, v) @@ -3731,7 +3660,7 @@ def kern( lmin = fkey(lmn) lmax = fkey(lmx) else: - # default: KPT scattered fkey ldg gathers, batch-then-fold (L1443-1450) + # default: KPT scattered fkey ldg gathers, batch-then-fold xs = [] for t in cutlass.range_constexpr(KPT): xv = cutlass.Float32(0.0) @@ -3746,13 +3675,13 @@ def kern( if uk > lmax: lmax = uk - # ---- block min/max in ONE barrier (L1452-1462); publishes hist clear + # ---- block min/max in ONE barrier; publishes hist clear lmin = warp_min_u32(lmin) lmax = warp_max_u32(lmax) if lane == cutlass.Int32(0): s_wmn[tid >> cutlass.Int32(5)] = lmin s_wmx[tid >> cutlass.Int32(5)] = lmax - cute.arch.barrier() # L1456 + cute.arch.barrier() # warp partials published a = cutlass.Uint32(0xFFFFFFFF) c = cutlass.Uint32(0) if lane < cutlass.Int32(NW): @@ -3763,7 +3692,7 @@ def kern( Tv = invkey(lmin) GMAX = invkey(lmax) - # ---- collapse guard, NaN-safe (L1464-1471) + # ---- collapse guard, NaN-safe okc = cutlass.Int32(0) if Tv < GMAX: if (GMAX - Tv) > cutlass.Float32(1e-30): @@ -3772,7 +3701,7 @@ def kern( Tv = cutlass.Float32(SENT_LO) GMAX = cutlass.Float32(SENT_HI) - # ---- bin transform constants (L1485-1511) + # ---- bin transform constants BRL = cutlass.const_expr(self.brl) # noqa: F841 OFFf = cutlass.Float32(1.0 if self.brl else 0.0) recip = 1.0 / float(self.nbh - (2 if self.brl else 0)) @@ -3781,33 +3710,33 @@ def kern( if WD > cutlass.Float32(0.0): wsel = WD # rcp.approx (single MUFU.RCP) — the CUDA arm's exact lowering of - # `1.0f / wsel`; the previous `1.0 / wsel` spelling emitted the IEEE - # div.rn Newton triple + slowpath CALL on the barrier-bounded chain + # `1.0f / wsel`; a plain `1.0 / wsel` would emit the IEEE div.rn + # Newton triple + slowpath CALL on the barrier-bounded chain # feeding all S classify FMULs. Output exactness is SC-invariant - # (any SC > 0 preserves the sign/monotonicity invariants, L1485-1511) - # and the WD > 0 arm is now bit-identical to CUDA's MUFU.RCP. + # (any SC > 0 preserves the sign/monotonicity invariants) and the + # WD > 0 arm is bit-identical to CUDA's MUFU.RCP. SC = cute.arch.rcp_approx(wsel) QCAPf = cutlass.Float32(float(self.nbh - 1)) CQ0 = OFFf - Tv * SC CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf(CQ0) + cutlass.Float32(1.0)) - # ---- histogram (L1513-1526) + # ---- histogram if cutlass.const_expr(self.brl): - # fix-2 P4 (GATED, removable as one hunk): A1 ported to the BRL - # classify arm — hist base pinned ONCE via the same - # _smem_addr_reg__reg identity-mov used in the !BRL arm below, and the - # result-discarded classify atomics spelled as resultless - # red.shared (_red_shared_add1__reg). Value-identical: same +1 to the - # same byte address (hb + 4*bn == &s_hist[bn]), same .relaxed.cta - # ordering; the q/bn computations are untouched so classify/emit + # BRL classify arm: hist base pinned ONCE via the same + # _smem_addr_reg__reg identity-mov used in the !BRL arm below, + # and the result-discarded classify atomics spelled as + # resultless red.shared (_red_shared_add1__reg). + # Value-identical: same +1 to the same byte address + # (hb + 4*bn == &s_hist[bn]), same .relaxed.cta ordering; the + # q/bn computations are untouched so classify/emit # bit-identity (BRL requirement) is preserved. Emit-path hist - # atomics (results used, L1630+) are NOT touched. + # atomics (results used) are NOT touched. hb = _smem_addr_reg__reg(sbase + cutlass.Int32(STATIC_WORDS * 4)) for s in cutlass.range_constexpr(S): q = _fmaf__reg(_val(frags, s), SC, CQ) bn = _umin_u32(f2u_rz(q), cutlass.Uint32(self.nbh - 1)) _red_shared_add1__reg(hb + (cutlass.Int32(bn) << cutlass.Int32(2))) - qt = _fmaf__reg(tval, SC, CQ) # unconditional (L1517) + qt = _fmaf__reg(tval, SC, CQ) # unconditional bnt = _umin_u32(f2u_rz(qt), cutlass.Uint32(self.nbh - 1)) _red_shared_add1__reg(hb + (cutlass.Int32(bnt) << cutlass.Int32(2))) else: @@ -3822,14 +3751,14 @@ def kern( qt = _submul_asm(tval, Tv, SC) if qt >= cutlass.Float32(0.0): _red_shared_add1__reg(hb + (f2s_rz(fmin_f32(qt, QCAPf)) << cutlass.Int32(2))) - cute.arch.barrier() # L1527 + cute.arch.barrier() # histogram done - # ---- crossing-bin find (L1528-1538) + # ---- crossing-bin find if cutlass.const_expr(self.cur or self.nbh > 1024): scan_cross_w(s_hist, s_ws, k, tid, s_res, blk=self.blk, nb=self.nbh) else: find_cross(s_hist, k, tid, s_res, nb=self.nbh) - cute.arch.barrier() # L1535 + cute.arch.barrier() # crossing published above = s_res[RES_ABOVE] m = s_res[RES_M] Bv = s_res[RES_B] @@ -3838,7 +3767,7 @@ def kern( if need >= m: whole = cutlass.Int32(1) - # ---- prod-fix ESCAPE (L1540-1617): 32-step key-space bisection + # ---- ESCAPE: 32-step key-space bisection esc = cutlass.Int32(0) if whole == cutlass.Int32(0): if m > cmp_: @@ -3847,17 +3776,15 @@ def kern( if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) s_cnt[1] = cutlass.Int32(0) - # DEVIATION (race fix, see notes): the CUDA zeroes s_o1/s_oc - # again between the nA read (L1581) and the emit (L1584) with - # only ONE barrier pair around both — a read/write race that - # nvcc's schedule happens to win and ptxas' does not (observed - # ~8% row corruption under CTA co-residency). We instead emit - # through the path-exclusive s_e1/s_e2 slots, zeroed HERE under - # the existing L1550 barrier; the racy mid-emit rezero is + # DEVIATION (race fix): the CUDA zeroes s_o1/s_oc again + # between the nA read and the emit with only ONE barrier + # pair around both — a read/write race. Emit instead + # through the path-exclusive s_e1/s_e2 slots, zeroed HERE + # under the existing barrier; the racy mid-emit rezero is # dropped. Barrier count unchanged. s_e12[0] = cutlass.Int32(0) s_e12[1] = cutlass.Int32(0) - cute.arch.barrier() # L1550 + cute.arch.barrier() # escape init klo = cutlass.Uint32(0) bit = cutlass.Int32(31) while bit >= cutlass.Int32(0): @@ -3877,13 +3804,13 @@ def kern( if lane == cutlass.Int32(0): if cnt != cutlass.Int32(0): atomic_add_cta(s_cnt.iterator, cnt) - cute.arch.barrier() # L1563 + cute.arch.barrier() # count published if s_cnt[0] >= k: klo = kt - cute.arch.barrier() # L1565 + cute.arch.barrier() # count consumed if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) - cute.arch.barrier() # L1567 + cute.arch.barrier() # count reset bit = bit - cutlass.Int32(1) ethr = cutlass.Int64(klo) # k-th largest key abv = cutlass.Int32(0) @@ -3901,11 +3828,11 @@ def kern( if lane == cutlass.Int32(0): if abv != cutlass.Int32(0): atomic_add_cta(s_cnt.iterator + 1, abv) - cute.arch.barrier() # L1580 + cute.arch.barrier() # above-count published nA = s_cnt[1] nT = k - nA # (rezero dropped — emit counters live in s_e12, see race-fix note) - cute.arch.barrier() # L1583 + cute.arch.barrier() # nA consumed lml = cutlass.Int32(cute.arch.lanemask_lt()) for s in cutlass.range_constexpr(S): ixv = ( @@ -3939,7 +3866,7 @@ def kern( if q2e == cutlass.Int32(1): if p2e < nT: out_row[nA + p2e] = ixv - # tail element (L1601-1615) + # tail element u64 = cutlass.Int64(-1) if tid < ntail: u64 = cutlass.Int64(fkey(tval)) @@ -3970,9 +3897,9 @@ def kern( out_row[nA + p2e] = tix # (CUDA returns here — everything below is the else-arm) else: - # ---- emit (L1619-1701) + # ---- emit if cutlass.const_expr(self.cur): - LOQ = cutlass.Float32(Bv) # int->float cvt (L1626) + LOQ = cutlass.Float32(Bv) # int->float cvt lim1 = above if whole == cutlass.Int32(1): lim1 = above + m @@ -3980,7 +3907,7 @@ def kern( if cutlass.const_expr(self.brl): q = _fmaf__reg(_val(frags, s), SC, CQ) # bit-identical to classify else: - q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) # L1630 emit spelling + q = _fmaf__reg(_val(frags, s) - Tv, SC, OFFf) # emit spelling idx = ( (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2) ) + cutlass.Int32(s % 4) @@ -3998,7 +3925,7 @@ def kern( if q2i < cmp_: # escape-made-safe guard ck[q2i] = fkey(_val(frags, s)) ci[q2i] = idx - # tail (L1640-1647) + # tail if cutlass.const_expr(self.brl): qt2 = _fmaf__reg(tval, SC, CQ) else: @@ -4016,7 +3943,7 @@ def kern( ck[q2i] = fkey(tval) ci[q2i] = tix else: - # two-mask ballot emit (L1650-1701) + # two-mask ballot emit HIf = cutlass.Float32(_POS_INF) LOf = cutlass.Float32(_POS_INF) if whole == cutlass.Int32(1): @@ -4050,7 +3977,7 @@ def kern( t2 = cutlass.Int32(1) c1 = popc(m1) + t1 c2 = popc(m2) + t2 - s1, s2 = warp_incl_scan_add2(c1, c2, lane) # L1669-1673 + s1, s2 = warp_incl_scan_add2(c1, c2, lane) b1 = cutlass.Int32(0) b2 = cutlass.Int32(0) if lane == cutlass.Int32(31): @@ -4078,7 +4005,7 @@ def kern( if p1 < lim1: out_row[p1] = tix p1 = p1 + cutlass.Int32(1) - if m2 != cutlass.Int32(0): # static-unrolled (L1689) + if m2 != cutlass.Int32(0): # static-unrolled for s in cutlass.range_constexpr(S): if (m2 & cutlass.Int32(1 << s)) != cutlass.Int32(0): idx = ( @@ -4096,7 +4023,7 @@ def kern( # ---- refine (skipped when whole — CUDA returned inside emit) if whole == cutlass.Int32(0): - cute.arch.barrier() # L1703 + cute.arch.barrier() # emit done if cutlass.const_expr(self.cur): mc = m if mc > cmp_: @@ -4110,7 +4037,7 @@ def kern( if mc <= qc: quad = cutlass.Int32(1) if quad == cutlass.Int32(1): - # O(mc^2) index-tie-broken rank (L1706-1718) + # O(mc^2) index-tie-broken rank i = tid while i < mc: uq = cutlass.Uint32(ck[i]) @@ -4130,18 +4057,18 @@ def kern( out_row[above + r] = ci[i] i = i + cutlass.Int32(BLK) else: - # ---- fallback: exact key-space narrowing (L1720-1775) + # ---- fallback: exact key-space narrowing if tid == cutlass.Int32(0): s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) s_kmm[1] = cutlass.Uint32(0) - cute.arch.barrier() # L1724 + cute.arch.barrier() # kmm init i = tid while i < mc: kv = cutlass.Uint32(ck[i]) atomic_min_cta(s_kmm.iterator, kv) atomic_max_cta(s_kmm.iterator + 1, kv) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1726 + cute.arch.barrier() # key range published rlo = cutlass.Uint32(s_kmm[0]) rhi = cutlass.Uint32(s_kmm[1]) ethr = cutlass.Int64(rlo) @@ -4151,16 +4078,16 @@ def kern( lev = cutlass.Int32(0) done = cutlass.Int32(0) while done == cutlass.Int32(0): - if needC == mm: # L1730 + if needC == mm: ethr = cutlass.Int64(rlo) - cutlass.Int64(1) aboveC = aboveC + mm needC = cutlass.Int32(0) done = cutlass.Int32(1) if done == cutlass.Int32(0): - if rlo >= rhi: # L1731 + if rlo >= rhi: ethr = cutlass.Int64(rlo) done = cutlass.Int32(1) - if lev >= cutlass.Int32(6): # L1732 + if lev >= cutlass.Int32(6): ethr = cutlass.Int64(rlo) done = cutlass.Int32(1) if done == cutlass.Int32(0): @@ -4173,7 +4100,7 @@ def kern( sh2 = b2w - cutlass.Int32(LNBH) for z in cutlass.range_constexpr(self.nbh // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L1737 + cute.arch.barrier() # level clear i = tid while i < mc: unar = cutlass.Uint32(ck[i]) @@ -4186,14 +4113,14 @@ def kern( cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L1743 + cute.arch.barrier() # level hist if cutlass.const_expr(self.nbh > 1024): scan_cross_w( s_hist, s_ws, needC, tid, s_res, blk=self.blk, nb=self.nbh ) else: find_cross(s_hist, needC, tid, s_res, nb=self.nbh) - cute.arch.barrier() # L1746 + cute.arch.barrier() # level scan aboveC = aboveC + s_res[RES_ABOVE] needC = needC - s_res[RES_ABOVE] mm = s_res[RES_M] @@ -4206,11 +4133,11 @@ def kern( ) rlo = nlo lev = lev + cutlass.Int32(1) - # final two-predicate ballot emit (L1752-1775) + # final two-predicate ballot emit if tid == cutlass.Int32(0): s_e12[0] = cutlass.Int32(0) s_e12[1] = cutlass.Int32(0) - cute.arch.barrier() # L1753 + cute.arch.barrier() # emit counters lml = cutlass.Int32(cute.arch.lanemask_lt()) it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) it = cutlass.Int32(0) @@ -4369,54 +4296,46 @@ def reg_topk(logits, pre_idx, n, out, rd=None): # =========================================================================== -# ==== ct_clus.py — family: clus =============================== +# ==== family: clus ============================================ # =========================================================================== -"""ct_clus.py — op46 gvr_clus (clustered streaming GVR) CuTeDSL port. +"""gvr_clus — clustered streaming self-sampling GVR; per-CTA stream mirrors +gvr_main. -Ground truth: src_cuda/kernel.cu L1793-2356 (frozen); phase contract, smem -map and barrier inventory per TRANSLATION_SPEC.md §5.3; DSL spellings pinned -by probes/PROBE_RESULTS.md (BINDING) + op43 lessons L1..L5; per-CTA stream -mirrors gvr_main (idioms reused from the proven src/ct_main.py). - -Ctor knobs (compile-time, mirror of the CUDA template params, spec §4c): +Ctor knobs (compile-time, mirror of the CUDA template params): BLK = 1024, U ∈ {1,2,4,8}, MINB = 1, NBS = 256, CS ∈ {2,4,8} (+ scap/cmp smem-extent knobs: every reachable route has 8192/2048 — SCAP/CMP stay LIVE runtime args for all value logic, ABI parity). -Derived: HB=NBS, STEPC=BLK*U, PFD=min(U,4) (kernel.cu L1905). +Derived: HB=NBS, STEPC=BLK*U, PFD=min(U,4). -Signature (ABI parity with kernel.cu L1797-1799; Q is dead in-kernel): +Signature (ABI parity with the CUDA form; Q is dead in-kernel): run__clus(logits[b,npad] f32, pre_idx[b,k] i32, out[b,k] i32) via kern(..., n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2) -Grid dim3(CS, b) native 2-D + cluster (CS,1,1) (probe P11); block 1024; -min_blocks_per_mp=1 (64-reg wall, probes P2/P15); smem one SmemAllocator -blob mirroring the CUDA dynamic map hist|cbuf|ck64c|mrg (L1802-1814), -dyn-equivalent bytes == host smc formula L3130 (asserted in run__clus()). +Grid dim3(CS, b) native 2-D + cluster (CS,1,1); block 1024; +min_blocks_per_mp=1 (64-register wall); smem one SmemAllocator blob +mirroring the CUDA dynamic map hist|cbuf|ck64c|mrg, dyn-equivalent bytes == +the host dispatch formula (asserted in run__clus()). -int2 staging convention (same as ct_main): int2(value bits, index) is ONE +int2 staging convention (same as gvr_main): int2(value bits, index) is ONE little-endian Uint64 = (idx << 32) | value_bits — single u64 smem ld/st. -Barrier / cluster-op inventory implemented (kernel.cu line cites, op43 L5): - L1902 (sample redux publish), L1948 (sample hist), L1956 (scan publish), +Barrier / cluster-op placement mirrors the CUDA source one-for-one: + sample redux publish, sample hist, scan publish, [degenerate sample: 2 inside gather_hint], - retry preamble: clus.sync L2016 + __syncthreads L2024, - clus.sync L2135 (merge), __syncthreads L2143 (merge publish), + retry preamble: clus.sync + __syncthreads, + clus.sync (merge), __syncthreads (merge publish), [ladder gather: 2 inside gather_hint], - clus.sync L2226 (EXIT RENDEZVOUS — the only one; rank!=0 falls through), - narrowing: L2255, L2257, per-level L2270/L2273, L2281, - degen: per-level L2321/L2325/L2327 (+1 INSIDE scan_cross), L2337. + clus.sync (EXIT RENDEZVOUS — the only one; rank!=0 falls through), + narrowing / degen per-level barriers (+1 INSIDE scan_cross). NO loop-tail ladder barriers (gvr_clus has none — unlike gvr_main). - All clus.sync = releasing aligned arrive+wait (risk R3, never relaxed). -Cluster ops: merge = _merge_scan0_local, a LOCAL patched copy of the frozen - ct_common.merge_scan0 that rematerializes mapa per (q, r) like the CUDA - L135-137 (register-pressure fix, see notes; DSMEM v4 spelling = probe - P5b via ct_common ops); ONE packed u64 st.shared::cluster candidate push - to rank-0 ck64c (L2185/L2205/L2219 — never split 4B, op43 L3); mapa of - ck64c to rank 0. PTX inventory audited: 3 arrive+3 wait (aligned), 3 - st.shared::cluster.u64 sites, 2*CS ld.shared::cluster.v4 sites, 4 - prefetch.global.L2 sites on U=8 only, zero griddepcontrol. + All clus.sync = releasing aligned arrive+wait, never relaxed. +Cluster ops: merge = _merge_scan0_local, a LOCAL patched copy of + merge_scan0 that rematerializes mapa per (q, r) like the CUDA + (register-pressure fix); ONE packed u64 st.shared::cluster candidate + push to rank-0 ck64c (never split into 4B stores); mapa of ck64c to + rank 0. Every rung/ladder decision is cluster-uniform by construction (identical sample locations on every rank; merged tot; block-uniform gather) — the -conditional retry clus.sync at L2016 cannot deadlock (spec §5.3). +conditional retry clus.sync cannot deadlock. """ @@ -4426,9 +4345,8 @@ def reg_topk(logits, pre_idx, n, out, rd=None): # --------------------------------------------------------------------------- -# single-rounding fma.rn.f32 (probe P6 emit spelling; local — ct_common has -# no fma export). Sites: T (L1957), Tk/T3 (L1966/1977/1992), HIC (L1981). -# (x-TF)*SC classify shapes stay plain sub+mul (P6: uncontractible). +# single-rounding fma.rn.f32, used at the T / Tk / T3 / HIC sites. +# (x-TF)*SC classify shapes stay plain sub+mul (uncontractible). # --------------------------------------------------------------------------- @dsl_user_op def _fmaf__clus(a, b, c, *, loc=None, ip=None): @@ -4445,7 +4363,7 @@ def _fmaf__clus(a, b, c, *, loc=None, ip=None): class GvrClusKernel: - """CuTeDSL port of gvr_clus (kernel.cu L1795).""" + """gvr_clus — clustered streaming GVR.""" def __init__( self, @@ -4460,9 +4378,9 @@ def __init__( next_n: int = 1, cr_shift: int = 0, ): - assert blk == 1024, "gvr_clus is always BLK=1024 (dispatch L3132)" - assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1) (L1796)" - assert nbs == 256, "SNB must stay 256 (kernel.cu L170-177)" + assert blk == 1024, "gvr_clus is always BLK=1024" + assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1)" + assert nbs == 256, "SNB must stay 256" assert u in (1, 2, 4, 8) and cs in (2, 4, 8) # per-row varlen mode (production heuristicTopKDecode contract, same # semantics as GvrMainKernel / GvrRegClusKernel): n and the sampling- @@ -4482,20 +4400,20 @@ def __init__( self.cs = cs self.scap = scap # smem extents only — self.cmp = cmp_ # value logic uses rt args - self.hb = nbs # L1801 - self.stepc = blk * u # L1836 - self.pfd = u if u < 4 else 4 # L1905 PFD=min(U,4) + self.hb = nbs + self.stepc = blk * u + self.pfd = u if u < 4 else 4 # PFD=min(U,4) self.lb = nbs.bit_length() - 1 # log2(NBS)=8 - # dynamic-region byte map (L1802-1814): hist | cbuf(int2) | ck64c | mrg + # dynamic-region byte map: hist | cbuf(int2) | ck64c | mrg self.cbuf_bytes = (scap + 4) * 8 assert self.cbuf_bytes % 16 == 0 self.ck_off = self.cbuf_bytes # inside the blob self.dyn_bytes = nbs * 4 + self.cbuf_bytes + cmp_ * 8 + nbs * 4 - # == host smc = SNB*8 + (SCAP+4)*8 + CMP*8 (L3130) + # == host smc = SNB*8 + (SCAP+4)*8 + CMP*8 # ------------------------------------------------------------------ - # GVR_EMITK (kernel.cu L2096-2104): classify+stage one survivor. - # bn via UNSIGNED saturating convert (f2u_rz, P4); staging store is ONE + # GVR_EMITK: classify+stage one survivor. + # bn via UNSIGNED saturating convert (f2u_rz); staging store is ONE # u64; branchless trash slot min(pos, SCAP) (runtime SCAP). Returns pos+1. # ------------------------------------------------------------------ @cute.jit @@ -4515,11 +4433,10 @@ def _emitk(self, xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2): return pos + cutlass.Int32(1) # ------------------------------------------------------------------ - # P5 emit step (kernel.cu L2178-2186 == L2196-2204 == L2211-2219): - # bn via SIGNED rz convert (__float2int_rz, L2180); bn>=B gate; LOCAL - # mrg atomicAdd whose result is a CLUSTER-GLOBAL position (prefix-biased - # cursors from merge_scan0); overflow -> ONE packed u64 DSMEM store to - # rank-0 ck64c (probe P5b, op43 L3 — never split 4B). + # P5 emit step: bn via SIGNED rz convert (__float2int_rz); bn>=B gate; + # LOCAL mrg atomicAdd whose result is a CLUSTER-GLOBAL position + # (prefix-biased cursors from merge_scan0); overflow -> ONE packed u64 + # DSMEM store to rank-0 ck64c (never split into 4B stores). # ------------------------------------------------------------------ @cute.jit def _p5_emit(self, xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, rk64): @@ -4542,15 +4459,14 @@ def _p5_emit(self, xv, idv, TF, SC, B, above, lim1, whole, CMP, s_mrg, out_row, ) # ------------------------------------------------------------------ - # LOCAL patched copy of ct_common.merge_scan0 (frozen sibling — not - # edited): rematerializes mapa per (q, r) exactly like the CUDA L135-137 - # instead of holding CS mapped base addresses across the whole merge. - # The hoisted-array form costs CS extra long-lived registers; with the - # U>=4 sixteen-register pf prime batch it tips ptxas into spilling the - # batch across the rung phase (reg-audit finding, notes). Semantics, - # the DSMEM v4 load spelling (probe P5b), the register accumulation and - # the prefix-biased STS.128 cursor write are IDENTICAL to ct_common. - # NO barrier inside (caller pays L2143). + # LOCAL patched copy of merge_scan0: rematerializes mapa per (q, r) + # exactly like the CUDA instead of holding CS mapped base addresses + # across the whole merge. The hoisted-array form costs CS extra + # long-lived registers; with the U>=4 sixteen-register pf prime batch + # it tips ptxas into spilling the batch across the rung phase. + # Semantics, the DSMEM v4 load spelling, the register accumulation and + # the prefix-biased STS.128 cursor write are IDENTICAL to merge_scan0. + # NO barrier inside (caller pays the merge-publish barrier). # ------------------------------------------------------------------ @cute.jit def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): @@ -4562,7 +4478,7 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): lane = tidx atom = C.smem_atom_i32_128() hbase = s_hist.iterator.toint() - # pass 1 (L131-143): remote v4 accumulation of tot/pre per vector + # pass 1: remote v4 accumulation of tot/pre per vector tot_r = [] pre_r = [] sm = cutlass.Int32(0) @@ -4573,13 +4489,13 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): for r in cutlass.range_constexpr(CS): mapped = C._mapa_shared_cluster_addr( hbase + boff, cutlass.Int32(r) - ) # per-use mapa (L136) + ) # per-use mapa v0, v1, v2, v3 = C._ld_shared_cluster_v4_u32(mapped) t[0] = t[0] + v0 t[1] = t[1] + v1 t[2] = t[2] + v2 t[3] = t[3] + v3 - if cutlass.Int32(r) < rank: # L140 predicated adds + if cutlass.Int32(r) < rank: # predicated adds p[0] = p[0] + v0 p[1] = p[1] + v1 p[2] = p[2] + v2 @@ -4587,7 +4503,7 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): tot_r.append(t) pre_r.append(p) sm = sm + t[0] + t[1] + t[2] + t[3] - # inclusive scan + totals (L144-148) + # inclusive scan + totals w = C.warp_incl_scan_add(sm, lane) tt = cute.arch.shuffle_sync(w, cutlass.Int32(31)) after = tt - w @@ -4595,7 +4511,7 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): s_res[C.RES_TOT] = tt base = lane * cutlass.Int32(BPT) # descending walk: crossing pin + prefix-biased cursors into mrg - for q in cutlass.range_constexpr(NV - 1, -1, -1): # L151-165 + for q in cutlass.range_constexpr(NV - 1, -1, -1): o4 = cute.make_rmem_tensor((4,), cutlass.Int32) for j in cutlass.range_constexpr(3, -1, -1): cq = tot_r[q][j] @@ -4616,8 +4532,8 @@ def _merge_scan0_local(self, s_hist, s_mrg, rank, target, tidx, s_res): C.sts128_i32(atom, o4, s_mrg.iterator.toint(), boff) # ------------------------------------------------------------------ - # two-predicate warp-ballot emit step (narrowing L2279-2301 and degen - # L2338-2354) — same helper as ct_main. s_scal[1]=s_o1, s_scal[2]=s_o2. + # two-predicate warp-ballot emit step (narrowing and degen emits) — + # same helper as the main family. s_scal[1]=s_o1, s_scal[2]=s_o2. # ------------------------------------------------------------------ @cute.jit def _ballot_pair_emit(self, p1, p2, idv, base1, cap1, base2, cap2, out_row, s_scal, lane): @@ -4673,7 +4589,7 @@ def kern( NW = BLK // 32 tidx, _, _ = cute.arch.thread_idx() - bx, by, _ = cute.arch.block_idx() # (rank, row) L1824-1825 + bx, by, _ = cute.arch.block_idx() # (rank, row) rank = bx row = by lane = tidx & cutlass.Int32(31) @@ -4714,17 +4630,16 @@ def kern( short = cutlass.Int32(1) if short == cutlass.Int32(0): n = nv - # ---- aim ladder (P2b cheap mirror, all-thread) ---- + # ---- aim ladder (cheap mirror, all-thread) ---- # Schedule quantities only (exactness is schedule-invariant, # same argument as the always-sample deviation above): divides # become MUFU.RCP multiplies and the isqrt fixup loops collapse # to single steps (f32 sqrt of an exactly-representable int # (6n <= 2^23) is within 1 of isqrt). All-thread on purpose: - # this family's mirror redundancy is small (inst x1.26-1.28) - # and a warp0+barrier hoist EXPOSES the chain's serial latency - # at ~1 CTA/SM — measured 1.14->1.31 tax regression on - # pro_1024k r32 — while the redundant form hides it across - # warps. Q (chunk ownership) keeps its exact shift form. + # this family's mirror redundancy is small, and a + # warp0+barrier hoist EXPOSES the chain's serial latency at + # ~1 CTA/SM, while the redundant form hides it across warps. + # Q (chunk ownership) keeps its exact shift form. x6 = cutlass.Int32(6) * nv ri = cutlass.Int32(cmath.sqrt(cutlass.Float32(x6))) if ri * ri > x6: @@ -4807,7 +4722,7 @@ def kern( ov = tidx out[row, tidx] = ov - # ---- shared memory (CUDA dynamic map order L1802-1814, then static allocs) ---- + # ---- shared memory (CUDA dynamic map order, then static allocs) ---- smem = SmemAllocator() s_hist = smem.allocate_tensor( # hist[NBS] @ blob start cutlass.Int32, cute.make_ordered_layout((self.hb,), order=(0,)), byte_alignment=128 @@ -4817,12 +4732,12 @@ def kern( cute.make_ordered_layout((self.cbuf_bytes + self.cmp * 8,), order=(0,)), byte_alignment=16, ) - s_mrg = smem.allocate_tensor( # mrg[NBS] (L1814) + s_mrg = smem.allocate_tensor( # mrg[NBS] cutlass.Int32, cute.make_ordered_layout((self.nbs,), order=(0,)), byte_alignment=16 ) - s_ws = smem.allocate_tensor( # L1817 (degen scan only) + s_ws = smem.allocate_tensor( # degen scan only cutlass.Int32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 - ) # Int32: ct_common + ) # scan_cross predeclares its second-stage partial as Int32 and reads # s_ws inside a dynamic if — a Uint32 ws tensor trips the DSL type- # stability check (frozen sibling; counts < 2^31 so Int32 is exact). @@ -4832,17 +4747,17 @@ def kern( s_wmx = smem.allocate_tensor( cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 ) - s_res = smem.allocate_tensor( # ct_common slot map + s_res = smem.allocate_tensor( # shared slot map cutlass.Int32, cute.make_ordered_layout((8,), order=(0,)), byte_alignment=16 ) - # scalar block: [0]=s_bufn [1]=s_o1 [2]=s_o2 (L1818) + # scalar block: [0]=s_bufn [1]=s_o1 [2]=s_o2 s_scal = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 ) - s_tsh = smem.allocate_tensor( # L1820 + s_tsh = smem.allocate_tensor( cutlass.Float32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 ) - s_kmm = smem.allocate_tensor( # L1821 [0]=kmin [1]=kmax + s_kmm = smem.allocate_tensor( # [0]=kmin [1]=kmax cutlass.Uint32, cute.make_ordered_layout((2,), order=(0,)), byte_alignment=8 ) sbase = blob.iterator.toint() @@ -4861,7 +4776,7 @@ def kern( # every cluster barrier below stays aligned; short rows already # emitted identity + -1 tail in the prologue) ---- if short == cutlass.Int32(0): - # ---- row bases (L1830-1833; pre_idx is request-level under varlen) ---- + # ---- row bases (pre_idx is request-level under varlen) ---- row64 = cutlass.Int64(row) x_addr = logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) p_addr = pre_idx.iterator.toint() + cutlass.Int64(prow) * cutlass.Int64( @@ -4869,7 +4784,7 @@ def kern( ) * cutlass.Int64(4) out_row = out[row, None] - # ---- interleaved chunk ownership (L1835-1846) ---- + # ---- interleaved chunk ownership ---- n4 = n >> cutlass.Int32(2) nCh = (n4 + cutlass.Int32(STEPC - 1)) // cutlass.Int32(STEPC) nFullG = n4 // cutlass.Int32(STEPC) @@ -4878,16 +4793,16 @@ def kern( if rank == cutlass.Int32(0): tailn = n - tail0 - if tidx == cutlass.Int32(0): # L1848 + if tidx == cutlass.Int32(0): s_res[C.RES_B2] = cutlass.Int32(-1) s_res[C.RES_B3] = cutlass.Int32(-1) s_scal[0] = cutlass.Int32(0) # s_bufn - if tidx < cutlass.Int32(self.hb): # L1849 (HB<=BLK) + if tidx < cutlass.Int32(self.hb): # HB<=BLK s_hist[tidx] = cutlass.Int32(0) - # ============ P1: QUAD sample (hint gather LAZY, L1851-1878) ======== + # ============ P1: QUAD sample (hint gather LAZY) ===================== # one 64B line = 4 float4 per location, TWO threads: tid takes the - # lower pair at p4, tid+SMP the upper pair at p4+2 (L1868-1869). + # lower pair at p4, tid+SMP the upper pair at p4+2. atom128 = C.g2r_atom_f32(128, invariant=True) fsa = cute.make_rmem_tensor((4,), cutlass.Float32) fsb = cute.make_rmem_tensor((4,), cutlass.Float32) @@ -4902,10 +4817,10 @@ def kern( C.ld_g_f32x4(atom128, x_addr, p4, fsa) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) - # ============ P2: quantile rung, redundant per CTA (L1883-1997) ===== + # ============ P2: quantile rung, redundant per CTA =================== smn = cutlass.Float32(float("inf")) smx = cutlass.Float32(float("-inf")) - if shas != cutlass.Int32(0): # L1928-1932 + if shas != cutlass.Int32(0): for t in cutlass.range_constexpr(4): smn = C.fmin_f32(smn, fsa[t]) smx = C.fmax_f32(smx, fsa[t]) @@ -4914,7 +4829,7 @@ def kern( smx = C.fmax_f32(smx, fsb[t]) fma_ = cute.make_rmem_tensor((4,), cutlass.Float32) # mop-up pair bufs fmb_ = cute.make_rmem_tensor((4,), cutlass.Float32) - j = tidx + cutlass.Int32(BLK) # mop-up L1891-1897 + j = tidx + cutlass.Int32(BLK) # mop-up while j < smp2: p4 = j * SS2 * cutlass.Int32(4) if j >= SMP: @@ -4928,23 +4843,23 @@ def kern( smn = C.fmin_f32(smn, fmb_[t]) smx = C.fmax_f32(smx, fmb_[t]) j = j + cutlass.Int32(BLK) - a0 = C.warp_min_u32(C.fkey(smn)) # L1898-1901 + a0 = C.warp_min_u32(C.fkey(smn)) c0m = C.warp_max_u32(C.fkey(smx)) if lane == cutlass.Int32(0): s_wmn[tidx >> cutlass.Int32(5)] = a0 s_wmx[tidx >> cutlass.Int32(5)] = c0m - cute.arch.barrier() # ---- barrier L1902 ---- + cute.arch.barrier() # ---- barrier (sample redux publish) ---- - # PRIME-LATE (L1903-1916): every rank's sample has landed; prime NOW. - lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) # L1904 + # PRIME-LATE: every rank's sample has landed; prime NOW. + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(PFD)] - for uu in cutlass.range_constexpr(PFD): # clamped prime L1906 + for uu in cutlass.range_constexpr(PFD): # clamped prime i_ = rank * cutlass.Int32(STEPC) + tidx + cutlass.Int32(uu * BLK) ic = i_ if ic >= n4: ic = lim4 C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - # asm prefetch gate (L1912-1916): DEEP rows only; empty for U<=PFD + # asm prefetch gate: DEEP rows only; empty for U<=PFD if cutlass.const_expr(U > PFD): gpp = cutlass.Int32(0) if n4 >= cutlass.Int32(32768): @@ -4960,7 +4875,7 @@ def kern( * cutlass.Int64(16) ) - # cross-warp sample reduce (L1917-1923) + # cross-warp sample reduce av = cutlass.Uint32(0xFFFFFFFF) cv = cutlass.Uint32(0) if lane < cutlass.Int32(NW): @@ -4969,20 +4884,20 @@ def kern( SMIN = C.invkey(C.warp_min_u32(av)) SMAX = C.invkey(C.warp_max_u32(cv)) - GMIN = cutlass.Float32(C.SENT_LO) # L1925-1926 + GMIN = cutlass.Float32(C.SENT_LO) # sentinels GMAX = cutlass.Float32(C.SENT_HI) T = cutlass.Float32(_NEG_INF__clus) HIC = cutlass.Float32(_NEG_INF__clus) w = cutlass.Float32(0.0) - sok = cutlass.Int32(0) # L1930 + sok = cutlass.Int32(0) if SMP > cutlass.Int32(0): if SMAX > SMIN: sok = cutlass.Int32(1) - if sok != cutlass.Int32(0): # L1932-1947 sample hist + if sok != cutlass.Int32(0): # sample hist w = (SMAX - SMIN) * cutlass.Float32(1.0 / 256.0) - # CUDA --use_fast_math lowers `1.0f / w` to a bare MUFU.RCP; the - # plain `/` spelling emits the IEEE div.rn rcp+Newton+CALL chain - # (fix-4). w > 0 by the sok guard; bucketing is SC-invariant. + # CUDA --use_fast_math lowers `1.0f / w` to a bare MUFU.RCP; + # a plain `/` would emit the IEEE div.rn rcp+Newton+CALL + # chain. w > 0 by the sok guard; bucketing is SC-invariant. sc_s = cute.arch.rcp_approx(w) if shas != cutlass.Int32(0): for t in cutlass.range_constexpr(4): @@ -5013,8 +4928,8 @@ def kern( bq = cutlass.Int32(NBS - 1) C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) j = j + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L1948 ---- - # triple-target ZERO scan (L1952-1955): TGT / TGT2 / 2*TGT + cute.arch.barrier() # ---- barrier (sample histogram) ---- + # triple-target ZERO scan: TGT / TGT2 / 2*TGT C.scan_cross0( s_hist, TGT, @@ -5028,14 +4943,14 @@ def kern( two=True, three=True, ) - cute.arch.barrier() # ---- barrier L1956 ---- + cute.arch.barrier() # ---- barrier (scan publish) ---- tot0 = s_res[C.RES_TOT] b1v = s_res[C.RES_B] - if sok != cutlass.Int32(0): # L1957 + if sok != cutlass.Int32(0): if tot0 >= TGT: T = _fmaf__clus(cutlass.Float32(b1v), w, SMIN) - needg = cutlass.Int32(1) # L1958-1963 + needg = cutlass.Int32(1) if T > cutlass.Float32(_NEG_INF__clus): needg = cutlass.Int32(0) if needg != cutlass.Int32(0): @@ -5044,7 +4959,7 @@ def kern( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 ) # 2 barriers T = GMIN - if sok != cutlass.Int32(0): # L1964-1982 HIC + if sok != cutlass.Int32(0): # HIC tighten if tot0 >= TGT: b2v = s_res[C.RES_B2] if b2v >= cutlass.Int32(0): @@ -5061,7 +4976,7 @@ def kern( _fmaf__clus(cutlass.Float32(4.0), up, T), _fmaf__clus(cutlass.Float32(8.0), w, T), ) - # r4 (a000) ladder floor in SHARED (64-reg wall, L1983-1996) + # ladder floor kept in SHARED (64-register wall) if tidx == cutlass.Int32(0): t5 = cutlass.Float32(_NEG_INF__clus) if sok != cutlass.Int32(0): @@ -5074,7 +4989,7 @@ def kern( t5 = T3 s_tsh[0] = t5 - # ============ attempt loop (L1999-2162) — MUST NOT unroll =========== + # ============ attempt loop — MUST NOT unroll =========== listN = cutlass.Int32(0) above = cutlass.Int32(0) m = cutlass.Int32(0) @@ -5087,13 +5002,13 @@ def kern( fr = [ cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(U - PFD) - ] # explicit batch (op43 L1) + ] # explicit batch # (empty for U<=PFD — every row-pass float4 then comes from pf[]) att = cutlass.Int32(0) running = cutlass.Int32(1) while running != cutlass.Int32(0): - if att > cutlass.Int32(0): # retry preamble L2005-2024 - # EXACTNESS: re-prime pf[] (stale roll data, L2006-2015) + if att > cutlass.Int32(0): # retry preamble + # exactness: re-prime pf[] (holds stale roll data) if rank < nFullG: for uu in cutlass.range_constexpr(PFD): C.ld_g_f32x4( @@ -5109,14 +5024,14 @@ def kern( if ic >= n4: ic = lim4 C.ld_g_f32x4(atom128, x_addr, ic, pf[uu]) - C._cluster_sync_aligned() # ==== clus.sync L2016 ==== - if tidx < cutlass.Int32(NBS): # L2021-2022 + C._cluster_sync_aligned() # ==== clus.sync (retry) ==== + if tidx < cutlass.Int32(NBS): s_hist[tidx] = cutlass.Int32(0) if tidx == cutlass.Int32(0): s_scal[0] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2024 ---- + cute.arch.barrier() # ---- barrier (retry reset) ---- - TF = T # window L2026-2031 + TF = T # window hi = C.fmax_f32(GMAX, T) if HIC > T: if HIC < hi: @@ -5127,11 +5042,11 @@ def kern( wdok = cutlass.Int32(1) if wdok == cutlass.Int32(0): WD = cutlass.Float32(1e-30) - # MUFU.RCP spelling (fix-4): WD >= 1e-30 finite by the wdok - # clamp; classify bucketing is SC-invariant for any SC > 0. + # MUFU.RCP spelling: WD >= 1e-30 finite by the wdok clamp; + # classify bucketing is SC-invariant for any SC > 0. SC = cute.arch.rcp_approx(WD) - # ---- P3 row pass over OWNED CHUNKS (L2033-2121) ---- + # ---- P3 row pass over OWNED CHUNKS ---- g = rank + cutlass.Int32(0) while g < nCh: i0 = g * cutlass.Int32(STEPC) + tidx @@ -5139,7 +5054,7 @@ def kern( isfull = cutlass.Int32(0) if g < nFullG: isfull = cutlass.Int32(1) - if isfull != cutlass.Int32(0): # full body L2042-2049 + if isfull != cutlass.Int32(0): # full body for uu in cutlass.range_constexpr(PFD, U): C.ld_g_f32x4( atom128, x_addr, i0 + cutlass.Int32(uu * BLK), fr[uu - PFD] @@ -5151,7 +5066,7 @@ def kern( vv = fr[uu - PFD] for q in cutlass.range_constexpr(4): M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) - else: # partial body L2051-2064 + else: # partial body for uu in cutlass.range_constexpr(PFD, U): i_ = i0 + cutlass.Int32(uu * BLK) ic = i_ @@ -5172,8 +5087,8 @@ def kern( M = M | ( cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q) ) - # ROLL THE PREFETCH FORWARD (L2066-2081): next OWNED chunk, - # issued before the reservation and the survivor walk. + # ROLL THE PREFETCH FORWARD: next OWNED chunk, issued + # before the reservation and the survivor walk. g2 = g + cutlass.Int32(CS) if g2 < nCh: j0 = g2 * cutlass.Int32(STEPC) + tidx @@ -5190,7 +5105,7 @@ def kern( if jc >= n4: jc = lim4 C.ld_g_f32x4(atom128, x_addr, jc, pf[uu]) - # warp-aggregated slot reservation (L2082-2095) + # warp-aggregated slot reservation cnt = cutlass.Int32(C.popc(M)) inc = C.warp_incl_scan_add(cnt, lane) bpos = cutlass.Int32(0) @@ -5198,7 +5113,7 @@ def kern( if inc != cutlass.Int32(0): bpos = C.atomic_add_cta(s_scal.iterator + 0, inc) pos = cute.arch.shuffle_sync(bpos, cutlass.Int32(31)) + (inc - cnt) - # survivor bit-walk, software-pipelined ONE deep (L2105-2119); + # survivor bit-walk, software-pipelined ONE deep; # reload X[idx] — never hold the U float4s across the walk if M != cutlass.Int32(0): bp = C.ffs_m1(M) @@ -5220,12 +5135,12 @@ def kern( xv = xv2 pos = self._emitk(xv, idx, pos, TF, SC, SCAP, s_hist, s_cbuf2) g = g + cutlass.Int32(CS) - # rank-0 scalar tail (L2122-2130): per-thread atomics, bound-check + # rank-0 scalar tail: per-thread atomics, bound-check i = tidx while i < tailn: x = C.ldg_f32(x_addr, tail0 + i) if x >= TF: - bq = C.f2s_rz((x - TF) * SC) # signed form L2125 + bq = C.f2s_rz((x - TF) * SC) # signed form if bq > cutlass.Int32(NBS - 1): bq = cutlass.Int32(NBS - 1) C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) @@ -5236,16 +5151,16 @@ def kern( ) | cutlass.Uint64(C.u32_of_f32(x)) i = i + cutlass.Int32(BLK) - # ---- cluster merge (L2132-2148) ---- - C._cluster_sync_aligned() # ==== clus.sync L2135 ==== - myn = s_scal[0] # L2140 + # ---- cluster merge ---- + C._cluster_sync_aligned() # ==== clus.sync (merge) ==== + myn = s_scal[0] self._merge_scan0_local(s_hist, s_mrg, rank, k, tidx, s_res) - cute.arch.barrier() # ---- barrier L2143 ---- + cute.arch.barrier() # ---- barrier (merge publish) ---- tot = s_res[C.RES_TOT] acc = cutlass.Int32(0) if tot >= k: acc = cutlass.Int32(1) - if acc != cutlass.Int32(0): # L2145-2148 accept + if acc != cutlass.Int32(0): # accept valid = cutlass.Int32(1) complete = cutlass.Int32(0) if myn <= SCAP: @@ -5257,10 +5172,10 @@ def kern( B = s_res[C.RES_B] running = cutlass.Int32(0) else: - if att == cutlass.Int32(2): # L2149 + if att == cutlass.Int32(2): # ladder exhausted running = cutlass.Int32(0) else: - # rung ladder (L2150-2161) — cluster-uniform on every arm + # rung ladder — cluster-uniform on every arm tshtaken = cutlass.Int32(0) if att == cutlass.Int32(0): T5 = s_tsh[0] @@ -5274,16 +5189,16 @@ def kern( GMIN, GMAX = C.gather_hint( x_addr, p_addr, k, n, tidx, s_wmn, s_wmx, blk=BLK, kpt=1 ) # 2 barriers inside - floorhit = cutlass.Int32(1) # L2160 + floorhit = cutlass.Int32(1) if T > GMIN: floorhit = cutlass.Int32(0) if floorhit != cutlass.Int32(0): running = cutlass.Int32(0) else: - T = GMIN # L2161 + T = GMIN att = att + cutlass.Int32(1) - # ============ classification (L2165-2173) ============ + # ============ classification ============ whole = cutlass.Int32(0) if valid != cutlass.Int32(0): if need >= m: @@ -5299,12 +5214,12 @@ def kern( mc = cutlass.Int32(0) if degen == cutlass.Int32(0): mc = m - # crossing candidates land in RANK 0's ck64c via DSMEM (L2173) + # crossing candidates land in RANK 0's ck64c via DSMEM rk64 = C._mapa_shared_cluster_addr(ck_addr, cutlass.Int32(0)) if degen == cutlass.Int32(0): if complete != cutlass.Int32(0): - # ---- P5 emit from staged cbuf (L2176-2187) ---- + # ---- P5 emit from staged cbuf ---- i = tidx while i < listN: pk64 = s_cbuf2[i] @@ -5316,8 +5231,7 @@ def kern( ) i = i + cutlass.Int32(BLK) else: - # ---- EXACTNESS re-sweep: OWNED CHUNKS + rank-0 true tail - # (L2188-2222) ---- + # ---- exactness re-sweep: OWNED CHUNKS + rank-0 true tail ---- g = rank + cutlass.Int32(0) while g < nCh: lo2 = (g * cutlass.Int32(STEPC)) << cutlass.Int32(2) @@ -5344,22 +5258,22 @@ def kern( ) t2 = t2 + cutlass.Int32(BLK) - # ============ EXIT RENDEZVOUS (L2226) ============ - # all DSMEM traffic retired; the ONLY exit rendezvous. rank!=0 falls - # through to the kernel end (post-barrier asymmetric exit L2227); + # ============ EXIT RENDEZVOUS ============ + # all DSMEM traffic retired; the ONLY exit rendezvous. rank!=0 + # falls through to the kernel end (post-barrier asymmetric exit); # NO later cluster barrier. - C._cluster_sync_aligned() # ==== clus.sync L2226 ==== + C._cluster_sync_aligned() # ==== clus.sync (exit) ==== if rank == cutlass.Int32(0): if degen == cutlass.Int32(0): if whole == cutlass.Int32(0): - # ---- P6 rank-0 refine (L2229-2303) ---- - if mc <= cutlass.Int32(QUADC_CLUS__clus): # L2232-2247 O(mc^2) + # ---- P6 rank-0 refine ---- + if mc <= cutlass.Int32(QUADC_CLUS__clus): # O(mc^2) mc2 = mc & cutlass.Int32(~1) i = tidx while i < mc: - # G1: re-assert Uint64 at every unsigned compare - # in/after dynamic loops (ct_main notes). + # re-assert Uint64 at every unsigned compare + # in/after dynamic loops. u64v = s_ck64[i] r_ = cutlass.Int32(0) jq = cutlass.Int32(0) @@ -5371,7 +5285,7 @@ def kern( + cutlass.Int32(vhi > cutlass.Uint64(u64v)) ) jq = jq + cutlass.Int32(2) - if mc2 < mc: # odd tail L2244 + if mc2 < mc: # odd tail r_ = r_ + cutlass.Int32( cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) ) @@ -5383,20 +5297,20 @@ def kern( ) i = i + cutlass.Int32(BLK) else: - # key-space narrowing over ck64c (L2249-2278) + # key-space narrowing over ck64c if tidx == cutlass.Int32(0): s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) s_kmm[1] = cutlass.Uint32(0) - if tidx < cutlass.Int32(NBS): # cleared ONCE L2251 + if tidx < cutlass.Int32(NBS): # cleared ONCE s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2255 ---- + cute.arch.barrier() # ---- barrier (narrowing init) ---- i = tidx while i < mc: kk = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) C.atomic_min_cta(s_kmm.iterator + 0, kk) C.atomic_max_cta(s_kmm.iterator + 1, kk) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2257 ---- + cute.arch.barrier() # ---- barrier (key range) ---- rlo = s_kmm[0] rhi = s_kmm[1] ethr = cutlass.Int64(cutlass.Uint32(rlo)) @@ -5405,7 +5319,7 @@ def kern( mm = mc brk = cutlass.Int32(0) lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L2258-2278 (<=6 lvl) + while brk == cutlass.Int32(0): # <=6 levels if needC == mm: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) aboveC = aboveC + mm @@ -5427,7 +5341,7 @@ def kern( sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) i = tidx - while i < mc: # L2266-2269 re-bin + while i < mc: # re-bin uq = cutlass.Uint32(s_ck64[i] >> cutlass.Uint64(32)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -5439,7 +5353,7 @@ def kern( cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2270 ---- + cute.arch.barrier() # ---- barrier (level hist) ---- C.scan_cross0( s_hist, needC, @@ -5451,25 +5365,25 @@ def kern( nb=NBS, zero=True, ) - cute.arch.barrier() # ---- barrier L2273 ---- + cute.arch.barrier() # ---- barrier (level scan) ---- aboveC = aboveC + s_res[C.RES_ABOVE] needC = needC - s_res[C.RES_ABOVE] mm = s_res[C.RES_M] sB = s_res[C.RES_B] nlo = cutlass.Uint32(rlo) + (cutlass.Uint32(sB) << sh2u) - if sB != cutlass.Int32(NBS - 1): # L2276 + if sB != cutlass.Int32(NBS - 1): rhi = nlo + ( (cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1) ) rlo = nlo lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L2279 + if tidx == cutlass.Int32(0): s_scal[1] = cutlass.Int32(0) s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2281 ---- + cute.arch.barrier() # ---- barrier (emit counters) ---- it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) it = cutlass.Int32(0) - while it < it2: # L2282-2301 ballot emit + while it < it2: # ballot emit i = it * cutlass.Int32(BLK) + tidx p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) @@ -5499,7 +5413,7 @@ def kern( it = it + cutlass.Int32(1) else: # ---- degen fallback: whole-row key-space narrowing - # (L2305-2355; per-level clear + scan_cross w/ ws) ---- + # (per-level clear + scan_cross w/ ws) ---- rlo = cutlass.Uint32(0) rhi = cutlass.Uint32(0xFFFFFFFF) above2 = cutlass.Int32(0) @@ -5509,7 +5423,7 @@ def kern( tie_m = cutlass.Int32(1) brk = cutlass.Int32(0) lev = cutlass.Int32(0) - while brk == cutlass.Int32(0): # L2312-2331 (<=8 lvl) + while brk == cutlass.Int32(0): # <=8 levels if need2 == m2: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) above2 = above2 + m2 @@ -5531,11 +5445,11 @@ def kern( if sh2 < cutlass.Int32(0): sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) - if tidx < cutlass.Int32(NBS): # per-level clear L2320 + if tidx < cutlass.Int32(NBS): # per-level clear s_hist[tidx] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2321 ---- + cute.arch.barrier() # ---- barrier (level clear) ---- i = tidx - while i < n: # L2322-2324 whole row + while i < n: # whole row uq = C.fkey(C.ldg_f32(x_addr, i)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -5546,9 +5460,9 @@ def kern( s_hist.iterator + cutlass.Int32(du), cutlass.Int32(1) ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # ---- barrier L2325 ---- - # block-parallel scan (ONE internal barrier; only use - # of ws in this kernel, L2326) + cute.arch.barrier() # ---- barrier (level hist) ---- + # block-parallel scan (ONE internal barrier; only + # use of ws in this kernel) C.scan_cross( s_hist, s_ws, @@ -5560,7 +5474,7 @@ def kern( nb=NBS, two=False, ) - cute.arch.barrier() # ---- barrier L2327 ---- + cute.arch.barrier() # ---- barrier (level scan) ---- above2 = above2 + s_res[C.RES_ABOVE] need2 = need2 - s_res[C.RES_ABOVE] m2 = s_res[C.RES_M] @@ -5570,18 +5484,18 @@ def kern( rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) rlo = nlo lev = lev + cutlass.Int32(1) - if tidx == cutlass.Int32(0): # L2336 + if tidx == cutlass.Int32(0): s_scal[1] = cutlass.Int32(0) s_scal[2] = cutlass.Int32(0) - cute.arch.barrier() # ---- barrier L2337 ---- - nA = k # L2338 + cute.arch.barrier() # ---- barrier (emit counters) ---- + nA = k nT = cutlass.Int32(0) if tie_m != cutlass.Int32(0): nA = above2 nT = need2 it2 = (n + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) it = cutlass.Int32(0) - while it < it2: # L2340-2354 + while it < it2: i = it * cutlass.Int32(BLK) + tidx p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) @@ -5599,8 +5513,8 @@ def kern( it = it + cutlass.Int32(1) # ------------------------------------------------------------------ - # host launcher: grid dim3(CS, b) + cluster (CS,1,1) (probe P11); - # min_blocks_per_mp=1 == __launch_bounds__(1024, 1) 64-reg wall. + # host launcher: grid dim3(CS, b) + cluster (CS,1,1); + # min_blocks_per_mp=1 == __launch_bounds__(1024, 1) 64-register wall. # ------------------------------------------------------------------ @cute.jit def __call__( @@ -5699,7 +5613,7 @@ def get_compiled__clus( def run__clus(logits, pre_idx, n: int, out): """torch-facing single-call entry: routes (b, n, k) through ct_dispatch, asserts the shape lands on gvr_clus, launches the matching variant. - gvr_clus takes NO workspace (spec §4c).""" + gvr_clus takes NO workspace.""" import torch # debug-entry only: module stays torch-free at import try: @@ -5761,17 +5675,13 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): # =========================================================================== -# ==== ct_regclus.py — family: regclus ========================= +# ==== family: regclus ========================================= # =========================================================================== -"""ct_regclus.py — op46 gvr_reg_clus CuTeDSL translation. - -CUDA ground truth: src_cuda/kernel.cu L2359-2648 (CLUSTERED register-resident -GVR: the register algorithm — T = GMIN directly, one float-space histogram, -one register sweep — run across a cluster of CS CTAs; per-CTA instruction -stream intentionally identical to the single-CTA reg path plus two hardware -cluster barriers and CS DSMEM reads per bin). Contract: TRANSLATION_SPEC.md -§5.4. Probe verdicts (probes/PROBE_RESULTS.md) binding; op43 lessons L1-L5 -applied; ct_main G1 signedness mitigation applied at every unsigned +"""gvr_reg_clus — CLUSTERED register-resident GVR: the register algorithm +(T = GMIN directly, one float-space histogram, one register sweep) run +across a cluster of CS CTAs; per-CTA instruction stream intentionally +identical to the single-CTA reg path plus two hardware cluster barriers and +CS DSMEM reads per bin. Signedness rule applied at every unsigned compare/shift in/after dynamic loops. Template knobs (CUDA `gvr_reg_clus`, all instantiations @@ -5785,35 +5695,35 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): __shared__ block folded into the first 512 B — byte-identical layout in every CTA, a mapa/DSMEM requirement): - [0..5] s_res (ct_common slot map RES_B/M/ABOVE/TOT/B2/B3) - [6..7] s_cnt (s_o1, s_o2 — kernel.cu L2392) - [8..9] s_kmm (s_kmin, s_kmax — Uint32, L2393) - [16..16+32) ws (scan_cross_w workspace, L2390) + [0..5] s_res (shared slot map RES_B/M/ABOVE/TOT/B2/B3) + [6..7] s_cnt (s_o1, s_o2) + [8..9] s_kmm (s_kmin, s_kmax — Uint32) + [16..16+32) ws (scan_cross_w workspace) [48..48+32) wmn (Uint32 warp min partials) [80..80+32) wmx (Uint32 warp max partials) - [128..1152) hist (this CTA's raw counts, L2384) + [128..1152) hist (this CTA's raw counts) [1152..2176) mrg (cluster totals -> per-CTA global write cursors) [2176..3200) hoff (this CTA's rank-exclusive bin offset) [3200..7296) ck (crossing keys, Uint32, CMPC=4096 slots) [7296..11392) ci (crossing indices, Int32, CMPC slots) Launch smem = 45,568 B (compile-time constant -> plain int at .launch(); -MINB==1 so the _build_kernel_attrs carveout path is not taken and ct_reg's -_no_carveout workaround is unnecessary here). +MINB==1 so the _build_kernel_attrs carveout path is not taken and the reg +family's _no_carveout workaround is unnecessary here). """ -# ---- constants (kernel.cu lines) ------------------------------------------- -NB__regclus = 1024 # L16 (histogram bins; == BLKC here) -LNB = 10 # L1267 log2(NB__regclus) — reg_clus narrowing shift -QUADC__regclus = 96 # L21 O(mc^2) rank gate (L2534) -CMPC = 4096 # L2372 crossing slots PER CTA (pow2) -LCMPC = 12 # L2373 log2(CMPC) -BLKC = 1024 # L2374 CTA size +# ---- constants -------------------------------------------------------------- +NB__regclus = 1024 # histogram bins; == BLKC here +LNB = 10 # log2(NB__regclus) — reg_clus narrowing shift +QUADC__regclus = 96 # O(mc^2) rank gate +CMPC = 4096 # crossing slots PER CTA (pow2) +LCMPC = 12 # log2(CMPC) +BLKC = 1024 # CTA size STATIC_WORDS__regclus = 128 # DSL smem prelude (static-__shared__ mirror) STATIC_BYTES__regclus = STATIC_WORDS__regclus * 4 -DYN_SMEM_BYTES = (3 * NB__regclus + 2 * CMPC) * 4 # 45,056 (L2926) +DYN_SMEM_BYTES = (3 * NB__regclus + 2 * CMPC) * 4 # 45,056 SMEM_BYTES = STATIC_BYTES__regclus + DYN_SMEM_BYTES # 45,568 # word offsets into the shared window (module docstring) @@ -5828,12 +5738,11 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): # --------------------------------------------------------------------------- -# module-local FP/util spellings (copied from frozen sibling ct_reg.py — -# probe P6 fma discipline; kept local so this module is self-contained) +# module-local FP/util spellings (same forms as the reg family) # --------------------------------------------------------------------------- @dsl_user_op def _fmaf__regclus(a, b, c, *, loc=None, ip=None): - """CUDA fmaf: single fma.rn.f32 (P6 spelling; classify == emit bit-exact).""" + """CUDA fmaf: single fma.rn.f32 (classify == emit bit-exact).""" return cutlass.Float32( mlir_math.fma( a.ir_value(loc=loc, ip=ip), @@ -5875,7 +5784,7 @@ def _val__regclus(frags, s: int): class GvrRegClusKernel: - """gvr_reg_clus (kernel.cu L2376-2379).""" + """gvr_reg_clus.""" def __init__( self, @@ -5887,7 +5796,7 @@ def __init__( next_n: int = 1, cr_shift: int = 0, ): - assert blk == BLKC, "all instantiations BLK=BLKC=1024 (spec §4b)" + assert blk == BLKC, "all instantiations BLK=BLKC=1024" assert vpt in (1, 2, 4) and cs in (2, 4, 8) self.blk = blk self.vpt = vpt @@ -5902,8 +5811,8 @@ def __init__( self.cr_shift = int(cr_shift) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) - self.S = vpt * 4 # L2381 - self.span = blk * vpt # L2382 (float4 per CTA) + self.S = vpt * 4 + self.span = blk * vpt # float4 per CTA # ------------------------------------------------------------------ @cute.kernel @@ -5922,10 +5831,10 @@ def kern( NW = cutlass.const_expr(self.blk // 32) if cutlass.const_expr(self.pdl): - cute.arch.griddepcontrol_wait() # L2380 (knob default off) + cute.arch.griddepcontrol_wait() # knob default off tid, _, _ = cute.arch.thread_idx() - rank, row, _ = cute.arch.block_idx() # L2396-2397 (P11: bx=rank) + rank, row, _ = cute.arch.block_idx() # bx=rank lane = tid & cutlass.Int32(31) # ================= per-row varlen prologue (varlen mode only) ========= @@ -5971,7 +5880,7 @@ def kern( # ------------------------------------------------------------------ # Predeclarations (DSL AST rule: every scalar (re)assigned under a # dynamic if/while must pre-exist with a stable type; constant inits - # are dead-coded — ct_reg precedent, reg-audited free). + # are dead-coded). # ------------------------------------------------------------------ i = cutlass.Int32(0) j = cutlass.Int32(0) @@ -6055,20 +5964,20 @@ def kern( ck_addr = sbase + cutlass.Int32(W_CK * 4) ci_addr = sbase + cutlass.Int32(W_CI * 4) - n4 = n >> cutlass.Int32(2) # L2405 - ntail = n - (n4 << cutlass.Int32(2)) # L2406 - base4 = rank * cutlass.Int32(self.span) # L2407 - tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` L2425 + n4 = n >> cutlass.Int32(2) + ntail = n - (n4 << cutlass.Int32(2)) + base4 = rank * cutlass.Int32(self.span) + tix = (n4 << cutlass.Int32(2)) + tid # CUDA `tidx` - # ---- P0: redundant hint gather, EVERY CTA (L2410-2413; k<=BLK by - # dispatch gate L2897). One coalesced word per thread, NO cluster - # barrier — GMIN/GMAX identical everywhere by construction. + # ---- P0: redundant hint gather, EVERY CTA (k<=BLK by dispatch + # gate). One coalesced word per thread, NO cluster barrier — + # GMIN/GMAX identical everywhere by construction. if tid < k: pv0 = ld_g_i32(p_addr, tid) - # ---- P1: row load — predicated flat float4[VPT] batch (L2415-2424; - # the CUDA has NO exact-fit peel here, guard is per-load). Issue all - # loads first (op43 L1), then -INFINITY-fill missed slots (op43 L2). + # ---- P1: row load — predicated flat float4[VPT] batch (the CUDA + # has NO exact-fit peel here, guard is per-load). Issue all loads + # first, then -INFINITY-fill missed slots. atom128 = g2r_atom_f32(128, invariant=True) frags = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(VPT)] for u in cutlass.range_constexpr(VPT): @@ -6077,26 +5986,26 @@ def kern( ld_g_f32x4(atom128, x_addr, i, frags[u]) for u in cutlass.range_constexpr(VPT): i = base4 + tid + cutlass.Int32(u * self.blk) - if i >= n4: # -INFINITY fill L2421 + if i >= n4: # -INFINITY fill for z in cutlass.range_constexpr(4): frags[u][z] = cutlass.Float32(_NEG_INF__regclus) - # tail element: rank 0 only (L2425-2426) + # tail element: rank 0 only if rank == cutlass.Int32(0): if tid < ntail: tval = ldg_f32(x_addr, tix) - # ---- P2: init (L2428-2429). NB__regclus == BLK -> single-pass hist clear. + # ---- P2: init. NB__regclus == BLK -> single-pass hist clear. if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) s_cnt[1] = cutlass.Int32(0) for z in cutlass.range_constexpr(NB__regclus // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - # ---- P3: GMIN/GMAX from the hint (L2431-2445), ONE barrier fold. + # ---- P3: GMIN/GMAX from the hint, ONE barrier fold. lmin = cutlass.Uint32(0xFFFFFFFF) lmax = cutlass.Uint32(0) if cutlass.Uint32(pv0) < cutlass.Uint32(n): - uk = fkey(ldg_f32(x_addr, pv0)) # __ldg(X+pv0) L2433 + uk = fkey(ldg_f32(x_addr, pv0)) # __ldg(X+pv0) lmin = uk lmax = uk lmin = warp_min_u32(lmin) @@ -6104,7 +6013,7 @@ def kern( if lane == cutlass.Int32(0): s_wmn[tid >> cutlass.Int32(5)] = lmin s_wmx[tid >> cutlass.Int32(5)] = lmax - cute.arch.barrier() # L2438 + cute.arch.barrier() # warp partials published a = cutlass.Uint32(0xFFFFFFFF) c = cutlass.Uint32(0) if lane < cutlass.Int32(NW): @@ -6115,7 +6024,7 @@ def kern( Tv = invkey(lmin) GMAX = invkey(lmax) - # ---- collapse guard, NaN-safe (L2446-2453) + # ---- collapse guard, NaN-safe okc = cutlass.Int32(0) if Tv < GMAX: if (GMAX - Tv) > cutlass.Float32(1e-30): @@ -6124,19 +6033,19 @@ def kern( Tv = cutlass.Float32(SENT_LO) GMAX = cutlass.Float32(SENT_HI) - # ---- bin transform constants (L2454-2467): branchless trash bin. + # ---- bin transform constants: branchless trash bin. WD = (GMAX - Tv) * cutlass.Float32(1.0 / float(NB__regclus - 2)) wsel = cutlass.Float32(1e-30) if WD > cutlass.Float32(0.0): wsel = WD - # MUFU.RCP spelling (fix-4, mirrors the reg family's L3668 site): + # MUFU.RCP spelling (mirrors the reg family's site): # wsel >= 1e-30 finite; bucketing is SC-invariant for any SC > 0. SC = cute.arch.rcp_approx(wsel) CQ0 = cutlass.Float32(1.0) - Tv * SC CQ = CQ0 + cutlass.Float32(1e-6) * (_fabsf__regclus(CQ0) + cutlass.Float32(1.0)) - # ---- P4: histogram (L2469-2472); tval add UNCONDITIONAL (trash bin - # swallows -INFINITY via the saturating cvt). + # ---- P4: histogram; tval add UNCONDITIONAL (trash bin swallows + # -INFINITY via the saturating cvt). for s in cutlass.range_constexpr(S): qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) @@ -6145,12 +6054,12 @@ def kern( bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) atomic_add_cta(s_hist.iterator + cutlass.Int32(bn), cutlass.Int32(1)) - # ---- P5: cluster merge (L2474-2484) - _cluster_sync_aligned() # L2474 + # ---- P5: cluster merge + _cluster_sync_aligned() # histograms complete on every rank for z in cutlass.range_constexpr(NB__regclus // self.blk): i = tid + cutlass.Int32(z * self.blk) - # CS-unrolled remote u32 loads: batch-issue, then fold (#pragma - # unroll L2477; one mapa per (i, r) exactly like map_shared_rank) + # CS-unrolled remote u32 loads: batch-issue, then fold + # (one mapa per (i, r) exactly like the CUDA) hvals = [] for r in cutlass.range_constexpr(CS): ma = _mapa_shared_cluster_addr( @@ -6166,10 +6075,10 @@ def kern( s_mrg[i] = tot_a s_hoff[i] = pre_a - # ---- P6: scan (L2485-2492) - cute.arch.barrier() # L2485 + # ---- P6: scan + cute.arch.barrier() # merge published scan_cross_w(s_mrg, s_ws, k, tid, s_res, blk=self.blk, nb=NB__regclus) - cute.arch.barrier() # L2487 + cute.arch.barrier() # scan published above = s_res[RES_ABOVE] m = s_res[RES_M] Bv = s_res[RES_B] @@ -6182,17 +6091,17 @@ def kern( degen = cutlass.Int32(1) for z in cutlass.range_constexpr(NB__regclus // self.blk): i = tid + cutlass.Int32(z * self.blk) - s_mrg[i] = s_mrg[i] + s_hoff[i] # L2491 global cursor - cute.arch.barrier() # L2492 + s_mrg[i] = s_mrg[i] + s_hoff[i] # global cursor + cute.arch.barrier() # cursors published - # ---- P7: register sweep emit (L2494-2527, !degen) + # ---- P7: register sweep emit (!degen) if degen == cutlass.Int32(0): - LOQ = cutlass.Float32(Bv) # L2495 + LOQ = cutlass.Float32(Bv) lim1 = above if whole == cutlass.Int32(1): - lim1 = above + m # L2496 + lim1 = above + m for s in cutlass.range_constexpr(S): - qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) # bit-identical L2499 + qv = _fmaf__regclus(_val__regclus(frags, s), SC, CQ) # bit-identical if qv >= LOQ: bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) p = atomic_add_cta(s_mrg.iterator + cutlass.Int32(bn), cutlass.Int32(1)) @@ -6204,7 +6113,7 @@ def kern( else: if whole == cutlass.Int32(0): # crossing overflow -> striped DSMEM slabs; TWO - # separate u32 remote stores (NOT packed, L2507-10) + # separate u32 remote stores (NOT packed) q2i = p - above rnk = q2i >> cutlass.Int32(LCMPC) j = (q2i & cutlass.Int32(CMPC - 1)) << cutlass.Int32(2) @@ -6215,7 +6124,7 @@ def kern( _st_shared_cluster_i32( _mapa_shared_cluster_addr(ci_addr + j, rnk), idx ) - # tail element (L2514-2526): tval == -INF fails q>=LOQ elsewhere + # tail element: tval == -INF fails q>=LOQ elsewhere qv = _fmaf__regclus(tval, SC, CQ) if qv >= LOQ: bn = _umin_u32__regclus(f2u_rz(qv), cutlass.Uint32(NB__regclus - 1)) @@ -6232,20 +6141,20 @@ def kern( ) _st_shared_cluster_i32(_mapa_shared_cluster_addr(ci_addr + j, rnk), tix) - # ---- P8 (L2529-2530): release staging to rank 0 - cute.arch.barrier() # L2529 - _cluster_sync_aligned() # L2530 + # ---- P8: release staging to rank 0 + cute.arch.barrier() + _cluster_sync_aligned() - # ---- P9: rank-0 selection (L2532-2647) + # ---- P9: rank-0 selection if rank == cutlass.Int32(0): if whole == cutlass.Int32(0): mc = m if degen == cutlass.Int32(1): - mc = cutlass.Int32(0) # L2533 + mc = cutlass.Int32(0) if degen == cutlass.Int32(0): if mc <= cutlass.Int32(QUADC__regclus): # (1) quad-96: all candidates LOCAL (96 < CMPC), - # O(mc^2) slot-order tie-broken rank (L2535-2543) + # O(mc^2) slot-order tie-broken rank i = tid while i < mc: uq = cutlass.Uint32(s_ck[i]) @@ -6265,12 +6174,12 @@ def kern( out_row[above + rnk] = s_ci[i] i = i + cutlass.Int32(BLK) else: - # (2) key-space narrowing over striped DSMEM slabs - # (L2544-2596): slot = i & (CMPC-1), rank = i >> LCMPC + # (2) key-space narrowing over striped DSMEM slabs: + # slot = i & (CMPC-1), rank = i >> LCMPC if tid == cutlass.Int32(0): s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) s_kmm[1] = cutlass.Uint32(0) - cute.arch.barrier() # L2546 + cute.arch.barrier() # kmm init i = tid while i < mc: kv = cutlass.Uint32( @@ -6285,7 +6194,7 @@ def kern( atomic_min_cta(s_kmm.iterator, kv) atomic_max_cta(s_kmm.iterator + 1, kv) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2551 + cute.arch.barrier() # key range published rlo = cutlass.Uint32(s_kmm[0]) rhi = cutlass.Uint32(s_kmm[1]) ethr = cutlass.Int64(cutlass.Uint32(rlo)) @@ -6294,8 +6203,8 @@ def kern( mm = mc lev = cutlass.Int32(0) done = cutlass.Int32(0) - while done == cutlass.Int32(0): # <=6 levels L2553 - if needC == mm: # L2554 + while done == cutlass.Int32(0): # <=6 levels + if needC == mm: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) aboveC = aboveC + mm needC = cutlass.Int32(0) @@ -6317,7 +6226,7 @@ def kern( sh2 = b2w - cutlass.Int32(LNB) for z in cutlass.range_constexpr(NB__regclus // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L2563 + cute.arch.barrier() # level clear i = tid while i < mc: unar = cutlass.Uint32( @@ -6345,9 +6254,9 @@ def kern( cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2568 + cute.arch.barrier() # level hist find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) - cute.arch.barrier() # L2570 + cute.arch.barrier() # level scan aboveC = aboveC + s_res[RES_ABOVE] needC = needC - s_res[RES_ABOVE] mm = s_res[RES_M] @@ -6362,7 +6271,7 @@ def kern( ) rlo = nlo lev = lev + cutlass.Int32(1) - cute.arch.barrier() # L2576 + cute.arch.barrier() # narrowing done # two-predicate ballot emit over the striped slabs lml = cutlass.Int32(cute.arch.lanemask_lt()) it2 = (mc + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) @@ -6419,9 +6328,9 @@ def kern( out_row[above + aboveC + p2e] = idv it = it + cutlass.Int32(1) else: - # (3) degen safety net (L2597-2645): crossing bin larger - # than the whole cluster buffer -> exact whole-row - # key-space narrowing by rank 0 alone, <=8 levels. + # (3) degen safety net: crossing bin larger than the + # whole cluster buffer -> exact whole-row key-space + # narrowing by rank 0 alone, <=8 levels. rlo = cutlass.Uint32(0) rhi = cutlass.Uint32(0xFFFFFFFF) aboveC = cutlass.Int32(0) # above2 @@ -6432,7 +6341,7 @@ def kern( lev = cutlass.Int32(0) done = cutlass.Int32(0) while done == cutlass.Int32(0): - if needC == mm: # L2603 + if needC == mm: ethr = cutlass.Int64(cutlass.Uint32(rlo)) - cutlass.Int64(1) aboveC = aboveC + mm needC = cutlass.Int32(0) @@ -6442,7 +6351,7 @@ def kern( if cutlass.Uint32(rlo) >= cutlass.Uint32(rhi): ethr = cutlass.Int64(cutlass.Uint32(rlo)) done = cutlass.Int32(1) - if lev >= cutlass.Int32(8): # L2605 + if lev >= cutlass.Int32(8): ethr = cutlass.Int64(cutlass.Uint32(rlo)) done = cutlass.Int32(1) if done == cutlass.Int32(0): @@ -6455,9 +6364,9 @@ def kern( sh2 = b2w - cutlass.Int32(LNB) for z in cutlass.range_constexpr(NB__regclus // self.blk): s_hist[tid + cutlass.Int32(z * self.blk)] = cutlass.Int32(0) - cute.arch.barrier() # L2612 + cute.arch.barrier() # level clear i = tid - while i < n: # whole-row bin L2613 + while i < n: # whole-row bin unar = fkey(ldg_f32(x_addr, i)) if cutlass.Uint32(unar) >= cutlass.Uint32(rlo): if cutlass.Uint32(unar) <= cutlass.Uint32(rhi): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index c761aca49f78..165048c8dcfe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -14,38 +14,34 @@ """Self-sampling GVR top-K decode — host side (dispatch, workspace, entry). -Companion to ``gvr_topk_decode_self_sampling.py`` (the merged device -module). Three sections, each a rename-only merge of the per-family source -(fork branch ``GVR-selfsampling-CuTeDSL``): +Companion to ``gvr_topk_decode_self_sampling.py`` (the device module). +Three sections: -1. dispatch — bit-exact transcription of the CUDA host dispatch, a pure - function ``route(b, n, npad, k)`` (cross-checked against an independent - second transcription by a 1,159,168-case boundary+fuzz sweep); +1. dispatch — the CUDA host dispatch as a pure function + ``route(b, n, npad, k)``; 2. workspace — one zero-initialised per-device slab (20,973,568 B) via the - torch caching allocator, keep-alive + double-checked locking semantics - mirrored from the CUDA binding; + torch caching allocator, with keep-alive + double-checked locking; 3. operator entry — ``run(logits, pre_idx, n_valid, indices)`` / - ``run_ws(..., workspace)`` DPS forms with the CUDA binding's hardening - battery, BIND-ONCE launch cache keyed on ``(b, n, npad, k)``. - -OPERATOR CONTRACT (standalone; not wired into the decode path): ``n_valid`` -is one host python int for the whole batch — every row shares the same -valid prefix, in COMPRESSED index space (the caller applies any -``compressRatio`` division). ``pre_idx`` is consumed AS-IS — raw prev-step -top-K indices, uniformly for DSv3.2 / DSv4 Flash / Pro. This deliberately -drops the +1 temporal shift ``heuristicTopKDecode.cu`` applies for cr==1: -hints only steer the sampling ladder (exactness never depends on them), and -on real V3.2 decode captures raw prev-step hints land on MORE of the current -top-K than +1-shifted ones (mean overlap 0.773 vs 0.536 across 15 cells x 14 -consecutive step-pairs, the gap widening with ISL), so one offset-free hint -convention serves all three models. The production per-row contract -(per-request ``kv_lens`` read on-device, per-row MTP offsets — sync-free and -CUDA-graph-replay safe with growing KV) is implemented by ``run_varlen``, -which is the entry the opt-in DSA dispatch seam calls. The batch-uniform -``run``/``run_ws`` entries keep the original standalone contract (one -host-side ``n_valid`` for the whole batch), are exercised for unit tests and -benchmarking only, and must not be substituted for the tiered path under -continuous batching, MTP (``next_n > 1``), or CUDA-graph capture. + ``run_ws(..., workspace)`` DPS forms with input hardening and a + bind-once launch cache keyed on ``(b, n, npad, k)``. + +OPERATOR CONTRACT (batch-uniform entries): ``n_valid`` is one host python +int for the whole batch — every row shares the same valid prefix, in +COMPRESSED index space (the caller applies any ``compressRatio`` division). +``pre_idx`` is consumed as-is — raw prev-step top-K indices, uniformly for +DSv3.2 / DSv4 Flash / Pro. The +1 temporal shift ``heuristicTopKDecode.cu`` +applies for cr==1 is deliberately dropped: hints only steer the sampling +ladder (exactness never depends on them), and raw prev-step hints overlap +the current top-K at least as well as +1-shifted ones on real decode data, +so one offset-free hint convention serves all three models. The production +per-row contract (per-request ``kv_lens`` read on-device, per-row MTP +offsets — sync-free and CUDA-graph-replay safe with growing KV) is +implemented by ``run_varlen``, which is the entry the opt-in DSA dispatch +seam calls. The batch-uniform ``run``/``run_ws`` entries keep the simpler +contract (one host-side ``n_valid`` for the whole batch), are exercised for +unit tests and benchmarking only, and must not be substituted for +``run_varlen`` under continuous batching, MTP (``next_n > 1``), or +CUDA-graph capture. """ import math @@ -72,62 +68,23 @@ def _device(): # =========================================================================== -# ==== dispatch (ct_dispatch.py) ============================================ +# ==== dispatch ============================================================= # =========================================================================== -"""Pure-Python transcription of the frozen GVR CUDA dispatch (gvr_topk_launch). - -Source of truth: ../src_cuda/kernel.cu (3197 lines). route(b, n, npad, k) is a -PURE function of its four ints -- no env knobs, no GPU, stdlib only. - -Branch map (kernel.cu line citations): - constants NB L16, QUADC L21, SNB L170, CMPC L2372, BLKC L2374 - reg-block prologue L2757-2822: wide=(b<=148) L2757; n4=n>>2 L2759; - CMP=min(n,2560) L2764; QC=(b>148?1024:QUADC) L2768; - CURE L2775; DEGE L2788; DEG widens CMP to n L2791; - NBSEL L2820; IMGOFF=NBSEL L2821; smem=(NBSEL+2*CMP)*4 L2822 - LAUNCH_REG2/DEG/REG macros L2823-2847 (KPT ladder 1/2/4; DEG forces KPT=1, - CUR=CURE both places); IMGW/smi/IMGE L2852-2854; - LAUNCH_REGIMG L2861-2863 -> gvr_topk_reg<...,KPT=1,CUR=true, - DEG=false,IMG=true,NBH=2*NB> via launch_regimg L2672-2686 - n4 rungs n4<=256 L2864; n4<=512 L2865; n4<=1024 wide/img/else L2866-2884 - clustered reg path L2897-2940: gate n4>4096 && n4<=8*BLKC*4 && k<=BLKC L2897; - av/amax L2898-2899; two-pass cs=8 co-residency veto - (pass==0 && c==8 && b>15 -> skip) L2917-2923; 64-bit product - (long long)c*BLKC*v < n4 L2919; smc=(3*NB+2*CMPC)*4 L2926; - grid dim3(cs,b) L2666 - wide 4k fallback n4<=4096 && wide -> LAUNCH_REG(1024,4,1,2*NB) L2945-2947 - streaming R L2959-2975 (b<=32: R=min(148/b, ((n>>2)+1023)/1024)); - r11 shallow split b<=74 && n4>=16384 && k<=1024 -> R=2 L2985; - cluster clamp R->pow2, useclus, only if 2<=R<=8 && k<=1024 L2994 - big/SCAP/CMP big=(b*R<=148) L2995; SCAP L3009-3010; CMP L3011 - aim L3039-3040; sqrt floor r=int(0.5+sqrt(6LL*n)) L3041-3042; - SFAC L3072-3073; amin L3079-3080; clamps L3081-3082 - sample geometry small_dense gate L3091 ((k>1024)&&!big&&n<=SCAP&&n>2*k); - PAIR form (sel>>3, half=n4s>>1, SMP*8) L3092-3109; - clus QUAD override (sel>>4, quarter=n4s>>2, SMP*16), - gated n>SCAP only, L3115-3128 - Q Q=(n4s+R-1)/R L3110 - clus launch smc=SNB*8+(SCAP+4)*8+CMP*8 L3130; U ladder per=Q>>10 - L3132-3142; CS=R in {2,4,8} L3143-3145; grid dim3(CS,b) L2704 - main launch smem=(SCAP+4)*((R>1||b<=296)?8:4)+(CMP+1)*8 L3149; - KPT ladder 1/2/4/8 L3150-3169; big: per=Q>>10 U ladder, - SPLIT=(R>1), grid dim3(R,b) L3173-3185 + L2750; - b<=296 -> (512,2,8,false) L3193; else (256,4,8,false) L3194 +"""Pure-Python mirror of the GVR CUDA host dispatch (gvr_topk_launch). + +route(b, n, npad, k) is a PURE function of its four ints -- no env knobs, no +GPU, stdlib only. It returns the kernel family, its compile-time template +tuple, the runtime scalar pack `rt`, grid/cluster/block geometry, smem size, +and whether the family needs the workspace. rt carries the FULL runtime scalar list each kernel receives, in signature -order, always starting with (n, npad, k) -- every launch site passes them -(L2666-2667 reg_clus, L2684-2685 regimg, L2704-2705 clus, L2726-2727 reg, -L2750-2751 main). [dispatch x-check 2026-08-13: rt previously omitted the -leading n/npad/k; fixed for full-ABI parity with the independent spec -transcription.] - -Dead ABI-parity args: gvr_main's 7th/8th params are declared `int SCAP_, int CMP_` -(kernel.cu L381) and are NEVER read by the kernel body -- it recomputes SCPB/CMPB -as constexprs of (BLK, SPLIT, KBIG) (L413-424) that mirror the host formulas -bit-identically. They are kept in rt under their source names 'SCAP_'/'CMP_' -purely for ABI parity. gvr_clus's SCAP/CMP (L1798) are LIVE runtime args. -`aim` and `SFAC` are host-side intermediates only (never cross the ABI), so they -do not appear in rt. +order, always starting with (n, npad, k). + +Dead ABI-parity args: gvr_main's `int SCAP_, int CMP_` params are NEVER read +by the kernel body -- it recomputes them as constexprs that mirror the host +formulas bit-identically. They are kept in rt purely for ABI parity. +gvr_clus's SCAP/CMP are LIVE runtime args. `aim` and `SFAC` are host-side +intermediates only (never cross the ABI), so they do not appear in rt. C-semantics notes encoded here: * every `/` on ints is C truncating division -> Python `//` (all operands @@ -137,41 +94,41 @@ def _device(): * `int r = (int)(0.5 + sqrt((double)(6LL*n)))` truncates toward zero after the +0.5 -> `int(0.5 + math.sqrt(float(6*n)))`; * `IMGW = (n + 3) & ~3` four-element float4 round-up; - * the reg-block CMP (possibly widened to n by DEGE) is scoped to the braces - at L2758-2949; the streaming path re-derives its own CMP. + * the reg-block CMP (possibly widened to n by DEGE) is scoped to the + register-resident block; the streaming path re-derives its own CMP. """ -# ---- constants lifted from kernel.cu --------------------------------------- -NB = 1024 # L16 register-path histogram bins -QUADC = 96 # L21 crossing-bin O(mc^2) rank gate (streaming/reg paths) -SNB = 256 # L170 streaming-path bin count -CMPC = 4096 # L2372 crossing-bin slots per CTA, clustered register path -BLKC = 1024 # L2374 CTA size of the clustered register path +# ---- dispatch constants (must match the device kernels) --------------------- +NB = 1024 # register-path histogram bins +QUADC = 96 # crossing-bin O(mc^2) rank gate (streaming/reg paths) +SNB = 256 # streaming-path bin count +CMPC = 4096 # crossing-bin slots per CTA, clustered register path +BLKC = 1024 # CTA size of the clustered register path def route(b: int, n: int, npad: int, k: int) -> dict[str, object]: - """Mirror of gvr_topk_launch (kernel.cu L2754-3197). Pure. See module doc.""" + """Mirror of the CUDA gvr_topk_launch dispatch. Pure. See module doc.""" if b < 1: raise RuntimeError(f"route requires b >= 1, got {b}") - wide = b <= 148 # L2757 - - # ================= register-resident block (L2758-2949) ================= - n4 = n >> 2 # L2759 - CMP = n if n < 2560 else 2560 # L2764 - QC = 1024 if b > 148 else QUADC # L2768 - CURE = not (n < 2 * k and b > 148) # L2775 - DEGE = (n <= 3 * k) or (n <= 4 * k + 64) # L2788 - if DEGE and CMP < n: # L2791 + wide = b <= 148 + + # ======================= register-resident block ======================== + n4 = n >> 2 + CMP = n if n < 2560 else 2560 + QC = 1024 if b > 148 else QUADC + CURE = not (n < 2 * k and b > 148) + DEGE = (n <= 3 * k) or (n <= 4 * k + 64) + if DEGE and CMP < n: CMP = n - NBSEL = (2 * NB) if (n4 > 512 and not (n4 <= 1024 and not wide)) else NB # L2820 - IMGOFF = NBSEL # L2821 - smem_reg = (NBSEL + 2 * CMP) * 4 # L2822 + NBSEL = (2 * NB) if (n4 > 512 and not (n4 <= 1024 and not wide)) else NB + IMGOFF = NBSEL + smem_reg = (NBSEL + 2 * CMP) * 4 def _reg(BLK, VPT, MINB, NBH): - # LAUNCH_REG (L2844-2847): DEG wins, else CUR flag; KPT ladder L2823-2834. + # DEG wins over the CUR flag; DEG forces KPT=1, else KPT ladder 1/2/4. if DEGE: - tpl = (BLK, VPT, MINB, 1, CURE, True, False, NBH) # LAUNCH_DEG L2836-2843 + tpl = (BLK, VPT, MINB, 1, CURE, True, False, NBH) else: kpt = 1 if k <= BLK else (2 if k <= 2 * BLK else 4) tpl = (BLK, VPT, MINB, kpt, CURE, False, False, NBH) @@ -181,7 +138,7 @@ def _reg(BLK, VPT, MINB, NBH): "rt": { "n": n, "npad": npad, - "k": k, # L2726-2727 full ABI + "k": k, # full ABI "CMP": CMP, "IMGOFF": IMGOFF, "QC": QC, @@ -193,26 +150,25 @@ def _reg(BLK, VPT, MINB, NBH): "ws": False, } - IMGW = (n + 3) & ~3 # L2852 - smi = (NBSEL + (2 * CMP if 2 * CMP > IMGW else IMGW)) * 4 # L2853 - IMGE = wide and (not DEGE) and k <= 1024 # L2854 + IMGW = (n + 3) & ~3 + smi = (NBSEL + (2 * CMP if 2 * CMP > IMGW else IMGW)) * 4 + IMGE = wide and (not DEGE) and k <= 1024 - if n4 <= 256: # L2864 + if n4 <= 256: return _reg(256, 1, 8, NB) - if n4 <= 512: # L2865 + if n4 <= 512: return _reg(512, 1, 4, NB) - if n4 <= 1024: # L2866-2884 + if n4 <= 1024: if wide: - if IMGE: # LAUNCH_REGIMG(1024,1,2) L2872 - # launch_regimg<1024,1,2,NBV=2*NB,KPTV=1> -> gvr_topk_reg - # <1024,1,2,1,true,false,true,2048> (L2672-2686) + if IMGE: + # regimg launch: gvr_topk_reg<1024,1,2,1,true,false,true,2048> return { "kernel": "regimg", "tpl": (1024, 1, 2, 1, True, False, True, 2 * NB), "rt": { "n": n, "npad": npad, - "k": k, # L2684-2685 full ABI + "k": k, # full ABI "CMP": CMP, "IMGOFF": IMGOFF, "QC": QC, @@ -223,38 +179,36 @@ def _reg(BLK, VPT, MINB, NBH): "smem": smi, "ws": False, } - return _reg(1024, 1, 2, 2 * NB) # L2872 else-arm - return _reg(512, 2, 4, NB) # L2883 + return _reg(1024, 1, 2, 2 * NB) + return _reg(512, 2, 4, NB) - # ---- clustered register-resident path (L2897-2940) ---- - if n4 > 4096 and n4 <= 8 * BLKC * 4 and k <= BLKC: # L2897 - av = 148 // (b if b > 0 else 1) # L2898 truncating - amax = 1 # L2899 + # ---- clustered register-resident path ---- + if n4 > 4096 and n4 <= 8 * BLKC * 4 and k <= BLKC: + av = 148 // (b if b > 0 else 1) # truncating + amax = 1 while (amax << 1) <= av and amax < 8: amax <<= 1 vsel = 0 cs = 0 - if amax >= 2: # L2901 - # knife5 (layer 9): UNCONDITIONAL cs=8 co-residency veto -- - # the L2w pass-1 rescue is deleted; 512k b>15 falls through to - # streaming, made retry-safe by TSH-floor staging (S1) and the - # gvr_clus veto (S2). + if amax >= 2: + # cs=8 co-residency veto: an 8-CTA cluster with b > 15 exceeds + # GPC packing; such shapes fall through to the streaming path. for v in (1, 2, 4): - c = 1 # 64-bit product + c = 1 # 64-bit product in C while c * BLKC * v < n4: c <<= 1 - if c == 8 and b > 15: # THE VETO + if c == 8 and b > 15: # the veto continue if c <= amax: vsel = v cs = c break - if vsel and cs >= 2: # L2925 - smc = (3 * NB + 2 * CMPC) * 4 # L2926 + if vsel and cs >= 2: + smc = (3 * NB + 2 * CMPC) * 4 return { "kernel": "reg_clus", "tpl": (BLKC, vsel, cs), - "rt": {"n": n, "npad": npad, "k": k}, # dims only, L2666-2667 + "rt": {"n": n, "npad": npad, "k": k}, # dims only "grid": (cs, b), "cluster": cs, "block": BLKC, @@ -262,129 +216,127 @@ def _reg(BLK, VPT, MINB, NBH): "ws": False, } - if n4 <= 4096 and wide: # L2945-2947 + if n4 <= 4096 and wide: return _reg(1024, 4, 1, 2 * NB) - # ================= streaming / collect path (L2950-3196) ================= - R = 1 # L2959 - if b <= 32: # L2960-2975 + # ====================== streaming / collect path ======================== + R = 1 + if b <= 32: r1 = 148 // b if r1 < 1: r1 = 1 - r2 = ((n >> 2) + 1023) // 1024 # L2972 + r2 = ((n >> 2) + 1023) // 1024 if r2 < 1: r2 = 1 R = r1 if r1 < r2 else r2 if R < 1: R = 1 - elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: # L2985 r11 split + elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: # shallow R=2 split R = 2 - useclus = False # L2993-2994 + useclus = False if 2 <= R <= 8 and k <= 1024: p2 = 1 while (p2 << 1) <= R: p2 <<= 1 - # knife5 (layer 8): gvr_clus cs=8 hits the same GPC packing wall as - # the clustered register path; same veto, same b>15 threshold. + # gvr_clus cs=8 hits the same GPC packing wall as the clustered + # register path; same veto, same b > 15 threshold. if p2 == 8 and b > 15: p2 = 4 R = p2 useclus = True - big = b * R <= 148 # L2995 - SCAP = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) # L3009-3010 - CMP = (4096 if k > 1024 else 2048) if big else 1024 # L3011 + big = b * R <= 148 + SCAP = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) + CMP = (4096 if k > 1024 else 2048) if big else 1024 aim = ( ((4 * k if k >= 1024 else 2 * k) if R == 1 else 2 * k) if big else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) - ) # L3039-3040 - q = 6 * n # L3041: 6LL * n - r = int(0.5 + math.sqrt(float(q))) # L3041 C cast trunc - if r > aim: # L3042 + ) + q = 6 * n # 6LL * n + r = int(0.5 + math.sqrt(float(q))) # C cast trunc + if r > aim: aim = r - SFAC = ( - (32 if R == 2 else (48 if k > 1024 else 16)) if R > 1 else (64 if k >= 1024 else 32) - ) # L3072-3073 - amin = 3 * k if R == 2 else (7 * k) // 2 # L3079 - if R > 1 and aim < amin: # L3080 + SFAC = (32 if R == 2 else (48 if k > 1024 else 16)) if R > 1 else (64 if k >= 1024 else 32) + amin = 3 * k if R == 2 else (7 * k) // 2 + if R > 1 and aim < amin: aim = amin - if aim > (SCAP >> 1): # L3081 + if aim > (SCAP >> 1): aim = SCAP >> 1 - if aim < k: # L3082 + if aim < k: aim = k - n4s = n >> 2 # L3084 - SMP, SS2, TGT, TGT2 = 0, 1, 0, 0 # L3085 - small_dense = (k > 1024) and (not big) and n <= SCAP and n > 2 * k # L3091 - if (n > SCAP or small_dense) and n4s >= 4: # L3092: PAIR sample - sel = SFAC * n // aim # L3095 64-bit - if sel < 256: # L3096 + n4s = n >> 2 + SMP, SS2, TGT, TGT2 = 0, 1, 0, 0 + small_dense = (k > 1024) and (not big) and n <= SCAP and n > 2 * k + if (n > SCAP or small_dense) and n4s >= 4: # PAIR sample + sel = SFAC * n // aim # 64-bit + if sel < 256: sel = 256 - if sel > n // 2: # L3097 + if sel > n // 2: sel = n // 2 - pairs = sel >> 3 # L3098 + pairs = sel >> 3 if pairs < 1: pairs = 1 - half = n4s >> 1 # L3099 + half = n4s >> 1 if half < 1: half = 1 - if pairs > half: # L3100 + if pairs > half: pairs = half - SS2 = half // pairs # L3101 + SS2 = half // pairs if SS2 < 1: SS2 = 1 - SMP = half // SS2 # L3102 + SMP = half // SS2 if SMP < 1: SMP = 1 - TGT = (aim * (SMP * 8)) // n # L3103 64-bit - if TGT < 1: # L3104 + TGT = (aim * (SMP * 8)) // n # 64-bit + if TGT < 1: TGT = 1 - TGT2 = (k * (SMP * 8)) // n # L3107 64-bit - if TGT2 < 1: # L3108 + TGT2 = (k * (SMP * 8)) // n # 64-bit + if TGT2 < 1: TGT2 = 1 - Q = (n4s + R - 1) // R # L3110 + Q = (n4s + R - 1) // R - if useclus: # L3111-3147 - if n > SCAP and n4s >= 4: # L3115: QUAD override - sel = SFAC * n // aim # L3116 + if useclus: + if n > SCAP and n4s >= 4: # QUAD override + sel = SFAC * n // aim if sel < 256: sel = 256 if sel > n // 2: sel = n // 2 - quads = sel >> 4 # L3119 + quads = sel >> 4 if quads < 1: quads = 1 - quarter = n4s >> 2 # L3120 + quarter = n4s >> 2 if quarter < 1: quarter = 1 - if quads > quarter: # L3121 + if quads > quarter: quads = quarter - SS2 = quarter // quads # L3122 + SS2 = quarter // quads if SS2 < 1: SS2 = 1 - SMP = quarter // SS2 # L3123 + SMP = quarter // SS2 if SMP < 1: SMP = 1 - TGT = (aim * (SMP * 16)) // n # L3124 + TGT = (aim * (SMP * 16)) // n if TGT < 1: TGT = 1 - TGT2 = (k * (SMP * 16)) // n # L3126 + TGT2 = (k * (SMP * 16)) // n if TGT2 < 1: TGT2 = 1 - smc = SNB * 8 + (SCAP + 4) * 8 + CMP * 8 # L3130 - per = Q >> 10 # L3131 - U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) # L3134-3141 - CS = 2 if R == 2 else (4 if R == 4 else 8) # L3143-3145 + smc = SNB * 8 + (SCAP + 4) * 8 + CMP * 8 + per = Q >> 10 + U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) + CS = 2 if R == 2 else (4 if R == 4 else 8) return { "kernel": "clus", "tpl": (1024, U, 1, SNB, CS), "rt": { "n": n, "npad": npad, - "k": k, # L2704-2705 ABI (live) + "k": k, # ABI (live) "SCAP": SCAP, "CMP": CMP, "SMP": SMP, @@ -400,26 +352,25 @@ def _reg(BLK, VPT, MINB, NBH): "ws": False, } - smem_main = (SCAP + 4) * (8 if (R > 1 or b <= 296) else 4) + (CMP + 1) * 8 # L3149 + smem_main = (SCAP + 4) * (8 if (R > 1 or b <= 296) else 4) + (CMP + 1) * 8 def _main(BLK, MINB, U, SPLIT): - # LAUNCH_MAIN KPT ladder 1/2/4/8 (L3150-3169); grid dim3(gx=R, gy=b) L2750. + # KPT ladder 1/2/4/8; grid = (R, b). kpt = 1 if k <= BLK else (2 if k <= 2 * BLK else (4 if k <= 4 * BLK else 8)) - # knife5 (layer 7) TSH-floor staging gate. CUDA form: grid-uniform - # RUNTIME gate gridDim.y > 15 && k <= 1024 && (n >> 2) <= 32768 with - # a dual scan-instantiation branch. Here: compile-time key -- the - # ungated variant IS the pre-knife5 kernel; per-launch semantics are - # identical because the gate is uniform over the grid. + # TSH-floor staging gate. The CUDA form is a grid-uniform RUNTIME + # gate (gridDim.y > 15 && k <= 1024 && (n >> 2) <= 32768); here it + # is a compile-time key -- per-launch semantics are identical + # because the gate is uniform over the grid. tshg = bool(SPLIT) and b > 15 and k <= 1024 and (n >> 2) <= 32768 return { "kernel": "main", "tpl": (BLK, U, MINB, SNB, kpt, SPLIT, tshg), - # SCAP_/CMP_ are DEAD ABI-parity args: gvr_main (L381) never reads - # them, it uses constexpr SCPB/CMPB (L413-424). Kept for ABI parity. + # SCAP_/CMP_ are dead ABI-parity args: gvr_main never reads them + # (it recomputes them as constexprs). "rt": { "n": n, "npad": npad, - "k": k, # L2750-2751 full ABI + "k": k, # full ABI "SCAP_": SCAP, "CMP_": CMP, "R": R, @@ -436,13 +387,13 @@ def _main(BLK, MINB, U, SPLIT): "ws": True, } - if big: # L3173-3185 - per = Q >> 10 # L3174 + if big: + per = Q >> 10 U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) return _main(1024, 1, U, R > 1) # SPLIT iff R>1 - if b <= 296: # L3193 + if b <= 296: return _main(512, 2, 8, False) - return _main(256, 4, 8, False) # L3194 + return _main(256, 4, 8, False) if __name__ == "__main__": @@ -454,14 +405,14 @@ def _main(BLK, MINB, U, SPLIT): (64, 4096, 4096, 512), # regimg wide !DEGE k<=1024 (64, 4096, 4096, 1024), # reg wide but DEGE (n<=4k+64) (8, 65536, 65536, 1024), # reg_clus (vsel=2, cs=8; b<=15 no veto) - (16, 131072, 131072, 512), # knife5: veto fall-through -> SPLIT slab, tshg=True + (16, 131072, 131072, 512), # main cs=8 veto fall-through -> SPLIT slab, tshg=True (64, 16384, 16384, 1024), # reg wide 4k fallback (1024,4,1) - (64, 262144, 262144, 1024), # clus r11 R=2 shallow cluster split + (64, 262144, 262144, 1024), # clus R=2 shallow cluster split (1, 1048576, 1048576, 1024), # main deep slab SPLIT R=148 (20, 262144, 262144, 2048), # main k>1024 split (no useclus) (512, 131072, 131072, 1024), # main b>296 BLK=256 (256, 6144, 6144, 2048), # main small_dense sample gate - (256, 262144, 262144, 2048), # main v32 KBIG-domain, BLK=512 KPT=4 + (256, 262144, 262144, 2048), # main KBIG-domain (k>1024), BLK=512 KPT=4 ] for shp in smoke: print(shp, "->", route(*shp)) @@ -479,10 +430,10 @@ def _main(BLK, MINB, U, SPLIT): # formulas): n, CMP (reg families), the sampling ladder # SMP/TGT/SS2/TGT2/Q (streaming families), and the reg-family smem # footprint. -# INVARIANT (fuzz-verified): merging route_dynamic back into route_static -# reproduces route() EXACTLY for every n. The capture-time policy of which n -# to freeze the static half at (e.g. max_seq_len) is a later, perf-only -# choice — this split only proves the factorization is lossless. +# INVARIANT: merging route_dynamic back into route_static reproduces +# route() EXACTLY for every n. The policy of which n to freeze the static +# half at (e.g. max_seq_len) is a perf-only choice — the factorization +# itself is lossless. _DYN_RT = { "reg": ("n", "CMP"), @@ -509,9 +460,8 @@ def route_static(b: int, n: int, npad: int, k: int) -> dict[str, object]: def route_dynamic(static: dict[str, object], n: int) -> tuple[dict[str, object], int]: """Recompute the redacted n-continuous scalars from (static, n). - Returns (rt_updates, smem). Transcribed independently from route() — - the factorization fuzz is the equivalence proof, and the device-side - per-row engine mirrors exactly these formulas.""" + Returns (rt_updates, smem). Must stay equivalent to route(); the + device-side per-row engine mirrors exactly these formulas.""" fam = static["kernel"] k = static["rt"]["k"] if fam in ("reg", "regimg"): @@ -603,10 +553,9 @@ def route_streaming( capture policy: per-row kernels must be picked from the families that are correct for ANY row length, so the register-resident specialists are skipped even when the envelope n would normally land on them. Where - route() itself lands on main/clus this is IDENTICAL to route() (fuzz: - 110,003/110,003 agreement). force_main additionally skips the clus - rounding (v1 varlen engine ships the gvr_main port first; the raw - min(r1, r2) R then matches the CUDA else-branch exactly).""" + route() itself lands on main/clus this is IDENTICAL to route(). + force_main additionally skips the clus rounding, so the raw + min(r1, r2) R matches the CUDA else-branch exactly.""" if b < 1: raise RuntimeError(f"route_streaming requires b >= 1, got {b}") R = 1 @@ -731,8 +680,8 @@ def _main(blk_, minb_, u_, split_): def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): - """Capture-time varlen plan + compiled launcher (v1 = the gvr_main port, - universally correct across the envelope; specialist tiers below). Every + """Capture-time varlen plan + compiled launcher. The gvr_main port is + the universally correct fallback; specialist family tiers below. Every choice here is a function of capture-stable quantities only — mirroring the in-tree runner's pick_tuning(graph_capture=...) discipline.""" key = (num_rows, npad, k, n_env, next_n, cr) @@ -743,12 +692,10 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): cr_shift = 0 if cr == 1 else 2 dev = _device() # ---- route() parity, family tier 1: clustered register-resident -------- - # Admit reg_clus exactly where the free route picks it (886-real-capture - # grid: this family recovers the deep-layer distribution tail and the - # large-N small-rows band; its whole admission window n4 <= 32768 fits - # capture-frozen envelopes). The choice is a pure function of this cache - # key, so CUDA-graph replay safety is unchanged; per-row n / short-row - # handling lives in-kernel (GvrMainKernel varlen discipline). + # Admit reg_clus exactly where the free route picks it; its whole + # admission window (n4 <= 32768) fits capture-frozen envelopes. The + # choice is a pure function of this cache key, so CUDA-graph replay + # safety is unchanged; per-row n / short-row handling lives in-kernel. plan_free = route(num_rows, n_eff, npad, k) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( @@ -763,7 +710,7 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # smem are envelope-derived launch constants -- in-kernel they are pure # capacity clamps (CMP), a fast-path threshold (QC) and the launch smem # size, all safe upper bounds for every per-row n <= envelope; per-row n - # / short-row handling lives in-kernel (GvrRegClusKernel discipline). + # / short-row handling lives in-kernel. if plan_free["kernel"] in ("reg", "regimg"): fn = dev.get_compiled__reg( tuple(plan_free["tpl"]), varlen=True, next_n=next_n, cr_shift=cr_shift @@ -781,8 +728,8 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): # large-N mid-rows band). SCAP/CMP are launch-stable (pure functions of # rows/CS/k — never of n) so the envelope values are the per-row values; # the sampling-ladder scalars (SMP/TGT/Q/SS2/TGT2) are dead launch slots, - # re-derived per row in-kernel by the route_dynamic clus mirror - # (GvrMainKernel discipline). Per-row n / short-row handling in-kernel. + # re-derived per row in-kernel by the route_dynamic clus mirror. + # Per-row n / short-row handling in-kernel. if plan_free["kernel"] == "clus": rt_f = plan_free["rt"] fn = dev.get_compiled__clus( @@ -821,12 +768,10 @@ def _varlen_launcher(num_rows, npad, k, n_env, next_n, cr): ) amin = 3 * k if r_const == 2 else (7 * k) // 2 sd_en = 1 if (k > 1024 and not big) else 0 - # TSH-floor staging: gate on SPLIT and K only. The old num_rows > 15 - # condition stranded small batches (rows <= 15) in SPLIT-main without - # the staged floor — a distribution-dependent 6x tail on real deep-layer - # captures (v4_pro_512k L46/L52, n4 = 32768, rows 1-8: 142-151 us vs - # 25 us with staging; healthy layers and n4 > 32768 rows unaffected — - # the kernel gates TSH per row at runtime anyway). + # TSH-floor staging: gate on SPLIT and K only. Gating additionally on + # num_rows > 15 would strand small batches in SPLIT-main without the + # staged floor (a distribution-dependent tail regression); the kernel + # gates TSH per row at runtime anyway. tsh_en = 1 if (tpl[5] and k <= 1024) else 0 pre = (0, npad, k, rt["SCAP_"], rt["CMP_"], r_const, 0, 0, 0, 0, 0) tail = (aim_base, sfac, amin, sd_en, tsh_en) @@ -859,57 +804,55 @@ def route_bands( # =========================================================================== -# ==== workspace (ct_workspace.py) ========================================== +# ==== workspace ============================================================ # =========================================================================== -"""op46 workspace mirror of src_cuda/main.cpp B2 (L15-37) + run_ws checks -(L107-114) and kernel.h workspace_bytes contract. +"""Per-device workspace slab for the multi-CTA SPLIT path. -B2 semantics mirrored exactly: +Semantics: * ONE zero-initialised slab workspace per device, lazily allocated through - the torch caching allocator (main.cpp:32-33 `at::zeros(..., kByte)`); - * keep-alive store (`ws_keep[GVR_MAX_DEV]`) -> module dict `_ws_keep` - (tensor refcount = keep-alive, same as the C static array); + the torch caching allocator; + * keep-alive store: module dict `_ws_keep` (tensor refcount = keep-alive); * double-checked locking: lock-free hot-path load (a GIL-atomic dict get - plays the `std::memory_order_acquire` load, main.cpp:26-27), slow path - re-checks under a mutex (main.cpp:28-31); - * device index bounds `0 <= d < GVR_MAX_DEV` (main.cpp:24-25) -- checked - BEFORE the CUDA-ness of the tensor, exactly like the C binding (run() - resolves the default workspace before run_impl's B1 checks, so a CPU - logits tensor dies here with "device index out of range: -1"). + plays an acquire load), slow path re-checks under a mutex; + * device index bounds `0 <= d < GVR_MAX_DEV` -- checked BEFORE the + CUDA-ness of the tensor (run() resolves the default workspace before the + input checks, so a CPU logits tensor dies here with "device index out of + range: -1"). Concurrent STREAMS on one device that may both take the multi-CTA SPLIT path -must pass their own workspace via run_ws() (main.cpp:16-17). +must pass their own workspace via run_ws(). -Size: gvr_topk_workspace_bytes() = GVR_WS_BUF_OFF + MAXC*GCAP*sizeof(int2) - = 2048 + 160*16384*8 = 20,973,568 B (kernel.cu L44-46). +Size: workspace_bytes() = GVR_WS_BUF_OFF + MAXC*GCAP*sizeof(int2) + = 2048 + 160*16384*8 = 20,973,568 B. -Kernel-facing view: ct_main's compiled signature takes the workspace as a -1-D contiguous int32 tensor (fake tensor dtype Int32, assumed_align=16 -- -torch caching-allocator bases are 256B-aligned so the default slab always -satisfies it). `kernel_view()` reproduces the C binding's raw -`workspace.data_ptr()` semantics for arbitrary user tensors by aliasing the -underlying storage at the tensor's byte offset. +Kernel-facing view: the compiled main-family signature takes the workspace +as a 1-D contiguous int32 tensor (fake tensor dtype Int32, assumed_align=16 +-- torch caching-allocator bases are 256B-aligned so the default slab always +satisfies it). `kernel_view()` reproduces raw `workspace.data_ptr()` +semantics for arbitrary user tensors by aliasing the underlying storage at +the tensor's byte offset. """ -GVR_MAX_DEV = 64 # kernel.cu L19 / main.cpp:19 -_MAXC = 160 # kernel.cu L17 -_GCAP = 16384 # kernel.cu L18 -_GVR_WS_BUF_OFF = 2048 # kernel.cu L43 -WS_BYTES = _GVR_WS_BUF_OFF + _MAXC * _GCAP * 8 # 20,973,568 (kernel.cu L44-46) +# workspace geometry constants -- must match the device kernels +GVR_MAX_DEV = 64 +_MAXC = 160 +_GCAP = 16384 +_GVR_WS_BUF_OFF = 2048 +WS_BYTES = _GVR_WS_BUF_OFF + _MAXC * _GCAP * 8 # 20,973,568 assert WS_BYTES == 20_973_568 -_mu = threading.Lock() # main.cpp:28 slow-path mutex +_mu = threading.Lock() # slow-path mutex _ws_keep = {} # device index -> keep-alive int32 view def workspace_bytes() -> int: - """kernel.h:12 gvr_topk_workspace_bytes().""" + """Workspace bytes required by the multi-CTA SPLIT path.""" return WS_BYTES def default_workspace(ref: torch.Tensor) -> torch.Tensor: - """main.cpp:23-37 default_workspace(ref) -> per-device cached slab. + """Per-device cached workspace slab. Returns the kernel-facing 1-D int32 view (zero-initialised on first use; the kernel restores the zeros it consumes, so one zeroing suffices for @@ -924,8 +867,8 @@ def default_workspace(ref: torch.Tensor) -> torch.Tensor: ws = _ws_keep.get(d) if ws is not None: return ws - # lazy zeros via the torch caching allocator (at::zeros kByte, - # main.cpp:32-33), viewed int32 for the DSL launch signature. + # lazy zeros via the torch caching allocator, viewed int32 for the + # DSL launch signature. buf = torch.zeros(WS_BYTES, dtype=torch.uint8, device=ref.device) ws = buf.view(torch.int32) _ws_keep[d] = ws # keep-alive (ws_keep[d] = tensor) @@ -933,7 +876,7 @@ def default_workspace(ref: torch.Tensor) -> torch.Tensor: def validate_run_ws(workspace: torch.Tensor, logits: torch.Tensor) -> None: - """main.cpp:107-114 run_ws() workspace hardening, same predicate order: + """run_ws() workspace hardening, in a fixed predicate order: CUDA + same device as logits; numel*element_size >= workspace_bytes(); base 8-byte aligned.""" if not (workspace.is_cuda and workspace.get_device() == logits.get_device()): @@ -945,14 +888,14 @@ def validate_run_ws(workspace: torch.Tensor, logits: torch.Tensor) -> None: def kernel_view(workspace: torch.Tensor) -> torch.Tensor: - """Raw-pointer semantics of the C binding (main.cpp:115 passes - workspace.data_ptr() and nothing else): alias the first WS_BYTES bytes at - the tensor's data_ptr() as int32[WS_BYTES/4], ignoring dtype/shape. + """Raw-pointer view of a user workspace tensor: alias the first WS_BYTES + bytes at the tensor's data_ptr() as int32[WS_BYTES/4], ignoring + dtype/shape. NOTE: the DSL-side fake tensor declares assumed_align=16; a workspace at - 8-but-not-16-byte alignment passes the C-contract check above but is + 8-but-not-16-byte alignment passes the validate_run_ws check but is rejected by the DSL at conversion -- surfaced as a launch failure with - shape context by ct_op (documented in notes/ct_op_NOTES.md).""" + shape context.""" if ( workspace.dtype is torch.int32 and workspace.dim() == 1 @@ -978,52 +921,41 @@ def _reset_for_tests() -> None: # =========================================================================== -# ==== operator entry (ct_op.py) ============================================ +# ==== operator entry ======================================================= # =========================================================================== -"""op46 operator entry: CuTeDSL mirror of src_cuda/main.cpp run()/run_ws()/ -workspace_bytes() (spec section 1). - -B1 hardening checks run in the SAME ORDER with the SAME PREDICATES as -main.cpp:43-88 (run_impl): - 1. all three tensors CUDA (main.cpp:43-44) - 2. dtypes: logits f32, pre_idx i32, indices i32 (45-47) - 3. all 2-D (48-49) - 4. all contiguous (50-51) - 5. n_valid unwrap (57-67): python-int fast path (strict integral cast, like - pybind cast); Tensor path checks - torch.cuda.is_current_stream_capturing() FIRST and fails loudly (B1d), - else .item() (the D2H sync) - 6. b/npad from logits, k = pre_idx.size(1) (68-70) - 7. b == 0 -> early no-op (71, B1f) - 8. npad % 4 == 0 (74-75, B1e float4 row loads) - 9. logits base 16-byte aligned (76-78) - 10. pre_idx/indices batch dims match (79-81) - 11. indices width >= k (84-85) - 12. n_valid >= 0 (86) - 13. n = min(nv, npad) clamped in unbounded ints BEFORE any narrowing (88) - -Dispatch: ct_dispatch.route(b, n, npad, k) -> compile-cache keyed on -(kernel family, constexpr tuple) inside each family module -> BIND-ONCE -launch cache keyed on the shape key (b, n, npad, k): caches the compiled -callable + the prebuilt runtime-scalar arg pack as plain Python ints (probe -P12: plain ints, never pre-wrapped cutlass.Int32; pre-binding removes only -route()/marshal-prep work -- the tvm-ffi per-argument cost is paid every -call). Hot enqueue target ~3-6 us (P12 arg-width tax); measured numbers in -notes/ct_op_NOTES.md. - -Error contract (spec 1.4): launch failures surface as exceptions WITH -(b, n, npad, k) context, mirroring main.cpp:94-95. - -All four family modules are imported LAZILY (first shape that routes to -them), so a missing/broken sibling only fails when actually routed to, with -(b, n, npad, k) context. Wired compiled ABIs (verified against each -module's __call__ signature): - ct_reg (logits, pre_idx, out, n, CMP, QC, smem_bytes) - ct_main (logits, pre_idx, out, ws, n, npad, k, SCAP_, CMP_, R, SMP, - TGT, Q, SS2, TGT2) [only family taking workspace] - ct_clus (logits, pre_idx, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, - SS2, TGT2) [get_compiled keyed +scap/cmp_] - ct_regclus (logits, pre_idx, out, n) +"""Operator entry: input hardening, dispatch, and bind-once launch cache. + +Hardening checks run in a fixed order with fixed predicates: + 1. all three tensors CUDA + 2. dtypes: logits f32, pre_idx i32, indices i32 + 3. all 2-D + 4. all contiguous + 5. n_valid unwrap: python-int fast path (strict integral cast); Tensor + path checks torch.cuda.is_current_stream_capturing() FIRST and fails + loudly, else .item() (the D2H sync) + 6. b/npad from logits, k = pre_idx.size(1) + 7. b == 0 -> early no-op + 8. npad % 4 == 0 (float4 row loads) + 9. logits base 16-byte aligned + 10. pre_idx/indices batch dims match + 11. indices width >= k + 12. n_valid >= 0 + 13. n = min(nv, npad) clamped in unbounded ints BEFORE any narrowing + +Dispatch: route(b, n, npad, k) -> compile cache keyed on (kernel family, +constexpr tuple) in the device module -> bind-once launch cache keyed on the +shape key (b, n, npad, k): caches the compiled callable + the prebuilt +runtime-scalar arg pack as plain Python ints (never pre-wrapped +cutlass.Int32 -- the FFI per-argument cost is paid every call regardless; +pre-binding removes only route()/marshal-prep work). + +Error contract: launch failures surface as exceptions WITH +(b, n, npad, k) context. + +The device module is imported LAZILY (first shape that routes to it), so a +missing/broken module only fails when actually reached, with (b, n, npad, k) +context. The per-family compiled ABIs are documented at each launcher +builder in _build_launcher; only the main family takes the workspace. """ @@ -1042,9 +974,8 @@ def _dummy_kv(dev_index, device): return t -# hot-path local bindings (each torch. lookup costs ~0.1 us; the B1 -# battery runs on EVERY call — mirror of main.cpp's "sub-100ns predicted -# branches" intent within Python's reach; measured in notes/ct_op_NOTES.md) +# hot-path local bindings: each torch. lookup costs ~0.1 us and the +# validation battery runs on EVERY call _F32 = torch.float32 _I32 = torch.int32 _TENSOR = torch.Tensor @@ -1068,7 +999,7 @@ def _build_launcher(b, n, npad, k): # compiled ABI: (logits, pre_idx, kv_lens, out, n, CMP, QC, # smem_total) -- kv_lens is the dead varlen slot in batch-uniform - # mode (dummy, gvr_main/reg_clus precedent) + # mode (cached dummy tensor) def fn(lg, pi, o, *a, _raw=raw): _raw(lg, pi, _dummy_kv(lg.get_device(), lg.device), o, *a) @@ -1104,8 +1035,8 @@ def fn(lg, pi, o, w, *a, _raw=raw): dev = _device() # compile key carries the smem-extent scalars (scap/cmp_); compiled # ABI: (logits, pre_idx, kv_lens, out, n, npad, k, SCAP, CMP, SMP, - # TGT, Q, SS2, TGT2) -- NO workspace (spec §4c); kv_lens is the - # dead varlen slot in batch-uniform mode (dummy, reg_clus precedent) + # TGT, Q, SS2, TGT2) -- NO workspace; kv_lens is the dead + # varlen slot in batch-uniform mode (cached dummy tensor) fn = dev.get_compiled__clus(tpl, scap=rt["SCAP"], cmp_=rt["CMP"]) args = ( rt["n"], @@ -1127,7 +1058,7 @@ def _call(lg, pi, idx, _fn=fn, _args=args): if fam == "reg_clus": dev = _device() # compiled ABI: (logits, pre_idx, kv_lens, out, n) -- kv_lens is the - # dead varlen slot in batch-uniform mode (dummy, gvr_main precedent); + # dead varlen slot in batch-uniform mode (cached dummy tensor); # smem/k derived in-module fn = dev.get_compiled__regclus(tpl) n_arg = rt["n"] @@ -1141,7 +1072,7 @@ def _call(lg, pi, idx, _fn=fn, _n=n_arg): # --------------------------------------------------------------------------- -# run_impl mirror (main.cpp:39-96) +# shared implementation of the batch-uniform entries # --------------------------------------------------------------------------- def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): @@ -1158,8 +1089,8 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): if not (logits.is_contiguous() and pre_idx.is_contiguous() and indices.is_contiguous()): raise RuntimeError("tensors must be contiguous") - # n_valid unwrap (main.cpp:57-67): tensor path = D2H sync, illegal under - # CUDA graph capture -- fail loudly instead of crashing the capture (B1d). + # n_valid unwrap: tensor path = D2H sync, illegal under CUDA graph + # capture -- fail loudly instead of crashing the capture. if isinstance(n_valid, _TENSOR): if _is_capturing(): raise RuntimeError( @@ -1168,12 +1099,12 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): ) nv = int(n_valid.item()) else: - # strict integral cast (pybind cast rejects floats/strings) + # strict integral cast (rejects floats/strings) nv = _index(n_valid) b, npad = lsh k = psh[1] - if b == 0: # empty batch: no-op (main.cpp:71, B1f) + if b == 0: # empty batch: no-op return if npad & 3: raise RuntimeError(f"npad (logits stride) must be a multiple of 4, got {npad}") @@ -1187,15 +1118,15 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): raise RuntimeError(f"indices width {ish[1]} < k={k} (k is pre_idx.size(1))") if nv < 0: raise RuntimeError(f"n_valid must be non-negative, got {nv}") - # clamp BEFORE any narrowing (main.cpp:87-88; python ints are unbounded, - # so min() is the exact 64-bit clamp) + # clamp BEFORE any narrowing (python ints are unbounded, so min() is the + # exact 64-bit clamp) n = nv if nv < npad else npad - # CUDA out-indexing mirror: every kernel derives O = out + row*k - # (kernel.cu L475/L1309 etc.) -- flat PACKED rows, ignoring the actual - # indices width. The DSL kernels index out[row, :] with the tensor's own - # row stride, so a wider `indices` must be re-viewed packed (pure view, - # no copy; contiguity already checked). + # CUDA out-indexing mirror: every kernel derives O = out + row*k -- + # flat PACKED rows, ignoring the actual indices width. The DSL kernels + # index out[row, :] with the tensor's own row stride, so a wider + # `indices` must be re-viewed packed (pure view, no copy; contiguity + # already checked). if ish[1] != k: indices = indices.reshape(-1)[: b * k].view(b, k) @@ -1219,7 +1150,7 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): if vsh[1] != k: values = values.reshape(-1)[: b * k].view(b, k) - # ---- n <= k short path (heuristicTopKDecode.cu:72-84) ------------------- + # ---- n <= k short path (heuristicTopKDecode.cu parity) ------------------ # Every valid position is in the top-K: emit identity indices and pad the # tail with -1 (the production pad convention; downstream treats -1 as # invalid). Order is contract-irrelevant — exactness is tie-interchangeable @@ -1258,7 +1189,7 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): # --------------------------------------------------------------------------- -# exports (main.cpp:98-124) +# exports # --------------------------------------------------------------------------- def run( logits: torch.Tensor, @@ -1271,15 +1202,14 @@ def run( device kv_lens; this entry assumes one batch-uniform host ``n_valid``, which real serving batches do not satisfy). - Fast 4-arg form: signature-identical to the original candidate. - ``values`` (optional DPS output, default None = OFF) mirrors the - production values writeback; see _run_impl. - Default per-device slab workspace resolved FIRST (main.cpp:99-102 -- - a CPU logits tensor therefore dies with 'device index out of range'). - Hot path inlines the C binding's check + atomic-load + cache-hit - (main.cpp:24-27); the slow path allocates under ct_workspace's lock.""" + Fast 4-arg form. ``values`` (optional DPS output, default None = OFF) + mirrors the production values writeback; see _run_impl. + The default per-device slab workspace is resolved FIRST (a CPU logits + tensor therefore dies with 'device index out of range'). + Hot path inlines the device check + atomic load + cache hit; the slow + path allocates under the workspace lock.""" d = logits.get_device() - if not 0 <= d < _GVR_MAX_DEV: # main.cpp:25, EVERY call + if not 0 <= d < _GVR_MAX_DEV: # checked on EVERY call raise RuntimeError(f"device index out of range: {d}") ws = _ws_hot.get(d) if ws is None: @@ -1297,7 +1227,7 @@ def run_ws( ) -> None: """TESTING/BENCH ONLY — production callers must use ``run_varlen(workspace=...)``. - Explicit-workspace form for multi-stream callers (main.cpp:105-116).""" + Explicit-workspace form for multi-stream callers.""" validate_run_ws(workspace, logits) _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) @@ -1341,14 +1271,13 @@ def run_varlen( KNOWN LIMITATION: on rows containing NaN logits the selected index SET can differ from ``heuristicTopKDecode.cu`` (both kernels order NaNs - implementation-specifically; inherited from the translation campaign's - probe battery). Finite inputs — including +/-inf and denormals — are - tie-aware exact. - - FULL-RANGE PRODUCTION CONTRACT: correct and dispatched for any - ``num_rows`` (BS 1..1024+ x next_n) and any envelope up to 1M kv tokens. - Family selection (streaming main / clustered register-resident) is a - pure function of the capture-stable launcher key. + implementation-specifically). Finite inputs — including +/-inf and + denormals — are tie-aware exact. + + CONTRACT: correct and dispatched for any ``num_rows`` + (BS 1..1024+ x next_n) and any envelope up to 1M kv tokens. Family + selection (streaming main / clustered register-resident) is a pure + function of the capture-stable launcher key. """ if logits.dtype is not torch.float32: raise RuntimeError( @@ -1396,10 +1325,11 @@ def run_varlen( if engine == "auto": # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- - # Full B1-style validation battery (the engine bypasses _run_impl — - # every check the legacy path enforces is replayed here; the batch-dim - # check is CRITICAL: the kernel grid comes from logits.shape[0], so a - # short indices/values tensor would be written out of bounds). + # Full validation battery (the engine bypasses _run_impl — every + # check the batch-uniform path enforces is replayed here; the + # batch-dim check is CRITICAL: the kernel grid comes from + # logits.shape[0], so a short indices/values tensor would be written + # out of bounds). if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): raise RuntimeError("all tensors must be CUDA") if logits.dtype is not _F32 or pre_idx.dtype is not _I32 or indices.dtype is not _I32: diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index a9d0a7987596..44fbb6396b0c 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -119,7 +119,7 @@ def test_selfsampling_topk_exactness(batch_size, top_k, n_valid): # n_valid <= top_k: every valid position is in the top-K. Production short -# path (heuristicTopKDecode.cu:72-84): identity indices + -1 tail padding. +# path convention: identity indices + -1 tail padding. _SHORT_CASES = [ (512, 256), (512, 511), @@ -199,7 +199,7 @@ def test_selfsampling_topk_degenerate_hints(hint_kind, top_k, n_valid): """Hints only steer the sampling ladder — exactness must survive the degenerate hint buffers production can produce: the all-zero cold start (dsa.py ``heuristic_prev_topk.zero_()`` init corners), fully duplicated - hints, and max-index hints (hint-robustness class of PR #17550).""" + hints, and max-index hints.""" logits, _, indices, ref_vals = _make_case(2, n_valid, top_k, seed=n_valid + top_k) if hint_kind == "all_zero": pre_idx = torch.zeros((2, top_k), dtype=torch.int32, device=_DEV) @@ -289,9 +289,8 @@ def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False, engine="aut ([131075, 32800, 2000], 1, 4, 512), # v4-style compressed index space ([9000, 5001], 2, 1, 512), # MTP: n varies per row within a request ([65540], 4, 4, 1024), # MTP: compressed-boundary-crossing rows - ([40000, 7003], 3, 4, 512), # MTP2 (next_n=3): no production config - # uses it today, but the window formula must generalize (gap noted - # in the PR's MTP-coverage section) + ([40000, 7003], 3, 4, 512), # MTP2 (next_n=3): unused in production + # today, but the window formula must generalize ], ids=["cr1_hetero_short", "cr4_hetero_short", "cr1_mtp2", "cr4_mtp4", "cr4_mtp3"], ) @@ -525,9 +524,9 @@ def test_selfsampling_topk_guards(): def test_selfsampling_route_large_batch_domain(): - """Production num_rows = max_batch_size * next_n can exceed the bench - grid's b<=1024 envelope (e.g. 1024 * 4 = 4096 rows). route() is a pure - function — assert the full domain stays well-formed up to 8192 rows.""" + """Production num_rows = max_batch_size * next_n can exceed the b<=1024 + range the plans were tuned on (e.g. 1024 * 4 = 4096 rows). route() is a + pure function — assert the full domain stays well-formed up to 8192 rows.""" for k in (512, 1024, 2048): for n in (k + 1, 4096, 65536, 262144): npad = (n + 63) // 64 * 64 From ff08b1d146bd5b1ddf8c7eb579d2ef8e9e75ad46 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:32:00 +0000 Subject: [PATCH 39/40] [TRTLLM-15293][chore] self-sampling top-K: comment slim tail Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index 543227e495fc..f6e5a59124a2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -722,27 +722,27 @@ def find_cross(s_hist, target, tidx, s_res, nb: cutlass.Constexpr): BPL = nb // 32 # python int at trace time if tidx < cutlass.Int32(32): lane = tidx - # per-lane span sum with rotated bank-skew indexing (L82-84) + # per-lane span sum with rotated bank-skew indexing part = cutlass.Int32(0) for j in cutlass.range_constexpr(BPL): idx = lane * cutlass.Int32(BPL) + ((cutlass.Int32(j) + lane) & cutlass.Int32(BPL - 1)) part = part + s_hist[idx] - # 5-step suffix scan (L85-89): v = sum of part over lanes >= lane + # 5-step suffix scan: v = sum of part over lanes >= lane v = warp_suffix_scan_add(part, lane) if lane == cutlass.Int32(0): s_res[RES_TOT] = v - # level 1: highest lane whose suffix still reaches target (L91-92) + # level 1: highest lane whose suffix still reaches target msk = ballot(v >= target) L = hi_bit_or_zero(msk) - aboveL = cute.arch.shuffle_sync(v - part, L) # L93 - # level 2: one bin per lane inside lane L's span (L94-100) + aboveL = cute.arch.shuffle_sync(v - part, L) + # level 2: one bin per lane inside lane L's span h = cutlass.Int32(0) if lane < cutlass.Int32(BPL): h = s_hist[L * cutlass.Int32(BPL) + lane] w = warp_suffix_scan_add(h, lane) msk2 = ballot((aboveL + w) >= target) J = hi_bit_or_zero(msk2) - if lane == J: # L103-107 + if lane == J: # pinning lane s_res[RES_B] = L * cutlass.Int32(BPL) + J s_res[RES_M] = h s_res[RES_ABOVE] = aboveL + (w - h) @@ -1057,7 +1057,7 @@ def gather_hint( ): NW = blk // 32 lane = tidx & cutlass.Int32(31) - # batch A: KPT coalesced pre_idx loads, predicated flat (L340) + # batch A: KPT coalesced pre_idx loads, predicated flat pvs = [] for t in cutlass.range_constexpr(kpt): pv = cutlass.Int32(-1) @@ -1065,14 +1065,14 @@ def gather_hint( if j < k: pv = ld_g_i32(p_addr, j) pvs.append(pv) - # batch B: KPT scattered read-only gathers, predicated flat (L341-343) + # batch B: KPT scattered read-only gathers, predicated flat xs = [] for t in cutlass.range_constexpr(kpt): xv = cutlass.Float32(0.0) if cutlass.Uint32(pvs[t]) < cutlass.Uint32(n): # (unsigned)p < (unsigned)n xv = ldg_f32(x_addr, pvs[t]) xs.append(xv) - # fold (L344-346) + # fold glmin = cutlass.Uint32(0xFFFFFFFF) glmax = cutlass.Uint32(0) for t in cutlass.range_constexpr(kpt): @@ -1082,14 +1082,14 @@ def gather_hint( glmin = u2 if u2 > glmax: glmax = u2 - # warp redux + staging (L347-348) + # warp redux + staging glmin = warp_min_u32(glmin) glmax = warp_max_u32(glmax) if lane == cutlass.Int32(0): s_wmn[tidx >> cutlass.Int32(5)] = glmin s_wmx[tidx >> cutlass.Int32(5)] = glmax - cute.arch.barrier() # L349 (barrier 1/2) - # cross-warp redux by EVERY thread — block-uniform outputs (L350-355) + cute.arch.barrier() # barrier 1/2 + # cross-warp redux by EVERY thread — block-uniform outputs a2 = cutlass.Uint32(0xFFFFFFFF) c2 = cutlass.Uint32(0) if lane < cutlass.Int32(NW): @@ -1097,14 +1097,14 @@ def gather_hint( c2 = s_wmx[lane] gm = invkey(warp_min_u32(a2)) gx = invkey(warp_max_u32(c2)) - # NaN-safe degeneracy guard (L356): !(GM < GX) — NaN compares false + # NaN-safe degeneracy guard: !(GM < GX) — NaN compares false ok = cutlass.Int32(0) if gm < gx: ok = cutlass.Int32(1) if ok == cutlass.Int32(0): gm = cutlass.Float32(SENT_LO) gx = cutlass.Float32(SENT_HI) - cute.arch.barrier() # L357 (barrier 2/2) + cute.arch.barrier() # barrier 2/2 return gm, gx @@ -1529,7 +1529,7 @@ def kern( n = nv n4v = n >> cutlass.Int32(2) # Ladder-scalar baselines only: the real SMP/SS2/TGT/TGT2 are - # derived by warp0 alone in the P2a block below (bit-identical + # derived by warp0 alone in the block below (bit-identical # formulas) and published through s_lad — every thread's local # copies here are overwritten by the post-barrier smem read. SMP = cutlass.Int32(0) @@ -4740,7 +4740,7 @@ def kern( ) # scan_cross predeclares its second-stage partial as Int32 and reads # s_ws inside a dynamic if — a Uint32 ws tensor trips the DSL type- - # stability check (frozen sibling; counts < 2^31 so Int32 is exact). + # stability check (counts < 2^31 so Int32 is exact). s_wmn = smem.allocate_tensor( cutlass.Uint32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 ) @@ -6381,9 +6381,9 @@ def kern( cutlass.Int32(1), ) i = i + cutlass.Int32(BLK) - cute.arch.barrier() # L2618 + cute.arch.barrier() # level hist find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) - cute.arch.barrier() # L2620 + cute.arch.barrier() # level scan aboveC = aboveC + s_res[RES_ABOVE] needC = needC - s_res[RES_ABOVE] mm = s_res[RES_M] @@ -6398,7 +6398,7 @@ def kern( ) rlo = nlo lev = lev + cutlass.Int32(1) - cute.arch.barrier() # L2626 + cute.arch.barrier() # narrowing done nA = k # tie_m ? above2 : k if tie_m == cutlass.Int32(1): nA = aboveC @@ -6408,7 +6408,7 @@ def kern( lml = cutlass.Int32(cute.arch.lanemask_lt()) it2 = (n + cutlass.Int32(self.blk - 1)) // cutlass.Int32(self.blk) it = cutlass.Int32(0) - while it < it2: # L2628-2645 + while it < it2: i = it * cutlass.Int32(BLK) + tid uke = cutlass.Uint32(0) if i < n: @@ -6442,7 +6442,7 @@ def kern( out_row[nA + p2e] = i it = it + cutlass.Int32(1) - # ---- P10: FINAL cluster rendezvous (L2648) — ALL ranks reach it; + # ---- P10: FINAL cluster rendezvous — ALL ranks reach it; # keeps peers resident until rank 0 has read their ck/ci. _cluster_sync_aligned() @@ -6534,7 +6534,7 @@ def regclus_topk(logits, pre_idx, n, out, rd=None): rd = route(logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1]) assert rd["kernel"] == "reg_clus", rd["kernel"] tpl = tuple(rd["tpl"]) - assert pre_idx.shape[1] <= tpl[0], "k <= BLK enforced by dispatch L2897" + assert pre_idx.shape[1] <= tpl[0], "k <= BLK enforced by dispatch" assert rd["smem"] == DYN_SMEM_BYTES compiled = get_compiled__regclus(tpl) compiled(logits, pre_idx, out, int(n)) From 62cfe5ddf2e841a563d9eefa2680d3021db666a2 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 03:24:59 +0000 Subject: [PATCH 40/40] [TRTLLM-15293][chore] self-sampling top-K: dedupe the per-row prologue contract comments Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 46 +++++++------------ 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index f6e5a59124a2..36c2d67ff965 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -4594,23 +4594,15 @@ def kern( row = by lane = tidx & cutlass.Int32(31) - # ================= per-row varlen prologue (varlen mode only) ========= - # Production contract (GvrMainKernel / GvrRegClusKernel discipline): - # row r serves request r // next_n with - # n = (kv_lens[req] - next_n + r % next_n + 1) >> cr_shift, clamped to - # the envelope launch arg n. The sampling-ladder scalars - # (SMP/TGT/Q/SS2/TGT2 — dead launch args in this mode) are re-derived - # from this row's n by the route_dynamic() clus formulas below, with - # ONE deviation: the QUAD sample geometry is computed for every - # non-short row instead of only n > SCAP — the host only ever launches - # this family with n > SCAP, so the SMP == 0 no-sample path is - # untested; short-of-envelope rows get a small valid schedule instead - # (sampling only steers the rung; exactness is schedule-invariant). - # Every quantity is a pure function of `row`, so all CS ranks of a - # row's cluster (and all threads) compute identical values — the - # whole-body guard below is cluster-uniform and the cluster barriers - # inside remain aligned. Short rows (n <= k) emit identity + (-1) - # tail from rank 0 here and SKIP the body entirely. + # ============ per-row varlen prologue — shared contract lives in ====== + # GvrMainKernel's prologue (per-row n from kv_lens; ladder scalars are + # dead launch args, re-derived here by the route_dynamic() clus + # formulas). Clus deviation: QUAD sample geometry runs for every + # non-short row (host never launches clus at n <= SCAP, so SMP == 0 is + # untested; sampling only steers the rung — exactness is + # schedule-invariant). All values are pure functions of `row`, so the + # whole-body guard and the cluster barriers stay cluster-uniform. + # Short rows (n <= k): rank 0 emits identity + (-1); body is SKIPped. short = cutlass.Int32(0) prow = row if cutlass.const_expr(self.varlen): @@ -5837,19 +5829,13 @@ def kern( rank, row, _ = cute.arch.block_idx() # bx=rank lane = tid & cutlass.Int32(31) - # ================= per-row varlen prologue (varlen mode only) ========= - # Same production contract as GvrMainKernel (see its prologue): row r - # serves request r // next_n with - # n = (kv_lens[req] - next_n + r % next_n + 1) >> cr_shift, clamped to - # the envelope launch arg n (the launcher admits this family only when - # the envelope fits its capacity window, so per-row n never exceeds - # capacity). Every quantity is a pure function of `row`, so all CS - # ranks of a row's cluster (and all threads) compute identical values - # -- the whole-body guard below is cluster-uniform and the cluster - # barriers inside remain aligned. Short rows (n <= k) emit - # identity + (-1) tail from rank 0 here and SKIP the body entirely: a - # zero-work pass would reach the degenerate crossing-overflow emitter - # and poison the output. + # ============ per-row varlen prologue — shared contract lives in ====== + # GvrMainKernel's prologue (per-row n from kv_lens, clamped to the + # envelope arg; the launcher admits this family only when the envelope + # fits its capacity window). Pure functions of `row` keep the body + # guard and cluster barriers cluster-uniform. Short rows (n <= k): + # rank 0 emits identity + (-1) and the body is SKIPped — a zero-work + # pass would reach the crossing-overflow emitter and poison the output. short = cutlass.Int32(0) prow = row if cutlass.const_expr(self.varlen):