diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index e32e806edb0e..177ba04fc1f2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -635,6 +635,23 @@ def __init__( self._enable_heuristic_topk = ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ) + # 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 + # 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) + ) self.mtp_index_share = sparse_params.mtp_index_share if self.use_cute_dsl_topk: @@ -647,6 +664,10 @@ def __init__( decode_top_k_implementation = TopKImplementation.CUDA_GVR else: decode_top_k_implementation = TopKImplementation.CUDA_RADIX + if self._use_self_sampling_topk and self._enable_heuristic_topk: + # 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, prefill_implementation=TopKImplementation.CUDA_RADIX, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 06ba0bc12f56..c647beaab653 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 @@ -16,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 ( @@ -336,6 +338,80 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: num_sms=self.num_sms, ) + 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``. The varlen launcher is keyed + by the exact row count AND the logits row stride: rows cover the + 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 + # same hardware gates as the dispatch flag (indexer __init__): never + # 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_gvr_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 + nn = int(next_n) + # 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 + 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 (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 + # 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), + 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 + # 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) -> None: # 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/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index 68e8f7bf8a45..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,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 run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -34,4 +35,5 @@ "DirectTopKKernel", "tiered_topk", "is_tiered_topk_supported", + "selfsampling_topk_run_varlen", ] 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..36c2d67ff965 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -0,0 +1,6537 @@ +# 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. + +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 +import sys + +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 +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. +C = sys.modules[__name__] + + +# =========================================================================== +# ==== shared device units ===================================== +# =========================================================================== +"""Device-helper library shared by the main / reg / clus / regclus families. + +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. 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; 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 +# --------------------------------------------------------------------------- +FULLM = 0xFFFFFFFF +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 +GVR_WS_BUF_OFF = 2048 # workspace g_buf byte offset + +# degenerate-hint sentinels (exact-equality flag values) +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 +# --------------------------------------------------------------------------- +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 — order-preserving float->u32 radix key. +# fkey: u ^ (((int32)u >> 31) | 0x80000000) [arithmetic-shift sign trick, +# 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. +# --------------------------------------------------------------------------- +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). x: dynamic Float32 -> Uint32 key.""" + return fkey_bits(u32_of_f32(x)) + + +def invkey_bits(K): + """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) + return K ^ m + + +def invkey(K): + """CUDA invkey(uint32). K: dynamic Uint32 key -> Float32.""" + return f32_of_u32(invkey_bits(K)) + + +# --------------------------------------------------------------------------- +# 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).""" + 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).""" + return cute.arch.warp_redux_sync(v, "add") + + +def fmin_f32(a, b): + """fminf -> native min.f32.""" + 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). 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`.""" + r = cutlass.Int32(0) + if msk != cutlass.Int32(0): + r = cutlass.Int32(31) - clz_i32(msk) + return r + + +# --------------------------------------------------------------------------- +# warp shfl scans (plus the TWO-interleaved variant gvr_topk_reg 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.""" + 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. + + 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.""" + for o in [1, 2, 4, 8, 16]: + val = _shfl_down_add(val, lane, o) + return val + + +# --------------------------------------------------------------------------- +# 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 + + (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). 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) +# --------------------------------------------------------------------------- +def threadfence_gpu(): + """__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 (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 (native ctors emit cvt.rzi.{u32,s32}.f32) +# --------------------------------------------------------------------------- +def f2u_rz(v): + """__float2uint_rz: saturating (neg/-inf -> 0, huge -> 0xffffffff, NaN -> 0). + + 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 +# --------------------------------------------------------------------------- +@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 +# --------------------------------------------------------------------------- +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 + ) + + +@dsl_user_op +def _ld_g_nc_v4_f32(gaddr, *, loc=None, ip=None): + """Pinned `ld.global.nc.v4.f32` (CUDA `__ldg(const float4*)`). + + 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)>") + 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]. + + 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 + frag[1] = v1 + frag[2] = v2 + frag[3] = v3 + + +def ldg_f32(base_addr, idx): + """__ldg(X + idx): scalar read-only 4B gather.""" + 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_rmem_tensor((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.""" + 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). + 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 (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 path): (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. 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): + """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.""" + 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 (never split into two 4B stores). + + 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 (cg::cluster.sync() == +# barrier.cluster.{arrive,wait}.aligned). Writers use the FULL (releasing) +# 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): + 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().""" + _cluster_arrive_aligned() + _cluster_wait_aligned() + + +# =========================================================================== +# 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))] — 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. +# =========================================================================== +@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 + 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: 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 + msk = ballot(v >= target) + L = hi_bit_or_zero(msk) + 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: # pinning lane + s_res[RES_B] = L * cutlass.Int32(BPL) + J + s_res[RES_M] = h + s_res[RES_ABOVE] = aboveL + (w - h) + + +# =========================================================================== +# 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( + 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 + 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 + 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 + 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 + for q in cutlass.range_constexpr(NV - 1, -1, -1): + if cutlass.const_expr(HOLD): + vv = frags[q] + else: + 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_rmem_tensor((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 + 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): + 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 +# 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; the caller pays its usual +# publish barrier after. Used by the gvr_clus whole-row degenerate path. +# =========================================================================== +@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): + c = s_hist[tidx] + w = warp_incl_scan_add(c, lane) + if lane == cutlass.Int32(31): + 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) + 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): + 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 +# 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. 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): + BPT = nb // blk + NW = blk // 32 + lane = tidx & cutlass.Int32(31) + wid = tidx >> cutlass.Int32(5) + loc = cute.make_rmem_tensor((BPT,), cutlass.Int32) + base = tidx * cutlass.Int32(BPT) + sm = cutlass.Int32(0) + 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) + if lane == cutlass.Int32(31): + s_ws[wid] = w + cute.arch.barrier() # the ONE internal barrier + vv = cutlass.Int32(0) + if lane < cutlass.Int32(NW): + 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)) + after = tot - (off + w) + if tidx == cutlass.Int32(0): + s_res[RES_TOT] = tot + 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) + 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 +# Warp-0-fused cluster merge + suffix scan: each lane reads its BPT-bin span +# from EVERY rank's hist via 16B DSMEM loads, 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) +# 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: 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. +# 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 + 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 + 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 + 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 + 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() # 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): + a2 = s_wmn[lane] + c2 = s_wmx[lane] + gm = invkey(warp_min_u32(a2)) + gx = invkey(warp_max_u32(c2)) + # 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() # barrier 2/2 + return gm, gx + + +# =========================================================================== +# ==== family: main ============================================ +# =========================================================================== +"""gvr_main — streaming self-sampling GVR top-K. + +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 (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 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; 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 (__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. +""" + + +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 + +_NEG_INF = float("-inf") + + +# --------------------------------------------------------------------------- +# 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): + 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, 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).""" + 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. + 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_=[ + 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.""" + 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. + + 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, + base_addr + cutlass.Int64(idx) * cutlass.Int64(sc4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + frag = cute.make_rmem_tensor((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. + + 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 (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(), + [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. + + `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 + 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: + """gvr_main — streaming self-sampling GVR.""" + + 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, + ): + 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 + self.u = u + self.minb = minb + 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 + # 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 (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)" + self.dyn_bytes = self.ck_off + (self.cmpb + 1) * 8 + self.lb = self.nbs.bit_length() - 1 # log2(NBS)=8 + + # ------------------------------------------------------------------ + # 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 + if bn_u > cutlass.Uint32(NBS - 1): + bn_u = cutlass.Uint32(NBS - 1) + bn = cutlass.Int32(bn_u) + if cutlass.const_expr(self.vstg): + # 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-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 + 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 + # (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 and both degen + # emits): q1 winners to out[base1+p] p> 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) + # "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 + n4v = n >> cutlass.Int32(2) + # Ladder-scalar baselines only: the real SMP/SS2/TGT/TGT2 are + # 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) + SS2 = cutlass.Int32(1) + TGT = cutlass.Int32(0) + TGT2 = cutlass.Int32(0) + 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) ---- + smem = SmemAllocator() + s_hist = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((self.hb,), order=(0,)), byte_alignment=128 + ) + 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( + 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 + ) + # 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 + s_scal = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 + ) + s_pk = smem.allocate_tensor( + cutlass.Int64, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=8 + ) + s_tsh = smem.allocate_tensor( + cutlass.Float32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 + ) + # 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( # [0]=kmin [1]=kmax + cutlass.Uint32, cute.make_ordered_layout((2,), order=(0,)), byte_alignment=8 + ) + if cutlass.const_expr(self.varlen): + # 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 region + cutlass.Int8, cute.make_ordered_layout((self.dyn_bytes,), order=(0,)), byte_alignment=16 + ) + sbase = blob.iterator.toint() + s_cbuf = cute.make_tensor( + cute.make_ptr(cutlass.Int32, sbase, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((SCPB + 4,)), + ) + s_cbuf2 = cute.make_tensor( + cute.make_ptr(cutlass.Uint64, sbase, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((SCPB + 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((CMPB + 1,)), + ) + + # 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()) + # 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 (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): + s_x4[0] = cutlass.Int32(4) + cute.arch.barrier() + x4_pin = s_x4[0] + + # ---- 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) + x_addr = _pin_i64(logits.iterator.toint() + row64 * cutlass.Int64(npad) * 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 # slab views + goff_addr = ws_addr + cutlass.Int64(C.GVR_WS_OFF_OFF) + gbuf_addr = ws_addr + cutlass.Int64(C.GVR_WS_BUF_OFF) + # 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). + 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) + c0 = cutlass.Int32(0) + c1 = n4 + if cutlass.const_expr(self.split): + c0 = part * Q + c1 = c0 + Q + if c1 > n4: + c1 = n4 + tail0 = n4 << cutlass.Int32(2) + tailn = cutlass.Int32(0) + if part == cutlass.Int32(0): + tailn = n - tail0 + + 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): # HB<=BLK always + s_hist[tidx] = cutlass.Int32(0) + + # ===== 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 (cheap mirror) ---- + # The ladder scalars steer the sampling rung only — + # exactness is schedule-invariant (retry/degen close every + # 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: + 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): 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) ================= + 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): + 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 ====================== + smn = cutlass.Float32(float("inf")) + smx = cutlass.Float32(float("-inf")) + 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]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsb[t]) + 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) # 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)) + 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 (sample redux publish) ---- + + # 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): # 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: # 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: 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): # 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): # prefetch site #3 + 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 + 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) # 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) + if SMP > cutlass.Int32(0): + if SMAX > SMIN: + sok = cutlass.Int32(1) + 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 + 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) # 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 (sample histogram) ---- + # triple-target ZERO scan: TGT / TGT2 / 2*TGT + # (THREE = SHD || gated-SPLIT) + 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 (scan publish) ---- + + tot0 = s_res[C.RES_TOT] + b1v = s_res[C.RES_B] + if sok != cutlass.Int32(0): + if tot0 >= TGT: + T = _fmaf(cutlass.Float32(b1v), w, SMIN) + Trung = T # snapshot + needg = cutlass.Int32(1) # 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): # 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) + 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) + ) + 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): + 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): + # 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] + # 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 — 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_rmem_tensor((4,), cutlass.Float32) for _ in range(max(U - PFD, 1)) + ] # 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 + if cutlass.const_expr(self.pf): + # exactness: re-prime pf[] (holds stale roll data) + 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): + s_hist[tidx] = cutlass.Int32(0) + if tidx == cutlass.Int32(0): + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (retry reset) ---- + + TF = T # window + 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) + # 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 ---- + span = c1 - c0 + step = cutlass.Int32(BLK * U) + nFull = cutlass.Int32(0) + rem = cutlass.Int32(0) + 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 + 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 + 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 + for uu in cutlass.range_constexpr(PFD, U): + i_ = i0 + cutlass.Int32(uu * BLK) + ic = i_ + if ic >= c1: + 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): + 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 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 + 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 + 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 + 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; + # 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)) + 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 + 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 (row pass) ---- + myn = s_scal[0] + + if cutlass.const_expr(self.split): + # ---- SLAB HAND-OFF; exactly ONE attempt ---- + if tidx == cutlass.Int32(0): + 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 (slab offset) ---- + base = s_scal[3] + 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: # overflow re-sweep + if tidx == cutlass.Int32(0): + s_scal[0] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (overflow reset) ---- + 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 + 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 (slab publish) ---- + if tidx == cutlass.Int32(0): # 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 (arrival word) ---- + pk = s_pk[0] + 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() # 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): # 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) + # 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 (slab histogram) ---- + 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 (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) # break (NATT==1) + else: + # ---- non-split verify + rung ladder ---- + 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 (verify scan) ---- + tot = s_res[C.RES_TOT] + acc = cutlass.Int32(0) + if tot >= k: + acc = cutlass.Int32(1) + if acc != cutlass.Int32(0): # 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): # ladder exhausted + running = cutlass.Int32(0) + else: + tshtaken = cutlass.Int32(0) # 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 (TSH retry) ---- + else: + # 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) + if T > GMIN: + floorhit = cutlass.Int32(0) + if floorhit != cutlass.Int32(0): + running = cutlass.Int32(0) + else: + T = GMIN + cute.arch.barrier() # ---- barrier (floor retry) ---- + att = att + cutlass.Int32(1) + + # ============ classification ============ + 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 ---- + if complete != cutlass.Int32(0): + 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) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + else: + 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: + 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 — + # 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) + 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 ---- + if whole == cutlass.Int32(0): + 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. + 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 + 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 + if tidx == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + if tidx < cutlass.Int32(NBS): # cleared ONCE + s_hist[tidx] = cutlass.Int32(0) + 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 (key range) ---- + 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): # <=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: # 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 (level hist) ---- + 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 (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): + rhi = nlo + ((cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1)) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (emit counters) ---- + it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # 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) # 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 ---- + 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): + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (degen A init) ---- + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # <=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: + 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 (level hist) ---- + 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 (level scan) ---- + 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): + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + 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: + 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 ---- + 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): + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (degen B init) ---- + brk = cutlass.Int32(0) + lev = cutlass.Int32(0) + while brk == cutlass.Int32(0): # <=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: # 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 (level hist) ---- + 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 (level scan) ---- + 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): + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + 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: + 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) + + # ---- varlen short-row epilogue (production heuristicTopKDecode + # 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); 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, + 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, + 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) + + +# --------------------------------------------------------------------------- +# 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) — 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 + 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 + ) + 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 + ) + 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, + out_fake, + ws_fake, + *([cutlass.Int32(0)] * 11), + kv_fake, + *([cutlass.Int32(0)] * 5), + 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"], + _legacy_dummy_kv(pre_idx), # dummy kv_lens (dead in legacy mode) + 0, + 0, + 0, + 0, + 0, + ) + 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 + + +# =========================================================================== +# ==== family: reg ============================================= +# =========================================================================== +"""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)` +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 (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 + [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 (IMGOFF==NBH) + +Launch smem = 512 + dispatch_smem_bytes (dynamic Int32). + +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 # 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 +# --------------------------------------------------------------------------- +@dsl_user_op +def _fmaf__reg(a, b, c, *, loc=None, ip=None): + """CUDA fmaf: single fma.rn.f32.""" + 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. + + 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( + 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. + + 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. + 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. 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: {} + # 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: + """gvr_topk_reg.""" + + def __init__( + self, + blk: int, + vpt: int, + minb: int, + kpt: int, + cur: bool, + deg: bool, + 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) + 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) + # 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 + 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, + kv_lens: 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() # knob default off + + tid, _, _ = cute.arch.thread_idx() + 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 + # 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) + + # 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) + lmin = cutlass.Uint32(0) + lmax = cutlass.Uint32(0) + esc = cutlass.Int32(0) + okc = cutlass.Int32(0) + whole = cutlass.Int32(0) + 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) + + 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 + cute.make_ptr(cutlass.Float32, ck_base, cute.AddressSpace.smem, assumed_align=16), + cute.make_layout((65536,)), + ) + 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` + + # ---- 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): + 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 + 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 + 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 + 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) + + # ---- 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 + 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() # image staged + 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() # img dies + elif cutlass.const_expr(self.use_bm): + nbw = (n + cutlass.Int32(31)) >> cutlass.Int32(5) + i = tid + while i < nbw: # bitmap clear + bm[i] = cutlass.Int32(0) + i = i + cutlass.Int32(BLK) + cute.arch.barrier() # bitmap cleared + 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() # bitmap set + 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 + 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): + 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 + 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; 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() # warp partials published + 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 + 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 + 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`; 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) 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 + if cutlass.const_expr(self.brl): + # 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) 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 + 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() # histogram done + + # ---- 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() # crossing published + 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) + + # ---- ESCAPE: 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): 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() # escape init + 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() # count published + if s_cnt[0] >= k: + klo = kt + cute.arch.barrier() # count consumed + if tid == cutlass.Int32(0): + s_cnt[0] = cutlass.Int32(0) + cute.arch.barrier() # count reset + 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() # 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() # nA consumed + 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 + 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 + if cutlass.const_expr(self.cur): + LOQ = cutlass.Float32(Bv) # int->float cvt + 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) # 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 + 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 + 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) + 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 + 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() # emit done + 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 + 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 + if tid == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + 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() # key range published + 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: + 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: + ethr = cutlass.Int64(rlo) + done = cutlass.Int32(1) + if lev >= cutlass.Int32(6): + 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() # level clear + 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() # 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() # level scan + 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 + if tid == cutlass.Int32(0): + s_e12[0] = cutlass.Int32(0) + s_e12[1] = cutlass.Int32(0) + 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) + 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, + kv_lens: 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, kv_lens, 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, 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), 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, + 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() + 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 + ) + 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}" + with _no_carveout(): + compiled = cute.compile( + kernel, + lg_fake, + pi_fake, + kv_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"] + 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 + + +# =========================================================================== +# ==== family: clus ============================================ +# =========================================================================== +"""gvr_clus — clustered streaming self-sampling GVR; per-CTA stream mirrors +gvr_main. + +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). + +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); 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 gvr_main): int2(value bits, index) is ONE +little-endian Uint64 = (idx << 32) | value_bits — single u64 smem ld/st. + +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 + __syncthreads, + clus.sync (merge), __syncthreads (merge publish), + [ladder gather: 2 inside gather_hint], + 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, 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 cannot deadlock. +""" + + +QUADC_CLUS__clus = C.QUADC_CLUS + +_NEG_INF__clus = float("-inf") + + +# --------------------------------------------------------------------------- +# 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): + 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: + """gvr_clus — clustered streaming GVR.""" + + def __init__( + 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" + 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- + # 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 + self.nbs = nbs + self.cs = cs + self.scap = scap # smem extents only — + self.cmp = cmp_ # value logic uses rt args + 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: 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 + + # ------------------------------------------------------------------ + # 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 + 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: 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): + 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 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): + 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: 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 + 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: # 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 + 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): + 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] + 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 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): + 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, + kv_lens: 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, + bigf: 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) + rank = bx + row = by + lane = tidx & cutlass.Int32(31) + + # ============ 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): + 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 (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, 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: + 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_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( + 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) + 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 = cutlass.Int32( + cutlass.Float32(quarter) * cute.arch.rcp_approx(cutlass.Float32(quads)) + ) + if SS2 < cutlass.Int32(1): + SS2 = cutlass.Int32(1) + 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) + 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) + TGT2 = cutlass.Int32(cutlass.Float32(k) * smp16f * rn_) + if TGT2 < cutlass.Int32(1): + TGT2 = cutlass.Int32(1) + # 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: + ov = cutlass.Int32(-1) + if tidx < nv: + ov = tidx + out[row, tidx] = ov + + # ---- 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 + ) + 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] + cutlass.Int32, cute.make_ordered_layout((self.nbs,), order=(0,)), byte_alignment=16 + ) + s_ws = smem.allocate_tensor( # degen scan only + cutlass.Int32, cute.make_ordered_layout((NW,), order=(0,)), byte_alignment=16 + ) + # 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 (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( # 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 + s_scal = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 + ) + s_tsh = smem.allocate_tensor( + cutlass.Float32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 + ) + 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() + 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,)), + ) + + # ---- 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 (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] + + # ---- interleaved chunk ownership ---- + 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): + 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): # HB<=BLK + s_hist[tidx] = cutlass.Int32(0) + + # ============ 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. + atom128 = C.g2r_atom_f32(128, invariant=True) + 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: + 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 =================== + smn = cutlass.Float32(float("inf")) + smx = cutlass.Float32(float("-inf")) + 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]) + for t in cutlass.range_constexpr(4): + smn = C.fmin_f32(smn, fsb[t]) + 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 + 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)) + 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 (sample redux publish) ---- + + # 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 + 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: 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 + 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) # 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) + if SMP > cutlass.Int32(0): + if SMAX > SMIN: + sok = cutlass.Int32(1) + 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; + # 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): + 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 (sample histogram) ---- + # triple-target ZERO scan: 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 (scan publish) ---- + + tot0 = s_res[C.RES_TOT] + b1v = s_res[C.RES_B] + if sok != cutlass.Int32(0): + if tot0 >= TGT: + T = _fmaf__clus(cutlass.Float32(b1v), w, SMIN) + needg = cutlass.Int32(1) + 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): # HIC tighten + 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), + ) + # ladder floor kept in SHARED (64-register wall) + 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 — 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_rmem_tensor((4,), cutlass.Float32) for _ in range(U - PFD) + ] # 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 + # exactness: re-prime pf[] (holds stale roll data) + 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 (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 (retry reset) ---- + + TF = T # window + 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: 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 ---- + 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 + 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 + 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: 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 + 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; + # 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: 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 + 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 ---- + 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 (merge publish) ---- + tot = s_res[C.RES_TOT] + acc = cutlass.Int32(0) + if tot >= k: + acc = cutlass.Int32(1) + if acc != cutlass.Int32(0): # 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): # ladder exhausted + running = cutlass.Int32(0) + else: + # rung ladder — 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) + if T > GMIN: + floorhit = cutlass.Int32(0) + if floorhit != cutlass.Int32(0): + running = cutlass.Int32(0) + else: + T = GMIN + att = att + cutlass.Int32(1) + + # ============ classification ============ + 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 + 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 ---- + 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 ---- + 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 ============ + # 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 (exit) ==== + + if rank == cutlass.Int32(0): + if degen == cutlass.Int32(0): + if whole == cutlass.Int32(0): + # ---- P6 rank-0 refine ---- + if mc <= cutlass.Int32(QUADC_CLUS__clus): # O(mc^2) + mc2 = mc & cutlass.Int32(~1) + i = tidx + while i < mc: + # re-assert Uint64 at every unsigned compare + # in/after dynamic loops. + 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 + 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 + if tidx == cutlass.Int32(0): + s_kmm[0] = cutlass.Uint32(0xFFFFFFFF) + s_kmm[1] = cutlass.Uint32(0) + if tidx < cutlass.Int32(NBS): # cleared ONCE + s_hist[tidx] = cutlass.Int32(0) + 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 (key range) ---- + 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): # <=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: # 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 (level hist) ---- + 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 (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): + rhi = nlo + ( + (cutlass.Uint32(1) << sh2u) - cutlass.Uint32(1) + ) + rlo = nlo + lev = lev + cutlass.Int32(1) + if tidx == cutlass.Int32(0): + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (emit counters) ---- + it2 = (mc + cutlass.Int32(BLK - 1)) // cutlass.Int32(BLK) + it = cutlass.Int32(0) + while it < it2: # 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 + # (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): # <=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) + if tidx < cutlass.Int32(NBS): # per-level clear + s_hist[tidx] = cutlass.Int32(0) + cute.arch.barrier() # ---- barrier (level clear) ---- + i = tidx + while i < n: # 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 (level hist) ---- + # block-parallel scan (ONE internal barrier; only + # use of ws in this kernel) + C.scan_cross( + s_hist, + s_ws, + need2, + tidx, + s_res, + cutlass.Int32(0), + blk=BLK, + nb=NBS, + two=False, + ) + 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] + 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): + s_scal[1] = cutlass.Int32(0) + s_scal[2] = cutlass.Int32(0) + 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: + 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); + # min_blocks_per_mp=1 == __launch_bounds__(1024, 1) 64-register wall. + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + logits: cute.Tensor, + pre_idx: cute.Tensor, + kv_lens: 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, + ): + # 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), + 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 = "", + 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, 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_, 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() + 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 + ) + 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, + 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.""" + 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: + 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"]) + dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot + fn( + logits, + pre_idx, + dkv, + 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).""" + 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( + logits, + pre_idx, + dkv, + out, + rt["n"], + rt["npad"], + rt["k"], + rt["SCAP"], + rt["CMP"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + + +# =========================================================================== +# ==== family: regclus ========================================= +# =========================================================================== +"""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 +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 (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) + [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 the reg +family's _no_carveout workaround is unnecessary here). +""" + + +# ---- 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 +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 (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 (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.""" + + 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" + 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 + self.span = blk * vpt # float4 per CTA + + # ------------------------------------------------------------------ + @cute.kernel + 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) + S = cutlass.const_expr(self.S) + NW = cutlass.const_expr(self.blk // 32) + + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_wait() # knob default off + + tid, _, _ = cute.arch.thread_idx() + rank, row, _ = cute.arch.block_idx() # bx=rank + lane = tid & cutlass.Int32(31) + + # ============ 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): + 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 + # are dead-coded). + # ------------------------------------------------------------------ + 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) + + 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) + 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 (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 (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): + 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 + for z in cutlass.range_constexpr(4): + frags[u][z] = cutlass.Float32(_NEG_INF__regclus) + # tail element: rank 0 only + if rank == cutlass.Int32(0): + if tid < ntail: + tval = ldg_f32(x_addr, tix) + + # ---- 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, 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) + 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() # warp partials published + 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 + 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: 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 (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; 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 + _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 + # (one mapa per (i, r) exactly like the CUDA) + 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 + 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() # scan published + 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] # global cursor + cute.arch.barrier() # cursors published + + # ---- P7: register sweep emit (!degen) + if degen == cutlass.Int32(0): + LOQ = cutlass.Float32(Bv) + lim1 = above + if whole == cutlass.Int32(1): + lim1 = above + m + for s in cutlass.range_constexpr(S): + 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)) + 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) + 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: 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: release staging to rank 0 + cute.arch.barrier() + _cluster_sync_aligned() + + # ---- 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) + 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 + 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: + # 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() # kmm init + 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() # key range published + 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 + if needC == mm: + 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() # level clear + 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() # level hist + find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) + cute.arch.barrier() # level scan + 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() # 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) + 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: 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: + 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): + 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() # level clear + i = tid + 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): + 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() # level hist + find_cross(s_hist, needC, tid, s_res, nb=NB__regclus) + cute.arch.barrier() # level scan + 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() # narrowing done + 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: + 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 — 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, + kv_lens: cute.Tensor, + out: cute.Tensor, + n: cutlass.Int32, + stream, + ): + 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), + 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, varlen=False, next_n=1, cr_shift=0): + """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS).""" + 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, 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() + 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 + ) + 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, + kv_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" + 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..165048c8dcfe --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -0,0 +1,1595 @@ +# 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 device module). +Three sections: + +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, with keep-alive + double-checked locking; +3. operator entry — ``run(logits, pre_idx, n_valid, indices)`` / + ``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 +import operator +import threading +from collections.abc import Sequence + +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 ============================================================= +# =========================================================================== +"""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). + +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 + 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 + register-resident block; the streaming path re-derives its own CMP. +""" + + +# ---- 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 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 + + # ======================= 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 + IMGOFF = NBSEL + smem_reg = (NBSEL + 2 * CMP) * 4 + + def _reg(BLK, VPT, MINB, NBH): + # 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) + 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, # full ABI + "CMP": CMP, + "IMGOFF": IMGOFF, + "QC": QC, + }, + "grid": (b, 1), + "cluster": 1, + "block": BLK, + "smem": smem_reg, + "ws": False, + } + + 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: + return _reg(256, 1, 8, NB) + if n4 <= 512: + return _reg(512, 1, 4, NB) + if n4 <= 1024: + if wide: + 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, # 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) + return _reg(512, 2, 4, NB) + + # ---- 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: + # 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 in C + 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: + smc = (3 * NB + 2 * CMPC) * 4 + return { + "kernel": "reg_clus", + "tpl": (BLKC, vsel, cs), + "rt": {"n": n, "npad": npad, "k": k}, # dims only + "grid": (cs, b), + "cluster": cs, + "block": BLKC, + "smem": smc, + "ws": False, + } + + if n4 <= 4096 and wide: + return _reg(1024, 4, 1, 2 * NB) + + # ====================== streaming / collect path ======================== + R = 1 + if b <= 32: + r1 = 148 // b + if r1 < 1: + r1 = 1 + 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: # shallow R=2 split + R = 2 + + useclus = False + if 2 <= R <= 8 and k <= 1024: + p2 = 1 + while (p2 << 1) <= R: + p2 <<= 1 + # 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 + 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) + ) + 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) + 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: # PAIR sample + sel = SFAC * n // aim # 64-bit + if sel < 256: + sel = 256 + if sel > n // 2: + sel = n // 2 + pairs = sel >> 3 + if pairs < 1: + pairs = 1 + half = n4s >> 1 + if half < 1: + half = 1 + if pairs > half: + pairs = half + SS2 = half // pairs + if SS2 < 1: + SS2 = 1 + SMP = half // SS2 + if SMP < 1: + SMP = 1 + TGT = (aim * (SMP * 8)) // n # 64-bit + if TGT < 1: + TGT = 1 + TGT2 = (k * (SMP * 8)) // n # 64-bit + if TGT2 < 1: + TGT2 = 1 + Q = (n4s + R - 1) // R + + 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 + if quads < 1: + quads = 1 + quarter = n4s >> 2 + if quarter < 1: + quarter = 1 + if quads > quarter: + quads = quarter + SS2 = quarter // quads + if SS2 < 1: + SS2 = 1 + SMP = quarter // SS2 + if SMP < 1: + SMP = 1 + TGT = (aim * (SMP * 16)) // n + if TGT < 1: + TGT = 1 + TGT2 = (k * (SMP * 16)) // n + if TGT2 < 1: + TGT2 = 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, # 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 + + def _main(BLK, MINB, U, SPLIT): + # 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)) + # 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 never reads them + # (it recomputes them as constexprs). + "rt": { + "n": n, + "npad": npad, + "k": k, # 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: + 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: + return _main(512, 2, 8, False) + return _main(256, 4, 8, False) + + +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), # 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 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 KBIG-domain (k>1024), BLK=512 KPT=4 + ] + 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: 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"), + "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: 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.""" + 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: dict[str, object], n: int) -> tuple[dict[str, object], int]: + """Recompute the redacted n-continuous scalars from (static, n). + 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"): + 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: 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) + 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_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 + skipped even when the envelope n would normally land on them. Where + 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 + 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. 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) + hit = _VARLEN_CACHE.get(key) + if hit is not None: + return hit + 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; 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( + 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 + # ---- 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. + 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 + # ---- 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. + # 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"] + r_const = rt["R"] + # 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) + 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-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) + lc = ("main", fn, pre, tail) + _VARLEN_CACHE[key] = lc + return lc + + +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 + 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 ============================================================ +# =========================================================================== +"""Per-device workspace slab for the multi-CTA SPLIT path. + +Semantics: + * ONE zero-initialised slab workspace per device, lazily allocated through + 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 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(). + +Size: workspace_bytes() = GVR_WS_BUF_OFF + MAXC*GCAP*sizeof(int2) + = 2048 + 160*16384*8 = 20,973,568 B. + +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. +""" + + +# 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() # slow-path mutex +_ws_keep = {} # device index -> keep-alive int32 view + + +def workspace_bytes() -> int: + """Workspace bytes required by the multi-CTA SPLIT path.""" + return WS_BYTES + + +def default_workspace(ref: torch.Tensor) -> torch.Tensor: + """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 + 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, 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: + """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()): + 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 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 validate_run_ws check but is + rejected by the DSL at conversion -- surfaced as a launch failure with + shape context.""" + 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 ======================================================= +# =========================================================================== +"""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. +""" + + +# 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 and the +# validation battery runs on EVERY call +_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() + 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 (cached dummy tensor) + 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": + 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) + # [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"], + 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, kv_lens, out, n, npad, k, SCAP, CMP, SMP, + # 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"], + rt["npad"], + rt["k"], + rt["SCAP"], + rt["CMP"], + rt["SMP"], + rt["TGT"], + rt["Q"], + rt["SS2"], + rt["TGT2"], + ) + + 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 + # 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"] + + 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}") + + +# --------------------------------------------------------------------------- +# 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): + 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: 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( + "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 (rejects floats/strings) + nv = _index(n_valid) + + b, npad = lsh + k = psh[1] + 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}") + 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 (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 -- + # 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) + + # ---- 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 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 + # 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 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) + 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 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 + if values is not None: + # 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 +# --------------------------------------------------------------------------- +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). + + 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: # checked on 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, values) + + +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.""" + validate_run_ws(workspace, logits) + _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) + + +def run_varlen( + 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 + ``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). + + 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. + + KNOWN LIMITATION: on rows containing NaN logits the selected index SET + can differ from ``heuristicTopKDecode.cu`` (both kernels order NaNs + 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( + 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: + 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)}" + ) + d = logits.get_device() + if not 0 <= d < _GVR_MAX_DEV: + raise RuntimeError(f"device index out of range: {d}") + 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) ---------------- + # 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: + 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 + 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 + # 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) + 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) + 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) + 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]) + 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) + if vals is not None: + idx64 = idx.to(torch.int64) + vals.copy_(lg.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" + ) + # 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): + # 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, + idx[r : r + 1], + ws, + None if vals is None else vals[r : r + 1], + ) + + +__all__ = [ + "route", + "route_static", + "route_dynamic", + "route_split", + "route_bands", + "run", + "run_ws", + "run_varlen", + "warmup_varlen", + "workspace_bytes", + "WS_BYTES", + "default_workspace", + "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: int, + max_seq_len: int, + 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. + + One tiny real launch per requested ``num_rows`` (compile keys do not + 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. + + ``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)) + # round each request down to a next_n multiple (min next_n) and dedup + 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"])) + elif plan_free["kernel"] in ("reg", "regimg"): + ekey = ("reg", 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)) + 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 + 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[: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 + 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/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 704ce11283bc..950b02c30f93 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -9,6 +9,8 @@ import torch import torch.nn as nn +from tensorrt_llm.logger import logger + from ..memory_buffer_utils import get_memory_buffers @@ -20,11 +22,13 @@ class TopKImplementation(str, Enum): CUTE_DSL_RADIX = "cute_dsl_radix" CUDA_GVR = "cuda_gvr" CUTE_DSL_GVR = "cute_dsl_gvr" + CUTE_DSL_GVR_V2 = "cute_dsl_gvr_v2" _GVR_IMPLEMENTATIONS = { TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, + TopKImplementation.CUTE_DSL_GVR_V2, } _MAX_RADIX_BLOCKS_PER_ROW = 10 @@ -258,7 +262,57 @@ def _forward_decode_gvr( gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: assert gvr_prior_indices is not None - if self.decode_implementation == TopKImplementation.CUDA_GVR: + if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR_V2: + assert max_seq_len is not None + if ( + # engine hardware-format gate (falls through otherwise): + # fp32 row-major scores 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). 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. + scores.dtype == torch.float32 + and scores.stride(1) == 1 + and scores.stride(0) % 4 == 0 + 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 + + logger.info_once( + "self-sampling GVR top-K engaged " + f"(K={self.top_k}, 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 constant + # (no host reads — CUDA-graph safe). Hints are consumed raw + # (offset-free contract). The module receives max_seq_len in + # COMPRESSED index space; run_varlen's max_seq_len is in + # kv-token space like sequence_lengths — multiply back. + selfsampling_topk_run_varlen( + scores, + gvr_prior_indices, + sequence_lengths, + output_indices, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_seq_len * self.compress_ratio, + ) + return output_indices + logger.warning_once( + "TRTLLM_GVR_SELF_SAMPLING=1 but the decode scores do not " + "satisfy the engine's hardware-format gate " + f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " + "falling through to the CUDA GVR top-K path.", + key="selfsampling_topk_fallthrough", + ) + if self.decode_implementation != TopKImplementation.CUTE_DSL_GVR: + # CUDA_GVR, or the V2 hardware-format fall-through above workspace = self._get_workspace( scores, (scores.shape[0], self.top_k), diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e5c096bae84d..25abd7ee435e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1683,6 +1683,9 @@ 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, 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 new file mode 100644 index 000000000000..44fbb6396b0c --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -0,0 +1,961 @@ +# 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() not in (100, 103): + pytest.skip( + "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 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) + + +# n_valid <= top_k: every valid position is in the top-K. Production short +# path convention: 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_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()) + + +@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"] +) +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.""" + 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) + 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( + 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 _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 + 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, + engine=engine, + ) + 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("engine", ["auto", "reference"]) +@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 + ([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"], +) +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 — 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, cr = 1024, 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(): + _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) + 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) + 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(): + """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 + 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_route_large_batch_domain(): + """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 + 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 + 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 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:]): + 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): + 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 + + +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) + + +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) + + +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) + + +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_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.""" + 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) + — 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 + + +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})"