[None][feat] GVR V2 decode top-k goes hint-free by default: the bracket comes from the current row

tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py

@@ -404,8 +404,9 @@ def warmup_selfsampling_topk(
             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)
+        # The helper takes max_seq_len in KV-token space;
+        # get_indexer_max_seq_len is compressed, so multiply it back as at
+        # the dispatch seam.
         try:
             _ss_host.warmup_varlen(
                 int(top_k),

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py

@@ -29,6 +29,7 @@
 
 import contextlib
 import sys
+from typing import Any
 
 import cutlass
 import cutlass.cute as cute
@@ -1326,7 +1327,8 @@ def __init__(
         next_n: int = 1,
         cr_shift: int = 0,
         r_const: int = 1,
-    ):
+        hint_free: bool = False,
+    ) -> None:
         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)
@@ -1346,6 +1348,8 @@ def __init__(
         self.next_n = int(next_n)
         self.cr_shift = int(cr_shift)
         self.r_const = int(r_const)
+        # hint-free: gather_hint sites compiled out (sentinel pass-through)
+        self.hint_free = bool(hint_free)
         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
@@ -1973,9 +1977,10 @@ def kern(
         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
+            if cutlass.const_expr(not self.hint_free):
+                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:
@@ -2373,9 +2378,10 @@ def kern(
                         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
-                                )
+                                if cutlass.const_expr(not self.hint_free):
+                                    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)
@@ -2960,19 +2966,21 @@ def __call__(
 _COMPILE_CACHE = {}
 
 
-def get_compiled(tpl, options_extra: str = ""):
+def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) -> Any:
     """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)
+    key = (tuple(tpl), options_extra, bool(hint_free))
     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))
+        kern = GvrMainKernel(
+            blk, u, minb, nbs, kpt, bool(split), bool(tshg), hint_free=bool(hint_free)
+        )
     else:
         blk, u, minb, nbs, kpt, split, tshg, next_n, cr_shift, r_const = tpl
         kern = GvrMainKernel(
@@ -2987,6 +2995,7 @@ def get_compiled(tpl, options_extra: str = ""):
             next_n=next_n,
             cr_shift=cr_shift,
             r_const=r_const,
+            hint_free=bool(hint_free),
         )
     r0, c0 = cute.sym_int(), cute.sym_int()
     r1, c1 = cute.sym_int(), cute.sym_int()
@@ -3309,7 +3318,8 @@ def __init__(
         varlen: bool = False,
         next_n: int = 1,
         cr_shift: int = 0,
-    ):
+        hint_free: bool = False,
+    ) -> None:
         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
@@ -3334,8 +3344,11 @@ def __init__(
         # 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
+        # hint-free: bracket = min/max fold of the first k row values
+        # (already in registers); the hint-gather bracket arms are forced off
+        self.hint_free = bool(hint_free)
+        self.use_bm = (not deg) and (not img) and kpt >= 2 and vpt == 1 and (not hint_free)
+        self.use_img = img and vpt == 1 and (not hint_free)
         self.brl = (minb * blk <= 1024) or (vpt == 1)
 
     # ------------------------------------------------------------------
@@ -3546,7 +3559,7 @@ def kern(
             # ---- hint prefetch: KPT coalesced pre_idx words BEFORE any
             # dependent gather; compiled out under DEG.
             pvs = []
-            if cutlass.const_expr(not self.deg):
+            if cutlass.const_expr(not (self.deg or self.hint_free)):
                 for t in cutlass.range_constexpr(KPT):
                     pv = cutlass.Int32(-1)
                     j = tid + cutlass.Int32(t * self.blk)
@@ -3641,6 +3654,19 @@ def kern(
                 lmin = fkey(lmn)
                 lmax = fkey(lmx)  # monotone
                 cute.arch.barrier()  # bm dies
+            elif cutlass.const_expr(self.hint_free and not self.deg):
+                lmn = cutlass.Float32(_POS_INF)
+                lmx = cutlass.Float32(_NEG_INF__reg)
+                for s in cutlass.range_constexpr(S):
+                    pos = (
+                        (tid + cutlass.Int32((s // 4) * self.blk)) << cutlass.Int32(2)
+                    ) + cutlass.Int32(s % 4)
+                    if pos < k:
+                        v = _val(frags, s)
+                        lmn = fmin_f32(lmn, v)
+                        lmx = fmax_f32(lmx, v)
+                lmin = fkey(lmn)
+                lmax = fkey(lmx)
             elif cutlass.const_expr(self.deg):
                 lmn = cutlass.Float32(_POS_INF)
                 lmx = cutlass.Float32(_NEG_INF__reg)
@@ -4259,10 +4285,18 @@ def __call__(
 _COMPILE_CACHE__reg: dict = {}
 
 
-def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0):
+def get_compiled__reg(
+    tpl: tuple,
+    dump_dir: str | None = None,
+    pdl: bool = False,
+    varlen: bool = False,
+    next_n: int = 1,
+    cr_shift: int = 0,
+    hint_free: bool = False,
+) -> Any:
     """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))
+    key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free))
     compiled = _COMPILE_CACHE__reg.get(key)
     if compiled is None:
         from cutlass.cute import runtime as _crt
@@ -4281,6 +4315,7 @@ def get_compiled__reg(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_
             varlen=varlen,
             next_n=next_n,
             cr_shift=cr_shift,
+            hint_free=hint_free,
         )
         nb_, nc_ = cute.sym_int(), cute.sym_int()
         nb2_, nc2_ = cute.sym_int(), cute.sym_int()
@@ -4429,7 +4464,8 @@ def __init__(
         varlen: bool = False,
         next_n: int = 1,
         cr_shift: int = 0,
-    ):
+        hint_free: bool = False,
+    ) -> None:
         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"
@@ -4444,6 +4480,7 @@ def __init__(
         self.cr_shift = int(cr_shift)
         if self.varlen:
             assert self.next_n >= 1 and self.cr_shift in (0, 2)
+        self.hint_free = bool(hint_free)  # hint-free: gather_hint sites compiled out
         self.lcs = cs.bit_length() - 1  # log2(CS) for the per-row Q shift
         self.blk = blk
         self.u = u
@@ -4999,9 +5036,10 @@ def kern(
                 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
+                if cutlass.const_expr(not self.hint_free):
+                    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:
@@ -5230,9 +5268,10 @@ def kern(
                         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
+                                if cutlass.const_expr(not self.hint_free):
+                                    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)
@@ -5604,24 +5643,44 @@ def __call__(
 
 
 def get_compiled__clus(
-    tpl,
+    tpl: tuple,
     scap: int = 8192,
     cmp_: int = 2048,
     options_extra: str = "",
     varlen: bool = False,
     next_n: int = 1,
     cr_shift: int = 0,
-):
+    hint_free: bool = False,
+) -> Any:
     """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))
+    key = (
+        tuple(tpl),
+        scap,
+        cmp_,
+        options_extra,
+        bool(varlen),
+        int(next_n),
+        int(cr_shift),
+        bool(hint_free),
+    )
     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
+        blk,
+        u,
+        minb,
+        nbs,
+        cs,
+        scap=scap,
+        cmp_=cmp_,
+        varlen=varlen,
+        next_n=next_n,
+        cr_shift=cr_shift,
+        hint_free=hint_free,
     )
     r0, c0 = cute.sym_int(), cute.sym_int()
     r1, c1 = cute.sym_int(), cute.sym_int()
@@ -5839,7 +5898,8 @@ def __init__(
         varlen: bool = False,
         next_n: int = 1,
         cr_shift: int = 0,
-    ):
+        hint_free: bool = False,
+    ) -> None:
         assert blk == BLKC, "all instantiations BLK=BLKC=1024"
         assert vpt in (1, 2, 4) and cs in (2, 4, 8)
         self.blk = blk
@@ -5855,6 +5915,8 @@ def __init__(
         self.cr_shift = int(cr_shift)
         if self.varlen:
             assert self.next_n >= 1 and self.cr_shift in (0, 2)
+        # hint-free: P0 samples the first k row elements (coalesced) instead of the hint
+        self.hint_free = bool(hint_free)
         self.S = vpt * 4
         self.span = blk * vpt  # float4 per CTA
 
@@ -6010,8 +6072,12 @@ def kern(
             # ---- 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)
+            if cutlass.const_expr(self.hint_free):
+                if tid < k:
+                    pv0 = tid
+            else:
+                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
@@ -6513,16 +6579,31 @@ def __call__(
 _COMPILE_CACHE__regclus: dict = {}
 
 
-def get_compiled__regclus(tpl, dump_dir=None, pdl=False, varlen=False, next_n=1, cr_shift=0):
+def get_compiled__regclus(
+    tpl: tuple,
+    dump_dir: str | None = None,
+    pdl: bool = False,
+    varlen: bool = False,
+    next_n: int = 1,
+    cr_shift: int = 0,
+    hint_free: bool = False,
+) -> Any:
     """Compile (or fetch) the variant for constexpr tuple (BLK, VPT, CS)."""
-    key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift))
+    key = (tuple(tpl), bool(pdl), bool(varlen), int(next_n), int(cr_shift), bool(hint_free))
     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
+            blk,
+            vpt,
+            cs,
+            pdl=pdl,
+            varlen=varlen,
+            next_n=next_n,
+            cr_shift=cr_shift,
+            hint_free=hint_free,
         )
         nb_, nc_ = cute.sym_int(), cute.sym_int()
         nb2_, nc2_ = cute.sym_int(), cute.sym_int()