diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 4038f452481a..f90ef8596940 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -720,6 +720,16 @@ def __init__( decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, ) + # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8 + # indexer epilogue emits candidates the GVR Top-K consumes (see + # gvr_emission / gvr_routing; state lives on the TopK module) + # only the FP4 scoring op accepts emission kwargs + self.use_gvr_emission = ( + os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1" + and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR + and self.use_cute_dsl_paged_mqa_logits + and self.use_fp4 + ) # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) @@ -1596,6 +1606,13 @@ def sparse_attn_indexer( gvr_prior_indices, request_offset=num_generations, ) + if self.use_gvr_emission: + # reused slots cold-start the emission closed loop; stale + # lines only mis-place cuts - counts are re-measured + # in-kernel, so exactness never rides on this reset + self.top_k.reset_gvr_emission_rows( + slice(num_generations, num_generations + num_contexts) + ) reuse_topk = ( self.mtp_index_share @@ -1688,6 +1705,25 @@ def sparse_attn_indexer( metadata.dsl_expand_factor > 1 and next_n == metadata.dsl_expand_factor * metadata.dsl_atom ) + gvr_emit_kwargs: dict = {} + # emitting for a step the Top-K cannot consume only churns + # the closed-loop state, so gate on the consumable shape + if ( + self.use_gvr_emission + and gvr_prior_indices is not None + and next_n == 1 + and not dsl_atom_split + and num_gen_tokens <= 256 + # ext tiers are single-CTA/sort-path only; row reordering + # routes the Top-K through order_row, which excludes them + and metadata.kv_lens_row_reorder is None + ): + gvr_emit_kwargs = self.top_k.prepare_gvr_emission( + num_generations, + indexer_max_seq_len, + torch.cuda.get_device_properties(q_decode.device).multi_processor_count, + gvr_prior_indices, + ) if self.use_fp4: # FP4 DSL signature splits DG's (q, sf_q) tuple into two # separate args and requires q.dtype == uint8 (q_decode @@ -1720,6 +1756,7 @@ def sparse_attn_indexer( dsl_block_table, dsl_schedule_meta, indexer_max_seq_len, + **gvr_emit_kwargs, ) else: # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. @@ -1747,6 +1784,7 @@ def sparse_attn_indexer( fp8_block_table, fp8_schedule_meta, indexer_max_seq_len, + **gvr_emit_kwargs, ) else: decode_q_scale = ( diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index face56e566c9..bbaa63c14f00 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7346,12 +7346,22 @@ def _compile( return_output_values: bool, cluster_size: int, seqlen_sorted: bool, + enable_block_skip: bool = False, + use_ext_counts: bool = False, + emit_xstate: bool = False, + use_ext_cand: bool = False, + ext_rungs: bool = False, + cand_cap: int = 5120, + accept_cap: Optional[int] = None, + kc_override: Optional[int] = None, ) -> tuple: key = (dtype, top_k, next_n, enable_unroll_4, enable_phase3_unroll, use_constant_hint, min_blocks_per_mp, use_256bit_load, num_threads_per_block, enable_warp_parallel_reduce, compress_ratio, return_output_values, cluster_size, - seqlen_sorted) + seqlen_sorted, enable_block_skip, use_ext_counts, + emit_xstate, use_ext_cand, ext_rungs, cand_cap, accept_cap, + kc_override) if key in cls.kernel_cache: return key n_rows = cute.sym_int() @@ -7384,6 +7394,31 @@ def _compile( order_row_fake = (cute.runtime.make_fake_compact_tensor( cutlass.Int32, (n_batch, ), stride_order=(0, )) if seqlen_sorted else None) + # emission-assisted tier fake tensors (list/counts/rungs) + block_max_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, cute.sym_int()), + stride_order=(1, 0), + assumed_align=16) if enable_block_skip else None) + seed_thr_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, 8 if use_ext_counts else 3), + stride_order=(1, 0), + assumed_align=4) if (use_ext_counts or ext_rungs) else None) + xstate_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, 8), + stride_order=(1, 0), + assumed_align=4) if emit_xstate else None) + cand_vals_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, cand_cap), + stride_order=(1, 0), + assumed_align=4) if use_ext_cand else None) + cand_idx_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (n_rows, cand_cap), + stride_order=(1, 0), + assumed_align=4) if use_ext_cand else None) + cand_ctl_fake = (cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (n_rows, 4), + stride_order=(1, 0), + assumed_align=8) if use_ext_cand else None) fake_stream = cute.runtime.make_fake_stream( use_tvm_ffi_env_stream=True) @@ -7402,6 +7437,18 @@ def _compile( return_output_values=return_output_values, cluster_size=cluster_size, seqlen_sorted=seqlen_sorted, + enable_block_skip=enable_block_skip, + use_ext_counts=use_ext_counts, + emit_xstate=emit_xstate, + use_ext_cand=use_ext_cand, + ext_rungs=ext_rungs, + cand_cap=cand_cap, + accept_cap=accept_cap, + kc_override=kc_override, + # ext modes need 3 rung slots (M_thr == 3); only the + # slot count matters, not the qfrac values (P1b skipped) + r0_qfracs=((0.85, 0.35) if + (use_ext_counts or ext_rungs) else None), ) cls.kernel_cache[key] = cute.compile( kernel, @@ -7412,6 +7459,13 @@ def _compile( out_indices_fake, order_row_fake, stream=fake_stream, + block_max=block_max_fake, + seed_thr=seed_thr_fake, + seed_counts=None, + xstate=xstate_fake, + cand_vals=cand_vals_fake, + cand_idx=cand_idx_fake, + cand_ctl=cand_ctl_fake, options="--enable-tvm-ffi", ) logger.debug(f"[compile cute_dsl gvr_topk_decode] {key}") @@ -7521,6 +7575,15 @@ def forward( order_row: Optional[torch.Tensor] = None, counters: Optional[torch.Tensor] = None, max_batch_size: Optional[int] = None, + seed_thr: Optional[torch.Tensor] = None, + xstate: Optional[torch.Tensor] = None, + cand_vals: Optional[torch.Tensor] = None, + cand_idx: Optional[torch.Tensor] = None, + cand_ctl: Optional[torch.Tensor] = None, + block_max: Optional[torch.Tensor] = None, + num_threads: Optional[int] = None, + accept_cap: Optional[int] = None, + kc_override: Optional[int] = None, ) -> None: """Three paths, picked by ``(counters, order_row)``: @@ -7543,9 +7606,10 @@ def forward( # Host-only guard — no device sync. The op signature and output # contract are unchanged (unordered int32 indices, -1 pad only # for degenerate rows). - if _is_tiered_topk_supported(logits, pre_idx, seq_lens, - output_indices, top_k, next_n, - compress_ratio, order_row, counters): + if (seed_thr is None and cand_vals is None and xstate is None + and block_max is None and _is_tiered_topk_supported( + logits, pre_idx, seq_lens, output_indices, top_k, + next_n, compress_ratio, order_row, counters)): _tiered_topk(logits, pre_idx, seq_lens, output_indices, top_k, next_n, compress_ratio) return @@ -7647,6 +7711,39 @@ def forward( ), ("order_row must be int32, CUDA, shape == seq_lens.shape " f"(={tuple(seq_lens.shape)}); got dtype={order_row.dtype} " f"shape={tuple(order_row.shape)}") + # emission-assisted tiers: mode from which ext tensors the + # caller handed in (see gvr_routing.plan_emission) + if seed_thr is not None: + assert seed_thr.shape[1] == 3 or seed_thr.shape[1] == 8, ( + "seed_thr must be [rows, 3] (rungs) or [rows, 8] " + "(packed lines + counts row, the width the kernel " + f"is compiled for); got width {seed_thr.shape[1]}") + use_ext_counts = seed_thr is not None and seed_thr.shape[1] == 8 + ext_rungs = seed_thr is not None and seed_thr.shape[1] == 3 + use_ext_cand = cand_vals is not None + enable_block_skip = block_max is not None + emit_xstate = xstate is not None + if use_ext_cand: + assert use_ext_counts, ( + "candidate list requires the packed seed row " + "([rows, 8]: lines + counts)") + assert (cand_idx is not None and cand_ctl is not None + and cand_vals.shape == cand_idx.shape + and cand_ctl.shape == (num_rows, 4)), ( + "list tier needs cand_vals/cand_idx same shape " + "+ cand_ctl [rows, 4]") + if (use_ext_counts or ext_rungs or use_ext_cand + or enable_block_skip): + assert not lb_mode and order_row is None, ( + "ext tiers are single-CTA/sort-path only") + assert logits.dtype == torch.float32 and next_n == 1, ( + "ext tiers are compiled for fp32 logits and next_n==1; " + f"got dtype={logits.dtype} next_n={next_n}") + if num_threads is not None: + tuning = dict(tuning, num_threads_per_block=num_threads) + elif use_ext_cand and top_k <= 512: + # small-K list rule: hit rows do O(list) work + tuning = dict(tuning, num_threads_per_block=512) key = cls._compile( cute_dtype, top_k, @@ -7655,17 +7752,29 @@ def forward( return_output_values=return_output_values, cluster_size=cluster_size, seqlen_sorted=seqlen_sorted, + enable_block_skip=enable_block_skip, + use_ext_counts=use_ext_counts, + emit_xstate=emit_xstate, + use_ext_cand=use_ext_cand, + ext_rungs=ext_rungs, + cand_cap=(cand_vals.shape[1] if use_ext_cand else 5120), + accept_cap=accept_cap, + kc_override=kc_override, **tuning, ) cls.kernel_cache[key](logits, pre_idx, seq_lens, None, - output_indices, order_row) + output_indices, order_row, block_max, + seed_thr, None, xstate, cand_vals, cand_idx, + cand_ctl) # TODO(dsa.py): wire ``order_row = argsort(seq_lens, descending=True)`` # (device-side, graph-safe) into the LJF row-reorder branch when # ``num_rows >= 2 * num_sms``. Physical meaning: wave-2 must fit a # full SM-row's worth of CTAs so the sort has long-vs-short rows to # swap. Below that threshold the win is noise / can regress a few - # percent (measured, N in {8K,16K,32K}). + # percent (B200 N∈{8K,16K,32K} sweep 2026-06-23). + # xstate is written by the kernel but stays out of mutates_args + # (optional-mutate None-default IndexError; see the fp4 op note) @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode", mutates_args=("output_indices", ), device_types="cuda") @@ -7682,6 +7791,15 @@ def cute_dsl_gvr_topk_decode( order_row: Optional[torch.Tensor] = None, counters: Optional[torch.Tensor] = None, max_batch_size: Optional[int] = None, + seed_thr: Optional[torch.Tensor] = None, + xstate: Optional[torch.Tensor] = None, + cand_vals: Optional[torch.Tensor] = None, + cand_idx: Optional[torch.Tensor] = None, + cand_ctl: Optional[torch.Tensor] = None, + block_max: Optional[torch.Tensor] = None, + num_threads: Optional[int] = None, + accept_cap: Optional[int] = None, + kc_override: Optional[int] = None, ) -> None: """CuTe DSL GVR (Guess-Verify-Refine) Top-K decode for Blackwell. @@ -7715,6 +7833,16 @@ def cute_dsl_gvr_topk_decode( max_batch_size: Required with ``counters``; ignored otherwise. Power of 2 in ``[64, 1024]``, must match the value passed to LB prepare. + + Of the optional hint tensors the kernel WRITES ``xstate`` (the + closed-loop publish); it cannot be declared in ``mutates_args``: + torch.library raises IndexError when a declared-mutable Optional + arg is None at call time (re-verified on the pinned torch), and + most calls pass no hints. Under torch.compile/functionalization + the undeclared write is invisible, so the hint path is eager / + CUDA-graph only. ``TRTLLM_GVR_EMISSION=1`` gates the + emission-assisted wiring that feeds these tensors (opt-in, + experimental). """ if not is_sm_100f(): raise ValueError( @@ -7752,6 +7880,15 @@ def cute_dsl_gvr_topk_decode( order_row=order_row, counters=counters, max_batch_size=max_batch_size, + seed_thr=seed_thr, + xstate=xstate, + cand_vals=cand_vals, + cand_idx=cand_idx, + cand_ctl=cand_ctl, + block_max=block_max, + num_threads=num_threads, + accept_cap=accept_cap, + kc_override=kc_override, ) @torch.library.register_fake("trtllm::cute_dsl_gvr_topk_decode") @@ -7768,6 +7905,15 @@ def _( order_row: Optional[torch.Tensor] = None, counters: Optional[torch.Tensor] = None, max_batch_size: Optional[int] = None, + seed_thr: Optional[torch.Tensor] = None, + xstate: Optional[torch.Tensor] = None, + cand_vals: Optional[torch.Tensor] = None, + cand_idx: Optional[torch.Tensor] = None, + cand_ctl: Optional[torch.Tensor] = None, + block_max: Optional[torch.Tensor] = None, + num_threads: Optional[int] = None, + accept_cap: Optional[int] = None, + kc_override: Optional[int] = None, ) -> None: return None @@ -8947,11 +9093,21 @@ def _compile(cls, num_epi_subtiles, epi_dtype, output_dtype, - remove_online_sf_transpose=False): + remove_online_sf_transpose=False, + emit_block_meta=False, + emit_hit_stats=True, + emit_seed_counts=False, + seed_packed=False, + emit_cand=False, + cand_cap=5120, + emit_cand_bucketed=False, + accept_cap=8192): """Compile kernel using fake tensors + TVM FFI.""" key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n, num_sms, num_epi_subtiles, epi_dtype, output_dtype, - remove_online_sf_transpose) + remove_online_sf_transpose, emit_block_meta, emit_hit_stats, + emit_seed_counts, seed_packed, emit_cand, cand_cap, + emit_cand_bucketed, accept_cap) if key in cls.kernel_cache: return @@ -9014,6 +9170,81 @@ def _compile(cls, (num_ctas, 2), stride_order=(1, 0)) + # Block-meta tensors (fused-GVR support): nb_pad*4 records. + block_max_fake = None + hit_stats_fake = None + hit_bitmap_fake = None + if emit_block_meta: + nb_sym = cute.sym_int() + block_max_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), nb_sym), + stride_order=(1, 0), + assumed_align=16) + if emit_hit_stats: + hit_stats_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 4), + stride_order=(1, 0), + assumed_align=16) + hit_bitmap_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_B, cute.sym_int()), + stride_order=(1, 0), + assumed_align=16) + cand_fake = None + cand_ctl_fake = None + if emit_cand: + cand_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), cand_cap * 2), + stride_order=(1, 0), + assumed_align=8) + cand_ctl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 2), + stride_order=(1, 0), + assumed_align=8) + cand_idx_fake = None + cand_cur_fake = None + if emit_cand_bucketed: + # bucketed SoA: cand slot reused as the fp32 VALUES tensor + wtot = 2 * accept_cap + cand_cap + cand_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), wtot), + stride_order=(1, 0), + assumed_align=4) + cand_idx_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), wtot), + stride_order=(1, 0), + assumed_align=4) + cand_ctl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 4), + stride_order=(1, 0), + assumed_align=4) + cand_cur_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 4), + stride_order=(1, 0), + assumed_align=4) + seed_thr_fake = None + seed_counts_fake = None + if emit_seed_counts: + if seed_packed: + # [rows, 8] fp32 packed seed row: lines at cols 0..2, + # counts at cols 3..5; same tensor bound to both params. + seed_thr_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 8), + stride_order=(1, 0), + assumed_align=4) + seed_counts_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 8), + stride_order=(1, 0), + assumed_align=4) + else: + seed_thr_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 3), + stride_order=(1, 0), + assumed_align=4) + seed_counts_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 3), + stride_order=(1, 0), + assumed_align=4) + fake_stream = cute.runtime.make_fake_stream( use_tvm_ffi_env_stream=True) @@ -9028,6 +9259,14 @@ def _compile(cls, epi_dtype=to_cutlass[epi_dtype], output_dtype=to_cutlass[output_dtype], remove_online_sf_transpose=remove_online_sf_transpose, + emit_block_meta=emit_block_meta, + emit_hit_stats=emit_hit_stats, + emit_seed_counts=emit_seed_counts, + seed_packed=seed_packed, + emit_cand=emit_cand, + cand_cap=cand_cap, + emit_cand_bucketed=emit_cand_bucketed, + accept_cap=accept_cap, ) compiled = cute.compile( @@ -9042,7 +9281,18 @@ def _compile(cls, sm_fake, cutlass.Int32(1), cutlass.Int32(1), + # keep __call__'s argument order: stream, then emission + # slots (the runtime call drops the stream, same sequence) fake_stream, + block_max_fake, + hit_stats_fake, + hit_bitmap_fake, + seed_thr=seed_thr_fake, + seed_counts=seed_counts_fake, + cand=cand_fake, + cand_ctl=cand_ctl_fake, + cand_idx_t=cand_idx_fake, + cand_cur=cand_cur_fake, options="--enable-tvm-ffi", ) cls.kernel_cache[key] = compiled @@ -9063,6 +9313,21 @@ def forward( epi_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, remove_online_sf_transpose: bool = False, + emit_block_meta: bool = False, + emit_hit_stats: bool = True, + hit_bitmap: Optional[torch.Tensor] = None, + block_max_out: Optional[torch.Tensor] = None, + hit_stats_out: Optional[torch.Tensor] = None, + emit_seed_counts: bool = False, + seed_thr: Optional[torch.Tensor] = None, + seed_counts_out: Optional[torch.Tensor] = None, + emit_cand: bool = False, + cand_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + emit_cand_bucketed: bool = False, + accept_cap: int = 8192, + cand_idx_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Execute FP4 paged MQA logits kernel. @@ -9078,8 +9343,25 @@ def forward( num_epi_subtiles: epilogue sub-tile count (1, 2, or 4) epi_dtype: epilogue compute dtype output_dtype: output logits dtype + emit_block_meta: also emit per-128-block metadata for the + fused GVR top-k. Requires ``hit_bitmap`` + [B, >= nb_pad*4] int32 (1 bit per compressed kv + position, request-level). ``block_max_out`` + [B*next_n, nb_pad*4] fp32 (4 warp-partial records per + block) is allocated when not supplied; + ``hit_stats_out`` [B*next_n, 4] fp32 must be supplied + when ``emit_hit_stats`` is set. Returns: - logits: [B*next_n, max_context_len] output_dtype + logits [B*next_n, max_context_len]; with emit_block_meta, + the tuple (logits, block_max, hit_stats). + + The optional emission tensors (``block_max_out`` / + ``seed_thr`` / ``cand_*``) are written by the kernel but + cannot be declared in ``mutates_args``: torch.library raises + IndexError when a declared-mutable Optional arg is None at + call time (re-verified on the pinned torch), and plain + logits-only calls pass none of them. Emission is therefore + eager / CUDA-graph only. """ B, next_n, H, half_D = q.shape N = next_n * H @@ -9132,10 +9414,154 @@ def forward( ) logits = logits[:, :max_context_len] + # Block-meta buffers (fused-GVR support). nb_pad mirrors the + # logits padding so WG1's odd-num_kv OOB tile lands in padding. + if emit_block_meta: + nb_pad = aligned_max_ctx // compute_block_kv + # 4 warp-partial records per block (see FP4MQALogitsKernel). + nrec = nb_pad * 4 + if emit_hit_stats: + assert ( + hit_bitmap is not None + and hit_bitmap.dtype == torch.int32 + and hit_bitmap.is_cuda and hit_bitmap.is_contiguous() + and hit_bitmap.dim() == 2 and hit_bitmap.shape[0] == B + and hit_bitmap.shape[1] >= nb_pad * 4 + ), (f"emit_hit_stats requires hit_bitmap int32 " + f"[{B}, >= {nb_pad * 4}]; got " + f"{None if hit_bitmap is None else (hit_bitmap.dtype, tuple(hit_bitmap.shape))}" + ) + # Per-row aggregate {enc_min, enc_max, sum, cnt}. The + # kernel only ATOMICALLY MERGES into this buffer — the + # caller must pre-initialize it to the identities + # {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} each step. + assert hit_stats_out is not None, ( + "emit_hit_stats requires a caller-initialized " + "hit_stats_out [B*next_n, 4] fp32 (identity-filled)") + assert (hit_stats_out.shape == (B * next_n, 4) + and hit_stats_out.is_contiguous()) + else: + hit_bitmap = None + hit_stats_out = None + if block_max_out is None: + block_max_out = torch.empty((B * next_n, nrec), + device=q.device, + dtype=torch.float32) + assert (block_max_out.shape == (B * next_n, nrec) + and block_max_out.is_contiguous()) + + seed_packed = False + if emit_seed_counts: + assert emit_block_meta, ( + "emit_seed_counts requires emit_block_meta") + if seed_counts_out is None: + # Packed contract: seed_thr IS the [rows, 8] fp32 seed + # row; lines at cols 0..2, counts accumulate as fp32 at + # cols 3..5. Caller zeroes cols 3..7 and writes lines + # each step. + seed_packed = True + assert ( + seed_thr is not None and seed_thr.dtype == torch.float32 + and seed_thr.is_cuda and seed_thr.is_contiguous() + and seed_thr.shape == (B * next_n, 8) + ), (f"packed emit_seed_counts requires seed_thr fp32 " + f"[{B * next_n}, 8]; got " + f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" + ) + seed_counts_out = seed_thr + else: + # Legacy split contract: 3 thresholds per row (fp32), + # counts accumulated with red.global.add.s32 into a + # caller-zeroed int32 [rows, 3]. + assert ( + seed_thr is not None and seed_thr.dtype == torch.float32 + and seed_thr.is_cuda and seed_thr.is_contiguous() + and seed_thr.shape == (B * next_n, 3) + ), (f"emit_seed_counts requires seed_thr fp32 " + f"[{B * next_n}, 3]; got " + f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" + ) + assert (seed_counts_out.dtype == torch.int32 + and seed_counts_out.is_cuda + and seed_counts_out.is_contiguous() + and seed_counts_out.shape == (B * next_n, 3)), ( + "emit_seed_counts requires a caller-zeroed " + "seed_counts_out int32 [B*next_n, 3]") + else: + seed_thr = None + seed_counts_out = None + + cand_cap = 0 + if emit_cand: + assert emit_seed_counts, ( + "emit_cand requires emit_seed_counts (t_0 threshold)") + # Unordered (value, index) pair scatter; the caller zeroes + # cand_ctl_out {claimed, void} each step. cand_out is + # [B*next_n, CAP*2] int32 (fp32 bits in even words). + assert ( + cand_out is not None and cand_out.dtype == torch.int32 + and cand_out.is_cuda and cand_out.is_contiguous() + and cand_out.dim() == 2 and cand_out.shape[0] == B * next_n + and cand_out.shape[1] % 2 == 0 and cand_out.shape[1] > 0), ( + "emit_cand requires cand_out int32 [B*next_n, CAP*2]") + assert (cand_ctl_out is not None + and cand_ctl_out.dtype == torch.int32 + and cand_ctl_out.is_cuda + and cand_ctl_out.is_contiguous() + and cand_ctl_out.shape == (B * next_n, 2)), ( + "emit_cand requires a caller-zeroed cand_ctl_out " + "int32 [B*next_n, 2]") + cand_cap = cand_out.shape[1] // 2 + elif emit_cand_bucketed: + assert emit_seed_counts, ( + "emit_cand_bucketed requires emit_seed_counts") + # SoA contract: cand_out = fp32 VALUES [rows, 2*segA+capC], + # cand_idx_out = int32 positions (same width), cand_cur_out = + # int32 [rows, 4] cursors (caller-zeroed), cand_ctl_out = + # int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed) + assert cand_out is not None, ( + "emit_cand_bucketed requires cand_out") + W = cand_out.shape[1] + assert ( + cand_out.dtype == torch.float32 and cand_out.is_cuda + and cand_out.is_contiguous() and cand_out.dim() == 2 + and cand_out.shape[0] == B * next_n + and W > 2 * accept_cap), ( + "bucketed requires cand_out fp32 [rows, 2*segA+capC]") + assert (cand_idx_out is not None + and cand_idx_out.dtype == torch.int32 + and cand_idx_out.is_cuda + and cand_idx_out.is_contiguous() + and cand_idx_out.shape == cand_out.shape), ( + "bucketed requires cand_idx_out int32, same shape") + assert (cand_ctl_out is not None + and cand_ctl_out.dtype == torch.int32 + and cand_ctl_out.is_cuda + and cand_ctl_out.is_contiguous() + and cand_ctl_out.shape == (B * next_n, 4)), ( + "bucketed requires caller-zeroed cand_ctl_out " + "int32 [rows, 4]") + assert (cand_cur_out is not None + and cand_cur_out.dtype == torch.int32 + and cand_cur_out.is_cuda + and cand_cur_out.is_contiguous() + and cand_cur_out.shape == (B * next_n, 4)), ( + "bucketed requires caller-zeroed cand_cur_out " + "int32 [rows, 4]") + cand_cap = W - 2 * accept_cap + else: + cand_out = None + cand_ctl_out = None + if not emit_cand_bucketed: + cand_idx_out = None + cand_cur_out = None + # Compile if needed (fake tensors, no real data required) key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms, num_epi_subtiles, epi_dtype, output_dtype, - remove_online_sf_transpose) + remove_online_sf_transpose, emit_block_meta, emit_hit_stats, + emit_seed_counts, seed_packed, emit_cand, cand_cap, + emit_cand_bucketed, accept_cap) if key not in cls.kernel_cache: cls._compile( compute_block_kv, @@ -9147,14 +9573,32 @@ def forward( num_epi_subtiles, epi_dtype, output_dtype, - remove_online_sf_transpose=remove_online_sf_transpose) + remove_online_sf_transpose=remove_online_sf_transpose, + emit_block_meta=emit_block_meta, + emit_hit_stats=emit_hit_stats, + emit_seed_counts=emit_seed_counts, + seed_packed=seed_packed, + emit_cand=emit_cand, + cand_cap=cand_cap, + emit_cand_bucketed=emit_cand_bucketed, + accept_cap=accept_cap) compiled = cls.kernel_cache[key] # TVM FFI: pass raw tensors, no dlpack/stream needed + if emit_block_meta: + compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table, + context_lens, schedule_meta, num_phys_blocks, B, + block_max_out, hit_stats_out, hit_bitmap, seed_thr, + seed_counts_out, cand_out, cand_ctl_out, cand_idx_out, + cand_cur_out) + return logits, block_max_out, hit_stats_out compiled(kv_flat, q_3d, sf_q_2d, w_2d, logits, block_table, - context_lens, schedule_meta, num_phys_blocks, B) + context_lens, schedule_meta, num_phys_blocks, B, None, + None, None, None, None, None, None, None, None) return logits + # NOTE: the optional emission tensors ARE written by the kernel but must + # stay out of mutates_args (torch.library IndexErrors on None defaults). @torch.library.custom_op("trtllm::cute_dsl_fp4_paged_mqa_logits", mutates_args=(), device_types="cuda") @@ -9171,6 +9615,13 @@ def cute_dsl_fp4_paged_mqa_logits( epi_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, remove_online_sf_transpose: bool = False, + block_max_out: Optional[torch.Tensor] = None, + seed_thr: Optional[torch.Tensor] = None, + cand_out: Optional[torch.Tensor] = None, + cand_idx_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, + accept_cap: int = 8192, ) -> torch.Tensor: if not is_sm_100f(): raise ValueError( @@ -9196,7 +9647,7 @@ def cute_dsl_fp4_paged_mqa_logits( f"epi_dtype={epi_dtype} output_dtype={output_dtype}", key="cute_dsl_fp4_paged_mqa_logits_inputs", ) - return CuteDSLFP4PagedMQALogitsRunner.forward( + ret = CuteDSLFP4PagedMQALogitsRunner.forward( q, sf_q, kv_fused, @@ -9208,7 +9659,21 @@ def cute_dsl_fp4_paged_mqa_logits( num_epi_subtiles=num_epi_subtiles, epi_dtype=epi_dtype, output_dtype=output_dtype, - remove_online_sf_transpose=remove_online_sf_transpose) + remove_online_sf_transpose=remove_online_sf_transpose, + emit_block_meta=block_max_out is not None, + emit_hit_stats=False, + block_max_out=block_max_out, + emit_seed_counts=seed_thr is not None, + seed_thr=seed_thr, + emit_cand_bucketed=cand_out is not None, + accept_cap=accept_cap, + cand_out=cand_out, + cand_idx_out=cand_idx_out, + cand_ctl_out=cand_ctl_out, + cand_cur_out=cand_cur_out) + # with emission on the runner returns a tuple; the op returns + # logits only (emission buffers are caller-owned) + return ret[0] if isinstance(ret, tuple) else ret @torch.library.register_fake("trtllm::cute_dsl_fp4_paged_mqa_logits") def _( @@ -9224,6 +9689,13 @@ def _( epi_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, remove_online_sf_transpose: bool = False, + block_max_out: Optional[torch.Tensor] = None, + seed_thr: Optional[torch.Tensor] = None, + cand_out: Optional[torch.Tensor] = None, + cand_idx_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, + accept_cap: int = 8192, ) -> torch.Tensor: B = q.shape[0] next_n = q.shape[1] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py index 1516eafc22ec..89f32856d77c 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py @@ -57,13 +57,111 @@ from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cutlass_dsl import dsl_user_op +from cutlass.cutlass_dsl import T, dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait # CuTe DSL CUDA 13 validates rounding modes as string literals. The string # form is also accepted by older wrappers, so keep it version-independent. _RND_RN = "rn" +# Reduction identities for emit_block_meta; must match the GVR kernel's +# FLT_MAX/NEG_FLT_MAX sentinels (gvr_topk_decode.py). +_META_FLT_MAX = 3.4028235e38 +_META_NEG_FLT_MAX = -3.4028235e38 + + +# Global-memory reductions for the per-row hit aggregate (emit_hit_stats). +# fp32 min/max use the order-preserving int encoding +# enc(f) = bits(f) >= 0 ? bits(f) : bits(f) ^ 0x7FFFFFFF (an involution) +# with red.global.{min,max}.s32; sum uses red.global.add.f32. +@dsl_user_op +def _red_global_fmin_ordered(addr_i64, fval, *, loc=None, ip=None): + llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b32 k;\n\t" + ".reg .pred p;\n\t" + "mov.b32 k, $1;\n\t" + "setp.lt.s32 p, k, 0;\n\t" + "@p xor.b32 k, k, 0x7FFFFFFF;\n\t" + "red.global.min.s32 [$0], k;\n\t" + "}", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _red_global_fmax_ordered(addr_i64, fval, *, loc=None, ip=None): + llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b32 k;\n\t" + ".reg .pred p;\n\t" + "mov.b32 k, $1;\n\t" + "setp.lt.s32 p, k, 0;\n\t" + "@p xor.b32 k, k, 0x7FFFFFFF;\n\t" + "red.global.max.s32 [$0], k;\n\t" + "}", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_global_add_s32(addr_i64, ival, *, loc=None, ip=None): + """atom.global.add.s32 returning the OLD value (warp batch-claim).""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [addr_i64.ir_value(loc=loc, ip=ip), ival.ir_value(loc=loc, ip=ip)], + "atom.global.add.s32 $0, [$1], $2;", + "=r,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_global_add_s32(addr_i64, ival, *, loc=None, ip=None): + llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), ival.ir_value(loc=loc, ip=ip)], + "red.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def _red_global_add_f32(addr_i64, fval, *, loc=None, ip=None): + llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), fval.ir_value(loc=loc, ip=ip)], + "red.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + @dsl_user_op def pack_f16x2( @@ -363,6 +461,14 @@ def __init__( output_dtype=cutlass.Float32, remove_online_sf_transpose: bool = False, use_batched_store: bool = True, + emit_block_meta: bool = False, + emit_hit_stats: bool = True, + emit_seed_counts: bool = False, + seed_packed: bool = False, + emit_cand: bool = False, + cand_cap: int = 5120, + emit_cand_bucketed: bool = False, + accept_cap: int = 8192, ): # Static FP4 invariants — see plan Sanity checklist. assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget" @@ -407,6 +513,71 @@ def __init__( # When True, defer per-t STG to register array and emit all STGs in # one contiguous LSU phase after the for-t loop (epilogue micro-opt). self.use_batched_store = use_batched_store + # When True, the epilogue additionally emits per-128-token-block + # metadata consumed by the fused GVR top-k (gvr_topk_decode.py). + # Emission is warp-autonomous: each of the WG's 4 warps writes one + # partial record per tile per t (record index = tile*4 + warp); + # the GVR consumer folds the 4 partials per block. No cross-warp + # barrier. + # block_max [num_rows, nb_pad*4] fp32 — warp-partial max of + # f32(stored logit) over valid positions (kv_pos < ctx), + # computed on the POST-conversion value so it bounds what GVR + # reads back bit-exactly. + # hit_agg [num_rows, 4] fp32 — per-row aggregate + # {enc_min, enc_max, sum, cnt} of stored logits at positions + # flagged in hit_bitmap; min/max slots hold the + # order-preserving int encoding (see _red_global_fmin_ordered). + # Buffer must be pre-initialized to + # {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step. + # hit_bitmap [batch, nb_pad*4] int32 — 1 bit per kv position + # (request-level; from the previous step's top-k). + # Hit accumulators are lane-local and flush once per q-transition + # via atomics (_flush_hit_agg). + # + # emit_hit_stats sub-knob (only meaningful with emit_block_meta): + # False emits block_max ONLY (no bitmap read, no hit aggregate). + self.emit_block_meta = emit_block_meta + self.emit_hit_stats = emit_hit_stats + # emit_seed_counts (requires emit_block_meta): per row, count + # stored logits >= each of the T=3 caller-provided thresholds. + # Counts are computed on the POST-conversion value over valid + # positions only. + if emit_seed_counts and not emit_block_meta: + raise ValueError("emit_seed_counts requires emit_block_meta") + self.emit_seed_counts = emit_seed_counts + # seed_packed: single [num_rows, 8] fp32 seed row per the top-k + # pre-packed contract - lines at cols 0..2, counts accumulated as + # floats at cols 3..5 (exact to 2^24; red.global.add.f32). The + # caller zeroes cols 3..7 and writes the lines each step. + if seed_packed and not emit_seed_counts: + raise ValueError("seed_packed requires emit_seed_counts") + self.seed_packed = seed_packed + # emit_cand: unordered pre-collect of all (value, index) pairs >= + # the t_0 seed threshold. claimed >= K certifies the candidate set + # covers the true top-K. + if emit_cand and not emit_seed_counts: + raise ValueError("emit_cand requires emit_seed_counts (t_0 source)") + self.emit_cand = emit_cand + self.cand_cap = cand_cap + # emit_cand_bucketed: three fixed SoA segments (A=[0,segA) holds + # >= t2, B=[segA,2segA) holds [t1,t2), C=[2segA,2segA+capC) holds + # [t0,t1)); a full segment spills to the next looser one. A/B use + # EXACT ballot claims (their prefixes must stay pad-free - the + # consumer's prefix math assumes it), C keeps the claim-window + # scheme (pads are legal there). Cursors live in caller-zeroed + # cand_cur [rows,4]; ctl [rows,4] carries {n0 incl C pads, void, + # n1, n2} with n1/n2 flushed from the seed counters. + if emit_cand_bucketed and not emit_seed_counts: + raise ValueError("emit_cand_bucketed requires emit_seed_counts") + if emit_cand_bucketed and emit_cand: + raise ValueError("emit_cand_bucketed and emit_cand are exclusive") + self.emit_cand_bucketed = emit_cand_bucketed + self.accept_cap = accept_cap + # Per-warp claim window: one atomic claims (hits + CAND_WIN) slots. + # The unconsumed tail is sentinel-filled (idx = -1) at + # q-transition/loop end, so `claimed` over-approximates the true + # count (counts[r][0] exact). + self.CAND_WIN = 8 # epi_bytes covers fp16 and bf16 (FP8 only handled fp16). self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4 # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy. @@ -628,6 +799,17 @@ def __call__( num_phys_blocks: cutlass.Int32, batch_size: cutlass.Int32, stream: cuda.CUstream, + # emission-only tensors; defaulted so the positional signature + # stays the one callers already use + block_max: cute.Tensor = None, # [num_rows, nb_pad*4] fp32 warp-partials + hit_stats: cute.Tensor = None, # [num_rows, 4] fp32 (emit_hit_stats) + hit_bitmap: cute.Tensor = None, # [batch, nb_pad*4] int32 (emit_block_meta) + seed_thr: cute.Tensor = None, # [num_rows, 3] fp32 (emit_seed_counts) + seed_counts: cute.Tensor = None, # [num_rows, 3] int32 out, caller-zeroed + cand: cute.Tensor = None, # [num_rows, CAP*2] int32 {val bits, idx} pairs + cand_ctl: cute.Tensor = None, # [num_rows, 2] int32 {claimed, void}, zeroed + cand_idx_t: cute.Tensor = None, # bucketed: [num_rows, 2*segA+capC] int32 SoA + cand_cur: cute.Tensor = None, # bucketed: [num_rows, 4] int32 cursors, zeroed ): # Derive KV data and SF views from the fused uint8 buffer. # Fused layout per phys block: [data half_head_dim*phys_block_kv bytes] @@ -826,6 +1008,15 @@ class SharedStorage: context_lens, schedule_meta, batch_size, + block_max, + hit_stats, + hit_bitmap, + seed_thr, + seed_counts, + cand, + cand_ctl, + cand_idx_t, + cand_cur, self.cluster_layout_vmnk, self.a_smem_layout_staged, self.b_smem_layout_staged, @@ -848,6 +1039,154 @@ class SharedStorage: stream=stream, ) + @cute.jit + def _flush_hit_agg( + self, + mHitAgg, # [num_rows, 4] fp32 {enc_min, enc_max, sum, cnt} + q_flush, # request whose accumulators are being flushed + hacc_min, + hacc_max, + hacc_sum, + hacc_cnt, + meta_lane, + ): + """Warp-reduce the per-lane hit accumulators and merge them into the + per-row global aggregate via one set of atomics (encoded-int + min/max + fp32 adds), then reset to identities. Called once per + q-transition per warp. The aggregate buffer must be pre-initialized + to {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} per step.""" + next_n = cutlass.const_expr(self.next_n) + base_addr = mHitAgg.iterator.toint() + for t in cutlass.range_constexpr(next_n): + w_min = cute.arch.warp_redux_sync(hacc_min[t], "fmin") + w_max = cute.arch.warp_redux_sync(hacc_max[t], "fmax") + w_sum = cute.arch.warp_reduction_sum(hacc_sum[t]) + w_cnt = cute.arch.warp_redux_sync(hacc_cnt[t], "add") + if meta_lane == cutlass.Int32(0): + if w_cnt > cutlass.Int32(0): + row = q_flush * cutlass.Int32(next_n) + cutlass.Int32(t) + row_addr = base_addr + cutlass.Int64(row) * cutlass.Int64(16) + _red_global_fmin_ordered(row_addr, w_min) + _red_global_fmax_ordered(row_addr + cutlass.Int64(4), w_max) + _red_global_add_f32(row_addr + cutlass.Int64(8), w_sum) + _red_global_add_f32(row_addr + cutlass.Int64(12), cutlass.Float32(w_cnt)) + hacc_min[t] = cutlass.Float32(_META_FLT_MAX) + hacc_max[t] = cutlass.Float32(_META_NEG_FLT_MAX) + hacc_sum[t] = cutlass.Float32(0.0) + hacc_cnt[t] = cutlass.Int32(0) + + @cute.jit + def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, cand_ctl=None): + """Warp-redux the lane-local seed counters and fire one lane-0 + red.global.add per (t, threshold). Caller zero-initializes the + count slots each step; cross-CTA totals accumulate atomically. + + seed_packed: mSeedCounts IS the [num_rows, 8] packed seed row - + counts land as fp32 at cols 3..5 (exact to 2^24).""" + next_n = cutlass.const_expr(self.next_n) + base_addr = mSeedCounts.iterator.toint() + for t in cutlass.range_constexpr(next_n): + for j in cutlass.range_constexpr(3): + w_cnt = cute.arch.warp_redux_sync(scnt[t * 3 + j], "add") + if meta_lane == cutlass.Int32(0): + row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + if cutlass.const_expr(self.seed_packed): + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(3 + j) + ) * cutlass.Int64(4) + _red_global_add_f32(addr, cutlass.Float32(w_cnt)) + else: + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j) + ) * cutlass.Int64(4) + _red_global_add_s32(addr, w_cnt) + if cutlass.const_expr(self.emit_cand_bucketed): + if j >= 1: + # consumer contract: ctl = {n0, void, n1, n2} + if meta_lane == cutlass.Int32(0): + row_b = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + ctl_a = cand_ctl.iterator.toint() + ( + cutlass.Int64(row_b) * cutlass.Int64(4) + cutlass.Int64(j + 1) + ) * cutlass.Int64(4) + _red_global_add_s32(ctl_a, w_cnt) + scnt[t * 3 + j] = cutlass.Int32(0) + if cutlass.const_expr(self.seed_packed and spass is not None): + # packed col 6: adaptive-skip pass count (lane0-accumulated) + for t in cutlass.range_constexpr(next_n): + w_bp = cute.arch.warp_redux_sync(spass[t], "add") + if meta_lane == cutlass.Int32(0): + row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(6) + ) * cutlass.Int64(4) + _red_global_add_f32(addr, cutlass.Float32(w_bp)) + spass[t] = cutlass.Int32(0) + + @cute.jit + def _flush_cand_window_bucketed(self, mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane): + """Sentinel-fill the unconsumed C-window tail in BOTH SoA columns + (score -inf, idx -1: the consumer pads by score) and invalidate + the window. Segment C sits at base 2*segA in each row.""" + next_n = cutlass.const_expr(self.next_n) + segA_f = cutlass.const_expr(self.accept_cap) + capC_f = cutlass.const_expr(self.cand_cap) + wtot_f = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) + vbase_f = mCand.iterator.toint() + ibase_f = mCandIdx.iterator.toint() + for t in cutlass.range_constexpr(next_n): + if cwleft[t] > cutlass.Int32(0): + row_f = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + sl_f = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and sl_f < cutlass.Int32(capC_f): + off_f = ( + cutlass.Int64(row_f) * cutlass.Int64(wtot_f) + + cutlass.Int64(2 * segA_f + sl_f) + ) * cutlass.Int64(4) + vp_f = cute.make_ptr( + cutlass.Float32, + vbase_f + off_f, + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp_f, cute.make_layout((1,)))[0] = cutlass.Float32( + _META_NEG_FLT_MAX + ) + ip_f = cute.make_ptr( + cutlass.Int32, + ibase_f + off_f, + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1) + cwbase[t] = cutlass.Int32(0) + cwleft[t] = cutlass.Int32(0) + + @cute.jit + def _flush_cand_window(self, mCand, q_idx, cwbase, cwleft, meta_lane): + """Sentinel-fill the unconsumed tail of each per-(warp, t) claim + window (idx word = -1; consumers skip sentinels) and invalidate the + window. wleft <= CAND_WIN + 31 always fits one lane round.""" + next_n = cutlass.const_expr(self.next_n) + CAP_C = cutlass.const_expr(self.cand_cap) + cand_base = mCand.iterator.toint() + for t in cutlass.range_constexpr(next_n): + if cwleft[t] > cutlass.Int32(0): + row_c = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + sl_f = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and sl_f < cutlass.Int32(CAP_C): + pair_f = cand_base + ( + cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + cutlass.Int64(sl_f) + ) * cutlass.Int64(8) + iptr_f = cute.make_ptr( + cutlass.Int32, + pair_f + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(iptr_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1) + cwbase[t] = cutlass.Int32(0) + cwleft[t] = cutlass.Int32(0) + @cute.kernel def kernel( self, @@ -867,6 +1206,15 @@ def kernel( mContextLens: cute.Tensor, # [batch_size] mScheduleMeta: cute.Tensor, # [num_sms+1, 2] int32 batch_size: cutlass.Int32, + mBlockMax: cute.Tensor, # [num_rows, nb_pad*4] fp32 warp-partials (or None) + mHitAgg: cute.Tensor, # [num_rows, 4] fp32 {enc_min, enc_max, sum, cnt} (or None) + mHitBitmap: cute.Tensor, # [batch, nb_pad*4] int32 (or None) + mSeedThr: cute.Tensor, # [num_rows, 3] fp32 seed thresholds (or None) + mSeedCounts: cute.Tensor, # [num_rows, 3] int32 counts out (or None) + mCand: cute.Tensor, # [num_rows, CAP*2] int32 pair scatter (or None) + mCandCtl: cute.Tensor, # [num_rows, 2] int32 {claimed, void} (or None) + mCandIdx: cute.Tensor, # bucketed: [num_rows, 2*segA+capC] int32 SoA (or None) + mCandCur: cute.Tensor, # bucketed: [num_rows, 4] int32 cursors (or None) cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, @@ -1115,6 +1463,9 @@ def kernel( layout=sf_kv_smem_layout_staged, byte_alignment=128, ) + # Block-meta emission is warp-autonomous: each warp's lane 0 writes + # its own warp-partial record straight to GMEM; the GVR side folds + # 4 partials per block. No SMEM scratch, no named barrier. a_mcast_mask = cpasync.create_tma_multicast_mask( cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 @@ -1661,7 +2012,6 @@ def kernel( ) // block_kv_val # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) - elif is_umma_warp_1: # UMMA warp for group 1 # Explicitly waits on Q pipeline — critical because TMA warp 1 @@ -1794,7 +2144,6 @@ def kernel( ) // block_kv_val # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) - elif is_math_warp: cute.arch.warpgroup_reg_alloc(240) @@ -1838,6 +2187,15 @@ def kernel( MAX_NUM_W_IN_REG = 64 else: # fp32, 4-byte weights MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64 + if cutlass.const_expr(self.emit_block_meta): + # Free ~8 registers for the meta accumulators/fragments; + # the epilogue's weight cache sits at the spill edge. + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 + if cutlass.const_expr(self.emit_hit_stats): + # Hit accumulators + bitmap word add ~6 more live + # registers across the tile loop. + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 + # emit_seed_counts needs no extra budget cut. NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) # Batched STG: hold reduced result per t in register; the @@ -1848,6 +2206,44 @@ def kernel( else: result_arr = None q_stage_local = cutlass.Int32(0) + if cutlass.const_expr(self.emit_block_meta): + ctx_cur = cutlass.Int32(0) + meta_warp = local_tidx // 32 + meta_lane = local_tidx % 32 + if cutlass.const_expr(self.emit_seed_counts): + sthr = cute.make_rmem_tensor(next_n * 3, cutlass.Float32) + scnt = cute.make_rmem_tensor(next_n * 3, cutlass.Int32) + spass = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + spass[_i] = cutlass.Int32(0) + for _i in cutlass.range_constexpr(next_n * 3): + sthr[_i] = cutlass.Float32(_META_FLT_MAX) + scnt[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed): + cwbase = cute.make_rmem_tensor(next_n, cutlass.Int32) + cwleft = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + cwbase[_i] = cutlass.Int32(0) + cwleft[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_hit_stats): + # Per-lane hit accumulators, carried across all + # tiles of the same q and flushed once per + # q-transition — no warp-wide ops per tile. + hacc_min = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_max = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_sum = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_cnt = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _t in cutlass.range_constexpr(next_n): + hacc_min[_t] = cutlass.Float32(_META_FLT_MAX) + hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX) + hacc_sum[_t] = cutlass.Float32(0.0) + hacc_cnt[_t] = cutlass.Int32(0) + # Batched bitmap read state: all 32 lanes of a warp + # need the SAME word per tile, so lane l loads the + # word for tile j+l once per 32 tiles and each tile + # takes its word via one shuffle. + meta_j = cutlass.Int32(0) + hitw_batch = cutlass.Int32(0) while has_work: # fetch_next_task: commit next → current @@ -1869,11 +2265,78 @@ def kernel( w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ (t_i * num_heads + w_j, q_stage_local) ] + if cutlass.const_expr(self.emit_block_meta): + # Flush the PREVIOUS request's hit accumulators + # before switching context. + if cutlass.const_expr(self.emit_hit_stats): + if q_idx_old < batch_size: + self._flush_hit_agg( + mHitAgg, + q_idx_old, + hacc_min, + hacc_max, + hacc_sum, + hacc_cnt, + meta_lane, + ) + # New bitmap row: invalidate the batched + # word cache (forces a reload). + meta_j = cutlass.Int32(0) + # Compressed-space context len; the meta valid + # mask (kv_pos < ctx_cur) keeps GEMM garbage in + # the aligned padding region out of block_max. + if cutlass.const_expr(self.emit_seed_counts): + if q_idx_old < batch_size: + self._flush_seed_counts( + mSeedCounts, + q_idx_old, + scnt, + meta_lane, + spass=spass, + cand_ctl=mCandCtl, + ) + # (re)load this q's thresholds - gated on + # emit_seed_counts, NOT emit_cand: counts- + # only mode needs them too (a stale + # FLT_MAX default zeroes every counter) + for _t in cutlass.range_constexpr(next_n): + for _j in cutlass.range_constexpr(3): + sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)] + if cutlass.const_expr(self.emit_cand): + if q_idx_old < batch_size: + self._flush_cand_window( + mCand, q_idx_old, cwbase, cwleft, meta_lane + ) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx_old < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane + ) + ctx_cur = mContextLens[q_idx] # Process KV block for group 0 (kv_idx + 0) # Unconditional Math: OOB results # written to aligned padding region in logits buffer. kv_pos = kv_idx * block_kv_val + m_coord + if cutlass.const_expr(self.emit_block_meta): + meta_kv_tile = kv_idx + meta_valid = kv_pos < ctx_cur + if cutlass.const_expr(self.emit_hit_stats): + # Warp-uniform reload once per 32 tiles: this + # WG's tile at counter j+l is kv_tile + 2*l, + # whose warp word index is (kv_tile+2*l)*4 + + # warp. Clamp keeps end-of-row lanes in + # bounds (their tiles are never consumed). + if (meta_j & cutlass.Int32(31)) == cutlass.Int32(0): + w_idx = ( + meta_kv_tile + cutlass.Int32(2) * meta_lane + ) * cutlass.Int32(4) + meta_warp + w_idx = min(w_idx, mHitBitmap.shape[1] - cutlass.Int32(1)) + hitw_batch = mHitBitmap[(q_idx, w_idx)] + hit_word = cute.arch.shuffle_sync( + hitw_batch, meta_j & cutlass.Int32(31) + ) + meta_j = meta_j + cutlass.Int32(1) # Step 5.7: drop kv_pipeline.consumer_wait/release and # scale_val LDS — UMMA owns KV+SF pipe; SF is baked into @@ -2054,11 +2517,388 @@ def kernel( else: result_t = s0x + s0y + s1x + s1y # Step 5.7: drop * scale_val (FP4 SF baked into acc). + stored_t = self.output_dtype(result_t) if cutlass.const_expr(self.use_batched_store): - result_arr[t] = self.output_dtype(result_t) + result_arr[t] = stored_t else: out_row = q_idx * next_n + t - mLogits[(out_row, kv_pos)] = self.output_dtype(result_t) + mLogits[(out_row, kv_pos)] = stored_t + if cutlass.const_expr(self.emit_block_meta): + # Meta reduction on the POST-conversion value so + # block_max bounds what GVR reads back bit-exactly. + f32_t = cutlass.Float32(stored_t) + bmax_v = cutlass.Float32(_META_NEG_FLT_MAX) + if meta_valid: + bmax_v = f32_t + r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax") + # Warp-autonomous store: record index = + # tile*4 + warp; the GVR consumer folds the + # 4 warp-partials per block. + if meta_lane == cutlass.Int32(0): + out_row_m = q_idx * next_n + t + rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp + mBlockMax[(out_row_m, rec_m)] = r_bmax + if cutlass.const_expr(self.emit_seed_counts): + # Seed-count accumulation: branchless 0/1 + # adds on the post-conversion value; the + # valid mask keeps aligned-padding garbage + # out (same contract as block_max). + valid_i1 = cutlass.Int32(meta_valid) + for _j in cutlass.range_constexpr(3): + ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j]) + scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1) + if cutlass.const_expr(self.seed_packed): + # adaptive-skip pass count: one record + # per (tile, warp); r_bmax is warp- + # uniform so lane0 alone accumulates + if meta_lane == cutlass.Int32(0): + spass[t] = spass[t] + cutlass.Int32( + r_bmax >= sthr[t * 3 + 0] + ) + if cutlass.const_expr(self.emit_cand): + # Candidate pre-collect at t_0 with per-warp + # claim windows: one atomic claims + # (hits + CAND_WIN) slots per refill. + # Unconsumed tail is sentinel-filled on + # flush; counts[r][0] stays the exact count. + # Gated on the warp-uniform 32-position + # bound: r_bmax < t_0 proves zero hits; + # bound >= t_0 guarantees a nonzero ballot + # (exact per-lane max, invalid -> -FLT_MAX). + if r_bmax >= sthr[t * 3 + 0]: + pred_c = cutlass.Int32(0) + if meta_valid: + if f32_t >= sthr[t * 3 + 0]: + pred_c = cutlass.Int32(1) + mask_c = cute.arch.vote_ballot_sync(pred_c != cutlass.Int32(0)) + row_c = q_idx * next_n + t + cnt_c = cutlass.Int32(cute.arch.popc(mask_c)) + lm_c = ( + cutlass.Uint32(1) << cutlass.Uint32(meta_lane) + ) - cutlass.Uint32(1) + off_c = cutlass.Int32(cute.arch.popc(mask_c & lm_c)) + CAP_C = cutlass.const_expr(self.cand_cap) + cand_b = mCand.iterator.toint() + if cnt_c > cwleft[t]: + # sentinel-fill the old tail, then + # refill: one atomic per window. + sl_o = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and sl_o < cutlass.Int32(CAP_C): + pair_o = cand_b + ( + cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + + cutlass.Int64(sl_o) + ) * cutlass.Int64(8) + iptr_o = cute.make_ptr( + cutlass.Int32, + pair_o + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(iptr_o, cute.make_layout((1,)))[0] = ( + cutlass.Int32(-1) + ) + m_c = cnt_c + cutlass.Int32(self.CAND_WIN) + ctl_addr = mCandCtl.iterator.toint() + ( + cutlass.Int64(row_c) * cutlass.Int64(8) + ) + nb_c = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0): + nb_c = _atom_global_add_s32(ctl_addr, m_c) + nb_c = cute.arch.shuffle_sync(nb_c, cutlass.Int32(0)) + cwbase[t] = nb_c + cwleft[t] = m_c + # one-shot void mark on crossing CAP + if meta_lane == cutlass.Int32(0): + if nb_c + m_c > cutlass.Int32( + CAP_C + ) and nb_c <= cutlass.Int32(CAP_C): + vdptr = cute.make_ptr( + cutlass.Int32, + ctl_addr + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vdptr, cute.make_layout((1,)))[ + 0 + ] = cutlass.Int32(1) + slot_c = cwbase[t] + off_c + if pred_c != cutlass.Int32(0) and slot_c < cutlass.Int32(CAP_C): + pair_addr = cand_b + ( + cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + + cutlass.Int64(slot_c) + ) * cutlass.Int64(8) + vptr_c = cute.make_ptr( + cutlass.Float32, + pair_addr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + cute.make_tensor(vptr_c, cute.make_layout((1,)))[0] = f32_t + iptr_c = cute.make_ptr( + cutlass.Int32, + pair_addr + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos + cwbase[t] = cwbase[t] + cnt_c + cwleft[t] = cwleft[t] - cnt_c + if cutlass.const_expr(self.emit_cand_bucketed): + # Bucketed SoA: A/B EXACT ballot claims + # (their prefixes must stay pad-free for + # the consumer's prefix math), C keeps the + # claim-window; a full segment spills to + # the next looser one. Every warp + # collective sits at the TOP level of this + # warp-uniform bound gate - no collectives + # inside nested dynamic branches (DSL). + if r_bmax >= sthr[t * 3 + 0]: + segA_k = cutlass.const_expr(self.accept_cap) + capC_k = cutlass.const_expr(self.cand_cap) + wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) + row_k = q_idx * next_n + t + vb_k = mCand.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) + ib_k = mCandIdx.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) + cur_k = mCandCur.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(16) + ctl_k = mCandCtl.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(16) + lmk_k = ( + cutlass.Uint32(1) << cutlass.Uint32(meta_lane) + ) - cutlass.Uint32(1) + # exclusive class predicates + pA_k = cutlass.Int32(0) + pB_k = cutlass.Int32(0) + pC_k = cutlass.Int32(0) + if meta_valid: + if f32_t >= sthr[t * 3 + 2]: + pA_k = cutlass.Int32(1) + if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0): + pB_k = cutlass.Int32(1) + if ( + f32_t >= sthr[t * 3 + 0] + and pA_k == cutlass.Int32(0) + and pB_k == cutlass.Int32(0) + ): + pC_k = cutlass.Int32(1) + # ---- A: exact claim ---- + mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0)) + cntA_k = cutlass.Int32(cute.arch.popc(mA_k)) + offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k)) + baseA_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0): + baseA_k = _atom_global_add_s32(cur_k, cntA_k) + baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0)) + slotA_k = baseA_k + offA_k + spA_k = cutlass.Int32(0) + if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32( + segA_k + ): + spA_k = cutlass.Int32(1) + if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k): + vp_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t + ip_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos + # ---- B: exact claim (native + A spill) ---- + pBe_k = cutlass.Int32(0) + if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0): + pBe_k = cutlass.Int32(1) + mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0)) + cntB_k = cutlass.Int32(cute.arch.popc(mB_k)) + offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k)) + baseB_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0): + baseB_k = _atom_global_add_s32( + cur_k + cutlass.Int64(4), cntB_k + ) + baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0)) + slotB_k = baseB_k + offB_k + spB_k = cutlass.Int32(0) + if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32( + segA_k + ): + spB_k = cutlass.Int32(1) + if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32( + segA_k + ): + vp2_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t + ip2_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos + # n0 += exact placements in A and B + plc_k = ( + cntA_k + - cutlass.Int32( + cute.arch.popc( + cute.arch.vote_ballot_sync( + spA_k != cutlass.Int32(0) + ) + ) + ) + ) + ( + cntB_k + - cutlass.Int32( + cute.arch.popc( + cute.arch.vote_ballot_sync( + spB_k != cutlass.Int32(0) + ) + ) + ) + ) + if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0): + _atom_global_add_s32(ctl_k, plc_k) + # ---- C: claim window (native + B spill) ---- + pCe_k = cutlass.Int32(0) + if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0): + pCe_k = cutlass.Int32(1) + mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0)) + cntC_k = cutlass.Int32(cute.arch.popc(mC_k)) + offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k)) + if cntC_k > cwleft[t]: + # sentinel-fill the old window tail + # (BOTH columns: the consumer pads + # by score -inf, idx -1) + slo_k = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k): + vpo_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(2 * segA_k + slo_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = ( + cutlass.Float32(_META_NEG_FLT_MAX) + ) + ipo_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(2 * segA_k + slo_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = ( + cutlass.Int32(-1) + ) + mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN) + nbC_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0): + nbC_k = _atom_global_add_s32( + cur_k + cutlass.Int64(8), mC2_k + ) + _atom_global_add_s32(ctl_k, mC2_k) + if nbC_k + mC2_k > cutlass.Int32( + capC_k + ) and nbC_k <= cutlass.Int32(capC_k): + vdp_k = cute.make_ptr( + cutlass.Int32, + ctl_k + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vdp_k, cute.make_layout((1,)))[ + 0 + ] = cutlass.Int32(1) + nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0)) + cwbase[t] = nbC_k + cwleft[t] = mC2_k + slotC_k = cwbase[t] + offC_k + if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32( + capC_k + ): + vpc_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(2 * segA_k + slotC_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t + ipc_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(2 * segA_k + slotC_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos + cwbase[t] = cwbase[t] + cntC_k + cwleft[t] = cwleft[t] - cntC_k + if cutlass.const_expr(self.emit_hit_stats): + # Lane-local accumulation; keep it branchless + # (`if meta_hit` compiles to real divergent + # branches). Bit-mask select is NaN-safe for + # OOB-tile garbage logits. No valid-mask + # needed: the bitmap contract only sets bits + # inside [0, ctx). + meta_hit = ( + hit_word >> (kv_pos & cutlass.Int32(31)) + ) & cutlass.Int32(1) + msk = cutlass.Int32(0) - meta_hit # 0 / ~0 + inv = cutlass.Int32(-1) - msk + fbits = cutlass.Int32( + llvm.bitcast(cutlass.Int32.mlir_type, f32_t.ir_value()) + ) + # bits(+FLT_MAX)=0x7F7FFFFF, + # bits(-FLT_MAX)=0xFF7FFFFF (as i32: neg). + selmin = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + ( + (fbits & msk) | (cutlass.Int32(0x7F7FFFFF) & inv) + ).ir_value(), + ) + ) + selmax = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + ( + (fbits & msk) | (cutlass.Int32(-8388609) & inv) + ).ir_value(), + ) + ) + seladd = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, (fbits & msk).ir_value() + ) + ) + hacc_min[t] = cutlass.min(hacc_min[t], selmin) + hacc_max[t] = cutlass.max(hacc_max[t], selmax) + hacc_sum[t] = hacc_sum[t] + seladd + hacc_cnt[t] = hacc_cnt[t] + meta_hit if cutlass.const_expr(self.use_batched_store): # Batched STG: all result_arr[t] → mLogits in one pass. @@ -2078,6 +2918,26 @@ def kernel( # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + # Flush the final request's hit accumulators (WG 0). + if cutlass.const_expr(self.emit_block_meta and self.emit_hit_stats): + if q_idx < batch_size: + self._flush_hit_agg( + mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane + ) + if cutlass.const_expr(self.emit_seed_counts): + if q_idx < batch_size: + self._flush_seed_counts( + mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl + ) + if cutlass.const_expr(self.emit_cand): + if q_idx < batch_size: + self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane + ) + # Release last Q stage (WG 0) if q_idx < batch_size: q_pipeline.consumer_release(q_cons_state) @@ -2101,6 +2961,15 @@ def kernel( MAX_NUM_W_IN_REG = 64 else: MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64 + if cutlass.const_expr(self.emit_block_meta): + # Free ~8 registers for the meta accumulators/fragments; + # the epilogue's weight cache sits at the spill edge. + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 + if cutlass.const_expr(self.emit_hit_stats): + # Hit accumulators + bitmap word add ~6 more live + # registers across the tile loop. + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 + # emit_seed_counts needs no extra budget cut. NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) # Batched STG: hold reduced result per t in register; the @@ -2111,6 +2980,44 @@ def kernel( else: result_arr = None q_stage_local = cutlass.Int32(0) + if cutlass.const_expr(self.emit_block_meta): + ctx_cur = cutlass.Int32(0) + meta_warp = local_tidx // 32 + meta_lane = local_tidx % 32 + if cutlass.const_expr(self.emit_seed_counts): + sthr = cute.make_rmem_tensor(next_n * 3, cutlass.Float32) + scnt = cute.make_rmem_tensor(next_n * 3, cutlass.Int32) + spass = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + spass[_i] = cutlass.Int32(0) + for _i in cutlass.range_constexpr(next_n * 3): + sthr[_i] = cutlass.Float32(_META_FLT_MAX) + scnt[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_cand or self.emit_cand_bucketed): + cwbase = cute.make_rmem_tensor(next_n, cutlass.Int32) + cwleft = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + cwbase[_i] = cutlass.Int32(0) + cwleft[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_hit_stats): + # Per-lane hit accumulators, carried across all + # tiles of the same q and flushed once per + # q-transition — no warp-wide ops per tile. + hacc_min = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_max = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_sum = cute.make_rmem_tensor(next_n, cutlass.Float32) + hacc_cnt = cute.make_rmem_tensor(next_n, cutlass.Int32) + for _t in cutlass.range_constexpr(next_n): + hacc_min[_t] = cutlass.Float32(_META_FLT_MAX) + hacc_max[_t] = cutlass.Float32(_META_NEG_FLT_MAX) + hacc_sum[_t] = cutlass.Float32(0.0) + hacc_cnt[_t] = cutlass.Int32(0) + # Batched bitmap read state: all 32 lanes of a warp + # need the SAME word per tile, so lane l loads the + # word for tile j+l once per 32 tiles and each tile + # takes its word via one shuffle. + meta_j = cutlass.Int32(0) + hitw_batch = cutlass.Int32(0) while has_work: # fetch_next_task: commit next → current @@ -2132,12 +3039,82 @@ def kernel( w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ (t_i * num_heads + w_j, q_stage_local) ] + if cutlass.const_expr(self.emit_block_meta): + # Flush the PREVIOUS request's hit accumulators + # before switching context. + if cutlass.const_expr(self.emit_hit_stats): + if q_idx_old < batch_size: + self._flush_hit_agg( + mHitAgg, + q_idx_old, + hacc_min, + hacc_max, + hacc_sum, + hacc_cnt, + meta_lane, + ) + # New bitmap row: invalidate the batched + # word cache (forces a reload). + meta_j = cutlass.Int32(0) + # Compressed-space context len; the meta valid + # mask (kv_pos < ctx_cur) keeps GEMM garbage in + # the aligned padding region out of block_max. + if cutlass.const_expr(self.emit_seed_counts): + if q_idx_old < batch_size: + self._flush_seed_counts( + mSeedCounts, + q_idx_old, + scnt, + meta_lane, + spass=spass, + cand_ctl=mCandCtl, + ) + # (re)load this q's thresholds - gated on + # emit_seed_counts, NOT emit_cand (see the + # WG0 twin above) + for _t in cutlass.range_constexpr(next_n): + for _j in cutlass.range_constexpr(3): + sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)] + if cutlass.const_expr(self.emit_cand): + if q_idx_old < batch_size: + self._flush_cand_window( + mCand, q_idx_old, cwbase, cwleft, meta_lane + ) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx_old < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane + ) + ctx_cur = mContextLens[q_idx] # Process KV block for group 1 (kv_idx + 1) # Unconditional Math kv_idx_1 = kv_idx + 1 kv_pos = kv_idx_1 * block_kv_val + m_coord + if cutlass.const_expr(self.emit_block_meta): + meta_kv_tile = kv_idx_1 + # See WG 0. For the odd-num_kv OOB tile + # (kv_idx_1 == num_kv) every lane has + # kv_pos >= ctx_cur, so identities land in the + # nb_pad padding slot — never read by GVR. + meta_valid = kv_pos < ctx_cur + if cutlass.const_expr(self.emit_hit_stats): + # Warp-uniform reload once per 32 tiles: this + # WG's tile at counter j+l is kv_tile + 2*l, + # whose warp word index is (kv_tile+2*l)*4 + + # warp. Clamp keeps end-of-row lanes in + # bounds (their tiles are never consumed). + if (meta_j & cutlass.Int32(31)) == cutlass.Int32(0): + w_idx = ( + meta_kv_tile + cutlass.Int32(2) * meta_lane + ) * cutlass.Int32(4) + meta_warp + w_idx = min(w_idx, mHitBitmap.shape[1] - cutlass.Int32(1)) + hitw_batch = mHitBitmap[(q_idx, w_idx)] + hit_word = cute.arch.shuffle_sync( + hitw_batch, meta_j & cutlass.Int32(31) + ) + meta_j = meta_j + cutlass.Int32(1) # Step 5.7: drop kv_pipeline.consumer_wait/release and # scale_val LDS — UMMA owns KV+SF pipe. @@ -2307,11 +3284,388 @@ def kernel( else: result_t = s0x + s0y + s1x + s1y # Step 5.7: drop * scale_val (FP4 SF baked into acc). + stored_t = self.output_dtype(result_t) if cutlass.const_expr(self.use_batched_store): - result_arr[t] = self.output_dtype(result_t) + result_arr[t] = stored_t else: out_row = q_idx * next_n + t - mLogits[(out_row, kv_pos)] = self.output_dtype(result_t) + mLogits[(out_row, kv_pos)] = stored_t + if cutlass.const_expr(self.emit_block_meta): + # Meta reduction on the POST-conversion value so + # block_max bounds what GVR reads back bit-exactly. + f32_t = cutlass.Float32(stored_t) + bmax_v = cutlass.Float32(_META_NEG_FLT_MAX) + if meta_valid: + bmax_v = f32_t + r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax") + # Warp-autonomous store: record index = + # tile*4 + warp; the GVR consumer folds the + # 4 warp-partials per block. + if meta_lane == cutlass.Int32(0): + out_row_m = q_idx * next_n + t + rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp + mBlockMax[(out_row_m, rec_m)] = r_bmax + if cutlass.const_expr(self.emit_seed_counts): + # Seed-count accumulation: branchless 0/1 + # adds on the post-conversion value; the + # valid mask keeps aligned-padding garbage + # out (same contract as block_max). + valid_i1 = cutlass.Int32(meta_valid) + for _j in cutlass.range_constexpr(3): + ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j]) + scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1) + if cutlass.const_expr(self.seed_packed): + # adaptive-skip pass count: one record + # per (tile, warp); r_bmax is warp- + # uniform so lane0 alone accumulates + if meta_lane == cutlass.Int32(0): + spass[t] = spass[t] + cutlass.Int32( + r_bmax >= sthr[t * 3 + 0] + ) + if cutlass.const_expr(self.emit_cand): + # Candidate pre-collect at t_0 with per-warp + # claim windows: one atomic claims + # (hits + CAND_WIN) slots per refill. + # Unconsumed tail is sentinel-filled on + # flush; counts[r][0] stays the exact count. + # Gated on the warp-uniform 32-position + # bound: r_bmax < t_0 proves zero hits; + # bound >= t_0 guarantees a nonzero ballot + # (exact per-lane max, invalid -> -FLT_MAX). + if r_bmax >= sthr[t * 3 + 0]: + pred_c = cutlass.Int32(0) + if meta_valid: + if f32_t >= sthr[t * 3 + 0]: + pred_c = cutlass.Int32(1) + mask_c = cute.arch.vote_ballot_sync(pred_c != cutlass.Int32(0)) + row_c = q_idx * next_n + t + cnt_c = cutlass.Int32(cute.arch.popc(mask_c)) + lm_c = ( + cutlass.Uint32(1) << cutlass.Uint32(meta_lane) + ) - cutlass.Uint32(1) + off_c = cutlass.Int32(cute.arch.popc(mask_c & lm_c)) + CAP_C = cutlass.const_expr(self.cand_cap) + cand_b = mCand.iterator.toint() + if cnt_c > cwleft[t]: + # sentinel-fill the old tail, then + # refill: one atomic per window. + sl_o = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and sl_o < cutlass.Int32(CAP_C): + pair_o = cand_b + ( + cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + + cutlass.Int64(sl_o) + ) * cutlass.Int64(8) + iptr_o = cute.make_ptr( + cutlass.Int32, + pair_o + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(iptr_o, cute.make_layout((1,)))[0] = ( + cutlass.Int32(-1) + ) + m_c = cnt_c + cutlass.Int32(self.CAND_WIN) + ctl_addr = mCandCtl.iterator.toint() + ( + cutlass.Int64(row_c) * cutlass.Int64(8) + ) + nb_c = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0): + nb_c = _atom_global_add_s32(ctl_addr, m_c) + nb_c = cute.arch.shuffle_sync(nb_c, cutlass.Int32(0)) + cwbase[t] = nb_c + cwleft[t] = m_c + # one-shot void mark on crossing CAP + if meta_lane == cutlass.Int32(0): + if nb_c + m_c > cutlass.Int32( + CAP_C + ) and nb_c <= cutlass.Int32(CAP_C): + vdptr = cute.make_ptr( + cutlass.Int32, + ctl_addr + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vdptr, cute.make_layout((1,)))[ + 0 + ] = cutlass.Int32(1) + slot_c = cwbase[t] + off_c + if pred_c != cutlass.Int32(0) and slot_c < cutlass.Int32(CAP_C): + pair_addr = cand_b + ( + cutlass.Int64(row_c) * cutlass.Int64(CAP_C) + + cutlass.Int64(slot_c) + ) * cutlass.Int64(8) + vptr_c = cute.make_ptr( + cutlass.Float32, + pair_addr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + cute.make_tensor(vptr_c, cute.make_layout((1,)))[0] = f32_t + iptr_c = cute.make_ptr( + cutlass.Int32, + pair_addr + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(iptr_c, cute.make_layout((1,)))[0] = kv_pos + cwbase[t] = cwbase[t] + cnt_c + cwleft[t] = cwleft[t] - cnt_c + if cutlass.const_expr(self.emit_cand_bucketed): + # Bucketed SoA: A/B EXACT ballot claims + # (their prefixes must stay pad-free for + # the consumer's prefix math), C keeps the + # claim-window; a full segment spills to + # the next looser one. Every warp + # collective sits at the TOP level of this + # warp-uniform bound gate - no collectives + # inside nested dynamic branches (DSL). + if r_bmax >= sthr[t * 3 + 0]: + segA_k = cutlass.const_expr(self.accept_cap) + capC_k = cutlass.const_expr(self.cand_cap) + wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) + row_k = q_idx * next_n + t + vb_k = mCand.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) + ib_k = mCandIdx.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) + cur_k = mCandCur.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(16) + ctl_k = mCandCtl.iterator.toint() + cutlass.Int64( + row_k + ) * cutlass.Int64(16) + lmk_k = ( + cutlass.Uint32(1) << cutlass.Uint32(meta_lane) + ) - cutlass.Uint32(1) + # exclusive class predicates + pA_k = cutlass.Int32(0) + pB_k = cutlass.Int32(0) + pC_k = cutlass.Int32(0) + if meta_valid: + if f32_t >= sthr[t * 3 + 2]: + pA_k = cutlass.Int32(1) + if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0): + pB_k = cutlass.Int32(1) + if ( + f32_t >= sthr[t * 3 + 0] + and pA_k == cutlass.Int32(0) + and pB_k == cutlass.Int32(0) + ): + pC_k = cutlass.Int32(1) + # ---- A: exact claim ---- + mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0)) + cntA_k = cutlass.Int32(cute.arch.popc(mA_k)) + offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k)) + baseA_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0): + baseA_k = _atom_global_add_s32(cur_k, cntA_k) + baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0)) + slotA_k = baseA_k + offA_k + spA_k = cutlass.Int32(0) + if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32( + segA_k + ): + spA_k = cutlass.Int32(1) + if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k): + vp_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t + ip_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos + # ---- B: exact claim (native + A spill) ---- + pBe_k = cutlass.Int32(0) + if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0): + pBe_k = cutlass.Int32(1) + mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0)) + cntB_k = cutlass.Int32(cute.arch.popc(mB_k)) + offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k)) + baseB_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0): + baseB_k = _atom_global_add_s32( + cur_k + cutlass.Int64(4), cntB_k + ) + baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0)) + slotB_k = baseB_k + offB_k + spB_k = cutlass.Int32(0) + if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32( + segA_k + ): + spB_k = cutlass.Int32(1) + if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32( + segA_k + ): + vp2_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t + ip2_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos + # n0 += exact placements in A and B + plc_k = ( + cntA_k + - cutlass.Int32( + cute.arch.popc( + cute.arch.vote_ballot_sync( + spA_k != cutlass.Int32(0) + ) + ) + ) + ) + ( + cntB_k + - cutlass.Int32( + cute.arch.popc( + cute.arch.vote_ballot_sync( + spB_k != cutlass.Int32(0) + ) + ) + ) + ) + if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0): + _atom_global_add_s32(ctl_k, plc_k) + # ---- C: claim window (native + B spill) ---- + pCe_k = cutlass.Int32(0) + if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0): + pCe_k = cutlass.Int32(1) + mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0)) + cntC_k = cutlass.Int32(cute.arch.popc(mC_k)) + offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k)) + if cntC_k > cwleft[t]: + # sentinel-fill the old window tail + # (BOTH columns: the consumer pads + # by score -inf, idx -1) + slo_k = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k): + vpo_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(2 * segA_k + slo_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = ( + cutlass.Float32(_META_NEG_FLT_MAX) + ) + ipo_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(2 * segA_k + slo_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = ( + cutlass.Int32(-1) + ) + mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN) + nbC_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0): + nbC_k = _atom_global_add_s32( + cur_k + cutlass.Int64(8), mC2_k + ) + _atom_global_add_s32(ctl_k, mC2_k) + if nbC_k + mC2_k > cutlass.Int32( + capC_k + ) and nbC_k <= cutlass.Int32(capC_k): + vdp_k = cute.make_ptr( + cutlass.Int32, + ctl_k + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vdp_k, cute.make_layout((1,)))[ + 0 + ] = cutlass.Int32(1) + nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0)) + cwbase[t] = nbC_k + cwleft[t] = mC2_k + slotC_k = cwbase[t] + offC_k + if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32( + capC_k + ): + vpc_k = cute.make_ptr( + cutlass.Float32, + vb_k + + cutlass.Int64(2 * segA_k + slotC_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t + ipc_k = cute.make_ptr( + cutlass.Int32, + ib_k + + cutlass.Int64(2 * segA_k + slotC_k) + * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos + cwbase[t] = cwbase[t] + cntC_k + cwleft[t] = cwleft[t] - cntC_k + if cutlass.const_expr(self.emit_hit_stats): + # Lane-local accumulation; keep it branchless + # (`if meta_hit` compiles to real divergent + # branches). Bit-mask select is NaN-safe for + # OOB-tile garbage logits. No valid-mask + # needed: the bitmap contract only sets bits + # inside [0, ctx). + meta_hit = ( + hit_word >> (kv_pos & cutlass.Int32(31)) + ) & cutlass.Int32(1) + msk = cutlass.Int32(0) - meta_hit # 0 / ~0 + inv = cutlass.Int32(-1) - msk + fbits = cutlass.Int32( + llvm.bitcast(cutlass.Int32.mlir_type, f32_t.ir_value()) + ) + # bits(+FLT_MAX)=0x7F7FFFFF, + # bits(-FLT_MAX)=0xFF7FFFFF (as i32: neg). + selmin = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + ( + (fbits & msk) | (cutlass.Int32(0x7F7FFFFF) & inv) + ).ir_value(), + ) + ) + selmax = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + ( + (fbits & msk) | (cutlass.Int32(-8388609) & inv) + ).ir_value(), + ) + ) + seladd = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, (fbits & msk).ir_value() + ) + ) + hacc_min[t] = cutlass.min(hacc_min[t], selmin) + hacc_max[t] = cutlass.max(hacc_max[t], selmax) + hacc_sum[t] = hacc_sum[t] + seladd + hacc_cnt[t] = hacc_cnt[t] + meta_hit if cutlass.const_expr(self.use_batched_store): # Batched STG: all result_arr[t] → mLogits in one pass. @@ -2331,6 +3685,26 @@ def kernel( # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + # Flush the final request's hit accumulators (WG 1). + if cutlass.const_expr(self.emit_block_meta and self.emit_hit_stats): + if q_idx < batch_size: + self._flush_hit_agg( + mHitAgg, q_idx, hacc_min, hacc_max, hacc_sum, hacc_cnt, meta_lane + ) + if cutlass.const_expr(self.emit_seed_counts): + if q_idx < batch_size: + self._flush_seed_counts( + mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl + ) + if cutlass.const_expr(self.emit_cand): + if q_idx < batch_size: + self._flush_cand_window(mCand, q_idx, cwbase, cwleft, meta_lane) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane + ) + # Release last Q stage (WG 1) if q_idx < batch_size: q_pipeline.consumer_release(q_cons_state) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py new file mode 100644 index 000000000000..fbaec11736c1 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Emission-assisted GVR top-k state for the DSA decode path. + +Owns the persistent (graph-address-stable) buffers the emission tiers +ride on, the device-side closed-loop seed-row update (pure tensor ops, +CUDA-graph capturable) and the per-step routing decision. Opt-in: +without the flag the DSA decode path is unchanged. + +Tier semantics (see gvr_routing): + * this step's TOP-K consumes what the PREVIOUS step's indexer + epilogue emitted; + * this step's INDEXER emits what the routing planned for the NEXT + step. +""" + +import math +from typing import Optional + +import torch + +from .gvr_routing import LIST_EMIT_MAX_B, LIST_EMIT_MIN_N, TopkRoute, pick_config, plan_emission + +# Bucketed candidate-list geometry: two tight segments of LIST_SEG_A +# entries plus a LIST_CAP_C-entry loose segment. +LIST_SEG_A = 8192 +LIST_CAP_C = 24576 +LIST_WIDTH = 2 * LIST_SEG_A + LIST_CAP_C + +__all__ = ["GvrEmissionState", "LIST_EMIT_MIN_N", "LIST_PARK_LINE"] + +# Closed-loop line placement: fit the slope of log2(count) vs threshold +# from the previous step's (lines, counts) and place the new lines at +# these K-relative target counts (t0 loosest .. t2 tightest). +LINE_TARGETS = (8.0, 5.0, 2.0) +LIST_T0_TARGET = 2.5 # list tier: single collect-line target (xK) +LIST_T0_COUNT_MAX = 6144.0 # keep n0 inside the [K+64, segA] admission band +SLOPE_MIN = 0.05 +SLOPE_MAX = 64.0 + +# No-fit fallback: multiplicative guards around the published k-th value +# (t1 hugs it from below; t0/t2 guard by GUARD_LO/GUARD_HI spans). +FALLBACK_REL = 2.0**0.125 - 1.0 +FALLBACK_ABS = 1e-3 +GUARD_LO = 2.0 +GUARD_HI = 0.5 + +# List tier only: park the two tight lines above any score so every +# admitted entry lands in the loosest segment. Any finite value above +# the score range works; the kernel's eligibility check only needs the +# three lines increasing and the loosest one finite. +LIST_PARK_LINE = 1.0e30 + + +class GvrEmissionState: + """Per-attention-backend emission state (persistent buffers).""" + + def __init__( + self, + max_rows: int, + top_k: int, + device: torch.device, + enable_list_tier: bool = True, + own_prior: bool = True, + ): + self.max_rows = max_rows + self.top_k = top_k + # packed seed row: lines at cols 0..2, counts (emission-filled) + # at 3..5, adaptive-skip pass count at 6 + self.seed_row = torch.zeros((max_rows, 8), dtype=torch.float32, device=device) + # contiguous alias of the three lines for the rungs tier (a + # [rows, 3] column view of the packed row is non-contiguous) + self.seed_rungs = torch.zeros((max_rows, 3), dtype=torch.float32, device=device) + self.xstate = torch.zeros((max_rows, 8), dtype=torch.float32, device=device) + self.cand_vals: Optional[torch.Tensor] = None + self.cand_idx: Optional[torch.Tensor] = None + self.cand_ctl: Optional[torch.Tensor] = None + self.cand_cur: Optional[torch.Tensor] = None + if enable_list_tier: + # the routing only ever picks the list tier at + # batch <= LIST_EMIT_MAX_B, so the wide candidate buffers + # need that many rows, not max_rows (~0.33 MB/row/layer) + cand_rows = min(max_rows, LIST_EMIT_MAX_B) + self.cand_vals = torch.zeros( + (cand_rows, LIST_WIDTH), dtype=torch.float32, device=device + ) + self.cand_idx = torch.zeros((cand_rows, LIST_WIDTH), dtype=torch.int32, device=device) + self.cand_ctl = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device) + self.cand_cur = torch.zeros((cand_rows, 4), dtype=torch.int32, device=device) + # previous-step top-k feedback (address-stable; zero-init -> + # first step's pre_idx points at index 0, a benign candidate). + # own_prior=False when the caller already keeps this state (the + # TopK module rides metadata's gvr_prior_indices). + self.prev_topk = ( + torch.zeros((max_rows, top_k), dtype=torch.int32, device=device) if own_prior else None + ) + # block_max prefix ([rows, nb_pad*4] fp32 warp-partials), + # allocated lazily once max_seq_len is known + self.block_max: Optional[torch.Tensor] = None + + def ensure_block_max(self, max_seq_len: int) -> torch.Tensor: + nb4 = ((max_seq_len + 255) // 256 * 256) // 128 * 4 + # exact width: the runner asserts shape == (rows, nrec), so a + # wider reused buffer would trip it + if self.block_max is None or self.block_max.shape[1] != nb4: + # allocating during CUDA graph capture would bake a dangling + # address into the graph, so fail loudly instead + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "GvrEmissionState.ensure_block_max: (re)allocation requested " + "during CUDA graph capture; the block_max buffer must be " + "created by a warmup step before capture" + ) + self.block_max = torch.zeros( + (self.max_rows, nb4), dtype=torch.float32, device=self.seed_row.device + ) + return self.block_max + + def plan( + self, batch: int, n_comp: int, num_sms: int, compress_ratio: int = 4 + ) -> tuple[str, TopkRoute]: + """Route this step: (tier the epilogue emits, launch knobs the + top-k consumes it with).""" + emit_tier = plan_emission( + batch, n_comp, self.top_k, have_epilogue=True, compress_ratio=compress_ratio + ) + if emit_tier == "list" and self.cand_vals is None: + # constructed with enable_list_tier=False: no candidate + # buffers to emit into, demote to the counts tier + emit_tier = "counts" + # emission and consumption happen inside the SAME forward (zero, + # emit, consume), so the consumer routes on this step's tier + route = pick_config(emit_tier, batch, n_comp, self.top_k, num_sms) + return emit_tier, route + + def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None: + """Device-side closed-loop line update from the last publish. + + Slope-fits log2(count) vs threshold from the previous step's + (lines, counts) and places the new lines at K-relative target + counts. Counts come from the packed row (counts/list emission) + or from the kernel's rung-count publish in xstate cols 4..6 + (rungs tier). Rows without a usable fit get multiplicative + guards around the published k-th value; rows with invalid + xstate (col 0 == 0, e.g. cold start) get non-finite lines, + which the kernel's validity guard routes to the stock path. + Pure tensor ops (graph-capturable). + """ + s = self.seed_row[:num_rows] + x = self.xstate[:num_rows] + valid = x[:, 0] > 0 + kth = x[:, 1] + anchor = x[:, 2] + t_prev0 = s[:, 0] + t_prev2 = s[:, 2] + cnts = x[:, 4:7] if emit_tier == "rungs" else s[:, 3:6] + k = float(self.top_k) + inf = torch.full_like(kth, float("inf")) + d_fb = kth.abs() * FALLBACK_REL + FALLBACK_ABS + if emit_tier == "list": + # two-point fit (t0_prev, n0) / (kth, K): kth is the exact + # k-th boundary on list rows + n0 = cnts[:, 0].clamp_min(1.0) + dthr = (kth - t_prev0).clamp_min(1e-3) + slope = ((torch.log2(n0) - math.log2(k)) / dthr).clamp(SLOPE_MIN, SLOPE_MAX) + tgt0 = min(LIST_T0_TARGET * k, LIST_T0_COUNT_MAX) + t0 = kth - math.log2(tgt0 / k) / slope + fit_ok = torch.isfinite(t_prev0) & (n0 > k) + t0 = torch.where(fit_ok, t0, kth - GUARD_LO * d_fb) + park = torch.full_like(kth, LIST_PARK_LINE) + new0 = torch.where(valid, t0, inf) + new1 = torch.where(valid, park, inf) + new2 = torch.where(valid, park + park, inf) + else: + c0 = cnts[:, 0].clamp_min(1.0) + c2 = cnts[:, 2].clamp_min(1.0) + dthr = (t_prev2 - t_prev0).clamp_min(1e-3) + slope = ((torch.log2(c0) - torch.log2(c2)) / dthr).clamp(SLOPE_MIN, SLOPE_MAX) + # anchor count estimate: slide the anchor onto the prev line fit + anch_c = (c2 * torch.exp2(-(anchor - t_prev2) * slope)).clamp(1.0, 1e6) + t0 = anchor + torch.log2(anch_c / (LINE_TARGETS[0] * k)) / slope + t1 = anchor + torch.log2(anch_c / (LINE_TARGETS[1] * k)) / slope + t2 = anchor + torch.log2(anch_c / (LINE_TARGETS[2] * k)) / slope + # t_prev2 < 1e29 also rejects a parked line left by a tier flip + fit_ok = torch.isfinite(t_prev0) & (t_prev2 < 1e29) & (c0 > c2) + t0 = torch.where(fit_ok, t0, kth - GUARD_LO * d_fb) + t1 = torch.where(fit_ok, t1, kth - 1e-6) + t2 = torch.where(fit_ok, t2, kth + GUARD_HI * d_fb) + # strictly ascending (kernel line-validity contract) + t1 = torch.maximum(t1, t0 + 1e-4) + t2 = torch.maximum(t2, t1 + 1e-4) + new0 = torch.where(valid, t0, inf) + new1 = torch.where(valid, t1, inf) + new2 = torch.where(valid, t2, inf) + s[:, 0] = new0 + s[:, 1] = new1 + s[:, 2] = new2 + s[:, 3:8] = 0.0 + rungs = self.seed_rungs[:num_rows] + rungs[:, 0] = new0 + rungs[:, 1] = new1 + rungs[:, 2] = new2 + if self.cand_ctl is not None: + nc = min(num_rows, self.cand_ctl.shape[0]) + self.cand_ctl[:nc].zero_() + self.cand_cur[:nc].zero_() + + def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict: + """kwargs for CuteDSLFP4PagedMQALogitsRunner.forward covering the + planned emission tier (caller merges into its call).""" + kw: dict = {} + if emit_tier in ("counts", "list"): + kw["seed_thr"] = self.seed_row[:num_rows] + if emit_tier == "list": + kw.update( + accept_cap=LIST_SEG_A, + cand_out=self.cand_vals[:num_rows], + cand_idx_out=self.cand_idx[:num_rows], + cand_ctl_out=self.cand_ctl[:num_rows], + cand_cur_out=self.cand_cur[:num_rows], + ) + return kw + + def topk_ext_kwargs( + self, route: TopkRoute, num_rows: int, block_max: Optional[torch.Tensor] + ) -> dict: + """kwargs for trtllm::cute_dsl_gvr_topk_decode consuming this + step's emission per the picked route.""" + kw: dict = { + "xstate": self.xstate[:num_rows], + "cluster_size": route.cluster_size, + } + if route.num_threads is not None: + kw["num_threads"] = route.num_threads + if route.tier in ("counts", "list"): + kw["seed_thr"] = self.seed_row[:num_rows] + elif route.tier == "rungs": + # [rows, 3] seed selects the op's ext_rungs variant + kw["seed_thr"] = self.seed_rungs[:num_rows] + if route.tier == "list": + # accept_cap must match the emitter's segment geometry: the + # buffers are laid out at bases 0 / LIST_SEG_A / 2*LIST_SEG_A, + # and the consumer derives the C capacity from the tensor + # width minus 2*accept_cap. + kw.update( + cand_vals=self.cand_vals[:num_rows], + cand_idx=self.cand_idx[:num_rows], + cand_ctl=self.cand_ctl[:num_rows], + accept_cap=LIST_SEG_A, + ) + if route.attach_block_max and block_max is not None: + kw["block_max"] = block_max + return kw diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py new file mode 100644 index 000000000000..2bc3992e0b15 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_routing.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host-side routing for the GVR top-k decode family. + +One kernel, three assist tiers selected by which emission inputs the +indexer epilogue produced for this step: + + * ``list`` - bucketed candidate list + packed seed row (v5): the + top-k pass never re-reads the row on a hit. + * ``counts`` - packed seed row only ([rows, 8]: three lines + three + counts): one filtered row pass, no in-kernel counting. + * ``rungs`` - closed-loop lines only (no emission): in-kernel + multi-line count, then the stock collect. + * ``none`` - stock kernel (v1 path). + +``plan_emission`` decides which tier the epilogue should emit for the +NEXT step (the emission tax is shape-dependent); ``pick_config`` maps +(tier, B, N, K) to concrete launch knobs for THIS step. + +All thresholds are B200 deployment defaults, kept in one place so +retuning is a constant edit, not a logic edit. +""" + +from dataclasses import dataclass +from typing import Optional + +# ---- thresholds (B200 deployment defaults) -------------------------------- +# +# UNITS: every threshold is fitted on KERNEL-ONLY time for both the +# indexer and the top-k; do not mix in wall-clock numbers when retuning. + +# Below this the stock kernel already runs at the fixed-cost floor. +ASSIST_MIN_N_COMP = 2048 + +# Block-skip prefix break-evens. Retune these from real captures only: +# the real block-max distribution decides the pass rate. +SKIP_MIN_N_COUNTS = 65536 # va: attach block_max from here up +SKIP_MIN_N_RUNGS_FLASH = 131072 # vb (flash): bm pays from here +SKIP_CS_MIN_N_RUNGS = 196608 # vb: cluster split from here up + +# Emission-cost model: counts emission is a batch-only latency chain +# (hides at large batch); list emission grows with batch and context. +LIST_EMIT_MIN_N = 65536 # shorter rows: the emission outweighs the saving +LIST_EMIT_MAX_B = 4 # past four rows the list stops repaying its emission +COUNTS_MIN_TOKENS = 524288 # B * raw length; below this the rungs tier wins +RUNGS_ONLY_MIN_N = 16384 # short-row band where rungs also beats counts +RUNGS_ONLY_MAX_N = 49152 + +# Mid-row weak band: the stock kernel's split grid fits one wave here +# and out-scans every assist tier, so the epilogue emits nothing. +ASSIST_WEAK_MIN_N = 49152 +# The band interior is unmeasured; it stays on the stock kernel. +ASSIST_WEAK_MAX_N = 65536 +ASSIST_WEAK_MAX_B = 8 + +# rungs block_max pays only at small K; at large K the prefix is overhead. +RUNGS_BM_MAX_K = 512 + +# 512-thread build wins for list-hit rows at small K (work is O(list)). +SMALL_K_LIST_THREADS = 512 +SMALL_K_MAX = 512 + +# GPC packing: cs=8 only while all row-clusters fit half the device; +# cs4/2 keep a 10% headroom. +CS8_HALF_DEVICE = 2 +CS_HEADROOM_NUM = 9 +CS_HEADROOM_DEN = 10 + + +@dataclass +class TopkRoute: + """Launch knobs for one decode step of the GVR top-k kernel.""" + + tier: str # list | counts | rungs | none + cluster_size: int = 1 + num_threads: Optional[int] = None # None = runner heuristic + attach_block_max: bool = False + + +def plan_emission( + batch: int, n_comp: int, k: int, have_epilogue: bool, compress_ratio: int = 4 +) -> str: + """Which assist tier the indexer epilogue should emit this step. + + ``n_comp``: compressed row length (post compress_ratio) - the + top-k kernel's N. Returns the tier name; the epilogue emits the + matching buffers and the next top-k launch routes on them. + """ + if n_comp < ASSIST_MIN_N_COMP: + # short rows: the stock kernel is already under the fixed cost + return "none" + if have_epilogue and n_comp >= LIST_EMIT_MIN_N and batch <= LIST_EMIT_MAX_B: + # must stay ahead of the weak-band gate: a list hit never scans the row + return "list" + if batch <= ASSIST_WEAK_MAX_B and ASSIST_WEAK_MIN_N <= n_comp < ASSIST_WEAK_MAX_N: + return "none" # stock's split grid wins this band outright + if ( + have_epilogue + and batch * n_comp * compress_ratio >= COUNTS_MIN_TOKENS + and not (RUNGS_ONLY_MIN_N <= n_comp < RUNGS_ONLY_MAX_N) + ): + return "counts" + return "rungs" # closed-loop lines cost nothing to carry + + +def pick_config(tier: str, batch: int, n_comp: int, k: int, num_sms: int) -> TopkRoute: + """Map (tier, B, N, K) to launch knobs. Pure function of shape.""" + r = TopkRoute(tier=tier) + if tier == "none": + return r + if tier == "list": + if k <= SMALL_K_MAX: + r.num_threads = SMALL_K_LIST_THREADS + # block_max: miss rows take a skip-walk instead of a dense re-scan + r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS + return r + if tier == "counts": + r.attach_block_max = n_comp >= SKIP_MIN_N_COUNTS + return r + # rungs (vb); the skip prefix and cluster split are mutually exclusive + if n_comp >= SKIP_CS_MIN_N_RUNGS: + if batch * 8 <= num_sms // CS8_HALF_DEVICE: + r.cluster_size = 8 + elif batch * 4 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN: + r.cluster_size = 4 + elif batch * 2 <= (num_sms * CS_HEADROOM_NUM) // CS_HEADROOM_DEN: + r.cluster_size = 2 + if r.cluster_size == 1 and k <= RUNGS_BM_MAX_K and n_comp >= SKIP_MIN_N_RUNGS_FLASH: + r.attach_block_max = True + return r diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py index 7f30511c6a1a..d877bdef6b3a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py @@ -27,6 +27,7 @@ """ import math +import os from dataclasses import dataclass from typing import Optional @@ -42,6 +43,47 @@ from .block_scan import warp_scan +def _env_flag(name: str) -> bool: + """Debug-knob parse that never raises at import time: any of + ""/0/false/off/no (case-insensitive) is False, everything else True.""" + return os.environ.get(name, "0").strip().lower() not in ("", "0", "false", "off", "no") + + +# Diagnostic knob: compile per-phase clock64 stamps of the list path +# into the spare xstate slots (harness-side analysis). Off by default; +# NEVER set in production. +_P4_TAIL_DBG = _env_flag("TRTLLM_GVR_P4_TAIL_DBG") +# P4 sub-phase clock64 breakdown -> xstate[1,2,4,5,6,7] (debug: clobbers +# the closed-loop thr/anch publish; single-shot cells only, not chains) +_P4_SUB_DBG = _env_flag("TRTLLM_GVR_P4_SUB_DBG") +# TRTLLM_GVR_P4_SUB_DBG=2: publish the P4 HEAD triple (minmax / histogram build / +# coarse search) instead of the tail triple, so the phase budget adds up. +_P4_SUB_HEAD = os.environ.get("TRTLLM_GVR_P4_SUB_DBG", "0").strip() == "2" +# Exact-tail small-class pair buffer. The scatter parks each member of the +# straddling tie class as (value bits, index) here, above the 256 digit bins +# the large-class radix zeroes and above its [256..258] scalars, so the repair +# never has to re-walk the candidates. The capacity is DERIVED from the bin +# count, never assumed: K=2048 with R0 shrinks the histogram to 512 bins (see +# the kNumBins override in __init__), where 260 + 2*128 would run 4 ints past +# the end of the allocation - into the per-thread count buffer that follows it. +# Rounded down to a power of two so the scatter can wrap the ordinal with a +# mask instead of a bounds branch; a class past the cap takes the large-class +# route, which does not use this buffer. +_PAIR_BASE = 260 +_PAIR_MAX = 128 + + +def _pair_cap_for(n_bins: int) -> int: + """Largest power-of-two pair count that fits [_PAIR_BASE, n_bins).""" + room = (n_bins - _PAIR_BASE) // 2 + if room < 1: + return 0 + return min(_PAIR_MAX, 1 << (room.bit_length() - 1)) + + +_SKIP_DBG = _env_flag("TRTLLM_GVR_SKIP_DBG") + + # --------------------------------------------------------------------------- # DSMEM primitives (inline PTX) # Adapted from single_pass_multi_cta_radix_topk_cluster.py. @@ -281,7 +323,14 @@ def __init__( smem_cache_elems: int = 32768, seqlen_sorted: bool = False, kc_diet: Optional[bool] = None, + pdl_wait_late: bool = False, + p4_fine_rangetest: Optional[bool] = None, + p4_scat_rangetest: bool = False, enable_r0: bool = True, + accept_cap: "int | None" = None, + kc_override: "int | None" = None, + self_scan: bool = False, + cap_c: "int | None" = None, r0_qfracs: Optional[tuple] = None, mt_unroll: int = 4, p1b_cache: Optional[bool] = None, @@ -292,8 +341,19 @@ def __init__( enable_p4_rank_scatter_exact: Optional[bool] = None, p4_exact_tail: Optional[bool] = None, p4_tail_fast: Optional[bool] = None, # [p4tt] + p4_tail_v3: Optional[bool] = None, + p4_no_fine: Optional[bool] = None, + p1r_rescue: bool = True, + num_bins: Optional[int] = None, p4_warp_redundant: bool = True, p2_warp_redundant: bool = True, + enable_block_skip: bool = False, + use_ext_counts: bool = False, + ext_rungs: bool = False, + use_ext_cand: bool = False, + cand_cap: int = 5120, + cand_rung: int = 0, + emit_xstate: bool = False, ): # Redundant-warp sync reduction: every warp replays the block # reduce + decision from the same staged SMEM partials in the @@ -403,7 +463,9 @@ def __init__( params = GvrParams.get(self._dtype_name, top_k, self.compress_ratio) self.kC = params.kC - self.kNumBins = params.kNumBins + # num_bins: the coarse histogram width. The bin search walks it + # per warp, so it sets the barrier count of the P4 head. + self.kNumBins = params.kNumBins if num_bins is None else int(num_bins) self.kFTarget = params.kFTarget # Kernel-wide constants. @@ -414,7 +476,7 @@ def __init__( self.FLT_MAX = 3.4028235e38 self.NEG_FLT_MAX = -self.FLT_MAX - # --- R0 histogram-ladder admission (default ON) --- + # --- op#26 R0 histogram-ladder admission (default ON) --- # enable_r0: replace the Phase-2 secant search with a single-pass # multi-threshold "rung ladder" admission seeded by a 256-bin # histogram over the prev-topK gathered values (P1b). @@ -422,20 +484,13 @@ def __init__( # workloads (25-cell seq-len scan) where R0 wins 24/25 vs the # secant baseline, geomean 1.33x (pro 128k 2.10x). Correctness is # value-set-exact vs torch.topk (186/186 across dtype/K/N/BS/cluster - # + tie plateaus). - # THE SECANT PATH IS NOT DEAD CODE AND MUST NOT BE DELETED. It has - # two distinct roles: - # (a) EXACT FALLBACK, live at the default enable_r0=True: when the - # rung ladder admits no candidate (R0 miss) the row falls - # through to phase2_secant_search, so this code runs in - # production on every hint-unrepresentative row; - # (b) DIFFERENTIAL ORACLE, via enable_r0=False: it is the classic - # baseline the R0 admission is checked against - # (test_..._r0_equivalence), and the direct-drive entry used to - # bisect admission-vs-baseline regressions. - # All R0 fields are const-foldable, so an enable_r0=False kernel is - # byte-identical to the pre-R0 upstream base and the disabled branch - # costs nothing at runtime. + # + tie plateaus). The secant path is retained verbatim and remains + # reachable via enable_r0=False; it is the exact fallback for the + # large-N / cold-hint (low preIdx hit-rate) regime where R0 can + # regress on the synthetic worst axis — a follow-up PR adds a + # data-driven dispatch guard to route between the two. All R0 fields + # are const-foldable, so an enable_r0=False kernel is byte-identical + # to the pre-R0 upstream base. # r0_qfracs: descending h-space quantile fractions defining the M # candidate rungs (ascending threshold values); None => no rungs. # r0_vseed: park P1's pmean (the secant init probe) as one extra @@ -454,9 +509,31 @@ def __init__( self.enable_r0 = bool(enable_r0) self.mt_unroll = int(mt_unroll) self.fb_fix = bool(fb_fix) - # C7 dispatch (host policy folded into the ctor; all gated on + # enable_block_skip: gate the R0 count pass and the Phase-3 + # stream-write on per-32-position upper bounds from the indexer + # epilogue (block_max [num_rows, nb_pad*4] fp32; record r = max + # over positions [r*32, r*32+32) of the stored logits). Lossless: + # the active list is built at the loosest rung, so a skipped + # block holds nothing >= any rung and every count equals its + # dense value. + self.enable_block_skip = bool(enable_block_skip) + self.SKIP_BLOCK = 32 + self.SKIP_BLOCK_LOG2 = 5 + self.SKIP_MAX_BLOCKS = 8192 # smem active-list budget (32KB of local + # ids); bounds block-skip to N_local <= 262144 at grain 32 - longer + # rows run the dense fallback and owe nothing to block-skip + self.SKIP_UNROLL = 2 + self.skip_order = "grouped" + if enable_block_skip and num_threads not in (512, 1024): + raise ValueError("enable_block_skip requires num_threads in {512, 1024}") + if enable_block_skip and not enable_r0: + # The compact machinery hangs off the R0 count pass and the + # phase-3 stream-write; without R0 the 16KB list SMEM would be + # allocated but the skip could never engage. + raise ValueError("enable_block_skip requires enable_r0") + # C7 dispatch (op#26 host policy folded into the ctor; all gated on # enable_r0 so an OFF kernel is byte-identical to the base): - # - qfracs default = M2D (0.85, 0.35): the shipped dispatch uses M2D for + # - qfracs default = M2D (0.85, 0.35): dispatch_r0_op26 ships M2D for # every (dtype, K, N); the M=2 pass is ~free and the R1 falsi shot # covers the 3-7% bracket misses. uh4 (M=4) was silicon-falsified # (mc geomean 0.956 — admission != latency). @@ -466,7 +543,7 @@ def __init__( # single-CTA path does NOT reproduce in the cluster kernel # (latency-bound, different SMEM budget). nsys cs=4: K1024 # ~1.01x / K2048 ~1.02x / K512 wash, 0 losses, exact. Matches - # the multi-CTA dispatch (unconditional ON). + # op26 dispatch_p1bc_mc (unconditional ON). # * cs=1 (single-CTA): (dtype != fp32). The gather-cache wins # +0.8-2.8% on 16-bit (random half-prec gather is the cost) but # is flat/negative on fp32 (occupancy at kC=6144), so OFF there. @@ -476,14 +553,14 @@ def __init__( if r0_vseed is None: r0_vseed = enable_r0 if enable_r0 and r0_qfracs is None: - # Per-K default (full-envelope audit, 2772 + # Per-K default (2026-07-16 vseed full-envelope audit, 2772 # cells): with the virtual seed rung on, pmean covers q.35's # admission region for K512/K1024 (2 count columns = zero # column tax); K2048 keeps q.35 (kC/K = 2.5 makes a fat admit # costlier than a slim 2-pass miss). Without vseed, q.35 must # stay for all K (it is the only slim rung). - # K2048 low rung 0.85 -> 0.6 (real-content rung recalibration, - # paired cold-L2 A/B): the shipped + # K2048 low rung 0.85 -> 0.6 (2026-07-19 real-content rung + # recalibration + paired nsys cold-L2 A/B, B200): the shipped # 0.85 rung's admission straddles [K, kC] on real V3.2 decode # captures (bracket on 86% of steps -> one extra falsi pass); # 0.6 lands the first pass. Measured: real V3.2 geomean @@ -506,12 +583,16 @@ def __init__( # kernel passes False for BOTH member instances so their SMEM layouts # stay byte-identical (the DSL sizes the launch from the last-traced # SmemAllocator only; see GvrTopKLBKernel). + # pdl_wait_late: move the PDL wait past the prologue so that work + # overlaps the producer's tail. Off by default: the entry-point + # wait is what upstream emits. + self.pdl_wait_late = bool(pdl_wait_late) if kc_diet is None: kc_diet = cluster_size == 1 if enable_r0 and top_k == 512 and kc_diet and self.kC > 3072: self.kC = 3072 - # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (paired - # cold-L2 A/B, all cells exact). The P4 zero / + # K2048 R0 Phase-4 histogram diet: 2048 -> 512 bins (2026-07-19 + # paired nsys cold-L2 A/B on B200, all cells exact). The P4 zero / # atomic build / serial scan all shrink 4x; the deeper boundary-bin # recursion costs less than the saved passes at kC=6144 candidates. # Measured vs this head: real V3.2 decode captures geomean +6.1% @@ -524,6 +605,22 @@ def __init__( # measured as a wash under the same protocol and stay stock. if enable_r0 and top_k == 2048 and self.kNumBins > 512: self.kNumBins = 512 + # p4_fine_rangetest: filter the fine recursion by value range + # instead of recomputing each candidate's bin. OFF - the two are + # not fp32-equivalent, and the scatter still classifies by bin + # recompute, so a candidate they disagree on is counted by one + # pass and placed by the other, leaving an output slot unwritten. + self.p4_fine_rangetest = False if p4_fine_rangetest is None else bool(p4_fine_rangetest) + # p4_scat_rangetest: the scatter classifies each candidate as + # above / inside / below the straddling bin and recomputed the + # bin index to do it; the same value-range compare the fine + # recursion uses answers it without the subtract, multiply and + # two clamps per candidate. + # DEFAULT OFF - same fp32 non-equivalence as p4_fine_rangetest, + # in the opposite direction: the range test admits a candidate + # the histogram binned elsewhere, so the scatter writes past the + # rank it reserved and out-of-range indices reach the output. + self.p4_scat_rangetest = bool(p4_scat_rangetest) self.r0_qfracs = tuple(float(q) for q in r0_qfracs) if r0_qfracs else () if self.r0_qfracs: assert all(0.0 < q < 1.0 for q in self.r0_qfracs), self.r0_qfracs @@ -531,7 +628,7 @@ def __init__( "r0_qfracs must be descending h (ascending threshold value)" ) self.M_thr = len(self.r0_qfracs) - # --- vseed: fold P1's pmean (the secant init + # --- vseed (2026-07-16): fold P1's pmean (the secant init # probe) into the M-ary R0 count pass as one extra "virtual rung". # Fixes the flash-1M fat-admission regression (the coarse q.85 rung # admits ~4400 candidates where pmean admits ~630 -> 7x P3/P4 cand @@ -545,6 +642,113 @@ def __init__( self.M_thr = self.M_qf + 1 # need[m] = ceil(q_m * K) prev-topK values >= rung m. self.qneeds = tuple(max(1, int(math.ceil(q * self.top_k))) for q in self.r0_qfracs) + # use_ext_counts (waterfall L1 admission): thresholds AND their + # exact counts arrive from the indexer epilogue (interface v2) — + # P1b and the M-ary count pass are skipped; an in-band rung is + # re-measured ONCE through the seeded refine (per-thread hand-off) + # and accepted; a miss seeds log-falsi with the external brackets. + # emit_xstate: write the per-row closed-loop state at Phase 4 exit + # (interface v2: [0] valid, [1] kth proxy, [2] accepted threshold, + # [3] cand_count). cs==1 only (leader-gather rows land later). + self.emit_xstate = bool(emit_xstate) + # use_ext_cand (waterfall L2 direct-to-P4): pre-collected (value, + # index) pairs from the epilogue land straight in smem_keys/vals — + # no P1, no counting, no P3 scan. Eligible when void==0, claimed + # <= cand_cap and the collect rung's exact count is in [K, kC]; + # ineligible rows fall through to the ext-counts path. + self.use_ext_cand = bool(use_ext_cand) + if kc_override is not None: + # physical candidate-buffer capacity override (B* search) + self.kC = int(kc_override) + # acceptance band top B*: a cut whose count fits [K, B*] goes + # straight to Phase 4. Physically bounded by kC. + self.accept_cap = int(accept_cap) if accept_cap is not None else self.kC + self.cand_cap = int(cand_cap) + # list path: the score column is staged into a DEDICATED smem + # region sized cand_cap fp32 (96KB at 24576). Budget note: this + # coexists with everything except a simultaneously-enabled big + # slice cache (128KB) - that combination exceeds the 227KB CTA + # limit and fails loudly at compile time. Long rows (the list + # path's target) cannot enable the slice cache anyway. + # v5 bucketed layout: tensor width = 2 * accept_cap + # (segments A, B) + segment-C capacity; the admission + # bounds the ENTRY COUNT by C's capacity (the only + # segment that can void). + self.list_cap = max(0, int(cand_cap) - 2 * self.accept_cap) + self.cand_rung = int(cand_rung) + # self_scan (fused self-contained mode): the kernel streams the + # row once against the three closed-loop lines, bucketing values + # into on-chip segments (A >= t2 / B [t1,t2) / C [t0,t1) at bases + # 0 / accept_cap / 2*accept_cap) and positions into cand_idx. A + # line cut compacts the winning segment to the smem_keys prefix + # and fills smem_vals with segment coordinates, after which the + # list consumer runs unchanged. Ineligible rows take the stock + # fallback. + self.self_scan = bool(self_scan) + if self.self_scan: + if not use_ext_counts: + raise ValueError("self_scan requires use_ext_counts") + if use_ext_cand: + raise ValueError("self_scan and use_ext_cand are exclusive") + if dtype != cutlass.Float32: + raise ValueError("self_scan is fp32-only (v1)") + if self.enable_smem_cache: + raise ValueError( + "self_scan and enable_smem_cache exceed the CTA smem budget together" + ) + # on-chip segment budget: values only, 4B/entry; C sized so + # keys(160KB) + vals(32KB) + hist + scratch stay under the + # 227KB CTA limit with room for the stage-2 skip list. + self.seg_total = 2 * self.accept_cap + int(cap_c if cap_c is not None else 16384) + self.cap_c = self.seg_total - 2 * self.accept_cap + # cp.async staging for the phase-0 dense scan: the pair-step + # pipeline keeps 2 pairs x 2 slots in flight per thread, so + # exactly 4 slot rows are required. The 64KB staging fits the + # CTA budget only with the C segment trimmed to <= 16384 + # (keys 128KB + staging 64KB + hist/scratch); larger C would + # silently overrun the alias, so reject it outright. Rows are + # 16B slots; never fewer than vals holds so the alias always + # covers it. + if self.cap_c > 16384: + raise ValueError("self_scan requires cap_c <= 16384 (staging budget)") + self.stage_slots = 4 + self.stage_rows = max(self.stage_slots * self.num_threads, self.kC // 4) + else: + self.seg_total = self.kC + self.cap_c = 0 + if use_ext_cand and not use_ext_counts: + raise ValueError("use_ext_cand requires use_ext_counts") + if (use_ext_counts or ext_rungs or use_ext_cand) and not enable_r0: + # the effective flags below are and-ed with enable_r0; reject + # instead of silently compiling the stock path + raise ValueError("ext tiers require enable_r0") + if use_ext_cand and self.list_cap <= 0: + raise ValueError( + f"cand_cap={cand_cap} leaves no C segment past 2*accept_cap={2 * self.accept_cap}" + ) + # ext_rungs (two-pass variant B): closed-loop rung THRESHOLDS come + # from the host (previous-step xstep lines); the kernel counts them + # itself via the stock R0 multi-count and admits the tightest rung + # in [K, kC]. Exclusive with use_ext_counts (which also imports + # the counts and skips nothing else). + self.ext_rungs = bool(ext_rungs) and bool(enable_r0) + if self.ext_rungs and bool(use_ext_counts): + raise ValueError("ext_rungs is exclusive with use_ext_counts") + self.use_ext_counts = bool(use_ext_counts) and bool(enable_r0) + if self.use_ext_counts: + if not self.fb_fix: + raise ValueError("use_ext_counts requires fb_fix") + if self.M_thr != 3: + raise ValueError("use_ext_counts expects exactly 3 seed rungs") + if self.ext_rungs: + if not self.fb_fix: + raise ValueError("ext_rungs requires fb_fix") + if self.M_thr != 3: + raise ValueError("ext_rungs expects exactly 3 seed rungs") + # cluster_size > 1 supported: the ext rungs/counts are + # per-row (identical across the cluster), the stock multi + # count pass cluster-merges as usual, and the L2 direct + # loader runs leader-only (peers contribute zero candidates). # R1 inline shot aim in log2-count space: geometric center of the # [K, kC] acceptance window. self.log2_r1aim = math.log2(math.sqrt(self.top_k * self.kC)) if self.r0_qfracs else 0.0 @@ -555,14 +759,14 @@ def __init__( else 0.0 ) - # --- P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- + # --- op#7 P4 fused rank-and-scatter (inert until enable_p4_rank_scatter) --- # Replaces phase4_histogram_snap's k-th-bin search + 2-pass writeback - # with a single rank-and-scatter pass (PR#15709), cutting Phase-4 + # with a single rank-and-scatter pass (op#7 PR#15709), cutting Phase-4 # barriers ~14 -> ~7. On a latency-bound kernel that is a whole-kernel # win (~1.078x, HW-invariant). enable_p4_rank_scatter_exact adds ONE # fine-histogram recursion on the straddling coarse bin so the result is # bit-exact vs torch.topk (adds a few barriers back but still < snap). - # Default ON with R0: measured over the 4k-1M BS=1 best/worst envelope + # Default ON with R0: nsys over the op22 4k-1M BS=1 best/worst envelope # gives geomean ~1.09x (K1024 1.12 / K2048 1.12 / K512 1.05) with NO # cell regressing >2%. Resolves to OFF when enable_r0 is False, so the # base kernel stays byte-identical to upstream. @@ -573,47 +777,62 @@ def __init__( self.enable_p4_rank_scatter = bool(enable_p4_rank_scatter) self.enable_p4_rank_scatter_exact = bool(enable_p4_rank_scatter_exact) # p4_exact_tail: ambiguity-gated exact tie-resolution for the fine - # straddling bin (see phase4_rank_scatter). The fine recursion - # resolves values to range/(kNumBins*256) — WINDOW-RELATIVE, so ANY - # dtype (fp32 or upconverted 16-bit) can leave distinct values in - # one fine bin whenever the Phase-2 bracket is wide relative to the - # boundary-local ULP (e.g. fp16 1.0 vs 1.25 under a [0, 65504] - # bracket); values straddling the kK boundary inside one fine bin - # were previously picked in arrival order (observed as |miss|=1 - # with |dv| ~ 3e-6 on real fp32 captures). The tail radix re-ranks - # on the full fp32 order key — candidate keys are ALWAYS fp32 - # (16-bit inputs are upcast injectively at collect), so the repair - # is exact for every supported dtype WHEN ENABLED. Default ON for - # fp32 only: on fp32 the gate fires rarely and the fix is ~free, - # but 16-bit quantization puts value plateaus at the boundary on - # virtually every input, so the gate fires constantly — measured - # B200 envelope cost (bf16, K512/K1024 x 16k-262k x BS 1-256, - # same-process paired) is gm 1.29-1.36x, worst 2.27x, while typical - # bf16 inputs (randn and quantized-tie, 48 paired runs) are already - # value-exact without it: 16-bit misses need an adversarially wide - # Phase-2 bracket (distinct 16-bit values inside one fine bin). - # 16-bit callers that need the guarantee opt in via the knob (see - # test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit). + # straddling bin (fp32 inputs only; see phase4_rank_scatter). The + # fine recursion resolves values to range/(kNumBins*256); two fp32 + # values closer than that straddling the kK boundary inside one fine + # bin were previously picked in arrival order (observed as |miss|=1 + # with |dv| ~ 3e-6 on real Pro 512k-ISL captures). Default ON for + # fp32 rank-scatter-exact kernels; 16-bit inputs keep the arrival + # fill (their upconverted keys are already fully resolved by the + # two-level histogram, and 16-bit tie plateaus are bitwise-equal, + # where arrival order is value-exact). if p4_exact_tail is None: p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32 self.p4_exact_tail = bool(p4_exact_tail) and self.enable_p4_rank_scatter_exact - # [p4tt] p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the - # exact-tail fire branch. When the (b*, sb*) tie class holds <= 128 - # entries (the real firing cells have 2), ONE candidate pass collects - # (value_bits, cand_idx) pairs into SMEM and thread0 selects the - # top-need exactly, replacing the 4 unconditional radix passes - # (~5.3us -> ~1 pass on pro/512k). Larger tie classes fall through to - # the existing radix select. Pure optimization (the radix backstop - # keeps exactness identical either way); False compiles the original - # text (byte-identical PTX modulo kernel name) for A/B. - # Default gate = p4_exact_tail AND top_k >= 1024: the non-firing - # codegen tax concentrates at K512 cs=1 mid-N (flash 64k/128k - # -6.6/-9.1%, cross-GPU reproducible) while the - # fire census (pro/512k bench + 9 per-layer fixture cells) contains - # NO K512 cell — so K512 keeps the original byte-identical kernel. + # p4_tail_fast: tiny-tie COLLECT+SELECT fast path inside the + # exact-tail fire branch: when the boundary tie class holds few + # enough entries to buffer, one candidate pass replaces the radix + # passes; larger classes fall through to the radix backstop, so + # exactness is identical either way. if p4_tail_fast is None: # [p4tt] - p4_tail_fast = self.p4_exact_tail and top_k >= 1024 + # default follows p4_exact_tail for every K (the compacted + # path pays for the boundary class only) + p4_tail_fast = self.p4_exact_tail self.p4_tail_fast = bool(p4_tail_fast) and self.p4_exact_tail # [p4tt] + # p4_tail_v3: compacted-class repair (block-parallel radix + + # pure-tie pre-check) in place of the stock thread0 serial + # select. Default follows the configuration: a stock kernel gets + # upstream's body, any emission-assisted one gets the rewrite. + if p4_tail_v3 is None: + p4_tail_v3 = bool( + use_ext_counts or use_ext_cand or ext_rungs or self_scan or enable_block_skip + ) + self.p4_tail_v3 = bool(p4_tail_v3) + # p4_no_fine: drop the 256-bin fine level from phase 4 and let the + # tail repair rank the whole straddling COARSE bin instead. A class + # past the tail's pair buffer falls into its radix, which handles + # any size. Gated to the same configuration as that small-class + # route; the stock kernel keeps the fine level so its codegen is + # untouched. + if p4_no_fine is None: + p4_no_fine = bool( + self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3 + ) and not (self.p4_fine_rangetest or self.p4_scat_rangetest) + self.p4_no_fine = bool(p4_no_fine) + if self.p4_no_fine: + if not (self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3): + raise ValueError("p4_no_fine requires the exact-tail repair chain") + if self.p4_fine_rangetest or self.p4_scat_rangetest: + raise ValueError("p4_no_fine is incompatible with the range-test arms") + if self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3: + if _pair_cap_for(self.kNumBins) < 1: + raise ValueError( + f"kNumBins={self.kNumBins} leaves no room for the tail pair buffer" + ) + # p1r_rescue: rebuild the refine bracket from the row when the + # seed bracket is degenerate (e.g. the zero-init prev_topk every + # request's first decode step feeds). ON by default. + self.p1r_rescue = bool(p1r_rescue) # ------------------------------------------------------------------ # SMEM slice cache loader. Streams this CTA's slice GMEM → SMEM via @@ -982,6 +1201,68 @@ def phase1_preidx_stats( s_iscalars[4] = cutlass.Int32(0) cute.arch.barrier() + # ------------------------------------------------------------------ + # P1r — degenerate-seed rescue: rebuild the refine bracket from the + # data itself. Runs only when the preIdx gather produced an unusable + # bracket (duplicate or invalid preIdx: cold-start zero-init slots, + # stale slots pointing past N, or an all-tied gather). A full-row + # min/max restores the P2 invariant count(>= v_lo) >= K, so the + # normal pipeline stays exact; the extra row scan is paid only by + # the (rare) degenerate rows. Bounds are clamped to +-FLT_MAX/2 so + # the secant range arithmetic stays finite against inf-laden rows. + # ------------------------------------------------------------------ + @cute.jit + def phase1r_data_reseed( + self, + input_row, # cute.Tensor [N] (row-major slice of the row) + N, # runtime row length + smem_wmin_f32, # cute.Tensor [NUM_WARPS] float32 (reused P1 buffer) + smem_wmax_f32, # cute.Tensor [NUM_WARPS] float32 (reused P1 buffer) + s_thr, # cute.Tensor [3] float32: [threshold, val_lo, val_hi] + s_iscalars, # [cand_count, done, cnt_lo, cnt_hi, out_count, ...] + s_mt_thr, # rung columns (r0_vseed parks the seed line here) + tidx, + warp_id, + lane, + ): + local_min = cutlass.Float32(self.FLT_MAX) + local_max = cutlass.Float32(self.NEG_FLT_MAX) + i = cutlass.Int32(tidx) + while i < N: + v = self._load_fp32(input_row, i) + local_max = cute.arch.fmax(local_max, v) + local_min = _fmin_f32_inline(local_min, v) + i = i + cutlass.Int32(self.num_threads) + wmin = self.warp_reduce_min_f32(local_min) + wmax = self.warp_reduce_max_f32(local_max) + if lane == 0: + smem_wmin_f32[warp_id] = wmin + smem_wmax_f32[warp_id] = wmax + cute.arch.barrier() + if tidx == 0: + rmin = cutlass.Float32(self.FLT_MAX) + rmax = cutlass.Float32(self.NEG_FLT_MAX) + for w in cutlass.range_constexpr(self.num_warps): + rmax = cute.arch.fmax(rmax, smem_wmax_f32[w]) + rmin = _fmin_f32_inline(rmin, smem_wmin_f32[w]) + # finite clamp keeps rng = val_hi - val_lo representable; rows + # with mass beyond +-FLT_MAX/2 are adversarial-only (production + # indexer scores are small finite values). + rmin = cute.arch.fmax(rmin, cutlass.Float32(self.NEG_FLT_MAX * 0.5)) + rmax = _fmin_f32_inline(rmax, cutlass.Float32(self.FLT_MAX * 0.5)) + mid = (rmin + rmax) * cutlass.Float32(0.5) + s_thr[0] = mid + s_thr[1] = rmin + s_thr[2] = rmax + if cutlass.const_expr(self.r0_vseed): + s_mt_thr[self.M_thr - 1] = mid + s_iscalars[0] = cutlass.Int32(0) # cand_count + s_iscalars[1] = cutlass.Int32(0) # done + s_iscalars[2] = N # cnt_lo: count(>= row min) = N, truthful + s_iscalars[3] = cutlass.Int32(1) # cnt_hi seed (same as P1) + s_iscalars[4] = cutlass.Int32(0) # out_count + cute.arch.barrier() + # ------------------------------------------------------------------ # P1b — 256-bin SMEM histogram over the prev-topK gathered values # (band [v_lo, v_hi] = P1's pmin/pmax = s_thr[1]/s_thr[2]), then M @@ -1155,6 +1436,427 @@ def phase1b_hspace_rungs_cached( # → warp reduce → block reduce → s_iscalars[0] = cand_count. # Optionally DSMEM-aggregates across the cluster. # ------------------------------------------------------------------ + @cute.jit + def phase0_scan_bucket( + self, + input_row, # cute.Tensor [N] dtype (fp32; full row, cs==1 only) + N, # int32 valid length (pad tail beyond N is never read) + seed_thr_row, # [3] fp32 closed-loop lines, ascending t0 < t1 < t2 + smem_keys, # [seg_total] fp32 value segments (A @0 / B @segA / C @2segA) + cand_idx_row, # [seg_total] int32 gmem POSITION column (write-only here) + block_max_row, # [nb_pad] fp32 per-32-position maxima, or None + smem_stage, # [stage_rows, 4] fp32 cp.async staging (aliases smem_vals) + s_seg, # [>=7] int32 scratch (reuses smem_wcnt_p1: P1 never runs + # on a row this phase succeeded on): [0..2] A/B/C claim + # cursors, [3] void, [4] n0, [5] n1, [6] n2 + tidx, + warp_id, + lane, + ): + """self_scan phase 0: ONE streaming pass buckets every element + >= t0 into the on-chip value segments (tightest line passed picks + the segment; a full segment spills to the next looser one) and + writes each entry's POSITION to the same coordinate of the gmem + column. The final cursor values ARE the line counts (attempts, + uncapped), so {n0, void, n1, n2} fall out for free — the same + contract the v5 emitter produced externally. + + The dense scan is a cp.async pipeline: each thread streams one + 16B vector per step into its private slot-major smem staging + slot, keeping ``stage_slots`` steps in flight; classification + reads the staged values and claims passers with per-element + direct smem atomics (they don't synchronize the warp). Segment + overflow is resolved in-claim by spilling to the next looser + segment (a segment overflows at most once per row).""" + num_threads = cutlass.const_expr(self.num_threads) + segA = cutlass.const_expr(self.accept_cap) + capC = cutlass.const_expr(self.cap_c) + elem_bytes = cutlass.const_expr(self.dtype.width // 8) + vec_align = cutlass.const_expr(self.vec_align_bytes) + p0ck0 = cutlass.Int64(0) + if cutlass.const_expr(_P4_SUB_DBG): + p0ck0 = cute.arch.clock64() + if tidx == cutlass.Int32(0): + s_seg[0] = cutlass.Int32(0) + s_seg[1] = cutlass.Int32(0) + s_seg[2] = cutlass.Int32(0) + cute.arch.barrier() + t0_s = seed_thr_row[0] + t1_s = seed_thr_row[1] + t2_s = seed_thr_row[2] + row_addr = input_row.iterator.toint() + # ---- stage-2 BLOCK-SKIP variant: the GEMM tail left per-32- + # position maxima; a block whose max < t0 contributes nothing to + # any count or segment, so it is never read. One warp per block: + # the bmax compare is warp-uniform (all lanes read the same + # scalar), a passing block is one coalesced 128B load, claims + # are the same non-synchronizing per-element atomics as the + # dense loop - so the block loop needs no uniform trip counts. + if cutlass.const_expr(self.enable_block_skip and block_max_row is not None): + # 8 blocks per warp iteration: every lane vector-loads the + # SAME 8 bmax values (L1 broadcast - the pass decisions are + # warp-uniform registers), then the passing blocks are + # loaded back-to-back as independent 128B coalesced loads. + bm_addr = block_max_row.iterator.toint() + nb0 = (N + cutlass.Int32(31)) >> cutlass.Int32(5) + # Two-pass skip: (1) DENSE-scan the bmax array itself (it + # is 1/32 of the row) with the tuned vector loop, compacting + # PASSING BLOCK IDS into the idle C segment (single-band mode + # never fills C; ids < 2^23 store exactly as floats); + # (2) walk the compact list, 8 blocks per warp round issued + # unguarded back-to-back. If the list overflows capC the row + # falls back to the dense full scan. + # pass-1 vectors: 128-bit (the bmax row base is only 16B + # aligned: nb_pad %% 4) + pass1_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=128, + ) + p1w = cutlass.const_expr(4) + frag_p = cute.make_rmem_tensor((p1w,), cutlass.Float32) + nfb0 = nb0 >> cutlass.const_expr((num_threads * 4).bit_length() - 1) + itp0 = cutlass.Int32(0) + while itp0 < nfb0: + ip0 = (itp0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(p1w) + pp0 = cute.make_ptr( + cutlass.Float32, + bm_addr + cutlass.Int64(ip0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=16, + ) + cute.copy( + pass1_atom, + cute.make_tensor(pp0, cute.make_layout((p1w,))), + frag_p, + ) + for _jp in cutlass.range_constexpr(p1w): + if cutlass.Float32(frag_p[_jp]) >= t2_s: + slp0 = atomicAdd(s_seg.iterator + cutlass.Int32(1), cutlass.Int32(1)) + if slp0 < cutlass.Int32(capC): + smem_keys[cutlass.Int32(2 * segA) + slp0] = cutlass.Float32( + ip0 + cutlass.Int32(_jp) + ) + itp0 = itp0 + cutlass.Int32(1) + ptb0 = nfb0 * cutlass.Int32(num_threads * 4) + tidx + while ptb0 < nb0: + bpt0 = cute.make_ptr( + cutlass.Float32, + bm_addr + cutlass.Int64(ptb0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + if cute.make_tensor(bpt0, cute.make_layout((1,)))[0] >= t2_s: + slp0 = atomicAdd(s_seg.iterator + cutlass.Int32(1), cutlass.Int32(1)) + if slp0 < cutlass.Int32(capC): + smem_keys[cutlass.Int32(2 * segA) + slp0] = cutlass.Float32(ptb0) + ptb0 = ptb0 + cutlass.Int32(num_threads) + cute.arch.barrier() + nlist0 = s_seg[1] + if nlist0 <= cutlass.Int32(capC): + # pass 2: 8 listed blocks per warp round + nwp = cutlass.const_expr(self.num_warps) + lb0 = warp_id * cutlass.Int32(8) + frag_v = cute.make_rmem_tensor((8,), cutlass.Float32) + while lb0 < nlist0: + # LOAD phase first: eight independent block loads in + # flight before any atomic (claims are memory-ordered + # and would serialize the blocks otherwise) + p2b0 = cutlass.Int32(0) + p2b1 = cutlass.Int32(0) + p2b2 = cutlass.Int32(0) + p2b3 = cutlass.Int32(0) + p2b4 = cutlass.Int32(0) + p2b5 = cutlass.Int32(0) + p2b6 = cutlass.Int32(0) + p2b7 = cutlass.Int32(0) + for _jb in cutlass.range_constexpr(8): + li0 = lb0 + cutlass.Int32(_jb) + bid0 = cutlass.Int32(-1) + vv0 = cutlass.Float32(self.NEG_FLT_MAX) + if li0 < nlist0: + bid0 = cutlass.Int32(smem_keys[cutlass.Int32(2 * segA) + li0]) + pos0 = (bid0 << cutlass.Int32(5)) + lane + if pos0 < N: + vp0 = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(pos0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + vv0 = cute.make_tensor(vp0, cute.make_layout((1,)))[0] + frag_v[_jb] = vv0 + if cutlass.const_expr(_jb == 0): + p2b0 = bid0 + elif cutlass.const_expr(_jb == 1): + p2b1 = bid0 + elif cutlass.const_expr(_jb == 2): + p2b2 = bid0 + elif cutlass.const_expr(_jb == 3): + p2b3 = bid0 + elif cutlass.const_expr(_jb == 4): + p2b4 = bid0 + elif cutlass.const_expr(_jb == 5): + p2b5 = bid0 + elif cutlass.const_expr(_jb == 6): + p2b6 = bid0 + else: + p2b7 = bid0 + # CLAIM phase + for _jb in cutlass.range_constexpr(8): + bidc = ( + p2b0 + if _jb == 0 + else p2b1 + if _jb == 1 + else p2b2 + if _jb == 2 + else p2b3 + if _jb == 3 + else p2b4 + if _jb == 4 + else p2b5 + if _jb == 5 + else p2b6 + if _jb == 6 + else p2b7 + ) + vvc = cutlass.Float32(frag_v[_jb]) + if bidc >= cutlass.Int32(0) and vvc >= t2_s: + posc = (bidc << cutlass.Int32(5)) + lane + if posc < N: + sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1)) + if sl0 < cutlass.Int32(segA): + smem_keys[sl0] = vvc + cand_idx_row[sl0] = posc + lb0 = lb0 + cutlass.Int32(nwp * 8) + if nlist0 > cutlass.Int32(capC): + # list overflow (pass rate too high for skip): dense full + # scan backup - nothing was read yet, plain re-run + itd0 = cutlass.Int32(0) + nfd0 = N >> cutlass.const_expr((num_threads * 4).bit_length() - 1) + while itd0 < nfd0: + idd0 = (itd0 * cutlass.Int32(num_threads) + tidx) * cutlass.Int32(p1w) + pd0 = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(idd0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + cute.copy( + pass1_atom, + cute.make_tensor(pd0, cute.make_layout((p1w,))), + frag_p, + ) + for _jd in cutlass.range_constexpr(p1w): + vd0 = cutlass.Float32(frag_p[_jd]) + if vd0 >= t2_s: + sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1)) + if sl0 < cutlass.Int32(segA): + smem_keys[sl0] = vd0 + cand_idx_row[sl0] = idd0 + cutlass.Int32(_jd) + itd0 = itd0 + cutlass.Int32(1) + ptd0 = nfd0 * cutlass.Int32(num_threads * 4) + tidx + while ptd0 < N: + pe0 = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(ptd0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + ve0 = cute.make_tensor(pe0, cute.make_layout((1,)))[0] + if ve0 >= t2_s: + sl0 = atomicAdd(s_seg.iterator, cutlass.Int32(1)) + if sl0 < cutlass.Int32(segA): + smem_keys[sl0] = ve0 + cand_idx_row[sl0] = ptd0 + ptd0 = ptd0 + cutlass.Int32(num_threads) + cpw = cutlass.const_expr(4) # cp.async caps at 16B per copy + step1 = cutlass.const_expr(num_threads * cpw) + st2log = cutlass.const_expr((2 * step1).bit_length() - 1) + # Pair-step cp.async pipeline: each step processes TWO 16B + # vectors per thread (2 pairs x 32B across the 4 staging slots). + # One commit group per pair; wait_group(1) pops the oldest pair. + # The staging buffer aliases smem_vals (only written after phase + # 0); every non-empty group is drained inside the loop, so + # nothing is in flight once the alias is read. FULL-pair steps + # only in the hot loop; the remainder takes the scalar tail below. + nfull = N >> st2log + if cutlass.const_expr(self.enable_block_skip and block_max_row is not None): + nfull = cutlass.Int32(0) + g2s_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cute.nvgpu.cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=128, + ) + stage_addr = smem_stage.iterator.toint() + for _p in cutlass.range_constexpr(2): + if cutlass.Int32(_p) < nfull: + for _v in cutlass.range_constexpr(2): + pr0 = cutlass.Int32(2 * _p + _v) * cutlass.Int32(num_threads) + tidx + pp0 = cutlass.Int32(2 * _p + _v) * cutlass.Int32(step1) + tidx * cutlass.Int32( + cpw + ) + pq0 = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(pp0) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=16, + ) + dq0 = cute.make_ptr( + cutlass.Float32, + stage_addr + cutlass.Int64(pr0 * cutlass.Int32(cpw)) * cutlass.Int64(4), + cute.AddressSpace.smem, + assumed_align=16, + ) + cute.copy( + g2s_atom, + cute.make_tensor(pq0, cute.make_layout((cpw,))), + cute.make_tensor(dq0, cute.make_layout((cpw,))), + ) + cute.arch.cp_async_commit_group() + it0 = cutlass.Int32(0) + while it0 < nfull: + cute.arch.cp_async_wait_group(1) + sp0 = (it0 & cutlass.Int32(1)) * cutlass.Int32(2 * num_threads) + tidx + ia0 = it0 * cutlass.Int32(2 * step1) + tidx * cutlass.Int32(cpw) + # v6: per-element DIRECT atomic claims for passers — smem + # atomics do NOT synchronize the warp and hide under the + # async copy stream. + for _jh in cutlass.range_constexpr(2): + sq0 = cute.make_ptr( + cutlass.Float32, + stage_addr + + cutlass.Int64( + (sp0 + cutlass.Int32(_jh) * cutlass.Int32(num_threads)) * cutlass.Int32(cpw) + ) + * cutlass.Int64(4), + cute.AddressSpace.smem, + assumed_align=16, + ) + srow = cute.make_tensor(sq0, cute.make_layout((cpw,))) + for _jv in cutlass.range_constexpr(cpw): + v0 = cutlass.Float32(srow[_jv]) + if v0 >= t0_s: + pos0 = ia0 + cutlass.Int32(_jh) * cutlass.Int32(step1) + cutlass.Int32(_jv) + c0 = cutlass.Int32(2) + if v0 >= t1_s: + c0 = cutlass.Int32(1) + if v0 >= t2_s: + c0 = cutlass.Int32(0) + while c0 >= cutlass.Int32(0) and c0 <= cutlass.Int32(2): + cap0 = cutlass.Int32(segA) + if c0 == cutlass.Int32(2): + cap0 = cutlass.Int32(capC) + sl0 = atomicAdd(s_seg.iterator + c0, cutlass.Int32(1)) + if sl0 < cap0: + cd0 = c0 * cutlass.Int32(segA) + sl0 + smem_keys[cd0] = v0 + cand_idx_row[cd0] = pos0 + c0 = cutlass.Int32(-1) + else: + c0 = c0 + cutlass.Int32(1) + # reissue the just-consumed slot pair for step it0 + 2 (the + # thread's own prior reads are ordered before the async write + # begins, so no fence is needed) + kn0 = it0 + cutlass.Int32(2) + if kn0 < nfull: + for _jh in cutlass.range_constexpr(2): + jr0 = sp0 + cutlass.Int32(_jh) * cutlass.Int32(num_threads) + jp0 = ( + kn0 * cutlass.Int32(2 * step1) + + cutlass.Int32(_jh) * cutlass.Int32(step1) + + tidx * cutlass.Int32(cpw) + ) + pq0 = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(jp0) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=16, + ) + dq0 = cute.make_ptr( + cutlass.Float32, + stage_addr + cutlass.Int64(jr0 * cutlass.Int32(cpw)) * cutlass.Int64(4), + cute.AddressSpace.smem, + assumed_align=16, + ) + cute.copy( + g2s_atom, + cute.make_tensor(pq0, cute.make_layout((cpw,))), + cute.make_tensor(dq0, cute.make_layout((cpw,))), + ) + cute.arch.cp_async_commit_group() + it0 = it0 + cutlass.Int32(1) + # scalar tail (< step1 elements): per-element DIRECT atomic + # claims — divergent-safe, no warp collectives + pt0 = (N >> st2log) * cutlass.Int32(2 * step1) + tidx + if cutlass.const_expr(self.enable_block_skip and block_max_row is not None): + pt0 = N + while pt0 < N: + spt = cute.make_ptr( + cutlass.Float32, + row_addr + cutlass.Int64(pt0) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + vt0 = cute.make_tensor(spt, cute.make_layout((1,)))[0] + if vt0 >= t0_s: + ct0 = cutlass.Int32(2) + if vt0 >= t1_s: + ct0 = cutlass.Int32(1) + if vt0 >= t2_s: + ct0 = cutlass.Int32(0) + while ct0 >= cutlass.Int32(0) and ct0 <= cutlass.Int32(2): + capt = cutlass.Int32(segA) + if ct0 == cutlass.Int32(2): + capt = cutlass.Int32(capC) + slt = atomicAdd(s_seg.iterator + ct0, cutlass.Int32(1)) + if slt < capt: + cdt = ct0 * cutlass.Int32(segA) + slt + smem_keys[cdt] = vt0 + cand_idx_row[cdt] = pt0 + ct0 = cutlass.Int32(-1) + else: + ct0 = ct0 + cutlass.Int32(1) + pt0 = pt0 + cutlass.Int32(num_threads) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + if cutlass.const_expr(self.enable_block_skip): + # single-band mode: every count is the t2 cursor; a cut + # can only land on t2 (in band), the sample-hist (over) + # or the fallback (under) - exactly the v5 state machine + # fed with n0 == n1 == n2 + curT0 = s_seg[0] + # claims past segA were dropped by the walk: honor the + # "void==0 means nothing was dropped" contract + s_seg[3] = cutlass.Int32(0) + if curT0 > cutlass.Int32(segA): + s_seg[3] = cutlass.Int32(1) + s_seg[4] = curT0 + s_seg[5] = curT0 + s_seg[6] = curT0 + if cutlass.const_expr(_P4_SUB_DBG): + s_seg[7] = cutlass.Int32(cute.arch.clock64() - p0ck0) + if cutlass.const_expr(not self.enable_block_skip): + curA0 = s_seg[0] + curB0 = s_seg[1] + curC0 = s_seg[2] + spA0 = curA0 - cutlass.Int32(segA) + if spA0 < cutlass.Int32(0): + spA0 = cutlass.Int32(0) + spB0 = curB0 - cutlass.Int32(segA) + if spB0 < cutlass.Int32(0): + spB0 = cutlass.Int32(0) + n1_0 = curA0 + curB0 - spA0 + n0_0 = n1_0 + curC0 - spB0 + s_seg[3] = cutlass.Int32(0) + if curC0 > cutlass.Int32(capC): + s_seg[3] = cutlass.Int32(1) + s_seg[4] = n0_0 + s_seg[5] = n1_0 + s_seg[6] = curA0 + cute.arch.barrier() + @cute.jit def block_count_ge( self, @@ -1411,7 +2113,7 @@ def block_count_ge( # path (same vec_w / 4-way-unroll / tail loops) with M static register # counters. Caches all M per-thread count columns in smem_ptcnt_multi so # the accepted rung's column seeds Phase 3 with zero rescan. This is the - # R0 admission primitive (multi-threshold lineage); it is only invoked + # R0 admission primitive (op#18 multithresh lineage); it is only invoked # from the enable_r0 path added in a later commit, so the base kernel is # unaffected. Slice + cluster form: each CTA scans [slice_start, # slice_end) and the M per-CTA totals are DSMEM all-reduced across the @@ -1420,6 +2122,109 @@ def block_count_ge( # totals are the answer. smem_ptcnt_multi holds slice-local per-thread # columns (the accepted rung's column seeds Phase 3 per CTA). # ------------------------------------------------------------------ + + # ---- block-skip machinery (enable_block_skip): grain 32, int16 + # list, grouped strided build, UN=2 scan ---- + + @cute.jit + def _list_ld(self, smem_active, idx): + return cutlass.Int32(smem_active[idx]) + + @cute.jit + def _list_st(self, smem_active, idx, val): + smem_active[idx] = cutlass.Int16(val) + + @cute.jit + def _block_bound(self, bm_addr, blk_id): + # grain 32: record blk_id IS the exact positional bound of + # [blk_id*32, blk_id*32+32) (indexer TMEM partition contract) — + # one 4B scalar load, no fold. + bm_ptr = cute.make_ptr( + cutlass.Float32, + bm_addr + cutlass.Int64(blk_id) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + return cute.make_tensor(bm_ptr, cute.make_layout((1,)))[0] + + @cute.jit + def _full_build_active( + self, + block_max_row, + slice_start, + slice_end, + threshold, + smem_wcnt, + smem_active, + s_active_cnt, + tidx, + warp_id, + lane, + ): + # Two-phase register-bitmask build. Each thread owns ids_per_thread + # block ids STRIDED (t, t+T, t+2T, ...) so each pass's bound loads + # coalesce warp-wide. Phase A flags+counts with no sync; phase B is + # ONE block-wide exclusive scan over per-thread counts; phase C + # scatters from the bitmask — 3 barriers TOTAL, independent of + # nb_slice. List order is thread-grouped, not ascending: fine, since + # the count scan and the Phase-3 stream-write walk the SAME list by + # position (determinism contract). + ids_per_thread = cutlass.const_expr(self.SKIP_MAX_BLOCKS // self.num_threads) + num_threads = cutlass.const_expr(self.num_threads) + # first FULL block: a block straddling slice_start would map list + # positions outside this CTA's slice; the sub-block head region + # [slice_start, blk_lo*32) is scanned separately by the callers. + blk_lo = (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32( + self.SKIP_BLOCK_LOG2 + ) + blk_hi = (slice_end + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32( + self.SKIP_BLOCK_LOG2 + ) + nb_slice = blk_hi - blk_lo + bm_addr = block_max_row.iterator.toint() + + # Phase A (no sync): flag my ids into a register bitmask. + bmask = cutlass.Int32(0) + cnt = cutlass.Int32(0) + for m in cutlass.range_constexpr(ids_per_thread): + ib = tidx + cutlass.Int32(m * num_threads) + if ib < nb_slice: + bound = self._block_bound(bm_addr, blk_lo + ib) + if bound >= threshold: + bmask = bmask | (cutlass.Int32(1) << cutlass.Int32(m)) + cnt = cnt + cutlass.Int32(1) + + # Phase B: one block-wide exclusive scan over per-thread counts. + tp = cnt + for off_i in cutlass.range_constexpr(5): + off_v = cutlass.const_expr(1 << off_i) + other = cute.arch.shuffle_sync_up(tp, off_v, mask_and_clamp=0) + if lane >= cutlass.Int32(off_v): + tp = tp + other + excl = tp - cnt + warp_total = cute.arch.shuffle_sync(tp, cutlass.Int32(self.WARP_SIZE - 1)) + if lane == 0: + smem_wcnt[warp_id] = warp_total + cute.arch.barrier() + if tidx == 0: + tot = cutlass.Int32(0) + for w in cutlass.range_constexpr(self.num_warps): + cw = smem_wcnt[w] + smem_wcnt[w] = tot + tot = tot + cw + s_active_cnt[0] = tot + cute.arch.barrier() + + # Phase C: scatter from the bitmask at deterministic offsets. + pos_out = smem_wcnt[warp_id] + excl + if bmask != cutlass.Int32(0): + for m in cutlass.range_constexpr(ids_per_thread): + if (bmask & (cutlass.Int32(1) << cutlass.Int32(m))) != cutlass.Int32(0): + ib_c = tidx + cutlass.Int32(m * num_threads) + self._list_st(smem_active, pos_out, blk_lo + ib_c) + pos_out = pos_out + cutlass.Int32(1) + cute.arch.barrier() + @cute.jit def block_count_ge_multi( self, @@ -1436,6 +2241,9 @@ def block_count_ge_multi( warp_id, lane, smem_ptcnt=None, # vseed: last column's per-thread counts land here + block_max_row=None, # block-skip: per-32-position upper bounds + smem_active=None, # block-skip: int16 active list + s_active_cnt=None, # block-skip: [0]=list length, [1]=list-current flag ): M = cutlass.const_expr(self.M_thr) num_threads = cutlass.const_expr(self.num_threads) @@ -1459,7 +2267,175 @@ def block_count_ge_multi( i = slice_start + tidx * cutlass.Int32(vec_w) step = cutlass.Int32(step_elem) - if self.enable_unroll_4: + # ---- block-skip compact iteration (lossless vs the dense path) ---- + # Build the active list at the LOOSEST threshold over all M columns: + # a skipped block bounds every element below min(t_m), so all M + # counts equal their dense values. The list covers FULL blocks only; + # the sub-block head region of an unaligned slice start is counted + # here separately (per-thread order contract: head elements FIRST, + # then list entries — Phase 3's compact write replays the same). + skip_ok = cutlass.Int32(0) + dense_ok = True # Python bool: no scf.if when block skip is off + if cutlass.const_expr(self.enable_block_skip and block_max_row is not None): + skip_ok = cutlass.Int32(1) + # Capacity/id-width guard: the active list holds at most + # SKIP_MAX_BLOCKS local ids and _list_st stores ABSOLUTE block + # ids as int16. A slice over 8192 full blocks (N_local > + # 262144) or reaching absolute id >= 32768 falls back to the + # dense walk (lossless; the list-current flag is never set, so + # phase3 stays dense too). + blk_lo_g = (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32( + self.SKIP_BLOCK_LOG2 + ) + blk_hi_g = (slice_end + cutlass.Int32(self.SKIP_BLOCK - 1)) >> cutlass.Int32( + self.SKIP_BLOCK_LOG2 + ) + if blk_hi_g - blk_lo_g > cutlass.Int32(self.SKIP_MAX_BLOCKS): + skip_ok = cutlass.Int32(0) + if blk_hi_g > cutlass.Int32(32767): + skip_ok = cutlass.Int32(0) + dense_ok = skip_ok == cutlass.Int32(0) + if cutlass.const_expr(self.enable_block_skip and block_max_row is not None): + if skip_ok == cutlass.Int32(1): + head_end = ( + (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) + >> cutlass.Int32(self.SKIP_BLOCK_LOG2) + ) << cutlass.Int32(self.SKIP_BLOCK_LOG2) + if head_end > slice_end: + head_end = slice_end + hh = slice_start + tidx + while hh < head_end: + vh = self._load_fp32(input_row, hh) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vh >= thr_frag[m]) + hh = hh + cutlass.Int32(num_threads) + # Rung-tightening build (cs==1): a rung whose active list + # exceeds CAP blocks cannot be accepted (count >= list + # length), so drop it and rebuild at the next tighter + # threshold. A dropped rung is only an unmeasured probe, + # recorded in the mask at s_active_cnt[2] and skipped by + # classify and by the fallback seeding. Bounded by M-1 + # extra builds. At cs>1 the per-CTA list lengths differ + # and the drop decision would diverge across the cluster, + # so keep the plain loosest-rung build there. + CAP_BLOCKS = cutlass.const_expr(3 * self.kC // 4) + if cutlass.const_expr(cluster_size == 1): + build_done = cutlass.Int32(0) + for _attempt in cutlass.range_constexpr(M): + if build_done == cutlass.Int32(0): + dmask = s_active_cnt[2] + tcur = cutlass.Float32(self.FLT_MAX) + mcur = cutlass.Int32(-1) + kept = cutlass.Int32(0) + for m in cutlass.range_constexpr(M): + if ( + dmask & (cutlass.Int32(1) << cutlass.Int32(m)) + ) == cutlass.Int32(0): + kept = kept + cutlass.Int32(1) + if thr_frag[m] < tcur: + tcur = thr_frag[m] + mcur = cutlass.Int32(m) + self._full_build_active( + block_max_row, + slice_start, + slice_end, + tcur, + smem_wcnt_multi, + smem_active, + s_active_cnt, + tidx, + warp_id, + lane, + ) + if s_active_cnt[0] <= cutlass.Int32( + CAP_BLOCKS + ) or kept <= cutlass.Int32(1): + build_done = cutlass.Int32(1) + else: + if tidx == 0: + s_active_cnt[2] = dmask | (cutlass.Int32(1) << mcur) + cute.arch.barrier() + else: + tmin = thr_frag[0] + for m in cutlass.range_constexpr(M): + tmin = _fmin_f32_inline(tmin, thr_frag[m]) + self._full_build_active( + block_max_row, + slice_start, + slice_end, + tmin, + smem_wcnt_multi, + smem_active, + s_active_cnt, + tidx, + warp_id, + lane, + ) + if tidx == 0: + s_active_cnt[1] = cutlass.Int32(1) # list-current flag + chunks_per_block = cutlass.const_expr( + self.SKIP_BLOCK // (self.vec_bits // self.dtype.width) + ) + tpb = cutlass.const_expr(chunks_per_block) + blocks_per_iter = cutlass.const_expr(self.num_threads // chunks_per_block) + UN = cutlass.const_expr(self.SKIP_UNROLL) + stride_un = cutlass.const_expr(blocks_per_iter * UN) + my_blk_slot = tidx // cutlass.Int32(tpb) + my_chunk0 = tidx % cutlass.Int32(tpb) + cnt_active = s_active_cnt[0] + frags = [cute.make_rmem_tensor((vec_w,), self.dtype) for _ in range(UN)] + li = my_blk_slot + while li < cnt_active: + poss = [] + valids = [] + for u in cutlass.range_constexpr(UN): + lu = li + cutlass.Int32(u * blocks_per_iter) + valid = lu < cnt_active + pos0 = cutlass.Int32(0) + if valid: + blk = self._list_ld(smem_active, lu) + pos0 = blk * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0 * cutlass.Int32( + vec_w + ) + # Vector-load only fully in-bounds chunks; the + # slice-end straddle re-reads scalars below (a + # tail block's chunks would otherwise read past + # the row/allocation when N % 32 != 0). + if pos0 + cutlass.Int32(vec_w) <= slice_end: + src_ptr_u = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(pos0) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + cute.copy( + copy_atom, + cute.make_tensor(src_ptr_u, cute.make_layout((vec_w,))), + frags[u], + ) + poss.append(pos0) + valids.append(valid) + for u in cutlass.range_constexpr(UN): + if valids[u]: + pos = poss[u] + if pos + cutlass.Int32(vec_w) <= slice_end: + for j in cutlass.range_constexpr(vec_w): + if cutlass.const_expr(self.dtype == cutlass.Float32): + vj = frags[u][j] + else: + vj = cutlass.Float32(frags[u][j]) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) + else: + jj = pos + while jj < slice_end: + vs = self._load_fp32(input_row, jj) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vs >= thr_frag[m]) + jj = jj + cutlass.Int32(1) + li = li + cutlass.Int32(stride_un) + + if self.enable_unroll_4 and dense_ok: rng_frag = cute.make_rmem_tensor((vec_w,), self.dtype) big_iters = cutlass.Int32(0) if slice_end > i + cutlass.Int32(vec_w - 1): @@ -1486,30 +2462,31 @@ def block_count_ge_multi( i = i + big_iters * cutlass.Int32(step_elem) tail_frag = cute.make_rmem_tensor((vec_w,), self.dtype) - while i + cutlass.Int32(vec_w - 1) < slice_end: - src_ptr = cute.make_ptr( - self.dtype, - row_addr + cutlass.Int64(i) * cutlass.Int64(elem_bytes), - cute.AddressSpace.gmem, - assumed_align=vec_align, - ) - src = cute.make_tensor(src_ptr, cute.make_layout((vec_w,))) - cute.copy(copy_atom, src, tail_frag) - for j in cutlass.range_constexpr(vec_w): - if cutlass.const_expr(self.dtype == cutlass.Float32): - vj = tail_frag[j] - else: - vj = cutlass.Float32(tail_frag[j]) - for m in cutlass.range_constexpr(M): - cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) - i = i + step + if dense_ok: + while i + cutlass.Int32(vec_w - 1) < slice_end: + src_ptr = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(i) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + src = cute.make_tensor(src_ptr, cute.make_layout((vec_w,))) + cute.copy(copy_atom, src, tail_frag) + for j in cutlass.range_constexpr(vec_w): + if cutlass.const_expr(self.dtype == cutlass.Float32): + vj = tail_frag[j] + else: + vj = cutlass.Float32(tail_frag[j]) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(vj >= thr_frag[m]) + i = i + step - it = n_aligned + tidx - while it < slice_end: - v = self._load_fp32(input_row, it) - for m in cutlass.range_constexpr(M): - cnt_frag[m] = cnt_frag[m] + cutlass.Int32(v >= thr_frag[m]) - it = it + cutlass.Int32(num_threads) + it = n_aligned + tidx + while it < slice_end: + v = self._load_fp32(input_row, it) + for m in cutlass.range_constexpr(M): + cnt_frag[m] = cnt_frag[m] + cutlass.Int32(v >= thr_frag[m]) + it = it + cutlass.Int32(num_threads) for m in cutlass.range_constexpr(M): if cutlass.const_expr(self.r0_vseed and m == self.M_qf): @@ -2027,6 +3004,8 @@ def phase3_collect_candidates( lane, do_cluster_sync, # bool: False = cs=1 / short-row degrade (skip cluster sync) smem_input=None, # optional SMEM-cached slice + smem_active=None, # block-skip: int16 active list (reused, not rebuilt) + s_active_cnt=None, # block-skip: [0]=list length, [1]=list-current flag ): """Retry-shrink (when P2 didn't converge) + prefix sum + stream-write. @@ -2045,6 +3024,20 @@ def phase3_collect_candidates( # always aggregates across the cluster, so every CTA sees the same # cluster-wide cand_count; cs=1 makes the aggregation a no-op. if s_iscalars[1] != cutlass.Int32(1): + # Any dense re-count invalidates the block-skip active list. The + # list is built from the loosest kept rung of the probe that set + # the flag, so it is a superset only at or above that rung, and + # the repair below anchors and bisects underneath it; separately, + # the compact stream-write replays the list walk against the + # smem_ptcnt of its matching compact pass, which a dense re-count + # overwrites. Cleared once here rather than per re-count: done==1 + # rows never reach this block, so the hot path keeps its compact + # write. + if cutlass.const_expr(self.enable_block_skip): + if tidx == cutlass.Int32(0): + s_active_cnt[1] = cutlass.Int32(0) + cute.arch.barrier() + # Re-count with current threshold (may already have stale cand_count) cur_thr = s_thr[0] self.block_count_ge( @@ -2249,6 +3242,89 @@ def phase3_collect_candidates( wc = my_write_pos step = cutlass.Int32(step_elem) + # ---- block-skip compact stream-write ---- + # Reuses the active list left by the R0 compact count pass (same + # ownership walk: my_blk_slot's list slots in ascending order, so + # each thread produces its candidates in the SAME per-thread order + # the count pass counted them — the prefix-sum positions match). + # Only taken when the list is CURRENT (s_active_cnt[1] == 1, set by + # the build; cleared on any dense fallback re-count). + skip_wr = cutlass.Int32(0) + park_cursors = False # Python bool: no scf.if when block skip is off + if cutlass.const_expr(self.enable_block_skip and smem_active is not None): + if s_active_cnt[1] == cutlass.Int32(1): + skip_wr = cutlass.Int32(1) + park_cursors = skip_wr == cutlass.Int32(1) + if cutlass.const_expr(self.enable_block_skip and smem_active is not None): + if skip_wr == cutlass.Int32(1): + # head region first — same per-thread order as the count pass + head_end_w = ( + (slice_start + cutlass.Int32(self.SKIP_BLOCK - 1)) + >> cutlass.Int32(self.SKIP_BLOCK_LOG2) + ) << cutlass.Int32(self.SKIP_BLOCK_LOG2) + if head_end_w > slice_end: + head_end_w = slice_end + hh_w = slice_start + tidx + while hh_w < head_end_w: + vh_w = self._load_fp32(input_row, hh_w) + if vh_w >= thr_final and wc < cutlass.Int32(kCC): + smem_keys[wc] = vh_w + smem_vals[wc] = hh_w + wc = wc + cutlass.Int32(1) + hh_w = hh_w + cutlass.Int32(num_threads) + chunks_per_block_w = cutlass.const_expr( + self.SKIP_BLOCK // (self.vec_bits // self.dtype.width) + ) + blocks_per_iter_w = cutlass.const_expr(self.num_threads // chunks_per_block_w) + my_blk_slot_w = tidx // cutlass.Int32(chunks_per_block_w) + my_chunk0_w = tidx % cutlass.Int32(chunks_per_block_w) + cnt_active_w = s_active_cnt[0] + wfrag = cute.make_rmem_tensor((vec_w,), self.dtype) + li_w = my_blk_slot_w + while li_w < cnt_active_w: + blk_w = self._list_ld(smem_active, li_w) + pos0_w = blk_w * cutlass.Int32(self.SKIP_BLOCK) + my_chunk0_w * cutlass.Int32( + vec_w + ) + if pos0_w + cutlass.Int32(vec_w) <= slice_end: + src_ptr_w = cute.make_ptr( + self.dtype, + row_addr + cutlass.Int64(pos0_w) * cutlass.Int64(elem_bytes), + cute.AddressSpace.gmem, + assumed_align=vec_align, + ) + cute.copy( + copy_atom, + cute.make_tensor(src_ptr_w, cute.make_layout((vec_w,))), + wfrag, + ) + for j in cutlass.range_constexpr(vec_w): + if cutlass.const_expr(self.dtype == cutlass.Float32): + vj = wfrag[j] + else: + vj = cutlass.Float32(wfrag[j]) + if vj >= thr_final and wc < cutlass.Int32(kCC): + smem_keys[wc] = vj + smem_vals[wc] = pos0_w + cutlass.Int32(j) + wc = wc + cutlass.Int32(1) + else: + jj_w = pos0_w + while jj_w < slice_end: + v_w = self._load_fp32(input_row, jj_w) + if v_w >= thr_final and wc < cutlass.Int32(kCC): + smem_keys[wc] = v_w + smem_vals[wc] = jj_w + wc = wc + cutlass.Int32(1) + jj_w = jj_w + cutlass.Int32(1) + li_w = li_w + cutlass.Int32(blocks_per_iter_w) + + # When the compact write ran, park the dense cursors at the end so + # all three dense loops below (4-way, vec tail, scalar tail) fall + # through without re-indenting them. + if park_cursors: + ic = N_local + n_aligned = N_local + # Phase3 unrolling: master gated by self.enable_phase3_unroll. # When OFF, only the tail 1-way loop runs (matches the pre-unroll # state of phase3_collect). When ON, the inner enable_unroll_4 @@ -2719,256 +3795,329 @@ def _kth_bin_search_rw(self, smem_hist, smem_wcnt, lo, binw, tidx, warp_id, lane return thr_out, sel_out # ------------------------------------------------------------------ - # Phase 4 (alt): fused rank-and-scatter (enable_p4_rank_scatter). - # Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py. + # _p4_coarse_rw - redundant-warp coarse bin search for the fused + # rank-and-scatter path: returns the straddling bin and the count + # strictly above it, resolved lane-parallel on every warp (an + # idx-shuffle scan + ballot locate the target slice, a second scan + # + the unique crossing test locate the bin inside it). Integer + # sums are associative, so every warp lands on the same answer + # bit-for-bit. Mirrors _kth_bin_search_rw (snap path). # ------------------------------------------------------------------ @cute.jit - def _p4_exact_tail_radix_select( - self, - kK: cutlass.Constexpr, - kBins: cutlass.Constexpr, - num_threads: cutlass.Constexpr, - num_warps: cutlass.Constexpr, - need0, - cand_count, - rank_above_fine, - b_star, - sb_star, - bmin_r, - f_lo, - finv, - fbins, - inv1, - tidx, - lane, - warp_id, - smem_hist, - smem_keys, - smem_vals, - smem_wcnt, - s_iscalars, - output_indices_row, - output_values_row, - ): - """MSB-first 4x8-bit exact radix select over the straddling - fine-bin tie set (``p4_exact_tail``) — single source shared by - the tiny-tie fast path's large-class fallback and the plain - exact-tail path (previously two verbatim copies; ``@cute.jit`` - helpers inline, so codegen is unchanged).""" - # Persistent scalars live above the 256 digit bins - # (kNumBins >= 512 always): [256] key prefix (chosen - # digits, remaining bits 0), [257] slots still to fill - # inside the current equal-prefix set, [258] ties - # strictly above the prefix (their slots precede it). - if tidx == cutlass.Int32(0): - smem_hist[256] = cutlass.Int32(0) - smem_hist[257] = need0 - smem_hist[258] = cutlass.Int32(0) - cute.arch.barrier() - for lvl in cutlass.range_constexpr(4): - shift = cutlass.const_expr(24 - 8 * lvl) - iz2 = tidx - while iz2 < cutlass.Int32(256): - smem_hist[iz2] = cutlass.Int32(0) - iz2 = iz2 + cutlass.Int32(num_threads) - cute.arch.barrier() - uthr_cur = smem_hist[256] - it2 = tidx - while it2 < cand_count: - vt = smem_keys[it2] - bt = cutlass.Int32((vt - bmin_r) * inv1) - if bt < cutlass.Int32(0): - bt = cutlass.Int32(0) - if bt > cutlass.Int32(kBins - 1): - bt = cutlass.Int32(kBins - 1) - if bt == b_star: - st2 = cutlass.Int32((vt - f_lo) * finv) - if st2 < cutlass.Int32(0): - st2 = cutlass.Int32(0) - if st2 > cutlass.Int32(fbins - 1): - st2 = cutlass.Int32(fbins - 1) - if st2 == sb_star: - uk = f32_order_key(vt) - pmatch = cutlass.Int32(1) - if cutlass.const_expr(lvl > 0): - if (uk >> cutlass.Int32(shift + 8)) != ( - uthr_cur >> cutlass.Int32(shift + 8) - ): - pmatch = cutlass.Int32(0) - if pmatch == cutlass.Int32(1): - dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) - atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) - it2 = it2 + cutlass.Int32(num_threads) - cute.arch.barrier() - # Two-stage descending digit scan (mirrors the - # fine 3-step search): per-warp partial sums, - # thread0 picks the target warp, its lane0 walks - # the warp's digit range — 2*num_warps serial - # steps instead of 256. - fdw = cutlass.const_expr(256 // self.num_warps) - wsum2 = cutlass.Int32(0) - for jd in cutlass.range_constexpr(fdw): - dix = cutlass.Int32(255) - warp_id * cutlass.Int32(fdw) - cutlass.Int32(jd) - wsum2 = wsum2 + smem_hist[dix] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = wsum2 - cute.arch.barrier() - if tidx == cutlass.Int32(0): - needl = smem_hist[257] - cw = cutlass.Int32(0) - tw3 = cutlass.Int32(num_warps - 1) - f3 = cutlass.Int32(0) - for w4 in cutlass.range_constexpr(self.num_warps): - cw = cw + smem_wcnt[w4] - if cw >= needl and f3 == cutlass.Int32(0): - tw3 = cutlass.Int32(w4) - f3 = cutlass.Int32(1) - pre3 = cutlass.Int32(0) - for w5 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w5) < tw3: - pre3 = pre3 + smem_wcnt[w5] - s_iscalars[4] = pre3 # prefix above target warp - s_iscalars[0] = tw3 # target warp - cute.arch.barrier() - pre4 = s_iscalars[4] - tw4 = s_iscalars[0] - if warp_id == tw4 and lane == cutlass.Int32(0): - needl2 = smem_hist[257] - base4 = pre4 - dstar = cutlass.Int32(0) - above_d = pre4 - sd4 = cutlass.Int32(0) - for jd2 in cutlass.range_constexpr(fdw): - dix2 = cutlass.Int32(255) - tw4 * cutlass.Int32(fdw) - cutlass.Int32(jd2) - ra4 = base4 - base4 = base4 + smem_hist[dix2] - if base4 >= needl2 and sd4 == cutlass.Int32(0): - dstar = dix2 - above_d = ra4 - sd4 = cutlass.Int32(1) - smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) - smem_hist[257] = needl2 - above_d - smem_hist[258] = smem_hist[258] + above_d - cute.arch.barrier() - # Rewrite the tie slot range: ties with key > u_thr - # first (there are exactly cnt_ab of them), then the - # first need_eq bitwise-equal-to-u_thr ties in arrival - # order (value-exact by construction). Signed compare - # needs the top bit flipped (unsigned-monotonic key). - u_thr = smem_hist[256] - cnt_ab = smem_hist[258] - need_eq = smem_hist[257] - ks_thr = u_thr ^ cutlass.Int32(-2147483648) - if tidx == cutlass.Int32(0): - s_iscalars[4] = cutlass.Int32(0) # above-writer ctr - s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr - cute.arch.barrier() - ir2 = tidx - while ir2 < cand_count: - vr = smem_keys[ir2] - br = cutlass.Int32((vr - bmin_r) * inv1) - if br < cutlass.Int32(0): - br = cutlass.Int32(0) - if br > cutlass.Int32(kBins - 1): - br = cutlass.Int32(kBins - 1) - if br == b_star: - sr = cutlass.Int32((vr - f_lo) * finv) - if sr < cutlass.Int32(0): - sr = cutlass.Int32(0) - if sr > cutlass.Int32(fbins - 1): - sr = cutlass.Int32(fbins - 1) - if sr == sb_star: - uk2 = f32_order_key(vr) - ks2 = uk2 ^ cutlass.Int32(-2147483648) - if ks2 > ks_thr: - o2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(4), - cutlass.Int32(1), - ) - pos = rank_above_fine + o2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - elif ks2 == ks_thr: - q2 = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), - cutlass.Int32(1), - ) - if q2 < need_eq: - pos = rank_above_fine + cnt_ab + q2 - if pos < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(vr) - output_indices_row[pos] = smem_vals[ir2] - ir2 = ir2 + cutlass.Int32(num_threads) - cute.arch.barrier() - - @cute.jit - def phase4_rank_scatter( - self, - smem_keys, - smem_vals, - smem_hist, - smem_wcnt, - s_thr, - s_iscalars, - output_values_row, - output_indices_row, - cand_count, - tidx, - warp_id, - lane, - ): + def _p4_coarse_rw(self, smem_hist, smem_wcnt, warp_id, lane): kK = cutlass.const_expr(self.top_k) kBins = cutlass.const_expr(self.kNumBins) - num_threads = cutlass.const_expr(self.num_threads) - num_warps = cutlass.const_expr(self.num_warps) bins_per_warp = cutlass.const_expr(kBins // self.num_warps) - if cand_count == cutlass.Int32(kK): - i4 = tidx - while i4 < cutlass.Int32(kK): - if cutlass.const_expr(self.return_output_values): - output_values_row[i4] = self.dtype(smem_keys[i4]) - output_indices_row[i4] = smem_vals[i4] - i4 = i4 + cutlass.Int32(num_threads) - elif cand_count > cutlass.Int32(kK): - # ---- block min/max over candidates ---- - local_cmin = cutlass.Float32(self.FLT_MAX) - local_cmax = cutlass.Float32(self.NEG_FLT_MAX) - i5 = tidx - while i5 < cand_count: - v = smem_keys[i5] - local_cmin = _fmin_f32_inline(local_cmin, v) - local_cmax = cute.arch.fmax(local_cmax, v) - i5 = i5 + cutlass.Int32(num_threads) - cmin = self.warp_reduce_min_f32(local_cmin) - cmax = self.warp_reduce_max_f32(local_cmax) - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = float_as_uint32(cmin) - smem_hist[warp_id] = float_as_uint32(cmax) - cute.arch.barrier() - bmin_r = cutlass.Float32(self.FLT_MAX) - bmax_r = cutlass.Float32(self.NEG_FLT_MAX) - for w in cutlass.range_constexpr(self.num_warps): - vmin = cutlass.Float32( - llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value()) + warp_bin_sum = cutlass.Int32(0) + if cutlass.const_expr(bins_per_warp % self.WARP_SIZE == 0): + for jm in cutlass.range_constexpr(bins_per_warp // self.WARP_SIZE): + bidx_s = ( + cutlass.Int32(kBins - 1) + - warp_id * cutlass.Int32(bins_per_warp) + - (lane + cutlass.Int32(jm * self.WARP_SIZE)) ) - vmax = cutlass.Float32( - llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[w].ir_value()) + warp_bin_sum = warp_bin_sum + smem_hist[bidx_s] + warp_bin_sum = self.warp_reduce_sum_i32(warp_bin_sum) + else: + for jb in cutlass.range_constexpr(bins_per_warp): + bidx_s = ( + cutlass.Int32(kBins - 1) + - warp_id * cutlass.Int32(bins_per_warp) + - cutlass.Int32(jb) ) - bmin_r = _fmin_f32_inline(bmin_r, vmin) - bmax_r = cute.arch.fmax(bmax_r, vmax) - if bmax_r <= bmin_r: - bmax_r = bmin_r + cutlass.Float32(1e-6) - cute.arch.barrier() - # ---- zero + build histogram ---- - i6 = tidx - while i6 < cutlass.Int32(kBins): - smem_hist[i6] = cutlass.Int32(0) - i6 = i6 + cutlass.Int32(num_threads) - cute.arch.barrier() + warp_bin_sum = warp_bin_sum + smem_hist[bidx_s] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = warp_bin_sum + cute.arch.barrier() + + # locate the target slice (lane w holds slot w) + v_s = cutlass.Int32(0) + if lane < cutlass.Int32(self.num_warps): + v_s = smem_wcnt[lane] + run2 = v_s + for d2 in cutlass.range_constexpr(5): + off2 = cutlass.const_expr(1 << d2) + src2 = lane - cutlass.Int32(off2) + if src2 < cutlass.Int32(0): + src2 = cutlass.Int32(0) + up2 = cute.arch.shuffle_sync(run2, src2) + if lane >= cutlass.Int32(off2): + run2 = run2 + up2 + m2 = cute.arch.vote_ballot_sync(run2 >= cutlass.Int32(kK)) + tw = cutlass.Int32(self.num_warps - 1) + if m2 != cutlass.Uint32(0): + low2 = m2 & (cutlass.Uint32(0) - m2) + tw = cutlass.Int32(cute.arch.popc(low2 - cutlass.Uint32(1))) + incl_tw = cute.arch.shuffle_sync(run2, tw) + slot_tw = cute.arch.shuffle_sync(v_s, tw) + prefix = incl_tw - slot_tw + + # locate the bin inside the target slice + ppl = cutlass.const_expr((bins_per_warp + self.WARP_SIZE - 1) // self.WARP_SIZE) + cnt_frag = cute.make_rmem_tensor((ppl,), cutlass.Int32) + my_sum = cutlass.Int32(0) + for j3 in cutlass.range_constexpr(ppl): + pos = lane * cutlass.Int32(ppl) + cutlass.Int32(j3) + cnt_j = cutlass.Int32(0) + if pos < cutlass.Int32(bins_per_warp): + bidx3 = cutlass.Int32(kBins - 1) - tw * cutlass.Int32(bins_per_warp) - pos + cnt_j = smem_hist[bidx3] + cnt_frag[j3] = cnt_j + my_sum = my_sum + cnt_j + run3 = my_sum + for d3 in cutlass.range_constexpr(5): + off3 = cutlass.const_expr(1 << d3) + src3 = lane - cutlass.Int32(off3) + if src3 < cutlass.Int32(0): + src3 = cutlass.Int32(0) + up3 = cute.arch.shuffle_sync(run3, src3) + if lane >= cutlass.Int32(off3): + run3 = run3 + up3 + base3 = prefix + (run3 - my_sum) + + b_loc = cutlass.Int32(kBins - 1) + ra_loc = prefix + hit = cutlass.Int32(0) + r3 = base3 + for j4 in cutlass.range_constexpr(ppl): + pos4 = lane * cutlass.Int32(ppl) + cutlass.Int32(j4) + cnt4 = cnt_frag[j4] + if ( + pos4 < cutlass.Int32(bins_per_warp) + and r3 < cutlass.Int32(kK) + and r3 + cnt4 >= cutlass.Int32(kK) + and hit == cutlass.Int32(0) + ): + b_loc = cutlass.Int32(kBins - 1) - tw * cutlass.Int32(bins_per_warp) - pos4 + ra_loc = r3 + hit = cutlass.Int32(1) + r3 = r3 + cnt4 + mask3 = cute.arch.vote_ballot_sync(hit != cutlass.Int32(0)) + b_out = cutlass.Int32(kBins - 1) + ra_out = prefix + if mask3 != cutlass.Uint32(0): + low = mask3 & (cutlass.Uint32(0) - mask3) + src = cutlass.Int32(cute.arch.popc(low - cutlass.Uint32(1))) + b_out = cute.arch.shuffle_sync(b_loc, src) + ra_out = cute.arch.shuffle_sync(ra_loc, src) + return b_out, ra_out + + # ------------------------------------------------------------------ + # _p4_fine_rw - redundant-warp variant of the fine sub-bin search, + # the same transformation _p4_coarse_rw applies one level up: every + # warp resolves it from the staged per-warp sums with an idx-shuffle + # scan and a ballot. Integer sums are associative: every warp lands + # on the same answer. + # ------------------------------------------------------------------ + @cute.jit + def _p4_fine_rw(self, smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane): + kK = cutlass.const_expr(self.top_k) + fbpw = cutlass.const_expr(fbins // self.num_warps) + + ws = cutlass.Int32(0) + if cutlass.const_expr(fbpw <= self.WARP_SIZE): + if lane < cutlass.Int32(fbpw): + bif = cutlass.Int32(fbins - 1) - warp_id * cutlass.Int32(fbpw) - lane + ws = smem_hist[bif] + ws = self.warp_reduce_sum_i32(ws) + else: + for jm in cutlass.range_constexpr(fbpw): + bif = cutlass.Int32(fbins - 1) - warp_id * cutlass.Int32(fbpw) - cutlass.Int32(jm) + ws = ws + smem_hist[bif] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = ws + cute.arch.barrier() + + v_s = cutlass.Int32(0) + if lane < cutlass.Int32(self.num_warps): + v_s = smem_wcnt[lane] + run2 = v_s + for d2 in cutlass.range_constexpr(5): + off2 = cutlass.const_expr(1 << d2) + src2 = lane - cutlass.Int32(off2) + if src2 < cutlass.Int32(0): + src2 = cutlass.Int32(0) + up2 = cute.arch.shuffle_sync(run2, src2) + if lane >= cutlass.Int32(off2): + run2 = run2 + up2 + m2 = cute.arch.vote_ballot_sync(rank_above + run2 >= cutlass.Int32(kK)) + tw = cutlass.Int32(self.num_warps - 1) + if m2 != cutlass.Uint32(0): + low2 = m2 & (cutlass.Uint32(0) - m2) + tw = cutlass.Int32(cute.arch.popc(low2 - cutlass.Uint32(1))) + incl_tw = cute.arch.shuffle_sync(run2, tw) + slot_tw = cute.arch.shuffle_sync(v_s, tw) + prefix = rank_above + (incl_tw - slot_tw) + + ppl = cutlass.const_expr((fbpw + self.WARP_SIZE - 1) // self.WARP_SIZE) + cnt_frag = cute.make_rmem_tensor((ppl,), cutlass.Int32) + my_sum = cutlass.Int32(0) + for j3 in cutlass.range_constexpr(ppl): + pos = lane * cutlass.Int32(ppl) + cutlass.Int32(j3) + cj = cutlass.Int32(0) + if pos < cutlass.Int32(fbpw): + sbi = cutlass.Int32(fbins - 1) - tw * cutlass.Int32(fbpw) - pos + cj = smem_hist[sbi] + cnt_frag[j3] = cj + my_sum = my_sum + cj + run3 = my_sum + for d3 in cutlass.range_constexpr(5): + off3 = cutlass.const_expr(1 << d3) + src3 = lane - cutlass.Int32(off3) + if src3 < cutlass.Int32(0): + src3 = cutlass.Int32(0) + up3 = cute.arch.shuffle_sync(run3, src3) + if lane >= cutlass.Int32(off3): + run3 = run3 + up3 + base3 = prefix + (run3 - my_sum) + + sb_loc = cutlass.Int32(fbins - 1) + ra_loc = prefix + hit = cutlass.Int32(0) + r3 = base3 + for j4 in cutlass.range_constexpr(ppl): + pos4 = lane * cutlass.Int32(ppl) + cutlass.Int32(j4) + c4 = cnt_frag[j4] + if ( + pos4 < cutlass.Int32(fbpw) + and r3 < cutlass.Int32(kK) + and r3 + c4 >= cutlass.Int32(kK) + and hit == cutlass.Int32(0) + ): + sb_loc = cutlass.Int32(fbins - 1) - tw * cutlass.Int32(fbpw) - pos4 + ra_loc = r3 + hit = cutlass.Int32(1) + r3 = r3 + c4 + mask3 = cute.arch.vote_ballot_sync(hit != cutlass.Int32(0)) + sb_out = cutlass.Int32(fbins - 1) + ra_out = prefix + if mask3 != cutlass.Uint32(0): + low = mask3 & (cutlass.Uint32(0) - mask3) + src = cutlass.Int32(cute.arch.popc(low - cutlass.Uint32(1))) + sb_out = cute.arch.shuffle_sync(sb_loc, src) + ra_out = cute.arch.shuffle_sync(ra_loc, src) + return sb_out, ra_out + + # ------------------------------------------------------------------ + # Phase 4 (alt): fused rank-and-scatter (enable_p4_rank_scatter). + # ------------------------------------------------------------------ + @cute.jit + def phase4_rank_scatter( + self, + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count, + tidx, + warp_id, + lane, + ext_range_flag=None, # list rows: walk pre-staged range + hist zero + ext_min=None, # list rows: cut line == exact candidate minimum + ): + kK = cutlass.const_expr(self.top_k) + kBins = cutlass.const_expr(self.kNumBins) + pair_cap = cutlass.const_expr(_pair_cap_for(self.kNumBins)) + num_threads = cutlass.const_expr(self.num_threads) + num_warps = cutlass.const_expr(self.num_warps) + bins_per_warp = cutlass.const_expr(kBins // self.num_warps) + + if cand_count == cutlass.Int32(kK): + i4 = tidx + while i4 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[i4] = self.dtype(smem_keys[i4]) + output_indices_row[i4] = smem_vals[i4] + i4 = i4 + cutlass.Int32(num_threads) + elif cand_count > cutlass.Int32(kK): + if cutlass.const_expr(_P4_SUB_DBG): + sc1 = cutlass.Int64(0) + sc2 = cutlass.Int64(0) + sc3 = cutlass.Int64(0) + sc4 = cutlass.Int64(0) + sc5 = cutlass.Int64(0) + sc6 = cutlass.Int64(0) + sc0 = cute.arch.clock64() + bmin_r = cutlass.Float32(self.FLT_MAX) + bmax_r = cutlass.Float32(self.NEG_FLT_MAX) + # a Python bool here, so with the feature off the stock body + # below traces straight-line instead of inside an scf.if whose + # predicate is a compile-time constant + run_stock_range = True + if cutlass.const_expr(ext_range_flag is not None): + use_ext_r = ext_range_flag + run_stock_range = use_ext_r == cutlass.Int32(0) + if use_ext_r == cutlass.Int32(1): + # list rows: the take walk pre-zeroed the hist and staged + # per-warp maxima in smem_wcnt (its end barrier orders + # them); min := cut line by construction. + if cutlass.const_expr(ext_min is not None): + bmin_r = ext_min + for w in cutlass.range_constexpr(self.num_warps): + vmax = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[w].ir_value()) + ) + bmax_r = cute.arch.fmax(bmax_r, vmax) + if bmax_r <= bmin_r: + bmax_r = bmin_r + cutlass.Float32(1e-6) + if run_stock_range: + # ---- block min/max over candidates ---- + # The accepted threshold is an EXACT lower bound on every + # candidate, so it can stand in for the min. Only the + # assist tiers have a published threshold, hence the gate. + use_thr_min = cutlass.const_expr(self.p4_no_fine) + local_cmin = cutlass.Float32(self.FLT_MAX) + local_cmax = cutlass.Float32(self.NEG_FLT_MAX) + i5 = tidx + while i5 < cand_count: + v = smem_keys[i5] + if cutlass.const_expr(not use_thr_min): + local_cmin = _fmin_f32_inline(local_cmin, v) + local_cmax = cute.arch.fmax(local_cmax, v) + i5 = i5 + cutlass.Int32(num_threads) + cmin = cutlass.Float32(0.0) + if cutlass.const_expr(not use_thr_min): + cmin = self.warp_reduce_min_f32(local_cmin) + cmax = self.warp_reduce_max_f32(local_cmax) + if lane == cutlass.Int32(0): + if cutlass.const_expr(not use_thr_min): + smem_wcnt[warp_id] = float_as_uint32(cmin) + smem_hist[warp_id] = float_as_uint32(cmax) + cute.arch.barrier() + # lane-parallel cross-warp fold: lane w holds slot w and one + # warp reduce settles it. min/max reassociate freely, so the + # result is bit-identical on every warp without a leader. + pmn = cutlass.Float32(self.FLT_MAX) + pmx = cutlass.Float32(self.NEG_FLT_MAX) + if lane < cutlass.Int32(self.num_warps): + if cutlass.const_expr(not use_thr_min): + pmn = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, smem_wcnt[lane].ir_value()) + ) + pmx = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, smem_hist[lane].ir_value()) + ) + if cutlass.const_expr(use_thr_min): + bmin_r = s_thr[0] + else: + bmin_r = _fmin_f32_inline(bmin_r, self.warp_reduce_min_f32(pmn)) + bmax_r = cute.arch.fmax(bmax_r, self.warp_reduce_max_f32(pmx)) + if bmax_r <= bmin_r: + bmax_r = bmin_r + cutlass.Float32(1e-6) + cute.arch.barrier() + if cutlass.const_expr(_P4_SUB_DBG): + sc1 = cute.arch.clock64() + # ---- zero + build histogram ---- + i6 = tidx + while i6 < cutlass.Int32(kBins): + smem_hist[i6] = cutlass.Int32(0) + i6 = i6 + cutlass.Int32(num_threads) + cute.arch.barrier() range1 = bmax_r - bmin_r inv1 = (cutlass.Float32(kBins - 1) + cutlass.Float32(0.99)) / range1 i7 = tidx @@ -2982,166 +4131,192 @@ def phase4_rank_scatter( atomicAdd(smem_hist.iterator + bin_i, cutlass.Int32(1)) i7 = i7 + cutlass.Int32(num_threads) cute.arch.barrier() - # ---- 3-step high→low bin search → straddling bin b* + rank_above ---- - warp_bin_sum = cutlass.Int32(0) - for jb in cutlass.range_constexpr(bins_per_warp): - bidx_s = ( - cutlass.Int32(kBins - 1) - - warp_id * cutlass.Int32(bins_per_warp) - - cutlass.Int32(jb) - ) - warp_bin_sum = warp_bin_sum + smem_hist[bidx_s] - if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = warp_bin_sum - cute.arch.barrier() - if tidx == cutlass.Int32(0): - cum = cutlass.Int32(0) - tw = cutlass.Int32(num_warps - 1) - found = cutlass.Int32(0) - for w2 in cutlass.range_constexpr(self.num_warps): - cum = cum + smem_wcnt[w2] - if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0): - tw = cutlass.Int32(w2) - found = cutlass.Int32(1) - cum2 = cutlass.Int32(0) - for w3 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w3) < tw: - cum2 = cum2 + smem_wcnt[w3] - s_iscalars[2] = cum2 # prefix-count before target warp - s_iscalars[3] = tw - cute.arch.barrier() - target_warp = s_iscalars[3] - cute.arch.barrier() - if warp_id == target_warp and lane == cutlass.Int32(0): - base_cum = s_iscalars[2] - b_star = cutlass.Int32(kBins - 1) - rank_above = base_cum - set_d = cutlass.Int32(0) - for jb2 in cutlass.range_constexpr(bins_per_warp): - bidx2 = ( - cutlass.Int32(kBins - 1) - - target_warp * cutlass.Int32(bins_per_warp) - - cutlass.Int32(jb2) - ) - ra_before = base_cum - base_cum = base_cum + smem_hist[bidx2] - if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0): - b_star = bidx2 - rank_above = ra_before # count in bins strictly above b* - set_d = cutlass.Int32(1) - s_iscalars[2] = rank_above - s_iscalars[3] = b_star - s_iscalars[4] = cutlass.Int32(0) # cnt_above - s_iscalars[1] = cutlass.Int32(0) # cnt_straddle - cute.arch.barrier() - b_star = s_iscalars[3] - rank_above = s_iscalars[2] - - # ---- EXACT: one fine-histogram recursion on the straddling bin b* ---- - if cutlass.const_expr(self.enable_p4_rank_scatter_exact): - # FIXED small fine-bin count (independent of kNumBins) — cuts the - # re-zero + 3-step cost (esp. K=2048 where kNumBins=2048); 256 - # sub-bins over bin b* gives kNumBins×256 effective resolution, - # enough to resolve the straddling bin to ≤1 distinct value. - fbins = cutlass.const_expr(256) - fbpw = cutlass.const_expr(256 // self.num_warps) - # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1) - f_lo = bmin_r + cutlass.Float32(b_star) / inv1 - finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1 - # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands - iz = tidx - while iz < cutlass.Int32(fbins): - smem_hist[iz] = cutlass.Int32(0) - iz = iz + cutlass.Int32(num_threads) - cute.arch.barrier() - ifb = tidx - while ifb < cand_count: - vf = smem_keys[ifb] - cb = cutlass.Int32((vf - bmin_r) * inv1) - if cb < cutlass.Int32(0): - cb = cutlass.Int32(0) - if cb > cutlass.Int32(kBins - 1): - cb = cutlass.Int32(kBins - 1) - if cb == b_star: - sb = cutlass.Int32((vf - f_lo) * finv) - if sb < cutlass.Int32(0): - sb = cutlass.Int32(0) - if sb > cutlass.Int32(fbins - 1): - sb = cutlass.Int32(fbins - 1) - atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1)) - ifb = ifb + cutlass.Int32(num_threads) + if cutlass.const_expr(_P4_SUB_DBG): + sc2 = cute.arch.clock64() + # ---- high→low bin search → straddling bin b* + rank_above ---- + if cutlass.const_expr(self.p4_warp_redundant): + b_star, rank_above = self._p4_coarse_rw(smem_hist, smem_wcnt, warp_id, lane) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # cnt_above + s_iscalars[1] = cutlass.Int32(0) # cnt_straddle cute.arch.barrier() - # fine 3-step search seeded at rank_above (over fbins bins) - fws = cutlass.Int32(0) - for jbf in cutlass.range_constexpr(fbpw): - bif = ( - cutlass.Int32(fbins - 1) - - warp_id * cutlass.Int32(fbpw) - - cutlass.Int32(jbf) + if cutlass.const_expr(_P4_SUB_DBG): + sc3 = cute.arch.clock64() + else: + warp_bin_sum = cutlass.Int32(0) + for jb in cutlass.range_constexpr(bins_per_warp): + bidx_s = ( + cutlass.Int32(kBins - 1) + - warp_id * cutlass.Int32(bins_per_warp) + - cutlass.Int32(jb) ) - fws = fws + smem_hist[bif] + warp_bin_sum = warp_bin_sum + smem_hist[bidx_s] if lane == cutlass.Int32(0): - smem_wcnt[warp_id] = fws + smem_wcnt[warp_id] = warp_bin_sum cute.arch.barrier() if tidx == cutlass.Int32(0): - cumf = rank_above - twf = cutlass.Int32(num_warps - 1) - fnd = cutlass.Int32(0) + cum = cutlass.Int32(0) + tw = cutlass.Int32(num_warps - 1) + found = cutlass.Int32(0) for w2 in cutlass.range_constexpr(self.num_warps): - cumf = cumf + smem_wcnt[w2] - if cumf >= cutlass.Int32(kK) and fnd == cutlass.Int32(0): - twf = cutlass.Int32(w2) - fnd = cutlass.Int32(1) - pre = rank_above + cum = cum + smem_wcnt[w2] + if cum >= cutlass.Int32(kK) and found == cutlass.Int32(0): + tw = cutlass.Int32(w2) + found = cutlass.Int32(1) + cum2 = cutlass.Int32(0) for w3 in cutlass.range_constexpr(self.num_warps): - if cutlass.Int32(w3) < twf: - pre = pre + smem_wcnt[w3] - # Stage prefix/target-warp metadata in spare s_iscalars - # slots, NOT smem_hist[0]/[1]: the last fine warp's reverse - # scan below walks fine bins down to 0/1, so reusing those - # histogram bins as scratch would corrupt sb_star/ra_fine - # when twf2 == num_warps-1. Slots [4]/[1] are dead here - # (re-zeroed at the cnt_above/cnt_strad reset below). - s_iscalars[4] = pre # prefix into target fine warp - s_iscalars[1] = twf # target fine warp + if cutlass.Int32(w3) < tw: + cum2 = cum2 + smem_wcnt[w3] + s_iscalars[2] = cum2 # prefix-count before target warp + s_iscalars[3] = tw cute.arch.barrier() - pre_f = s_iscalars[4] - twf2 = s_iscalars[1] - if warp_id == twf2 and lane == cutlass.Int32(0): - base_f = pre_f - sb_star = cutlass.Int32(fbins - 1) - ra_fine = base_f - sd = cutlass.Int32(0) - for jb3 in cutlass.range_constexpr(fbpw): - sbi = ( - cutlass.Int32(fbins - 1) - - twf2 * cutlass.Int32(fbpw) - - cutlass.Int32(jb3) + target_warp = s_iscalars[3] + cute.arch.barrier() + if warp_id == target_warp and lane == cutlass.Int32(0): + base_cum = s_iscalars[2] + b_star_s = cutlass.Int32(kBins - 1) + rank_above_s = base_cum + set_d = cutlass.Int32(0) + for jb2 in cutlass.range_constexpr(bins_per_warp): + bidx2 = ( + cutlass.Int32(kBins - 1) + - target_warp * cutlass.Int32(bins_per_warp) + - cutlass.Int32(jb2) ) - ra_b = base_f - base_f = base_f + smem_hist[sbi] - if base_f >= cutlass.Int32(kK) and sd == cutlass.Int32(0): - sb_star = sbi - ra_fine = ra_b - sd = cutlass.Int32(1) - smem_hist[2] = sb_star - smem_hist[3] = ra_fine + ra_before = base_cum + base_cum = base_cum + smem_hist[bidx2] + if base_cum >= cutlass.Int32(kK) and set_d == cutlass.Int32(0): + b_star_s = bidx2 + rank_above_s = ra_before # count in bins strictly above b* + set_d = cutlass.Int32(1) + s_iscalars[2] = rank_above_s + s_iscalars[3] = b_star_s + s_iscalars[4] = cutlass.Int32(0) # cnt_above + s_iscalars[1] = cutlass.Int32(0) # cnt_straddle cute.arch.barrier() + if cutlass.const_expr(_P4_SUB_DBG): + sc3 = cute.arch.clock64() + b_star = s_iscalars[3] + rank_above = s_iscalars[2] + + # ---- EXACT: one fine-histogram recursion on the straddling bin b* ---- + if cutlass.const_expr(self.enable_p4_rank_scatter_exact): + # FIXED small fine-bin count (independent of kNumBins): 256 + # sub-bins over bin b* gives kNumBins×256 effective resolution, + # enough to resolve the straddling bin to ≤1 distinct value. + fbins = cutlass.const_expr(256) + # bin b* value range under the inv1 binning: [f_lo, f_lo + 1/inv1) + f_lo = cutlass.Float32(0.0) + finv = cutlass.Float32(0.0) + if cutlass.const_expr(not self.p4_no_fine): + f_lo = bmin_r + cutlass.Float32(b_star) / inv1 + finv = (cutlass.Float32(fbins - 1) + cutlass.Float32(0.99)) * inv1 + if cutlass.const_expr(self.p4_no_fine): + # Sub-binning collapsed: finv 0 sends every member of the + # straddling coarse bin to sub-bin 0 == sb*, so the + # scatter parks the whole coarse class and the tail ranks + # it. MUST stay a compile-time removal (the re-zero, build + # and search below untraced), not a runtime branch. + finv = cutlass.Float32(0.0) + sb_star = cutlass.Int32(0) + rank_above_fine = rank_above + # bin b* spans [f_lo, f_hi); the clamped ends fold + # out-of-range values into bin 0 and bin kBins-1, so those + # two drop the matching side. Only the range-test arms read + # these, so upstream's build must not compute them. + if cutlass.const_expr(self.p4_fine_rangetest or self.p4_scat_rangetest): + f_hi = f_lo + cutlass.Float32(1.0) / inv1 + lo_edge = b_star == cutlass.Int32(0) + hi_edge = b_star == cutlass.Int32(kBins - 1) + if cutlass.const_expr(not self.p4_no_fine): + # re-zero (only fbins slots) + build fine sub-hist of bin-b* cands + iz = tidx + while iz < cutlass.Int32(fbins): + smem_hist[iz] = cutlass.Int32(0) + iz = iz + cutlass.Int32(num_threads) + cute.arch.barrier() + if cutlass.const_expr(self.p4_fine_rangetest): + # A candidate belongs to bin b* exactly when its value + # lies in [f_lo, f_hi). The clamped ends of the binning + # fold out-of-range values INTO bin 0 and bin kBins-1, + # so those two bins must drop the matching side of the + # range test to stay bit-identical. + ifb = tidx + while ifb < cand_count: + vf = smem_keys[ifb] + inb = vf >= f_lo and vf < f_hi + if lo_edge: + inb = vf < f_hi + if hi_edge: + inb = vf >= f_lo + if inb: + sb = cutlass.Int32((vf - f_lo) * finv) + if sb < cutlass.Int32(0): + sb = cutlass.Int32(0) + if sb > cutlass.Int32(fbins - 1): + sb = cutlass.Int32(fbins - 1) + atomicAdd(smem_hist.iterator + sb, cutlass.Int32(1)) + ifb = ifb + cutlass.Int32(num_threads) + else: + ifb = tidx + while ifb < cand_count: + vfo = smem_keys[ifb] + cbo = cutlass.Int32((vfo - bmin_r) * inv1) + if cbo < cutlass.Int32(0): + cbo = cutlass.Int32(0) + if cbo > cutlass.Int32(kBins - 1): + cbo = cutlass.Int32(kBins - 1) + if cbo == b_star: + sbo = cutlass.Int32((vfo - f_lo) * finv) + if sbo < cutlass.Int32(0): + sbo = cutlass.Int32(0) + if sbo > cutlass.Int32(fbins - 1): + sbo = cutlass.Int32(fbins - 1) + atomicAdd(smem_hist.iterator + sbo, cutlass.Int32(1)) + ifb = ifb + cutlass.Int32(num_threads) + cute.arch.barrier() + # fine sub-bin search, resolved lane-parallel on every + # warp (see _p4_fine_rw); the answer comes back in + # registers. + sb_star, rank_above_fine = self._p4_fine_rw( + smem_hist, smem_wcnt, fbins, rank_above, warp_id, lane + ) if tidx == cutlass.Int32(0): s_iscalars[4] = cutlass.Int32(0) # cnt_above s_iscalars[0] = cutlass.Int32(0) # cnt_mid (b*, sub>sb*) s_iscalars[1] = cutlass.Int32(0) # cnt_strad (b*, sub==sb*) cute.arch.barrier() - sb_star = smem_hist[2] - rank_above_fine = smem_hist[3] + if cutlass.const_expr(_P4_SUB_DBG): + sc4 = cute.arch.clock64() isc = tidx while isc < cand_count: v = smem_keys[isc] - bin_i = cutlass.Int32((v - bmin_r) * inv1) - if bin_i < cutlass.Int32(0): - bin_i = cutlass.Int32(0) - if bin_i > cutlass.Int32(kBins - 1): - bin_i = cutlass.Int32(kBins - 1) + if cutlass.const_expr(self.p4_scat_rangetest): + # same three-way split as the bin recompute, by value: + # above b* <=> v >= f_hi (impossible at the top bin, + # which absorbs everything higher), inside b* <=> v in + # [f_lo, f_hi) with the edge bins dropping their + # absorbed side. bin_i is only ever compared against + # b*, so encoding the class as b*-1 / b* / b*+1 keeps + # the branches below unchanged. + abv = v >= f_hi + inb2 = v >= f_lo and v < f_hi + if lo_edge: + inb2 = v < f_hi + if hi_edge: + abv = False + inb2 = v >= f_lo + bin_i = b_star - cutlass.Int32(1) + if abv: + bin_i = b_star + cutlass.Int32(1) + if inb2: + bin_i = b_star + else: + bin_i = cutlass.Int32((v - bmin_r) * inv1) + if bin_i < cutlass.Int32(0): + bin_i = cutlass.Int32(0) + if bin_i > cutlass.Int32(kBins - 1): + bin_i = cutlass.Int32(kBins - 1) if bin_i > b_star: pos = atomicAdd(s_iscalars.iterator + cutlass.Int32(4), cutlass.Int32(1)) if pos < cutlass.Int32(kK): @@ -3149,11 +4324,19 @@ def phase4_rank_scatter( output_values_row[pos] = self.dtype(v) output_indices_row[pos] = smem_vals[isc] elif bin_i == b_star: - sb = cutlass.Int32((v - f_lo) * finv) - if sb < cutlass.Int32(0): - sb = cutlass.Int32(0) - if sb > cutlass.Int32(fbins - 1): - sb = cutlass.Int32(fbins - 1) + # With the fine level compiled out the sub-bin is a + # constant 0 == sb*, so the whole three-way split + # below collapses to the park arm and the per- + # candidate recompute (a subtract, a multiply and two + # clamps) is dead. Keep it a compile-time constant so + # the scatter does not carry it. + sb = cutlass.Int32(0) + if cutlass.const_expr(not self.p4_no_fine): + sb = cutlass.Int32((v - f_lo) * finv) + if sb < cutlass.Int32(0): + sb = cutlass.Int32(0) + if sb > cutlass.Int32(fbins - 1): + sb = cutlass.Int32(fbins - 1) if sb > sb_star: o = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1)) pos = rank_above + o @@ -3163,6 +4346,21 @@ def phase4_rank_scatter( output_indices_row[pos] = smem_vals[isc] elif sb == sb_star: o = atomicAdd(s_iscalars.iterator + cutlass.Int32(1), cutlass.Int32(1)) + if cutlass.const_expr( + self.p4_exact_tail and self.p4_tail_fast and self.p4_tail_v3 + ): + # Park (value bits, index) so the exact-tail + # repair never re-walks the candidates. Pairs + # live above the digit bins the radix route + # zeroes; small classes are the only consumer. + # Unconditional, index wrapped: the buffer is + # only ever READ when the class fits it, and + # then the wrap is a no-op. + ow = (o & cutlass.Int32(pair_cap - 1)) * cutlass.Int32(2) + smem_hist[cutlass.Int32(_PAIR_BASE) + ow] = float_as_int32(v) + smem_hist[cutlass.Int32(_PAIR_BASE) + ow + cutlass.Int32(1)] = ( + smem_vals[isc] + ) pos = rank_above_fine + o if pos < cutlass.Int32(kK): if cutlass.const_expr(self.return_output_values): @@ -3170,6 +4368,8 @@ def phase4_rank_scatter( output_indices_row[pos] = smem_vals[isc] isc = isc + cutlass.Int32(num_threads) cute.arch.barrier() + if cutlass.const_expr(_P4_SUB_DBG): + sc5 = cute.arch.clock64() cnt_strad = s_iscalars[1] filled = rank_above_fine + cnt_strad if filled > cutlass.Int32(kK): @@ -3193,148 +4393,704 @@ def phase4_rank_scatter( # slot range [rank_above_fine, kK). Unambiguous rows (the # overwhelming majority) pay two scalar compares; the counters # and the fine histogram are reused, so SMEM does not grow. - # [p4tt] tiny-tie fast path: when the exact-tail gate fires - # with a small (b*, sb*) tie class (cnt_strad <= 128 — the - # real firing cells hold 2), ONE candidate pass collects the - # class and thread0 selects the top-need exactly, replacing - # the 4 unconditional radix passes. Larger classes take the - # UNMODIFIED radix select below (verbatim copy). + # boundary-class repair: pure-tie classes (one key value) + # exit on a warp-reduce precheck; classes parked in the + # pair buffer rank block-parallel over the parked pairs; + # anything larger takes the full-candidate radix below. if cutlass.const_expr(self.p4_exact_tail and self.p4_tail_fast): # [p4tt] - need0 = cutlass.Int32(kK) - rank_above_fine - if cnt_strad > need0 and need0 > cutlass.Int32(0): - if cnt_strad <= cutlass.Int32(128): - # [p4tt] SMEM: (value_bits, cand_idx) pairs at - # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255). - # The 256 digit bins are dead here (the fast path - # replaces the radix levels that used them); the - # sb_star/ra staging in slots 2/3 was read by - # every thread before the pre-scatter barrier. - # Persistent radix scalars [256..258] untouched. - # Collect counter = s_iscalars[0] (dead after the - # scatter; same reuse as the radix rewrite pass). + if cutlass.const_expr(self.p4_tail_v3): + need0 = cutlass.Int32(kK) - rank_above_fine + # per-thread compact buffers, bounded by the + # strided trip count over the candidate array + nbuf7 = cutlass.const_expr( + (self.kC + self.num_threads - 1) // self.num_threads + ) + rv7 = cute.make_rmem_tensor((nbuf7,), cutlass.Float32) + ri7 = cute.make_rmem_tensor((nbuf7,), cutlass.Int32) + if cutlass.const_expr(_P4_TAIL_DBG or _P4_SUB_DBG): if tidx == cutlass.Int32(0): - s_iscalars[0] = cutlass.Int32(0) - cute.arch.barrier() - itc = tidx - while itc < cand_count: - tv = smem_keys[itc] - tb = cutlass.Int32((tv - bmin_r) * inv1) - if tb < cutlass.Int32(0): - tb = cutlass.Int32(0) - if tb > cutlass.Int32(kBins - 1): - tb = cutlass.Int32(kBins - 1) - if tb == b_star: - ts = cutlass.Int32((tv - f_lo) * finv) - if ts < cutlass.Int32(0): - ts = cutlass.Int32(0) - if ts > cutlass.Int32(fbins - 1): - ts = cutlass.Int32(fbins - 1) - if ts == sb_star: - to = atomicAdd( - s_iscalars.iterator + cutlass.Int32(0), cutlass.Int32(1) + s_thr[1] = cutlass.Float32(cnt_strad) + s_thr[2] = cutlass.Float32(need0) + fast_done = cutlass.Int32(1) + if cnt_strad > need0 and need0 > cutlass.Int32(0): + fast_done = cutlass.Int32(0) + if cnt_strad <= cutlass.Int32(pair_cap): + # Small mixed class: the scatter already parked + # every member as a (value bits, index) pair, + # so there is nothing to collect - go straight + # to the rank. Rank with the WHOLE BLOCK: each + # warp owns a stride of the class and its 32 + # lanes split the comparisons. + e9 = warp_id + while e9 < cnt_strad: + be9 = smem_hist[_PAIR_BASE + e9 + e9] + ke9 = f32_order_key( + cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, be9.ir_value()) ) - if to < cutlass.Int32(128): - smem_hist[to + to] = float_as_int32(tv) - smem_hist[to + to + cutlass.Int32(1)] = smem_vals[itc] - itc = itc + cutlass.Int32(num_threads) - cute.arch.barrier() - # [p4tt] thread0 exact top-need0 select rewriting - # positions [rank_above_fine, kK). Consumed flag = - # the cand_idx slot set to -1 (indices are always - # >= 0), so a genuine -FLT_MAX value in the class - # remains selectable (no value sentinel). Ties - # (bit-equal values) pick arbitrarily: value-set - # exact. - if tidx == cutlass.Int32(0): - tj = cutlass.Int32(0) - while tj < need0: - tbv = cutlass.Float32(self.NEG_FLT_MAX) - tbi = cutlass.Int32(-1) - ti = cutlass.Int32(0) - while ti < cnt_strad: - tvi = smem_hist[ti + ti + cutlass.Int32(1)] - if tvi >= cutlass.Int32(0): - tvb = smem_hist[ti + ti] - tvv = cutlass.Float32( + ) ^ cutlass.Int32(-2147483648) + c9 = cutlass.Int32(0) + j9 = lane + while j9 < cnt_strad: + bj9 = smem_hist[_PAIR_BASE + j9 + j9] + kj9 = f32_order_key( + cutlass.Float32( llvm.bitcast( - cutlass.Float32.mlir_type, - tvb.ir_value(), + cutlass.Float32.mlir_type, bj9.ir_value() ) ) - take = cutlass.Int32(0) - if tbi < cutlass.Int32(0): - take = cutlass.Int32(1) - elif tvv > tbv: - take = cutlass.Int32(1) - if take == cutlass.Int32(1): - tbv = tvv - tbi = ti - ti = ti + cutlass.Int32(1) - pos = rank_above_fine + tj - if cutlass.const_expr(self.return_output_values): - output_values_row[pos] = self.dtype(tbv) - output_indices_row[pos] = smem_hist[ - tbi + tbi + cutlass.Int32(1) - ] - smem_hist[tbi + tbi + cutlass.Int32(1)] = cutlass.Int32(-1) - tj = tj + cutlass.Int32(1) - cute.arch.barrier() - else: - self._p4_exact_tail_radix_select( - kK, - kBins, - num_threads, - num_warps, - need0, - cand_count, - rank_above_fine, - b_star, - sb_star, - bmin_r, - f_lo, - finv, - fbins, - inv1, - tidx, - lane, - warp_id, - smem_hist, - smem_keys, - smem_vals, - smem_wcnt, - s_iscalars, - output_indices_row, - output_values_row, - ) - elif cutlass.const_expr(self.p4_exact_tail): # [p4tt] if->elif only + ) ^ cutlass.Int32(-2147483648) + if kj9 > ke9: + c9 = c9 + cutlass.Int32(1) + elif kj9 == ke9 and j9 < e9: + c9 = c9 + cutlass.Int32(1) + j9 = j9 + cutlass.Int32(32) + r9 = self.warp_reduce_sum_i32(c9) + if lane == cutlass.Int32(0) and r9 < need0: + pos9 = rank_above_fine + r9 + if pos9 < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos9] = self.dtype( + cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + be9.ir_value(), + ) + ) + ) + output_indices_row[pos9] = smem_hist[ + _PAIR_BASE + e9 + e9 + cutlass.Int32(1) + ] + e9 = e9 + cutlass.Int32(num_warps) + cute.arch.barrier() + else: + # Large class only: block-wide pure-tie check + # (min/max order key over the (b*, sb*) class) + # decides whether the radix can be skipped. + # A pure-tie class needs NO repair — the scatter's + # arrival fill of bit-equal values is already + # value-set exact. + # Staging mirrors the head min/max (wcnt + hist + # slots [0..31], both dead here; pairs live at + # 260+). + if tidx == cutlass.Int32(0): + s_iscalars[0] = cutlass.Int32(0) + nh7 = cutlass.Int32(0) + kmn6 = cutlass.Int32(2147483647) + kmx6 = cutlass.Int32(-2147483648) + it6 = tidx + # The pure-tie pre-check SKIPs the repair when + # the class is bit-uniform; only the large-class + # route runs it (the small route is already + # value-exact on a pure tie). + while it6 < cand_count: + v6 = smem_keys[it6] + b6 = cutlass.Int32((v6 - bmin_r) * inv1) + if b6 < cutlass.Int32(0): + b6 = cutlass.Int32(0) + if b6 > cutlass.Int32(kBins - 1): + b6 = cutlass.Int32(kBins - 1) + if b6 == b_star: + # with the fine level compiled out the + # whole coarse bin IS the class + s6 = cutlass.Int32(0) + if cutlass.const_expr(not self.p4_no_fine): + s6 = cutlass.Int32((v6 - f_lo) * finv) + if s6 < cutlass.Int32(0): + s6 = cutlass.Int32(0) + if s6 > cutlass.Int32(fbins - 1): + s6 = cutlass.Int32(fbins - 1) + if s6 == sb_star: + k6 = f32_order_key(v6) ^ cutlass.Int32(-2147483648) + if k6 < kmn6: + kmn6 = k6 + if k6 > kmx6: + kmx6 = k6 + # Buffer the member so the compaction + # below needs no walk of its own; the + # array is nbuf7 = kC/num_threads + # deep (unrolled dynamic-index store). + for sl7 in cutlass.range_constexpr(nbuf7): + if cutlass.Int32(sl7) == nh7: + rv7[sl7] = v6 + ri7[sl7] = smem_vals[it6] + nh7 = nh7 + cutlass.Int32(1) + it6 = it6 + cutlass.Int32(num_threads) + kmn6 = cute.arch.warp_redux_sync(kmn6, "min") + kmx6 = cute.arch.warp_redux_sync(kmx6, "max") + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = kmn6 + smem_hist[warp_id] = kmx6 + cute.arch.barrier() + # lane-parallel cross-warp fold: lane w holds + # slot w and one warp reduce settles it, instead + # of every thread walking all num_warps slots of + # two arrays with dependent SMEM reads. Same + # inputs in the same order on every warp, so the + # result stays bit-identical and leaderless. + pa8 = cutlass.Int32(2147483647) + pb8 = cutlass.Int32(-2147483648) + if lane < cutlass.Int32(self.num_warps): + pa8 = smem_wcnt[lane] + pb8 = smem_hist[lane] + kmn7 = cute.arch.warp_redux_sync(pa8, "min") + kmx7 = cute.arch.warp_redux_sync(pb8, "max") + if kmn7 == kmx7: + fast_done = cutlass.Int32(1) + if fast_done == cutlass.Int32(0): + # mixed class: compact the members buffered by + # the pure-tie pass above into + # smem_keys/vals[0..cnt_strad). The staging + # barrier above orders the buffered reads + # before these writes. + # warp-aggregated claim: intra-warp exclusive + # prefix via shfl scan + ONE atomic per warp + # (same-address claims would serialize) + pf7 = nh7 + for so3 in cutlass.range_constexpr(5): + oth3 = cute.arch.shuffle_sync_up( + pf7, cutlass.Int32(1 << so3), mask_and_clamp=0 + ) + if lane >= cutlass.Int32(1 << so3): + pf7 = pf7 + oth3 + tot7 = cute.arch.shuffle_sync(pf7, cutlass.Int32(31)) + wb7 = cutlass.Int32(0) + if lane == cutlass.Int32(31): + if tot7 > cutlass.Int32(0): + wb7 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), tot7 + ) + wb7 = cute.arch.shuffle_sync(wb7, cutlass.Int32(31)) + bs7 = wb7 + pf7 - nh7 + for sl8 in cutlass.range_constexpr(nbuf7): + if cutlass.Int32(sl8) < nh7: + smem_keys[bs7 + cutlass.Int32(sl8)] = rv7[sl8] + smem_vals[bs7 + cutlass.Int32(sl8)] = ri7[sl8] + cute.arch.barrier() + # Large class only (the small one never + # reaches here): block-parallel 4-level + # MSB radix over the compacted class (scans + # touch class pairs only; warp0 shuffle-scan + # digit search). + if tidx == cutlass.Int32(0): + smem_hist[256] = cutlass.Int32(0) + smem_hist[257] = need0 + smem_hist[258] = cutlass.Int32(0) + cute.arch.barrier() + for lvl2 in cutlass.range_constexpr(4): + shift2 = cutlass.const_expr(24 - 8 * lvl2) + iz3 = tidx + while iz3 < cutlass.Int32(256): + smem_hist[iz3] = cutlass.Int32(0) + iz3 = iz3 + cutlass.Int32(num_threads) + cute.arch.barrier() + uthr_c2 = smem_hist[256] + ic2 = tidx + while ic2 < cnt_strad: + uk3 = f32_order_key(smem_keys[ic2]) + pm2 = cutlass.Int32(1) + if cutlass.const_expr(lvl2 > 0): + if (uk3 >> cutlass.Int32(shift2 + 8)) != ( + uthr_c2 >> cutlass.Int32(shift2 + 8) + ): + pm2 = cutlass.Int32(0) + if pm2 == cutlass.Int32(1): + dg2 = ( + uk3 >> cutlass.Int32(shift2) + ) & cutlass.Int32(0xFF) + atomicAdd( + smem_hist.iterator + dg2, cutlass.Int32(1) + ) + ic2 = ic2 + cutlass.Int32(num_threads) + cute.arch.barrier() + if warp_id == cutlass.Int32(0): + ws3 = cutlass.Int32(0) + for jd3 in cutlass.range_constexpr(8): + di3 = ( + cutlass.Int32(255) + - lane * cutlass.Int32(8) + - cutlass.Int32(jd3) + ) + ws3 = ws3 + smem_hist[di3] + pre6 = ws3 + for so2 in cutlass.range_constexpr(5): + oth2 = cute.arch.shuffle_sync_up( + pre6, + cutlass.Int32(1 << so2), + mask_and_clamp=0, + ) + if lane >= cutlass.Int32(1 << so2): + pre6 = pre6 + oth2 + needl3 = smem_hist[257] + if pre6 >= needl3 and (pre6 - ws3) < needl3: + base5 = pre6 - ws3 + dstar2 = cutlass.Int32(0) + above5 = base5 + sd5 = cutlass.Int32(0) + for jd4 in cutlass.range_constexpr(8): + di4 = ( + cutlass.Int32(255) + - lane * cutlass.Int32(8) + - cutlass.Int32(jd4) + ) + ra5 = base5 + base5 = base5 + smem_hist[di4] + if base5 >= needl3 and sd5 == cutlass.Int32(0): + dstar2 = di4 + above5 = ra5 + sd5 = cutlass.Int32(1) + smem_hist[256] = uthr_c2 | ( + dstar2 << cutlass.Int32(shift2) + ) + smem_hist[257] = needl3 - above5 + smem_hist[258] = smem_hist[258] + above5 + cute.arch.barrier() + u_thr2 = smem_hist[256] + cnt_ab2 = smem_hist[258] + need_eq2 = smem_hist[257] + kthr2 = u_thr2 ^ cutlass.Int32(-2147483648) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) + s_iscalars[0] = cutlass.Int32(0) + cute.arch.barrier() + ir3 = tidx + while ir3 < cnt_strad: + vv3 = smem_keys[ir3] + uk4 = f32_order_key(vv3) + ks4 = uk4 ^ cutlass.Int32(-2147483648) + if ks4 > kthr2: + o4 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(4), + cutlass.Int32(1), + ) + pos = rank_above_fine + o4 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vv3) + output_indices_row[pos] = smem_vals[ir3] + elif ks4 == kthr2: + q4 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if q4 < need_eq2: + pos = rank_above_fine + cnt_ab2 + q4 + if pos < cutlass.Int32(kK): + if cutlass.const_expr( + self.return_output_values + ): + output_values_row[pos] = self.dtype(vv3) + output_indices_row[pos] = smem_vals[ir3] + ir3 = ir3 + cutlass.Int32(num_threads) + cute.arch.barrier() + else: + need0_s = cutlass.Int32(kK) - rank_above_fine + if cnt_strad > need0_s and need0_s > cutlass.Int32(0): + if cnt_strad <= cutlass.Int32(128): + # SMEM: (value_bits, cand_idx) pairs at + # smem_hist[2*o]/[2*o+1], o < 128 (slots 0..255). + # The 256 digit bins are dead here (the fast path + # replaces the radix levels that used them); the + # sb_star/ra staging in slots 2/3 was read by + # every thread before the pre-scatter barrier. + # Persistent radix scalars [256..258] untouched. + # Collect counter = s_iscalars[0] (dead after the + # scatter; same reuse as the radix rewrite pass). + if tidx == cutlass.Int32(0): + s_iscalars[0] = cutlass.Int32(0) + cute.arch.barrier() + itc = tidx + while itc < cand_count: + tv = smem_keys[itc] + tb = cutlass.Int32((tv - bmin_r) * inv1) + if tb < cutlass.Int32(0): + tb = cutlass.Int32(0) + if tb > cutlass.Int32(kBins - 1): + tb = cutlass.Int32(kBins - 1) + if tb == b_star: + ts = cutlass.Int32((tv - f_lo) * finv) + if ts < cutlass.Int32(0): + ts = cutlass.Int32(0) + if ts > cutlass.Int32(fbins - 1): + ts = cutlass.Int32(fbins - 1) + if ts == sb_star: + to = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if to < cutlass.Int32(128): + smem_hist[to + to] = float_as_int32(tv) + smem_hist[to + to + cutlass.Int32(1)] = smem_vals[ + itc + ] + itc = itc + cutlass.Int32(num_threads) + cute.arch.barrier() + # thread0 exact top-need0_s select rewriting + # positions [rank_above_fine, kK). Consumed flag = + # the cand_idx slot set to -1 (indices are always + # >= 0), so a genuine -FLT_MAX value in the class + # remains selectable (no value sentinel). Ties + # (bit-equal values) pick arbitrarily: value-set + # exact. + if tidx == cutlass.Int32(0): + tj = cutlass.Int32(0) + while tj < need0_s: + tbv = cutlass.Float32(self.NEG_FLT_MAX) + tbi = cutlass.Int32(-1) + ti = cutlass.Int32(0) + while ti < cnt_strad: + tvi = smem_hist[ti + ti + cutlass.Int32(1)] + if tvi >= cutlass.Int32(0): + tvb = smem_hist[ti + ti] + tvv = cutlass.Float32( + llvm.bitcast( + cutlass.Float32.mlir_type, + tvb.ir_value(), + ) + ) + take = cutlass.Int32(0) + if tbi < cutlass.Int32(0): + take = cutlass.Int32(1) + elif tvv > tbv: + take = cutlass.Int32(1) + if take == cutlass.Int32(1): + tbv = tvv + tbi = ti + ti = ti + cutlass.Int32(1) + pos_s = rank_above_fine + tj + if cutlass.const_expr(self.return_output_values): + output_values_row[pos_s] = self.dtype(tbv) + output_indices_row[pos_s] = smem_hist[ + tbi + tbi + cutlass.Int32(1) + ] + smem_hist[tbi + tbi + cutlass.Int32(1)] = cutlass.Int32(-1) + tj = tj + cutlass.Int32(1) + cute.arch.barrier() + else: + # Persistent scalars live above the 256 digit bins + # (kNumBins >= 512 always): [256] key prefix (chosen + # digits, remaining bits 0), [257] slots still to fill + # inside the current equal-prefix set, [258] ties + # strictly above the prefix (their slots precede it). + if tidx == cutlass.Int32(0): + smem_hist[256] = cutlass.Int32(0) + smem_hist[257] = need0_s + smem_hist[258] = cutlass.Int32(0) + cute.arch.barrier() + for lvl in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * lvl) + iz2 = tidx + while iz2 < cutlass.Int32(256): + smem_hist[iz2] = cutlass.Int32(0) + iz2 = iz2 + cutlass.Int32(num_threads) + cute.arch.barrier() + uthr_cur = smem_hist[256] + it2 = tidx + while it2 < cand_count: + vt = smem_keys[it2] + bt = cutlass.Int32((vt - bmin_r) * inv1) + if bt < cutlass.Int32(0): + bt = cutlass.Int32(0) + if bt > cutlass.Int32(kBins - 1): + bt = cutlass.Int32(kBins - 1) + if bt == b_star: + st2 = cutlass.Int32((vt - f_lo) * finv) + if st2 < cutlass.Int32(0): + st2 = cutlass.Int32(0) + if st2 > cutlass.Int32(fbins - 1): + st2 = cutlass.Int32(fbins - 1) + if st2 == sb_star: + uk = f32_order_key(vt) + pmatch = cutlass.Int32(1) + if cutlass.const_expr(lvl > 0): + if (uk >> cutlass.Int32(shift + 8)) != ( + uthr_cur >> cutlass.Int32(shift + 8) + ): + pmatch = cutlass.Int32(0) + if pmatch == cutlass.Int32(1): + dg = ( + uk >> cutlass.Int32(shift) + ) & cutlass.Int32(0xFF) + atomicAdd( + smem_hist.iterator + dg, cutlass.Int32(1) + ) + it2 = it2 + cutlass.Int32(num_threads) + cute.arch.barrier() + # Two-stage descending digit scan (mirrors the + # fine 3-step search): per-warp partial sums, + # thread0 picks the target warp, its lane0 walks + # the warp's digit range. + fdw = cutlass.const_expr(256 // self.num_warps) + wsum2 = cutlass.Int32(0) + for jd in cutlass.range_constexpr(fdw): + dix = ( + cutlass.Int32(255) + - warp_id * cutlass.Int32(fdw) + - cutlass.Int32(jd) + ) + wsum2 = wsum2 + smem_hist[dix] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = wsum2 + cute.arch.barrier() + if tidx == cutlass.Int32(0): + needl = smem_hist[257] + cw = cutlass.Int32(0) + tw3 = cutlass.Int32(num_warps - 1) + f3 = cutlass.Int32(0) + for w4 in cutlass.range_constexpr(self.num_warps): + cw = cw + smem_wcnt[w4] + if cw >= needl and f3 == cutlass.Int32(0): + tw3 = cutlass.Int32(w4) + f3 = cutlass.Int32(1) + pre3 = cutlass.Int32(0) + for w5 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w5) < tw3: + pre3 = pre3 + smem_wcnt[w5] + s_iscalars[4] = pre3 # prefix above target warp + s_iscalars[0] = tw3 # target warp + cute.arch.barrier() + pre4 = s_iscalars[4] + tw4 = s_iscalars[0] + if warp_id == tw4 and lane == cutlass.Int32(0): + needl2 = smem_hist[257] + base4 = pre4 + dstar = cutlass.Int32(0) + above_d = pre4 + sd4 = cutlass.Int32(0) + for jd2 in cutlass.range_constexpr(fdw): + dix2 = ( + cutlass.Int32(255) + - tw4 * cutlass.Int32(fdw) + - cutlass.Int32(jd2) + ) + ra4 = base4 + base4 = base4 + smem_hist[dix2] + if base4 >= needl2 and sd4 == cutlass.Int32(0): + dstar = dix2 + above_d = ra4 + sd4 = cutlass.Int32(1) + smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) + smem_hist[257] = needl2 - above_d + smem_hist[258] = smem_hist[258] + above_d + cute.arch.barrier() + # Rewrite the tie slot range: ties with key > u_thr + # first (there are exactly cnt_ab of them), then the + # first need_eq bitwise-equal-to-u_thr ties in arrival + # order (value-exact by construction). Signed compare + # needs the top bit flipped (unsigned-monotonic key). + u_thr = smem_hist[256] + cnt_ab = smem_hist[258] + need_eq = smem_hist[257] + ks_thr = u_thr ^ cutlass.Int32(-2147483648) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # above-writer ctr + s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr + cute.arch.barrier() + ir2 = tidx + while ir2 < cand_count: + vr = smem_keys[ir2] + br = cutlass.Int32((vr - bmin_r) * inv1) + if br < cutlass.Int32(0): + br = cutlass.Int32(0) + if br > cutlass.Int32(kBins - 1): + br = cutlass.Int32(kBins - 1) + if br == b_star: + sr = cutlass.Int32((vr - f_lo) * finv) + if sr < cutlass.Int32(0): + sr = cutlass.Int32(0) + if sr > cutlass.Int32(fbins - 1): + sr = cutlass.Int32(fbins - 1) + if sr == sb_star: + uk2 = f32_order_key(vr) + ks2 = uk2 ^ cutlass.Int32(-2147483648) + if ks2 > ks_thr: + o2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(4), + cutlass.Int32(1), + ) + pos_s = rank_above_fine + o2 + if pos_s < cutlass.Int32(kK): + if cutlass.const_expr( + self.return_output_values + ): + output_values_row[pos_s] = self.dtype(vr) + output_indices_row[pos_s] = smem_vals[ir2] + elif ks2 == ks_thr: + q2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if q2 < need_eq: + pos_s = rank_above_fine + cnt_ab + q2 + if pos_s < cutlass.Int32(kK): + if cutlass.const_expr( + self.return_output_values + ): + output_values_row[pos_s] = self.dtype( + vr + ) + output_indices_row[pos_s] = smem_vals[ir2] + ir2 = ir2 + cutlass.Int32(num_threads) + cute.arch.barrier() + elif cutlass.const_expr(self.p4_exact_tail): # if->elif only need0 = cutlass.Int32(kK) - rank_above_fine if cnt_strad > need0 and need0 > cutlass.Int32(0): - self._p4_exact_tail_radix_select( - kK, - kBins, - num_threads, - num_warps, - need0, - cand_count, - rank_above_fine, - b_star, - sb_star, - bmin_r, - f_lo, - finv, - fbins, - inv1, - tidx, - lane, - warp_id, - smem_hist, - smem_keys, - smem_vals, - smem_wcnt, - s_iscalars, - output_indices_row, - output_values_row, - ) + # Persistent scalars live above the 256 digit bins + # (kNumBins >= 512 always): [256] key prefix (chosen + # digits, remaining bits 0), [257] slots still to fill + # inside the current equal-prefix set, [258] ties + # strictly above the prefix (their slots precede it). + if tidx == cutlass.Int32(0): + smem_hist[256] = cutlass.Int32(0) + smem_hist[257] = need0 + smem_hist[258] = cutlass.Int32(0) + cute.arch.barrier() + for lvl in cutlass.range_constexpr(4): + shift = cutlass.const_expr(24 - 8 * lvl) + iz2 = tidx + while iz2 < cutlass.Int32(256): + smem_hist[iz2] = cutlass.Int32(0) + iz2 = iz2 + cutlass.Int32(num_threads) + cute.arch.barrier() + uthr_cur = smem_hist[256] + it2 = tidx + while it2 < cand_count: + vt = smem_keys[it2] + bt = cutlass.Int32((vt - bmin_r) * inv1) + if bt < cutlass.Int32(0): + bt = cutlass.Int32(0) + if bt > cutlass.Int32(kBins - 1): + bt = cutlass.Int32(kBins - 1) + if bt == b_star: + st2 = cutlass.Int32((vt - f_lo) * finv) + if st2 < cutlass.Int32(0): + st2 = cutlass.Int32(0) + if st2 > cutlass.Int32(fbins - 1): + st2 = cutlass.Int32(fbins - 1) + if st2 == sb_star: + uk = f32_order_key(vt) + pmatch = cutlass.Int32(1) + if cutlass.const_expr(lvl > 0): + if (uk >> cutlass.Int32(shift + 8)) != ( + uthr_cur >> cutlass.Int32(shift + 8) + ): + pmatch = cutlass.Int32(0) + if pmatch == cutlass.Int32(1): + dg = (uk >> cutlass.Int32(shift)) & cutlass.Int32(0xFF) + atomicAdd(smem_hist.iterator + dg, cutlass.Int32(1)) + it2 = it2 + cutlass.Int32(num_threads) + cute.arch.barrier() + # Two-stage descending digit scan (mirrors the + # fine 3-step search): per-warp partial sums, + # thread0 picks the target warp, its lane0 walks + # the warp's digit range — 2*num_warps serial + # steps instead of 256. + fdw = cutlass.const_expr(256 // self.num_warps) + wsum2 = cutlass.Int32(0) + for jd in cutlass.range_constexpr(fdw): + dix = ( + cutlass.Int32(255) + - warp_id * cutlass.Int32(fdw) + - cutlass.Int32(jd) + ) + wsum2 = wsum2 + smem_hist[dix] + if lane == cutlass.Int32(0): + smem_wcnt[warp_id] = wsum2 + cute.arch.barrier() + if tidx == cutlass.Int32(0): + needl = smem_hist[257] + cw = cutlass.Int32(0) + tw3 = cutlass.Int32(num_warps - 1) + f3 = cutlass.Int32(0) + for w4 in cutlass.range_constexpr(self.num_warps): + cw = cw + smem_wcnt[w4] + if cw >= needl and f3 == cutlass.Int32(0): + tw3 = cutlass.Int32(w4) + f3 = cutlass.Int32(1) + pre3 = cutlass.Int32(0) + for w5 in cutlass.range_constexpr(self.num_warps): + if cutlass.Int32(w5) < tw3: + pre3 = pre3 + smem_wcnt[w5] + s_iscalars[4] = pre3 # prefix above target warp + s_iscalars[0] = tw3 # target warp + cute.arch.barrier() + pre4 = s_iscalars[4] + tw4 = s_iscalars[0] + if warp_id == tw4 and lane == cutlass.Int32(0): + needl2 = smem_hist[257] + base4 = pre4 + dstar = cutlass.Int32(0) + above_d = pre4 + sd4 = cutlass.Int32(0) + for jd2 in cutlass.range_constexpr(fdw): + dix2 = ( + cutlass.Int32(255) + - tw4 * cutlass.Int32(fdw) + - cutlass.Int32(jd2) + ) + ra4 = base4 + base4 = base4 + smem_hist[dix2] + if base4 >= needl2 and sd4 == cutlass.Int32(0): + dstar = dix2 + above_d = ra4 + sd4 = cutlass.Int32(1) + smem_hist[256] = uthr_cur | (dstar << cutlass.Int32(shift)) + smem_hist[257] = needl2 - above_d + smem_hist[258] = smem_hist[258] + above_d + cute.arch.barrier() + # Rewrite the tie slot range: ties with key > u_thr + # first (there are exactly cnt_ab of them), then the + # first need_eq bitwise-equal-to-u_thr ties in arrival + # order (value-exact by construction). Signed compare + # needs the top bit flipped (unsigned-monotonic key). + u_thr = smem_hist[256] + cnt_ab = smem_hist[258] + need_eq = smem_hist[257] + ks_thr = u_thr ^ cutlass.Int32(-2147483648) + if tidx == cutlass.Int32(0): + s_iscalars[4] = cutlass.Int32(0) # above-writer ctr + s_iscalars[0] = cutlass.Int32(0) # equal-writer ctr + cute.arch.barrier() + ir2 = tidx + while ir2 < cand_count: + vr = smem_keys[ir2] + br = cutlass.Int32((vr - bmin_r) * inv1) + if br < cutlass.Int32(0): + br = cutlass.Int32(0) + if br > cutlass.Int32(kBins - 1): + br = cutlass.Int32(kBins - 1) + if br == b_star: + sr = cutlass.Int32((vr - f_lo) * finv) + if sr < cutlass.Int32(0): + sr = cutlass.Int32(0) + if sr > cutlass.Int32(fbins - 1): + sr = cutlass.Int32(fbins - 1) + if sr == sb_star: + uk2 = f32_order_key(vr) + ks2 = uk2 ^ cutlass.Int32(-2147483648) + if ks2 > ks_thr: + o2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(4), + cutlass.Int32(1), + ) + pos = rank_above_fine + o2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + elif ks2 == ks_thr: + q2 = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if q2 < need_eq: + pos = rank_above_fine + cnt_ab + q2 + if pos < cutlass.Int32(kK): + if cutlass.const_expr(self.return_output_values): + output_values_row[pos] = self.dtype(vr) + output_indices_row[pos] = smem_vals[ir2] + ir2 = ir2 + cutlass.Int32(num_threads) + cute.arch.barrier() else: # ---- APPROX rank-and-scatter (single pass), arbitrary straddling order ---- isc = tidx @@ -3370,6 +5126,17 @@ def phase4_rank_scatter( output_values_row[ipad] = self.dtype(self.NEG_FLT_MAX) output_indices_row[ipad] = cutlass.Int32(-1) ipad = ipad + cutlass.Int32(num_threads) + if cutlass.const_expr(_P4_SUB_DBG): + # smem_wcnt slots [8..13] are dead after the last warp-sum + # use above; the take-block publish copies them to xstate. + sc6 = cute.arch.clock64() + if tidx == cutlass.Int32(0): + smem_wcnt[8] = cutlass.Int32(sc1 - sc0) # minmax + smem_wcnt[9] = cutlass.Int32(sc2 - sc1) # hist build + smem_wcnt[10] = cutlass.Int32(sc3 - sc2) # coarse search + smem_wcnt[11] = cutlass.Int32(sc4 - sc3) # fine recursion + smem_wcnt[12] = cutlass.Int32(sc5 - sc4) # scatter + smem_wcnt[13] = cutlass.Int32(sc6 - sc5) # tail repair+pad else: i10 = tidx while i10 < cand_count: @@ -3910,6 +5677,13 @@ def gvr_topk_kernel( output_values: cute.Tensor, # [numRows, top_k] dtype output_indices: cute.Tensor, # [numRows, top_k] int32 order_row: cute.Tensor, # [batch_size] int32 (or None when seqlen_sorted=False) + block_max: cute.Tensor, # [numRows, nb_pad*4] fp32 (or None: no block-skip) + seed_thr: cute.Tensor, # [numRows, 3] fp32 (or None: no ext counts) + seed_counts: cute.Tensor, # [numRows, 3] int32 (or None) + xstate: cute.Tensor, # [numRows, 8] fp32 closed-loop state (or None) + cand_vals: cute.Tensor, # [numRows, CAP] fp32 scores (or None) + cand_idx: cute.Tensor, # [numRows, CAP] int32 positions (or None) + cand_ctl: cute.Tensor, # [numRows, 4] int32 {n0, void, n1, n2} (or None) ): """Thin entry: bidx → row_idx → run_one_row. @@ -3961,6 +5735,13 @@ def gvr_topk_kernel( seq_lens, output_values, output_indices, + block_max=block_max, + seed_thr=seed_thr, + seed_counts=seed_counts, + xstate=xstate, + cand_vals=cand_vals, + cand_idx=cand_idx, + cand_ctl=cand_ctl, ) @cute.jit @@ -3972,6 +5753,13 @@ def run_one_row( seq_lens: cute.Tensor, # [numRows / next_n] int32 output_values: cute.Tensor, # [numRows, top_k] dtype, optional output_indices: cute.Tensor, # [numRows, top_k] int32 + block_max: cute.Tensor = None, # [numRows, nb_pad*4] fp32 + seed_thr: cute.Tensor = None, # [numRows, 3] fp32 (ext counts) + seed_counts: cute.Tensor = None, # [numRows, 3] int32 (ext counts) + xstate: cute.Tensor = None, # [numRows, 8] fp32 (emit_xstate) + cand_vals: cute.Tensor = None, # [numRows, CAP] fp32 (ext cand) + cand_idx: cute.Tensor = None, # [numRows, CAP] int32 (ext cand) + cand_ctl: cute.Tensor = None, # [numRows, 4] int32 (ext cand) ): """Dispatch: compute per-row slice + cluster sync mode, call _run_phases. @@ -4028,6 +5816,56 @@ def run_one_row( # Slice per-row views. input_row = input_data[row_idx, None] pre_idx_row = pre_idx[pre_idx_row_idx, None] + # trace-time contract checks: a feature flag without its tensor + # must fail HERE, not as a NoneType subscript deep in the phases + if cutlass.const_expr((self.use_ext_counts or self.ext_rungs) and seed_thr is None): + raise ValueError("use_ext_counts/ext_rungs kernels require seed_thr") + if cutlass.const_expr(self.emit_xstate and xstate is None): + raise ValueError("emit_xstate kernels require xstate") + if cutlass.const_expr( + self.use_ext_cand and (cand_vals is None or cand_idx is None or cand_ctl is None) + ): + raise ValueError("use_ext_cand kernels require cand_vals/cand_idx/cand_ctl") + if cutlass.const_expr(self.enable_block_skip and block_max is None): + raise ValueError("enable_block_skip kernels require block_max") + if cutlass.const_expr(self.enable_block_skip and block_max is not None): + block_max_row = block_max[row_idx, None] + else: + block_max_row = None + if cutlass.const_expr(self.use_ext_counts and seed_thr is not None): + # packed seed row [>=6] fp32: [0..2] lines, [3..5] counts as + # floats (exact to 2^24) - ONE 32B sector serves both + seed_thr_row = seed_thr[row_idx, None] + seed_counts_row = None + elif cutlass.const_expr(self.ext_rungs and seed_thr is not None): + seed_thr_row = seed_thr[row_idx, None] + seed_counts_row = None + else: + seed_thr_row = None + seed_counts_row = None + if cutlass.const_expr(self.emit_xstate and xstate is not None): + xstate_row = xstate[row_idx, None] + else: + xstate_row = None + if cutlass.const_expr( + self.use_ext_cand + and cand_vals is not None + and cand_idx is not None + and cand_ctl is not None + ): + cand_vals_row = cand_vals[row_idx, None] + cand_idx_row = cand_idx[row_idx, None] + cand_ctl_row = cand_ctl[row_idx, None] + elif cutlass.const_expr(self.self_scan and cand_idx is not None): + # self_scan: only the POSITION column exists (values live in + # smem from birth; counts come from the phase-0 cursors) + cand_vals_row = None + cand_idx_row = cand_idx[row_idx, None] + cand_ctl_row = None + else: + cand_vals_row = None + cand_idx_row = None + cand_ctl_row = None # When return_output_values=False, ``output_values`` is None at # launch and the gated writes below are compiled out; slicing into # None would crash so we keep the view None as well. @@ -4038,23 +5876,50 @@ def run_one_row( output_indices_row = output_indices[row_idx, None] pre_idx_count = pre_idx.shape[1] - griddepcontrol_wait() + if cutlass.const_expr(not self.pdl_wait_late): + griddepcontrol_wait() # ---- Shared memory allocation ---- smem = SmemAllocator() # keys[kC] fp32 (P3 candidate values; smem keys always fp32 even for half-prec) # Use fp32 even for half-prec to make secant search algorithm keep the accuracy/precision and converge faster. + # self_scan: enlarged to seg_total (three value segments at bases + # 0 / accept_cap / 2*accept_cap); every later consumer only ever + # touches a <= kC prefix after cut compaction. smem_keys = smem.allocate_tensor( element_type=cutlass.Float32, - layout=cute.make_ordered_layout((kC,), order=(0,)), - byte_alignment=128, - ) - # vals[kC] int32 (P3 candidate indices) - smem_vals = smem.allocate_tensor( - element_type=cutlass.Int32, - layout=cute.make_ordered_layout((kC,), order=(0,)), + layout=cute.make_ordered_layout((cutlass.const_expr(self.seg_total),), order=(0,)), byte_alignment=128, ) + # vals[kC] int32 (P3 candidate indices). self_scan: holds the + # SEGMENT COORDINATE of each compacted candidate (identity for the + # deferred position gather via cand_idx[coord]) — every consumer + # (P4, tail repair, gather) works unchanged. + if cutlass.const_expr(self.self_scan): + # phase-0 cp.async staging, ALIASED over vals: vals is only + # written after phase 0 completes and every in-flight group + # is drained inside the dense loop, so the lifetimes never + # overlap. Slot-major (slot s of thread t at row + # s*num_threads + t): a warp's 16B reads/writes land on + # consecutive banks, conflict-free. + smem_stage = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=cute.make_ordered_layout( + (cutlass.const_expr(self.stage_rows), 4), order=(1, 0) + ), + byte_alignment=128, + ) + smem_vals = cute.make_tensor( + cute.recast_ptr(smem_stage.iterator, dtype=cutlass.Int32), + cute.make_ordered_layout((kC,), order=(0,)), + ) + else: + smem_stage = None + smem_vals = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((kC,), order=(0,)), + byte_alignment=128, + ) # histogram[kNumBins] int32 (P4 only) smem_hist = smem.allocate_tensor( element_type=cutlass.Int32, @@ -4067,6 +5932,22 @@ def run_one_row( layout=cute.make_ordered_layout((num_threads,), order=(0,)), byte_alignment=128, ) + # block-skip: int16 active list (16KB at 8192 entries) + control + # ([0] list length, [1] list-current flag for the Phase-3 reuse). + if cutlass.const_expr(self.enable_block_skip): + smem_active = smem.allocate_tensor( + element_type=cutlass.Int16, + layout=cute.make_ordered_layout((self.SKIP_MAX_BLOCKS,), order=(0,)), + byte_alignment=128, + ) + s_active_cnt = smem.allocate_tensor( + element_type=cutlass.Int32, + layout=cute.make_ordered_layout((4,), order=(0,)), + byte_alignment=16, + ) + else: + smem_active = None + s_active_cnt = None # warp_counts[NUM_WARPS] int32 (P3 prefix-sum scratch) # p2_warp_redundant parity-banks the Phase-2 staging (a warp one # round ahead writes the other half) — costs num_warps*4 bytes. @@ -4159,7 +6040,7 @@ def run_one_row( else: smem_input = None - # R0 admission scratch (single-CTA fast path). Allocated only + # op#26 R0 admission scratch (single-CTA fast path). Allocated only # when enable_r0; None otherwise so the base SMEM layout is byte-for- # byte unchanged and these propagate harmlessly through _run_phases' # const_expr(enable_r0)-gated branch (same idiom as s_cluster_partial @@ -4227,6 +6108,14 @@ def run_one_row( s_cluster_partial_m = None smem_gath = None + # PDL wait placed as late as possible so the prologue overlaps + # the producer indexer's tail. INVARIANT: nothing above may read + # producer-written data - seq_lens/pre_idx come from host-side + # metadata and the feedback buffer; logits / block_max / + # seed_thr / cand are first touched below. + if cutlass.const_expr(self.pdl_wait_late): + griddepcontrol_wait() + # ---- Per-row dispatch ---- # Three branches: # 1. Degenerate (N <= top_k): no GVR work, leader emits identity. @@ -4311,6 +6200,16 @@ def run_one_row( tidx, warp_id, lane, + block_max_row=block_max_row, + seed_thr_row=seed_thr_row, + seed_counts_row=seed_counts_row, + xstate_row=xstate_row, + cand_vals_row=cand_vals_row, + cand_idx_row=cand_idx_row, + cand_ctl_row=cand_ctl_row, + smem_active=smem_active, + s_active_cnt=s_active_cnt, + smem_stage=smem_stage, ) else: # Short row: only CTA 0 scans the full row; the other @@ -4352,6 +6251,16 @@ def run_one_row( tidx, warp_id, lane, + block_max_row=block_max_row, + seed_thr_row=seed_thr_row, + seed_counts_row=seed_counts_row, + xstate_row=xstate_row, + cand_vals_row=cand_vals_row, + cand_idx_row=cand_idx_row, + cand_ctl_row=cand_ctl_row, + smem_active=smem_active, + s_active_cnt=s_active_cnt, + smem_stage=smem_stage, ) else: # cs=1: one CTA per row, no cluster sync. @@ -4390,6 +6299,16 @@ def run_one_row( tidx, warp_id, lane, + block_max_row=block_max_row, + seed_thr_row=seed_thr_row, + seed_counts_row=seed_counts_row, + xstate_row=xstate_row, + cand_vals_row=cand_vals_row, + cand_idx_row=cand_idx_row, + cand_ctl_row=cand_ctl_row, + smem_active=smem_active, + s_active_cnt=s_active_cnt, + smem_stage=smem_stage, ) griddepcontrol_launch_dependents() @@ -4431,6 +6350,16 @@ def _run_phases( tidx, warp_id, lane, + block_max_row=None, # block-skip: this row's per-32-position bounds + seed_thr_row=None, # ext counts: this row's 3 seed thresholds (fp32) + seed_counts_row=None, # ext counts: this row's 3 exact counts (int32) + xstate_row=None, # emit_xstate: this row's [8] fp32 state slot + cand_vals_row=None, # ext cand: this row's [CAP] fp32 scores + cand_idx_row=None, # ext cand: this row's [CAP] int32 positions + cand_ctl_row=None, # ext cand: this row's [2] int32 {claimed, void} + smem_active=None, + s_active_cnt=None, + smem_stage=None, # self_scan: [stage_rows, 4] fp32 cp.async staging ): """Run Phase 1-4 + final cluster barrier on a given row slice. @@ -4445,302 +6374,1255 @@ def _run_phases( cluster_size = cutlass.const_expr(self.cluster_size) is_leader = cta_in_cluster == cutlass.Int32(0) - # ---- Phase 1: preIdx Min/Max/Mean ---- - self.phase1_preidx_stats( - input_row, - N, - pre_idx_row, - pre_idx_count, - pre_idx_offset, - smem_wmin, - smem_wmax, - smem_wsum, - smem_wcnt_p1, - s_thr, - s_iscalars, - tidx, - warp_id, - lane, - smem_gath=smem_gath, # p1b_cache: stash gathered values (None-op OFF) - s_mt_thr=s_mt_thr, # r0_vseed: park pmean in the last rung column - ) - - # Degenerate hint (all gathered values identical or out of range): - # reset to a synthetic bracket and fall through instead of emitting - # row[0:K]; the Phase-3 repair guarantees the answer, so the hint only - # affects speed. cnt_hi is seeded with top_k (not 0) so Phase 2's - # budget-collapse guard cannot fire on this unmeasured bracket. - v_lo = s_thr[1] - v_hi = s_thr[2] - # All threads read the bracket before tid0's degenerate-hint rewrite. - cute.arch.barrier() - if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi: - if tidx == 0: - s_thr[0] = cutlass.Float32(0.0) - s_thr[1] = cutlass.Float32(-1.0) - s_thr[2] = cutlass.Float32(1.0) - s_iscalars[2] = N # cnt_lo - s_iscalars[3] = cutlass.Int32(self.top_k) # cnt_hi - cute.arch.barrier() - - # Stage this CTA's slice into SMEM once before Phase 2's - # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses - # scatter-loads OUTSIDE this slice, so it stays on GMEM. - if cutlass.const_expr(self.enable_smem_cache): - self.load_slice_to_smem( - input_row, - slice_start, - slice_end, - smem_input, - tidx, - ) + # block-skip: the list-current flag starts INVALID every row; only + # the R0 compact pass's build sets it. Ordered ahead of all readers + # by Phase 1's internal barriers. + if cutlass.const_expr(self.enable_block_skip): + if tidx == cutlass.Int32(0): + s_active_cnt[1] = cutlass.Int32(0) + s_active_cnt[2] = cutlass.Int32(0) # dropped-rung mask + + # ---- Per-row dynamic routing (ext counts) ---- + # Use the epilogue rungs ONLY when the row is valid (finite t_0, + # xstate contract) AND some rung count already lies in [K, kC]. + # A miss/invalid row runs the full stock path (P1 + P1b + vseed + + # count). All threads read the same control words, so the + # predicate is CTA-uniform and the branches below stay convergent. + if cutlass.const_expr(_P4_TAIL_DBG): + ck0 = cutlass.Int64(0) + ck1 = cutlass.Int64(0) + ckE = cutlass.Int64(0) + ckE = cute.arch.clock64() # row-phase entry (device-residency ref) + ext_row = cutlass.Int32(0) + if cutlass.const_expr(self.use_ext_counts): + # line validity mirrors ext_rungs: ALL THREE lines must be + # finite and strictly ascending (a NaN in t1/t2 must not + # reach the refine brackets); invalid rows fall to the stock + # path, so exactness never rides on the host loop's line + # quality. + if ( + seed_thr_row[0] < cutlass.Float32(1e37) + and seed_thr_row[0] > cutlass.Float32(-1e37) + and seed_thr_row[1] > seed_thr_row[0] + and seed_thr_row[2] > seed_thr_row[1] + ): + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + cm_e = cutlass.Int32(seed_thr_row[3 + m]) + if cm_e >= cutlass.Int32(self.top_k) and cm_e <= cutlass.Int32(self.kC): + ext_row = cutlass.Int32(1) + # list path preview: when the SoA candidate list will be taken + # (count-only admission), Phase 1's gather buys nothing either + # - reuse the same skip (the seed rungs still provide the + # [t_0, t_2] bracket the degenerate check wants). + if cutlass.const_expr( + self.use_ext_cand + and self.use_ext_counts + and cluster_size == 1 + and self.dtype == cutlass.Float32 + ): + claimed_p = cutlass.Int32(cand_ctl_row[0]) + void_p = cutlass.Int32(cand_ctl_row[1]) + # real (non-parked) lines only: the skip stages the raw + # lines into the threshold scratch, and a parked line + # (1e30) would poison every later bracket. Parked rows + # keep Phase 1; the list take below is independent. + if ( + void_p == cutlass.Int32(0) + and claimed_p >= cutlass.Int32(self.top_k + 64) + and claimed_p <= cutlass.Int32(self.list_cap) + and seed_thr_row[0] < cutlass.Float32(1e37) + and seed_thr_row[0] > cutlass.Float32(-1e37) + and seed_thr_row[1] > seed_thr_row[0] + and seed_thr_row[2] > seed_thr_row[1] + and seed_thr_row[2] < cutlass.Float32(1e29) + ): + ext_row = cutlass.Int32(1) + # ---- self_scan phase 0: fused scan-bucket ---- + # The kernel streams the row itself (no external emitter); + # eligibility mirrors the list contract: nothing dropped + # (void == 0) and the loosest line provably covers the top-K + # (n0 >= K, exact counts - no sentinel slack needed). The + # seed-thr finite guard keeps the branch CTA-uniform, so the + # barriers/ballots inside phase 0 stay convergent. Cursor + # scratch = smem_wcnt_p1 (P1 only runs when this row is NOT + # taken, so the reuse never overlaps live data). + if cutlass.const_expr( + self.self_scan and cluster_size == 1 and self.dtype == cutlass.Float32 + ): + if seed_thr_row[0] < cutlass.Float32(1e37): + self.phase0_scan_bucket( + input_row, + N, + seed_thr_row, + smem_keys, + cand_idx_row, + block_max_row, + smem_stage, + smem_wcnt_p1, + tidx, + warp_id, + lane, + ) + if smem_wcnt_p1[3] == cutlass.Int32(0) and smem_wcnt_p1[4] >= cutlass.Int32( + self.top_k + ): + ext_row = cutlass.Int32(1) - # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast - # path) or the secant threshold search ---- - # R0 covers every cluster size: at cs>1 each CTA scans its own - # slice and block_count_ge_multi cluster-merges the rung counts - # (the P1b rungs are per-CTA identical because the preIdx stats are - # full-row). The secant search below is the exact fallback taken - # when the ladder admits nothing, plus the enable_r0=False - # differential-oracle entry. - if cutlass.const_expr(self.enable_r0): - # P1b rung placement -> ONE M-ary R0 count pass -> accept the - # tightest rung with count in [K, kC]. On a miss, fall back to - # the inline log-falsi R1 shot / fb_fix refine. At cs>1 each - # CTA scans its slice and block_count_ge_multi cluster-merges - # the rung counts (phase1b rungs are per-CTA identical since - # preIdx stats are full-row). - if cutlass.const_expr(self.p1b_cache): - # rungs from the SMEM gather-cache P1 stashed (no 2nd - # GMEM gather); 16-bit only. - self.phase1b_hspace_rungs_cached( - pre_idx_count, smem_gath, smem_hist, s_thr, s_mt_thr, tidx, warp_id, lane + # ---- Phase 1: preIdx Min/Max/Mean ---- + # ext counts: P1's only surviving products are the [v_lo, v_hi] + # outer bracket and the scalar state init — the ext rungs provide + # the bracket directly (host contract: t_0 < t_2, finite, all rows + # valid), so the preIdx gather is skipped wholesale. A miss whose + # target lies outside [t_0, t_2] recovers via the refine loop's + # 8x bracket expansion (same fail-soft as the stock path). + rungs_ok = cutlass.Int32(0) + if cutlass.const_expr(self.ext_rungs): + # runtime validity: finite AND strictly ascending; anything + # else (cold start, dropped row, NaN from the host loop) + # falls back to the stock seed path below - exactness never + # rides on the host's line quality + if ( + seed_thr_row[0] < cutlass.Float32(1e37) + and seed_thr_row[0] > cutlass.Float32(-1e37) + and seed_thr_row[1] > seed_thr_row[0] + and seed_thr_row[2] > seed_thr_row[1] + ): + rungs_ok = cutlass.Int32(1) + if cutlass.const_expr(self.ext_rungs): + if rungs_ok == cutlass.Int32(1): + # variant B: the rungs carry the bracket, so P1's gather + # buys nothing - same seed-line init as the ext-counts + # hit path. + if tidx == cutlass.Int32(0): + s_thr[0] = seed_thr_row[1] + s_thr[1] = seed_thr_row[0] + s_thr[2] = seed_thr_row[2] + s_iscalars[0] = cutlass.Int32(0) # cand_count + s_iscalars[1] = cutlass.Int32(0) # done + s_iscalars[2] = cutlass.Int32(-1) # cnt_lo (fb owns) + s_iscalars[3] = cutlass.Int32(-1) # cnt_hi + s_iscalars[4] = cutlass.Int32(0) # out_count + cute.arch.barrier() + if rungs_ok == cutlass.Int32(0): + self.phase1_preidx_stats( + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_wmin, + smem_wmax, + smem_wsum, + smem_wcnt_p1, + s_thr, + s_iscalars, + tidx, + warp_id, + lane, + smem_gath=smem_gath, + s_mt_thr=s_mt_thr, ) - else: - self.phase1b_hspace_rungs( + if cutlass.const_expr(self.use_ext_counts): + if ext_row == cutlass.Int32(1): + if tidx == cutlass.Int32(0): + s_thr[0] = seed_thr_row[1] + s_thr[1] = seed_thr_row[0] + s_thr[2] = seed_thr_row[2] + s_iscalars[0] = cutlass.Int32(0) # cand_count + s_iscalars[1] = cutlass.Int32(0) # done + s_iscalars[2] = cutlass.Int32(-1) # cnt_lo (fb seeding owns) + s_iscalars[3] = cutlass.Int32(-1) # cnt_hi + s_iscalars[4] = cutlass.Int32(0) # out_count + # rung parking folded into the SAME thread0 block + # (was a second thread0 block + barrier in the R0 + # region): tightest in-band count picks the admitted + # line; single-column builds stage it in s_thr[0] and + # pre-mark the rung column (M_qf = accepted, -2 = + # defensive M-ary rerun). + bx_m = cutlass.Int32(-1) + bx_c = cutlass.Int32(2147483647) + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + cx = cutlass.Int32(seed_thr_row[3 + m]) + if ( + cx >= cutlass.Int32(self.top_k) + and cx <= cutlass.Int32(self.kC) + and cx < bx_c + ): + bx_m = cutlass.Int32(m) + bx_c = cx + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + if bx_m >= cutlass.Int32(0): + s_mt_thr[m] = seed_thr_row[bx_m] + else: + s_mt_thr[m] = seed_thr_row[m] + if cutlass.const_expr(not self.enable_block_skip): + if bx_m >= cutlass.Int32(0): + s_thr[0] = seed_thr_row[bx_m] + s_r0col[0] = cutlass.Int32(self.M_qf) + else: + s_r0col[0] = cutlass.Int32(-2) + cute.arch.barrier() + if ext_row == cutlass.Int32(0): + self.phase1_preidx_stats( input_row, N, pre_idx_row, pre_idx_count, pre_idx_offset, - smem_hist, + smem_wmin, + smem_wmax, + smem_wsum, + smem_wcnt_p1, s_thr, - s_mt_thr, + s_iscalars, tidx, warp_id, lane, + smem_gath=smem_gath, # p1b_cache: stash gathered values (None-op OFF) + s_mt_thr=s_mt_thr, # r0_vseed: park pmean in the last rung column ) - self.block_count_ge_multi( + if cutlass.const_expr(not (self.use_ext_counts or self.ext_rungs)): + self.phase1_preidx_stats( input_row, - slice_start, - slice_end, - s_mt_thr, - smem_ptcnt_multi, - smem_wcnt_multi, - s_mt_cnt, - s_cluster_partial_m, - do_cluster_sync, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_wmin, + smem_wmax, + smem_wsum, + smem_wcnt_p1, + s_thr, + s_iscalars, tidx, warp_id, lane, - smem_ptcnt=smem_ptcnt, + smem_gath=smem_gath, # p1b_cache: stash gathered values (None-op OFF) + s_mt_thr=s_mt_thr, # r0_vseed: park pmean in the last rung column ) - cute.arch.barrier() - if tidx == 0: - # tightest admissible rung = SMALLEST count in [K, kC]. - # (Explicit argmin: with r0_vseed the pmean column is not - # sorted into the rung order; for sorted rungs this is - # equivalent to the old "last m in window" rule.) - best_m = cutlass.Int32(-1) - best_c = cutlass.Int32(2147483647) - for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): - cm = s_mt_cnt[m] - if ( - cm >= cutlass.Int32(self.top_k) - and cm <= cutlass.Int32(self.kC) - and cm < best_c - ): - best_m = cutlass.Int32(m) - best_c = cm - s_r0col[0] = best_m - if best_m >= cutlass.Int32(0): - s_thr[0] = s_mt_thr[best_m] - s_iscalars[0] = s_mt_cnt[best_m] - # done=1: the threshold is admitted, so Phase 3 must - # SKIP its retry-shrink and honor s_thr[0]. (block_count - # _ge / secant leave done via their own path; the R0 - # admission must set it explicitly or Phase 3 re-searches - # and the cluster collect diverges -> wrong output.) - s_iscalars[1] = cutlass.Int32(1) - # Snapshot this CTA's LOCAL slice count for the chosen - # rung into s_iscalars[5] — the per-CTA cand_count that - # Phase 3/4's cluster gather consumes (block_count_ge - # sets it too; the R0 admission must match). Without it - # the cluster collect under-counts -> wrong output. - if cutlass.const_expr(cluster_size > 1): - s_iscalars[5] = s_cluster_partial_m[best_m] - cute.arch.barrier() - bc = s_r0col[0] - if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf): - # accepted rung column: copy its cached per-thread counts - # into the secant hand-off buffer (zero rescan). The vseed - # column (bc == M_qf) is ALREADY in smem_ptcnt (v3 reuse). - smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx] - cute.arch.barrier() - # ---- R0 miss: SEEDED bounded log-falsi refine ---- - # At large N the M2D rungs straddle [K, kC]; the refine must - # find a threshold with count in [K, kC] between the measured - # rungs. SEED the loop with the rung bracket AND its known - # counts (clo/chi) so it does log-count regula-falsi from - # iter 0 with no re-measure and no separate R1 shot -> ~2-3 - # count passes instead of ~6. done=1 on - # accept so Phase 3 skips its retry-shrink. - if bc < cutlass.Int32(0): - if cutlass.const_expr(self.fb_fix): + + # Degenerate threshold init: val_hi <= -self.FLT_MAX or val_lo >= val_hi. + # A duplicate/invalid preIdx gather (cold-start zero-init slots, stale + # slots pointing past N, an all-tied gather) produces an unusable + # bracket. When N > K real selection work remains, so rebuild the + # bracket from the data itself (P1r) and run the normal pipeline; + # an identity shortcut here would NOT be the top-K. If the bracket + # is STILL degenerate after the rescue, every in-range value is + # identical (or N <= K), and identity output is then exact — keep + # the shortcut for exactly those rows. + if cutlass.const_expr(self.p1r_rescue): + v_lo = s_thr[1] + v_hi = s_thr[2] + if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi: + if N > cutlass.Int32(self.top_k): + self.phase1r_data_reseed( + input_row, + N, + smem_wmin, + smem_wmax, + s_thr, + s_iscalars, + s_mt_thr, + tidx, + warp_id, + lane, + ) + v_lo = s_thr[1] + v_hi = s_thr[2] + cute.arch.barrier() + if v_hi <= cutlass.Float32(self.NEG_FLT_MAX) or v_lo >= v_hi: + if cutlass.const_expr(cluster_size == 1): + if tidx == 0: + top_k = cutlass.const_expr(self.top_k) + # Emit identity output (first min(top_k, N) indices) + emit_count = cutlass.Int32(top_k) if cutlass.Int32(top_k) < N else N + je = cutlass.Int32(0) + while je < emit_count: + output_indices_row[je] = je + if cutlass.const_expr(self.return_output_values): + output_values_row[je] = input_row[je] + je = je + cutlass.Int32(1) + if cutlass.const_expr(self.emit_xstate): + xstate_row[0] = cutlass.Float32(0.0) # degenerate + else: + # cs>1: all cluster CTAs enter _run_phases; only leader writes. + if is_leader & (tidx == cutlass.Int32(0)): + top_k = cutlass.const_expr(self.top_k) + # Emit identity output (first min(top_k, N) indices) + emit_count = cutlass.Int32(top_k) if cutlass.Int32(top_k) < N else N + je = cutlass.Int32(0) + while je < emit_count: + output_indices_row[je] = je + if cutlass.const_expr(self.return_output_values): + output_values_row[je] = input_row[je] + je = je + cutlass.Int32(1) + if cutlass.const_expr(self.emit_xstate): + xstate_row[0] = cutlass.Float32(0.0) # degenerate + else: + # ---- List path: known-counts admission ---- + # The emitter wrote an SoA list collected at t0 = seed_thr[0] + # and counted the two tighter lines on the way out, so the + # control words carry {n0, void, n1, n2} with n_i = #(>= t_i). + # Admission is then a scalar lookup: + # 1. some n_i in [K, B*] -> cut at the tightest, one pass + # 2. every line straddles/overshoots -> histogram over the + # list clamped between the two known bracket lines + # 3. void, or n0 < K + 64 (emitter sentinel bound) -> fall + # back. The slack also lets accepted cuts load without + # an overflow net: count and load are the same compare. + # cs>1 and 16-bit dtypes keep the plain fallback. + take_cand = cutlass.Int32(0) + # Python bool: with every list/scan feature off, the stock + # Phase 2 + Phase 3 below trace straight-line rather than as + # one 470-line scf.if region with a large yield list. + run_stock_p23 = True + list_used = cutlass.Int32(0) # list path taken (xstate publish) + claimed_c = cutlass.Int32(0) + if cutlass.const_expr( + self.use_ext_cand + and self.use_ext_counts + and cluster_size == 1 + and self.dtype == cutlass.Float32 + ): + # ---- List path: bucketed segments ---- + # The emitter classifies each entry by the tightest line + # it passes and appends into one of three fixed segments + # (A = [0, segA) holds >= t2, B = [segA, 2*segA) holds + # [t1, t2), C = [2*segA, ...) holds [t0, t1)); a full + # segment spills to the next looser one. Segment caps = + # B*, so "line in the acceptance band" <=> "its segment + # prefix group is complete" - the same condition the cut + # selection already checks. A LINE cut therefore loads a + # dense prefix of known length: a pure mapped copy, no + # filtering, no ballots, no atomics, zero wasted reads. + # Histogram fallbacks value-scan the mapped extents. + claimed_c = cutlass.Int32(cand_ctl_row[0]) + void_c = cutlass.Int32(cand_ctl_row[1]) + n1_c = cutlass.Int32(cand_ctl_row[2]) + n2_c = cutlass.Int32(cand_ctl_row[3]) + # segment bases/extents follow the EMITTER geometry + # (accept_cap); the admission bound is additionally + # clamped by the physical candidate capacity kC. + segA = cutlass.const_expr(self.accept_cap) + bstar = cutlass.const_expr(min(self.accept_cap, self.kC)) + if cutlass.const_expr(_P4_TAIL_DBG): + ck0 = cute.arch.clock64() + ck1 = ck0 + # segment extents (pads live at C's tail: sentinel score + # -inf slots, harmless to copy, never rank) + lenA = n2_c + if lenA > cutlass.Int32(segA): + lenA = cutlass.Int32(segA) + spillA = n2_c - lenA + lenB = n1_c - n2_c + spillA + if lenB > cutlass.Int32(segA): + lenB = cutlass.Int32(segA) + lenC = claimed_c - lenA - lenB + total_l = claimed_c + usable = cutlass.Int32(0) + if ( + void_c == cutlass.Int32(0) + and claimed_c >= cutlass.Int32(self.top_k + 64) + and claimed_c <= cutlass.Int32(self.list_cap) + ): + usable = cutlass.Int32(1) + kK_l = cutlass.Int32(self.top_k) + bs_l = cutlass.Int32(bstar) + cut_t = cutlass.Float32(0.0) + cut_n = cutlass.Int32(0) + have = cutlass.Int32(0) + line_cut = cutlass.Int32(0) + anch_t = cutlass.Float32(0.0) + vbase = cutlass.Int64(0) + ibase = cutlass.Int64(0) + if cutlass.const_expr(True): + vbase = cand_vals_row.iterator.toint() + ibase = cand_idx_row.iterator.toint() + if usable == cutlass.Int32(1): + # cut = tightest line in [K, B*]; anchor = loosest. + if n2_c >= kK_l and n2_c <= bs_l: + cut_t = seed_thr_row[2] + cut_n = n2_c + anch_t = seed_thr_row[2] + have = cutlass.Int32(1) + line_cut = cutlass.Int32(1) + if n1_c >= kK_l and n1_c <= bs_l: + if have == cutlass.Int32(0): + cut_t = seed_thr_row[1] + cut_n = n1_c + anch_t = seed_thr_row[1] + have = cutlass.Int32(1) + line_cut = cutlass.Int32(1) + if claimed_c <= bs_l: + if have == cutlass.Int32(0): + cut_t = seed_thr_row[0] + cut_n = claimed_c + anch_t = seed_thr_row[0] + have = cutlass.Int32(1) + line_cut = cutlass.Int32(1) + if have == cutlass.Int32(0): + # ---- clamped-histogram fallback over the mapped + # extents (bracket between two known lines; the + # all-above case takes one max pass first) ---- + # histogram source = the bracket's own SEGMENT + # prefix: if the segment is full it is a value- + # blind (unbiased) SAMPLE of the band - scale the + # targets by band/segment and let the post-load + # count net verify; if not full it IS the exact + # band. Either way the scan shrinks from the + # whole list to <= one segment. + b_lo = seed_thr_row[0] + b_hi = seed_thr_row[1] + base_c = n1_c + hs_base = cutlass.Int32(2 * segA) # segment C + hs_len = lenC + hs_band = claimed_c - n1_c + if n1_c > bs_l: + b_lo = seed_thr_row[1] + b_hi = seed_thr_row[2] + base_c = n2_c + hs_base = cutlass.Int32(segA) # segment B + hs_len = lenB + hs_band = n1_c - n2_c + need_max = cutlass.Int32(0) + if n2_c > bs_l: + b_lo = seed_thr_row[2] + base_c = cutlass.Int32(0) + hs_base = cutlass.Int32(0) # segment A + hs_len = lenA + hs_band = n2_c + need_max = cutlass.Int32(1) + if b_hi >= cutlass.Float32(1e29): + # parked upper line: bracket by the segment max + need_max = cutlass.Int32(1) + if hs_band < cutlass.Int32(1): + hs_band = cutlass.Int32(1) + samp_f = (cutlass.Float32(1.0) * hs_len) / hs_band + if need_max == cutlass.Int32(1): + lmax = cutlass.Float32(self.NEG_FLT_MAX) + i_m = tidx + while i_m < hs_len: + for _ju in cutlass.range_constexpr(4): + j_m = i_m + cutlass.Int32(_ju * num_threads) + if j_m < hs_len: + src_m = hs_base + j_m + vp_m = cute.make_ptr( + cutlass.Float32, + vbase + cutlass.Int64(src_m) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + lmax = cute.arch.fmax( + lmax, cute.make_tensor(vp_m, cute.make_layout((1,)))[0] + ) + i_m = i_m + cutlass.Int32(4 * num_threads) + wmax_l = self.warp_reduce_max_f32(lmax) + if lane == cutlass.Int32(0): + smem_wmax[warp_id] = wmax_l + cute.arch.barrier() + vmax_l = cutlass.Float32(self.NEG_FLT_MAX) + for _wr in cutlass.range_constexpr(self.num_warps): + vmax_l = cute.arch.fmax(vmax_l, smem_wmax[_wr]) + b_hi = vmax_l + cutlass.Float32(1e-3) + NBL = cutlass.const_expr(self.kNumBins) + r_lo = b_lo + r_w = (b_hi - b_lo) / cutlass.Float32(NBL) + if r_w <= cutlass.Float32(0.0): + r_w = cutlass.Float32(1e-6) + # sample-unit targets (population targets scaled + # by segment/band); the descend base stays in + # sample units too - the post-load exact-count + # net absorbs the sampling error. + # fire target = 1.25x the K-need (headroom over + # the sampling noise) + kneedS = cutlass.Int32( + (cutlass.Float32(1.25) * (kK_l - base_c)) * samp_f + + cutlass.Float32(0.5) + ) + if kneedS < cutlass.Int32(1): + kneedS = cutlass.Int32(1) + kfitS = cutlass.Int32((cutlass.Float32(1.0) * (bs_l - base_c)) * samp_f) + sbase = cutlass.Int32(0) + searching = cutlass.Int32(1) + for _rd in cutlass.range_constexpr(3): + if searching == cutlass.Int32(1): + jz_l = tidx + while jz_l < cutlass.Int32(NBL): + smem_hist[jz_l] = cutlass.Int32(0) + jz_l = jz_l + cutlass.Int32(num_threads) + if tidx == cutlass.Int32(0): + s_iscalars[2] = cutlass.Int32(0) + cute.arch.barrier() + inv_wr = cutlass.Float32(1.0) / r_w + r_hi = r_lo + r_w * cutlass.Float32(NBL) + i_h = tidx + while i_h < hs_len: + for _ju in cutlass.range_constexpr(4): + j_h = i_h + cutlass.Int32(_ju * num_threads) + if j_h < hs_len: + src_h = hs_base + j_h + vp_h = cute.make_ptr( + cutlass.Float32, + vbase + cutlass.Int64(src_h) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + v_h = cute.make_tensor(vp_h, cute.make_layout((1,)))[0] + if v_h >= r_lo and v_h < r_hi: + b_h = cutlass.Int32((v_h - r_lo) * inv_wr) + if b_h < cutlass.Int32(0): + b_h = cutlass.Int32(0) + if b_h > cutlass.Int32(NBL - 1): + b_h = cutlass.Int32(NBL - 1) + atomicAdd( + smem_hist.iterator + b_h, cutlass.Int32(1) + ) + i_h = i_h + cutlass.Int32(4 * num_threads) + cute.arch.barrier() + if warp_id == cutlass.Int32(0): + SEGL = cutlass.const_expr(NBL // self.WARP_SIZE) + top_l = cutlass.Int32(NBL - 1) - lane * cutlass.Int32(SEGL) + seg_l = cute.make_rmem_tensor((SEGL,), cutlass.Int32) + part_l = cutlass.Int32(0) + for _js in cutlass.range_constexpr(SEGL): + v8_l = smem_hist[top_l - cutlass.Int32(_js)] + seg_l[_js] = v8_l + part_l = part_l + v8_l + tp_l = part_l + for _os in cutlass.range_constexpr(5): + ov_l = cutlass.const_expr(1 << _os) + oth_l = cute.arch.shuffle_sync_up( + tp_l, ov_l, mask_and_clamp=0 + ) + if lane >= cutlass.Int32(ov_l): + tp_l = tp_l + oth_l + excl_l = tp_l - part_l + kneed = kneedS - sbase + kfit = kfitS - sbase + run_l = cutlass.Int32(0) + for _js in cutlass.range_constexpr(SEGL): + run_l = run_l + seg_l[_js] + cum_at = excl_l + run_l + cum_bef = cum_at - seg_l[_js] + if cum_bef < kneed and cum_at >= kneed: + s_iscalars[3] = top_l - cutlass.Int32(_js) + smem_wcnt[0] = cum_bef + if cum_at <= kfit: + s_iscalars[2] = cutlass.Int32(1) + else: + s_iscalars[2] = cutlass.Int32(2) + cute.arch.barrier() + st_l = s_iscalars[2] + if st_l == cutlass.Int32(1): + cut_t = r_lo + cutlass.Float32(s_iscalars[3]) * r_w + anch_t = cut_t + have = cutlass.Int32(1) + searching = cutlass.Int32(0) + if st_l == cutlass.Int32(2): + sbase = sbase + smem_wcnt[0] + r_lo = r_lo + cutlass.Float32(s_iscalars[3]) * r_w + r_w = r_w / cutlass.Float32(NBL) + if r_w <= cutlass.Float32(0.0): + searching = cutlass.Int32(0) + if st_l == cutlass.Int32(0): + searching = cutlass.Int32(0) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_iscalars[2] = cutlass.Int32(-1) + s_iscalars[3] = cutlass.Int32(-1) + cute.arch.barrier() + if usable == cutlass.Int32(1) and have == cutlass.Int32(1): + take_cand = cutlass.Int32(1) + list_used = cutlass.Int32(1) if tidx == cutlass.Int32(0): - M = cutlass.const_expr(self.M_thr) - blo = v_lo - bhi = v_hi - clo = cutlass.Int32(-1) - chi = cutlass.Int32(-1) - for m in cutlass.range_constexpr(M): - cm = s_mt_cnt[m] - tm = s_mt_thr[m] - if cm > cutlass.Int32(self.kC) and (clo < cutlass.Int32(0) or tm > blo): - blo = tm - clo = cm - if cm < cutlass.Int32(self.top_k) and ( - chi < cutlass.Int32(0) or tm < bhi - ): - bhi = tm - chi = cm - s_thr[1] = blo - s_thr[2] = bhi - s_iscalars[2] = clo # SEED known rung counts - s_iscalars[3] = chi - s_iscalars[1] = cutlass.Int32(0) # done=0 - cand = (blo + bhi) * cutlass.Float32(0.5) - if clo > cutlass.Int32(0) and chi >= cutlass.Int32(0): - chic = chi - if chic < cutlass.Int32(1): - chic = cutlass.Int32(1) - l_lo = cmath.log2(cutlass.Float32(clo), fastmath=True) - l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) - den = l_lo - l_hi - if den > cutlass.Float32(0.0): - t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den - cnd3 = bhi + t3 * (blo - bhi) - if cnd3 > blo and cnd3 < bhi: - cand = cnd3 - elif chi < cutlass.Int32(0): - cand = bhi - elif clo < cutlass.Int32(0): - cand = blo - s_thr[0] = cand + s_iscalars[0] = cutlass.Int32(0) + s_iscalars[2] = cutlass.Int32(0) # non-sentinel count + s_thr[0] = anch_t + s_iscalars[1] = cutlass.Int32(1) # done cute.arch.barrier() - rs = cutlass.Int32(0) - while rs < cutlass.Int32(8) and s_iscalars[1] == cutlass.Int32(0): - if rs > cutlass.Int32(0): + lane_c = tidx & cutlass.Int32(self.WARP_SIZE - 1) + # fused P4 prologue: zero the coarse hist here and + # accumulate the candidate max INSIDE the cut walk + # (min := cut line by construction); the staging + # rides this walk's own end barrier. + izh_c = tidx + while izh_c < cutlass.Int32(self.kNumBins): + smem_hist[izh_c] = cutlass.Int32(0) + izh_c = izh_c + cutlass.Int32(num_threads) + wmax_acc = cutlass.Float32(self.NEG_FLT_MAX) + if line_cut == cutlass.Int32(1): + # ---- LINE cut: dense mapped-prefix COPY of + # exactly cut_n entries. No filter, no ballots, + # no atomics - every read is a winner candidate. + if tidx == cutlass.Int32(0): + s_iscalars[0] = cut_n + nreal_c = cutlass.Int32(0) + i_c = tidx + while i_c < cut_n: + for _ju in cutlass.range_constexpr(4): + j_c = i_c + cutlass.Int32(_ju * num_threads) + if j_c < cut_n: + src_c = j_c + if j_c >= lenA: + src_c = cutlass.Int32(segA) + j_c - lenA + if j_c >= lenA + lenB: + src_c = cutlass.Int32(2 * segA) + j_c - lenA - lenB + vp_c = cute.make_ptr( + cutlass.Float32, + vbase + cutlass.Int64(src_c) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + pv_c = cute.make_tensor(vp_c, cute.make_layout((1,)))[0] + smem_keys[j_c] = pv_c + wmax_acc = cute.arch.fmax(wmax_acc, pv_c) + # eager position fetch: the idx column + # rides the same ILP batch as the value + # read, so vals hold TRUE positions and + # the post-P4 slot swap disappears + ip_c = cute.make_ptr( + cutlass.Int32, + ibase + cutlass.Int64(src_c) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + iv_c = cute.make_tensor(ip_c, cute.make_layout((1,)))[0] + smem_vals[j_c] = iv_c + # sentinel pads carry idx -1; cut_n + # (= claimed n0) counts them, so the + # REAL candidate count must be + # re-measured during the copy + if iv_c >= cutlass.Int32(0): + nreal_c = nreal_c + cutlass.Int32(1) + i_c = i_c + cutlass.Int32(4 * num_threads) + wsum_c = self.warp_reduce_sum_i32(nreal_c) + if lane_c == cutlass.Int32(0): + atomicAdd(s_iscalars.iterator + cutlass.Int32(2), wsum_c) + wmax_w = self.warp_reduce_max_f32(wmax_acc) + if lane_c == cutlass.Int32(0): + smem_wcnt[tidx // cutlass.Int32(32)] = float_as_uint32(wmax_w) + cute.arch.barrier() + # demote when the pad-inflated claim admitted a + # list that holds fewer than K real candidates; + # the stock path below recovers exactly + real_c = s_iscalars[2] + if real_c < cutlass.Int32(self.top_k): + take_cand = cutlass.Int32(0) + list_used = cutlass.Int32(0) if tidx == cutlass.Int32(0): - lo3 = s_thr[1] - hi3 = s_thr[2] - clo3 = s_iscalars[2] - chi3 = s_iscalars[3] - cand = (lo3 + hi3) * cutlass.Float32(0.5) - if chi3 < cutlass.Int32(0): - cand = hi3 - elif clo3 < cutlass.Int32(0): - cand = lo3 - else: - chic = chi3 - if chic < cutlass.Int32(1): - chic = cutlass.Int32(1) - l_lo = cmath.log2(cutlass.Float32(clo3), fastmath=True) - l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) - den3 = l_lo - l_hi - if den3 > cutlass.Float32(0.0): - t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den3 - cnd3 = hi3 + t3 * (lo3 - hi3) - if cnd3 > lo3 and cnd3 < hi3: - cand = cnd3 - s_thr[0] = cand + s_iscalars[1] = cutlass.Int32(0) + cute.arch.barrier() + if line_cut == cutlass.Int32(0): + # ---- histogram-edge cut: value-filtered mapped + # walk with merged-ballot claims (float edges + # round independently -> demote net below). + i_c = tidx + while (i_c - lane_c) < total_l: + pvals = [] + pidxs = [] + keeps = [] + for _ju in cutlass.range_constexpr(4): + j_c = i_c + cutlass.Int32(_ju * num_threads) + pval = cutlass.Float32(self.NEG_FLT_MAX) + pidx = cutlass.Int32(-1) + keep = cutlass.Int32(0) + if j_c < total_l: + src_c = j_c + if j_c >= lenA: + src_c = cutlass.Int32(segA) + j_c - lenA + if j_c >= lenA + lenB: + src_c = cutlass.Int32(2 * segA) + j_c - lenA - lenB + vp_c = cute.make_ptr( + cutlass.Float32, + vbase + cutlass.Int64(src_c) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + pval = cute.make_tensor(vp_c, cute.make_layout((1,)))[0] + # eager position fetch (unconditional: + # keeps the load independent of the + # value compare, same ILP batch) + ip_c = cute.make_ptr( + cutlass.Int32, + ibase + cutlass.Int64(src_c) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + pidx = cute.make_tensor(ip_c, cute.make_layout((1,)))[0] + if pval >= cut_t: + keep = cutlass.Int32(1) + wmax_acc = cute.arch.fmax(wmax_acc, pval) + pvals.append(pval) + pidxs.append(pidx) + keeps.append(keep) + m0 = cute.arch.vote_ballot_sync(keeps[0] != cutlass.Int32(0)) + m1 = cute.arch.vote_ballot_sync(keeps[1] != cutlass.Int32(0)) + m2 = cute.arch.vote_ballot_sync(keeps[2] != cutlass.Int32(0)) + m3 = cute.arch.vote_ballot_sync(keeps[3] != cutlass.Int32(0)) + nk = cutlass.Int32( + cute.arch.popc(m0) + + cute.arch.popc(m1) + + cute.arch.popc(m2) + + cute.arch.popc(m3) + ) + bk = cutlass.Int32(0) + if nk > cutlass.Int32(0): + if lane_c == cutlass.Int32(0): + bk = atomicAdd(s_iscalars.iterator + cutlass.Int32(0), nk) + bk = cute.arch.shuffle_sync(bk, cutlass.Int32(0)) + lmk = ( + cutlass.Uint32(1) << cutlass.Uint32(lane_c) + ) - cutlass.Uint32(1) + off = bk + for _ju in cutlass.range_constexpr(4): + mj = ( + m0 + if _ju == 0 + else m1 + if _ju == 1 + else m2 + if _ju == 2 + else m3 + ) + if keeps[_ju] != cutlass.Int32(0): + wpos = off + cutlass.Int32(cute.arch.popc(mj & lmk)) + if wpos < cutlass.Int32(self.kC): + smem_keys[wpos] = pvals[_ju] + smem_vals[wpos] = pidxs[_ju] + off = off + cutlass.Int32(cute.arch.popc(mj)) + i_c = i_c + cutlass.Int32(4 * num_threads) + wmax_w2 = self.warp_reduce_max_f32(wmax_acc) + if lane_c == cutlass.Int32(0): + smem_wcnt[tidx // cutlass.Int32(32)] = float_as_uint32(wmax_w2) + cute.arch.barrier() + cnt_l = s_iscalars[0] + if cnt_l < cutlass.Int32(self.top_k) or cnt_l > cutlass.Int32(self.kC): + take_cand = cutlass.Int32(0) + list_used = cutlass.Int32(0) + if tidx == cutlass.Int32(0): + s_iscalars[1] = cutlass.Int32(0) + cute.arch.barrier() + if cutlass.const_expr(_P4_TAIL_DBG): + ck1 = cute.arch.clock64() + + # ---- parked count-free take (ext counts, no list): the + # emission already measured the admitted line's exact count, + # so the count pass exists ONLY to build P3's placement + # prefix. Claim-collect instead (v5 edge-cut walk shape: + # 4-way strided reads, ballot-merged claims): ONE cold pass + # replaces count(cold) + collect(L2-hot). P4 is candidate- + # order agnostic (the v5 list path feeds it emission-claim + # order already). The claim total re-measures the count; a + # mismatch with the parked band (stale host state) leaves + # take_cand=0 and the stock single-count path below recovers. + if cutlass.const_expr( + self.use_ext_counts + and not self.use_ext_cand + and not self.self_scan + and not self.enable_block_skip + and cluster_size == 1 + and self.dtype == cutlass.Float32 + ): + if ext_row == cutlass.Int32(1) and N < cutlass.Int32(16384): + # short rows only (< 16k): fused claim-collect; + # longer rows keep the two-pass count-then-place path + cut_p = s_thr[0] # parked line (staged at P1-init) + rbase = input_row.iterator.toint() + vw_p = cutlass.const_expr(self.vec_bits // self.dtype.width) + va_p = cutlass.const_expr(self.vec_align_bytes) + eb_p = cutlass.const_expr(self.dtype.width // 8) + cp_atom = self._make_load_copy_atom() + frag0 = cute.make_rmem_tensor((vw_p,), self.dtype) + frag1 = cute.make_rmem_tensor((vw_p,), self.dtype) + frag2 = cute.make_rmem_tensor((vw_p,), self.dtype) + frag3 = cute.make_rmem_tensor((vw_p,), self.dtype) + step4_p = cutlass.const_expr(4 * num_threads * vw_p) + nfull_p = (N // cutlass.Int32(step4_p)) * cutlass.Int32(step4_p) + it_p = tidx * cutlass.Int32(vw_p) + # 4 chunks in flight per iter (the count primitive's + # ILP shape); per-element DIRECT smem atomics for + # passers - no warp sync, claims hide under the reads + while it_p < nfull_p: + for _jf in cutlass.range_constexpr(4): + gp_p = cute.make_ptr( + self.dtype, + rbase + + cutlass.Int64(it_p + cutlass.Int32(_jf * num_threads * vw_p)) + * cutlass.Int64(eb_p), + cute.AddressSpace.gmem, + assumed_align=va_p, + ) + cute.copy( + cp_atom, + cute.make_tensor(gp_p, cute.make_layout((vw_p,))), + frag0 + if _jf == 0 + else frag1 + if _jf == 1 + else frag2 + if _jf == 2 + else frag3, + ) + for _jf in cutlass.range_constexpr(4): + for _jv in cutlass.range_constexpr(vw_p): + v_p = cutlass.Float32( + ( + frag0 + if _jf == 0 + else frag1 + if _jf == 1 + else frag2 + if _jf == 2 + else frag3 + )[_jv] + ) + if v_p >= cut_p: + sl_p = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if sl_p < cutlass.Int32(self.kC): + smem_keys[sl_p] = v_p + smem_vals[sl_p] = ( + it_p + + cutlass.Int32(_jf * num_threads * vw_p) + + cutlass.Int32(_jv) + ) + it_p = it_p + cutlass.Int32(step4_p) + i_p = nfull_p + tidx + while i_p < N: + vp_p = cute.make_ptr( + cutlass.Float32, + rbase + cutlass.Int64(i_p) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + pval = cute.make_tensor(vp_p, cute.make_layout((1,)))[0] + if pval >= cut_p: + sl_p = atomicAdd( + s_iscalars.iterator + cutlass.Int32(0), + cutlass.Int32(1), + ) + if sl_p < cutlass.Int32(self.kC): + smem_keys[sl_p] = pval + smem_vals[sl_p] = i_p + i_p = i_p + cutlass.Int32(num_threads) + cute.arch.barrier() + cnt_p = s_iscalars[0] + if cnt_p >= cutlass.Int32(self.top_k) and cnt_p <= cutlass.Int32(self.kC): + take_cand = cutlass.Int32(1) + if tidx == cutlass.Int32(0): + s_iscalars[1] = cutlass.Int32(1) # done + cute.arch.barrier() + if take_cand == cutlass.Int32(0): + # stale host counts: reset the claim counter so + # the stock single-count path re-measures cleanly + if tidx == cutlass.Int32(0): + s_iscalars[0] = cutlass.Int32(0) + cute.arch.barrier() + + # ---- self_scan take: cut straight from the phase-0 cursors ---- + # Same admission state machine as the v5 list (tightest line + # whose count fits [K, B*] wins; anchor = loosest in-band + # line), but the candidates already LIVE in smem: a line cut + # is a same-buffer run compaction (sources at >= segA, + # destinations below it - disjoint) plus the segment-coordinate + # fill of smem_vals that the unchanged P4 / tail repair / + # deferred position gather consume. Straddle / overshoot rows + # (no line in band) fall through to the stock fallback (v1). + if cutlass.const_expr( + self.self_scan and cluster_size == 1 and self.dtype == cutlass.Float32 + ): + segA_f = cutlass.const_expr(self.accept_cap) + if cutlass.const_expr(_P4_TAIL_DBG): + ck0 = cute.arch.clock64() + ck1 = ck0 + n0_f = smem_wcnt_p1[4] + n1_f = smem_wcnt_p1[5] + n2_f = smem_wcnt_p1[6] + kK_f = cutlass.Int32(self.top_k) + bs_f = cutlass.Int32(segA_f) + cut_n = cutlass.Int32(0) + have_f = cutlass.Int32(0) + anch_f = cutlass.Float32(0.0) + if ext_row == cutlass.Int32(1): + if n2_f >= kK_f and n2_f <= bs_f: + if have_f == cutlass.Int32(0): + cut_n = n2_f + anch_f = seed_thr_row[2] + have_f = cutlass.Int32(1) + if n1_f >= kK_f and n1_f <= bs_f: + if have_f == cutlass.Int32(0): + cut_n = n1_f + anch_f = seed_thr_row[1] + have_f = cutlass.Int32(1) + if n0_f <= bs_f: + if have_f == cutlass.Int32(0): + cut_n = n0_f + anch_f = seed_thr_row[0] + have_f = cutlass.Int32(1) + if have_f == cutlass.Int32(1): + take_cand = cutlass.Int32(1) + list_used = cutlass.Int32(1) + if tidx == cutlass.Int32(0): + s_iscalars[0] = cut_n + s_thr[0] = anch_f + s_iscalars[1] = cutlass.Int32(1) # done + cute.arch.barrier() + # line cut => no segment spilled (n_cut <= B* bounds + # every tighter count too) => lenA = n2, lenB = n1-n2 + j_f = tidx + while j_f < cut_n: + for _jw in cutlass.range_constexpr(4): + q_f = j_f + cutlass.Int32(_jw * num_threads) + if q_f < cut_n: + src_f = q_f + if q_f >= n2_f: + src_f = cutlass.Int32(segA_f) + q_f - n2_f + if q_f >= n1_f: + src_f = cutlass.Int32(2 * segA_f) + q_f - n1_f + if src_f != q_f: + smem_keys[q_f] = smem_keys[src_f] + smem_vals[q_f] = src_f + j_f = j_f + cutlass.Int32(4 * num_threads) + cute.arch.barrier() + if cutlass.const_expr(_P4_TAIL_DBG): + ck1 = cute.arch.clock64() + + if cutlass.const_expr(self.use_ext_cand or self.use_ext_counts or self.self_scan): + run_stock_p23 = take_cand == cutlass.Int32(0) + if run_stock_p23: + # Stage this CTA's slice into SMEM once before Phase 2's + # 6-10 secant iters re-scan it. Phase 1 (preIdx) uses + # scatter-loads OUTSIDE this slice, so it stays on GMEM. + if cutlass.const_expr(self.enable_smem_cache): + self.load_slice_to_smem( + input_row, + slice_start, + slice_end, + smem_input, + tidx, + ) + + # ---- Phase 2: R0 histogram-ladder admission (single-CTA fast + # path) or the secant threshold search ---- + # enable_r0 gates to cluster_size==1: R0 scans the full row in + # one CTA; cs>1 keeps the secant path. + if cutlass.const_expr(self.enable_r0): + # P1b rung placement -> ONE M-ary R0 count pass -> accept the + # tightest rung with count in [K, kC]. On a miss, fall back to + # the inline log-falsi R1 shot / fb_fix refine. At cs>1 each + # CTA scans its slice and block_count_ge_multi cluster-merges + # the rung counts (phase1b rungs are per-CTA identical since + # preIdx stats are full-row). + if cutlass.const_expr(self.use_ext_counts): + if ext_row == cutlass.Int32(1): + # ---- Waterfall L1 admission (ext rungs) ---- + # Rung thresholds arrive from the indexer epilogue, + # so only P1b is skipped: the stock M-ary count pass + # runs on the ext rungs and the block-skip list + # build, rung tightening, hand-off and classify all + # compose unchanged. When an ext count already lies + # in [K, kC] the admitted threshold is parked in all + # rung slots, degenerating the pass to one compact + # single-threshold count; a full miss keeps the + # three distinct rungs as brackets for the refine. + # Parking and staging happen in the P1-init thread0 + # block, one barrier for the whole prologue. + if cutlass.const_expr(not self.enable_block_skip): + # parked admission: count the parked threshold + # ONCE with the refine primitive - same + # per-thread ptcnt cache and cluster merge P3 + # consumes - and accept in place. + if s_r0col[0] == cutlass.Int32(self.M_qf): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + cpar = s_iscalars[0] + if cpar >= cutlass.Int32( + self.top_k + ) and cpar <= cutlass.Int32(self.kC): + s_iscalars[1] = cutlass.Int32(1) + else: + # emission counts disagree with + # the measured count (stale + # host state): rerun the full + # M-ary machinery on the three + # distinct seed lines. + s_r0col[0] = cutlass.Int32(-2) + for m in cutlass.range_constexpr( + cutlass.const_expr(self.M_thr) + ): + s_mt_thr[m] = seed_thr_row[m] + cute.arch.barrier() + if ext_row == cutlass.Int32(0): + if cutlass.const_expr(self.p1b_cache): + # rungs from the SMEM gather-cache P1 stashed (no 2nd + # GMEM gather); 16-bit only. + self.phase1b_hspace_rungs_cached( + pre_idx_count, + smem_gath, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + else: + self.phase1b_hspace_rungs( + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + if cutlass.const_expr(self.ext_rungs): + # variant B: rung thresholds = the closed-loop seed + # lines verbatim; the stock multi-count measures + # them and the argmin admission below picks the + # tightest one in [K, kC]. Invalid lines fall back + # to the stock P1b quantile rungs (P1 stats ran on + # this row in that case). + if rungs_ok == cutlass.Int32(1): + if tidx == cutlass.Int32(0): + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + s_mt_thr[m] = seed_thr_row[m] cute.arch.barrier() - self.block_count_ge( + if rungs_ok == cutlass.Int32(0): + self.phase1b_hspace_rungs( + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + if cutlass.const_expr(not (self.use_ext_counts or self.ext_rungs)): + if cutlass.const_expr(self.p1b_cache): + # rungs from the SMEM gather-cache P1 stashed (no 2nd + # GMEM gather); 16-bit only. + self.phase1b_hspace_rungs_cached( + pre_idx_count, + smem_gath, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + else: + self.phase1b_hspace_rungs( + input_row, + N, + pre_idx_row, + pre_idx_count, + pre_idx_offset, + smem_hist, + s_thr, + s_mt_thr, + tidx, + warp_id, + lane, + ) + r0_par = cutlass.Int32(0) + run_mary = True # Python bool: no scf.if without ext counts + if cutlass.const_expr(self.use_ext_counts and not self.enable_block_skip): + # single-column fast path accepted: the parked count + # is done and admitted; the M-ary pass, argmin, + # handoff and miss machinery all stand down + # (s_r0col == M_qf skips the copy and the refine). + if s_r0col[0] == cutlass.Int32(self.M_qf) and s_iscalars[ + 1 + ] == cutlass.Int32(1): + r0_par = cutlass.Int32(1) + run_mary = r0_par == cutlass.Int32(0) + if run_mary: + self.block_count_ge_multi( input_row, slice_start, slice_end, - s_thr[0], - smem_ptcnt, - smem_wcnt, - s_iscalars, - s_cluster_partial, + s_mt_thr, + smem_ptcnt_multi, + smem_wcnt_multi, + s_mt_cnt, + s_cluster_partial_m, + do_cluster_sync, tidx, warp_id, lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, + smem_ptcnt=smem_ptcnt, + block_max_row=block_max_row, + smem_active=smem_active, + s_active_cnt=s_active_cnt, ) - cute.arch.barrier() - if tidx == cutlass.Int32(0): - c3 = s_iscalars[0] - t3v = s_thr[0] - if c3 >= cutlass.Int32(self.top_k) and c3 <= cutlass.Int32(self.kC): - s_iscalars[1] = cutlass.Int32(1) # accept - elif c3 > cutlass.Int32(self.kC): - s_thr[1] = t3v - s_iscalars[2] = c3 - if t3v >= s_thr[2]: - rng3 = s_thr[2] - s_thr[1] - if rng3 < cutlass.Float32(1.0): - rng3 = cutlass.Float32(1.0) - s_thr[2] = s_thr[2] + rng3 * cutlass.Float32(8.0) - s_iscalars[3] = cutlass.Int32(-1) - else: - s_thr[2] = t3v - s_iscalars[3] = c3 - if t3v <= s_thr[1]: - rng3 = s_thr[2] - s_thr[1] - if rng3 < cutlass.Float32(1.0): - rng3 = cutlass.Float32(1.0) - s_thr[1] = s_thr[1] - rng3 * cutlass.Float32(8.0) - s_iscalars[2] = cutlass.Int32(-1) - cute.arch.barrier() - rs = rs + cutlass.Int32(1) - if s_iscalars[1] != cutlass.Int32(1): - # The retry budget could not land in [K, kC]. - # ONLY the coherent undershoot-overflow corner - # (count(>= lo) > kC AND 0 <= count(>= hi) < K, - # both counts CURRENT — the retry's bracket - # widening marks a side stale with -1 and thus - # fails this guard) collapses the bracket by - # pure bisection to ADJACENT floats, where the - # plateau terminal (done = 3, threshold = hi) - # is exact: Phase 4 emits the sure winners and - # the plateau fill completes the row from the - # tie class. A mid-collapse count landing in - # [K, kC] converges normally; anything else - # (incl. an exhausted collapse budget) falls - # through to the fail-soft terminal below. - it4 = cutlass.Int32(0) - if ( - s_iscalars[2] <= cutlass.Int32(self.kC) - or s_iscalars[3] < cutlass.Int32(0) - or s_iscalars[3] >= cutlass.Int32(self.top_k) - ): - it4 = cutlass.Int32(40) # guard: skip collapse - while it4 < cutlass.Int32(40) and s_iscalars[1] == cutlass.Int32(0): - cute.arch.barrier() + cute.arch.barrier() + if run_mary and tidx == 0: + # tightest admissible rung = SMALLEST count in [K, kC] + # (explicit argmin: with r0_vseed the pmean column is not + # sorted into the rung order). Dropped rungs (block-skip + # rung tightening) hold PARTIAL counts — never admissible. + dmask_c = cutlass.Int32(0) + if cutlass.const_expr(self.enable_block_skip): + dmask_c = s_active_cnt[2] + best_m = cutlass.Int32(-1) + best_c = cutlass.Int32(2147483647) + for m in cutlass.range_constexpr(cutlass.const_expr(self.M_thr)): + cm = s_mt_cnt[m] + if ( + cm >= cutlass.Int32(self.top_k) + and cm <= cutlass.Int32(self.kC) + and cm < best_c + and (dmask_c & (cutlass.Int32(1) << cutlass.Int32(m))) + == cutlass.Int32(0) + ): + best_m = cutlass.Int32(m) + best_c = cm + s_r0col[0] = best_m + if best_m >= cutlass.Int32(0): + s_thr[0] = s_mt_thr[best_m] + s_iscalars[0] = s_mt_cnt[best_m] + # done=1: the threshold is admitted, so Phase 3 must + # SKIP its retry-shrink and honor s_thr[0]. (block_count + # _ge / secant leave done via their own path; the R0 + # admission must set it explicitly or Phase 3 re-searches + # and the cluster collect diverges -> wrong output.) + s_iscalars[1] = cutlass.Int32(1) + # Snapshot this CTA's LOCAL slice count for the chosen + # rung into s_iscalars[5] — the per-CTA cand_count that + # Phase 3/4's cluster gather consumes (block_count_ge + # sets it too; the R0 admission must match). Without it + # the cluster collect under-counts -> wrong output. + if cutlass.const_expr(cluster_size > 1): + s_iscalars[5] = s_cluster_partial_m[best_m] + + cute.arch.barrier() + bc = s_r0col[0] + if bc >= cutlass.Int32(0) and bc < cutlass.Int32(self.M_qf): + # accepted rung column: copy its cached per-thread counts + # into the secant hand-off buffer (zero rescan). The vseed + # column (bc == M_qf) is ALREADY in smem_ptcnt (v3 reuse). + smem_ptcnt[tidx] = smem_ptcnt_multi[bc * cutlass.Int32(num_threads) + tidx] + cute.arch.barrier() + # ---- R0 miss: SEEDED bounded log-falsi refine ---- + # The refine must find a threshold with count in [K, kC] + # between the measured rungs. SEED the loop with the rung + # bracket AND its known counts (clo/chi) so it does + # log-count regula-falsi from iter 0 with no re-measure. + # done=1 on accept so Phase 3 skips its retry-shrink. + if bc < cutlass.Int32(0): + if cutlass.const_expr(self.enable_block_skip): + if tidx == cutlass.Int32(0): + s_active_cnt[1] = cutlass.Int32(0) + if cutlass.const_expr(self.fb_fix): if tidx == cutlass.Int32(0): - lo4 = s_thr[1] - hi4 = s_thr[2] - mid4 = (lo4 + hi4) * cutlass.Float32(0.5) - if mid4 == lo4 or mid4 == hi4: - s_thr[0] = hi4 - s_iscalars[1] = cutlass.Int32(3) - else: - s_thr[0] = mid4 + M = cutlass.const_expr(self.M_thr) + blo = v_lo + bhi = v_hi + clo = cutlass.Int32(-1) + chi = cutlass.Int32(-1) + dmask_f = cutlass.Int32(0) + if cutlass.const_expr(self.enable_block_skip): + dmask_f = s_active_cnt[2] + for m in cutlass.range_constexpr(M): + cm = s_mt_cnt[m] + tm = s_mt_thr[m] + m_ok = ( + dmask_f & (cutlass.Int32(1) << cutlass.Int32(m)) + ) == cutlass.Int32(0) + if ( + m_ok + and cm > cutlass.Int32(self.kC) + and (clo < cutlass.Int32(0) or tm > blo) + ): + blo = tm + clo = cm + if ( + m_ok + and cm < cutlass.Int32(self.top_k) + and (chi < cutlass.Int32(0) or tm < bhi) + ): + bhi = tm + chi = cm + s_thr[1] = blo + s_thr[2] = bhi + s_iscalars[2] = clo # SEED known rung counts + s_iscalars[3] = chi + s_iscalars[1] = cutlass.Int32(0) # done=0 + cand = (blo + bhi) * cutlass.Float32(0.5) + if clo > cutlass.Int32(0) and chi >= cutlass.Int32(0): + chic = chi + if chic < cutlass.Int32(1): + chic = cutlass.Int32(1) + l_lo = cmath.log2(cutlass.Float32(clo), fastmath=True) + l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) + den = l_lo - l_hi + if den > cutlass.Float32(0.0): + t3 = (cutlass.Float32(self.log2_mstar) - l_hi) / den + cnd3 = bhi + t3 * (blo - bhi) + if cnd3 > blo and cnd3 < bhi: + cand = cnd3 + elif chi < cutlass.Int32(0): + cand = bhi + elif clo < cutlass.Int32(0): + cand = blo + s_thr[0] = cand cute.arch.barrier() - if s_iscalars[1] == cutlass.Int32(0): + rs = cutlass.Int32(0) + while rs < cutlass.Int32(8) and s_iscalars[1] == cutlass.Int32(0): + if rs > cutlass.Int32(0): + if tidx == cutlass.Int32(0): + lo3 = s_thr[1] + hi3 = s_thr[2] + clo3 = s_iscalars[2] + chi3 = s_iscalars[3] + cand = (lo3 + hi3) * cutlass.Float32(0.5) + if chi3 < cutlass.Int32(0): + cand = hi3 + elif clo3 < cutlass.Int32(0): + cand = lo3 + else: + chic = chi3 + if chic < cutlass.Int32(1): + chic = cutlass.Int32(1) + l_lo = cmath.log2(cutlass.Float32(clo3), fastmath=True) + l_hi = cmath.log2(cutlass.Float32(chic), fastmath=True) + den3 = l_lo - l_hi + if den3 > cutlass.Float32(0.0): + t3 = ( + cutlass.Float32(self.log2_mstar) - l_hi + ) / den3 + cnd3 = hi3 + t3 * (lo3 - hi3) + if cnd3 > lo3 and cnd3 < hi3: + cand = cnd3 + s_thr[0] = cand + cute.arch.barrier() self.block_count_ge( input_row, slice_start, @@ -4758,31 +7640,142 @@ def _run_phases( ) cute.arch.barrier() if tidx == cutlass.Int32(0): - c4 = s_iscalars[0] - t4 = s_thr[0] - if c4 >= cutlass.Int32(self.top_k) and c4 <= cutlass.Int32( + c3 = s_iscalars[0] + t3v = s_thr[0] + if c3 >= cutlass.Int32(self.top_k) and c3 <= cutlass.Int32( self.kC ): - s_iscalars[1] = cutlass.Int32(1) - elif c4 > cutlass.Int32(self.kC): - s_thr[1] = t4 - s_iscalars[2] = c4 + s_iscalars[1] = cutlass.Int32(1) # accept + elif c3 > cutlass.Int32(self.kC): + s_thr[1] = t3v + s_iscalars[2] = c3 + if t3v >= s_thr[2]: + rng3 = s_thr[2] - s_thr[1] + if rng3 < cutlass.Float32(1.0): + rng3 = cutlass.Float32(1.0) + s_thr[2] = s_thr[2] + rng3 * cutlass.Float32(8.0) + s_iscalars[3] = cutlass.Int32(-1) else: - s_thr[2] = t4 - s_iscalars[3] = c4 + s_thr[2] = t3v + s_iscalars[3] = c3 + if t3v <= s_thr[1]: + rng3 = s_thr[2] - s_thr[1] + if rng3 < cutlass.Float32(1.0): + rng3 = cutlass.Float32(1.0) + s_thr[1] = s_thr[1] - rng3 * cutlass.Float32(8.0) + s_iscalars[2] = cutlass.Int32(-1) cute.arch.barrier() - it4 = it4 + cutlass.Int32(1) - if s_iscalars[1] == cutlass.Int32(3): - # recount at the terminal threshold so P3's - # cached per-thread counts describe the - # sure-winner set the fill completes. - self.block_count_ge( + rs = rs + cutlass.Int32(1) + if s_iscalars[1] != cutlass.Int32(1): + # The retry budget could not land in [K, kC]. + # ONLY the coherent undershoot-overflow corner + # (count(>= lo) > kC AND 0 <= count(>= hi) < K, + # both counts CURRENT — the retry's bracket + # widening marks a side stale with -1 and thus + # fails this guard) collapses the bracket by + # pure bisection to ADJACENT floats, where the + # plateau terminal (done = 3, threshold = hi) + # is exact: Phase 4 emits the sure winners and + # the plateau fill completes the row from the + # tie class. A mid-collapse count landing in + # [K, kC] converges normally; anything else + # (incl. an exhausted collapse budget) falls + # through to the fail-soft terminal below. + it4 = cutlass.Int32(0) + if ( + s_iscalars[2] <= cutlass.Int32(self.kC) + or s_iscalars[3] < cutlass.Int32(0) + or s_iscalars[3] >= cutlass.Int32(self.top_k) + ): + it4 = cutlass.Int32(40) # guard: skip collapse + while it4 < cutlass.Int32(40) and s_iscalars[1] == cutlass.Int32(0): + cute.arch.barrier() + if tidx == cutlass.Int32(0): + lo4 = s_thr[1] + hi4 = s_thr[2] + mid4 = (lo4 + hi4) * cutlass.Float32(0.5) + if mid4 == lo4 or mid4 == hi4: + s_thr[0] = hi4 + s_iscalars[1] = cutlass.Int32(3) + else: + s_thr[0] = mid4 + cute.arch.barrier() + if s_iscalars[1] == cutlass.Int32(0): + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + if tidx == cutlass.Int32(0): + c4 = s_iscalars[0] + t4 = s_thr[0] + if c4 >= cutlass.Int32(self.top_k) and c4 <= cutlass.Int32( + self.kC + ): + s_iscalars[1] = cutlass.Int32(1) + elif c4 > cutlass.Int32(self.kC): + s_thr[1] = t4 + s_iscalars[2] = c4 + else: + s_thr[2] = t4 + s_iscalars[3] = c4 + cute.arch.barrier() + it4 = it4 + cutlass.Int32(1) + if s_iscalars[1] == cutlass.Int32(3): + # recount at the terminal threshold so P3's + # cached per-thread counts describe the + # sure-winner set the fill completes. + self.block_count_ge( + input_row, + slice_start, + slice_end, + s_thr[0], + smem_ptcnt, + smem_wcnt, + s_iscalars, + s_cluster_partial, + tidx, + warp_id, + lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + ) + cute.arch.barrier() + elif s_iscalars[1] != cutlass.Int32(1): + # Non-converged terminal on the leader path. + # This used to recount at the undershoot side + # and stamp done = 1, shipping a -1-padded row + # as a documented "non-convergence encoding" - + # which also hid the row from Phase 3, since + # done == 1 never enters the repair. Stamp + # done = 2 and let Phase 3's two-sided + # bisection own it; the recount is dropped + # because that bisection measures anyway. + cute.arch.barrier() + if tidx == cutlass.Int32(0): + s_thr[0] = s_thr[2] + s_iscalars[1] = cutlass.Int32(2) + cute.arch.barrier() + else: + self.phase2_secant_search( input_row, + N, slice_start, slice_end, - s_thr[0], smem_ptcnt, smem_wcnt, + s_thr, s_iscalars, s_cluster_partial, tidx, @@ -4791,18 +7784,6 @@ def _run_phases( do_cluster_sync=do_cluster_sync, smem_input=smem_input, ) - cute.arch.barrier() - elif s_iscalars[1] != cutlass.Int32(1): - # Non-converged terminal. This used to stamp - # done = 1 on an undershooting threshold and - # ship a -1-padded row (e.g. ReLU-sparse rows - # with a 0.0 plateau wider than kC); done = 2 - # routes it into the Phase-3 repair instead. - cute.arch.barrier() - if tidx == cutlass.Int32(0): - s_thr[0] = s_thr[2] - s_iscalars[1] = cutlass.Int32(2) - cute.arch.barrier() else: self.phase2_secant_search( input_row, @@ -4820,183 +7801,57 @@ def _run_phases( do_cluster_sync=do_cluster_sync, smem_input=smem_input, ) - else: - self.phase2_secant_search( - input_row, - N, - slice_start, - slice_end, - smem_ptcnt, - smem_wcnt, - s_thr, - s_iscalars, - s_cluster_partial, - tidx, - warp_id, - lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, - ) - - # Cluster handoff #1 (end of Phase 2). Skipped when - # do_cluster_sync is False (cs=1 or short-row degrade). - if cutlass.const_expr(cluster_size > 1): - if do_cluster_sync: - cute.arch.cluster_arrive_relaxed() - cute.arch.cluster_wait() - - # ---- Phase 3: cluster-parallel candidate collect ---- - self.phase3_collect_candidates( - input_row, - N, - slice_start, - slice_end, - smem_keys, - smem_vals, - smem_ptcnt, - smem_wcnt, - s_thr, - s_iscalars, - s_cluster_partial, - tidx, - warp_id, - lane, - do_cluster_sync=do_cluster_sync, - smem_input=smem_input, - ) - # Cluster handoff #2: leader's DSMEM gather of peer - # smem_keys/smem_vals. Skipped at do_cluster_sync=False. - if cutlass.const_expr(cluster_size > 1): - if do_cluster_sync: - cute.arch.cluster_arrive() - cute.arch.cluster_wait() + # Cluster handoff #1 (end of Phase 2). Skipped when + # do_cluster_sync is False (cs=1 or short-row degrade). + if cutlass.const_expr(cluster_size > 1): + if do_cluster_sync: + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() - # Phase 4 runs on the leader only. const_expr (compile- - # time eliminated) split from runtime so cs=1 gets a flat - # code path with no leader/sync checks. - # Pre-init cand_count_p4 so CuTe DSL sees a stable Int32 type - # across the runtime ``if is_leader:`` branch in cs>1 mode - # (DSL forbids first-assigning a variable inside a dynamic if). - cand_count_p4 = cutlass.Int32(0) - if cutlass.const_expr(cluster_size == 1): - # cs=1: the single CTA per row IS the leader. - # Capture the P2 terminal BEFORE Phase 4: P4 reuses - # s_iscalars[1] as radix scratch. - if tidx == cutlass.Int32(0): - s_iscalars[6] = cutlass.Int32(0) - if s_iscalars[1] == cutlass.Int32(3): - s_iscalars[6] = cutlass.Int32(1) - cute.arch.barrier() - cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) - if cutlass.const_expr(self.enable_p4_rank_scatter): - self.phase4_rank_scatter( - smem_keys, - smem_vals, - smem_hist, - smem_wcnt, - s_thr, - s_iscalars, - output_values_row, - output_indices_row, - cand_count_p4, - tidx, - warp_id, - lane, - ) - else: - self.phase4_histogram_snap( + # ---- Phase 3: cluster-parallel candidate collect ---- + self.phase3_collect_candidates( + input_row, + N, + slice_start, + slice_end, smem_keys, smem_vals, - smem_hist, + smem_ptcnt, smem_wcnt, s_thr, s_iscalars, - output_values_row, - output_indices_row, - cand_count_p4, + s_cluster_partial, tidx, warp_id, lane, + do_cluster_sync=do_cluster_sync, + smem_input=smem_input, + smem_active=smem_active, + s_active_cnt=s_active_cnt, ) - # ---- plateau fill (done == 3): complete the row from the - # bitwise-equal plateau class. The terminal is only set on an - # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) - # is bitwise-equal; Phase 4 has already emitted the - # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of - # the tie class is a valid tie-aware completion. Ticket counter - # lives in the DEDICATED s_iscalars[7]. - if s_iscalars[6] == cutlass.Int32(1): - pv_lo = s_thr[1] - pv_hi = s_thr[0] - if tidx == cutlass.Int32(0): - # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] - # is radix scratch by now (same hazard as the flag). - s_iscalars[7] = cand_count_p4 - cute.arch.barrier() - ifp = tidx - while ifp < N: - vfp = cutlass.Float32(0.0) - if cutlass.const_expr(self.dtype == cutlass.Float32): - vfp = input_row[ifp] - else: - vfp = cutlass.Float32(input_row[ifp]) - if vfp >= pv_lo and vfp < pv_hi: - pfill = atomicAdd(s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1)) - if pfill < cutlass.Int32(self.top_k): - if cutlass.const_expr(self.return_output_values): - output_values_row[pfill] = self.dtype(vfp) - output_indices_row[pfill] = ifp - ifp = ifp + cutlass.Int32(self.num_threads) - cute.arch.barrier() - else: - # cs>1: only the leader (CTA 0 in cluster) runs Phase 4. - if is_leader: - if do_cluster_sync: - # DSMEM-gather peer candidates into the leader's - # smem_keys/smem_vals. Layout: leader's chunk goes - # to [0 .. leader_local_cnt); each peer r's chunk - # appends the next peer_r_local_cnt entries. - local_cnt_self = s_iscalars[5] - local_iscalars_ptr = s_iscalars.iterator + cutlass.Int32(5) - smem_keys_iter = smem_keys.iterator - smem_vals_iter = smem_vals.iterator - base_offset = local_cnt_self - for peer in cutlass.range_constexpr(1, cluster_size): - peer_iscalars_addr = mapa_shared_cluster( - local_iscalars_ptr, cutlass.Int32(peer) - ) - peer_cnt = ld_shared_cluster_i32(peer_iscalars_addr) - # Cap to kC (defense-in-depth vs. the - # done==2 bracket-exhaustion path). - peer_cnt = min(peer_cnt, cutlass.Int32(self.kC)) - i_gather = tidx - while i_gather < peer_cnt: - peer_key_addr = mapa_shared_cluster( - smem_keys_iter + i_gather, cutlass.Int32(peer) - ) - peer_val_addr = mapa_shared_cluster( - smem_vals_iter + i_gather, cutlass.Int32(peer) - ) - k_val = ld_shared_cluster_f32(peer_key_addr) - v_val = ld_shared_cluster_i32(peer_val_addr) - dst = base_offset + i_gather - if dst < cutlass.Int32(self.kC): - smem_keys[dst] = k_val - smem_vals[dst] = v_val - i_gather = i_gather + cutlass.Int32(num_threads) - base_offset = base_offset + peer_cnt - # Reset s_iscalars[0] to cluster-wide cand_count. - if tidx == cutlass.Int32(0): - s_iscalars[0] = base_offset - cute.arch.barrier() - # else: short-row degrade — leader (CTA 0) already - # holds the full row's candidates in its own - # smem_keys/smem_vals (no peers to gather from). - # ---- Phase 4: histogram snap + writeback ---- - # Capture the P2 terminal BEFORE Phase 4: P4 - # reuses s_iscalars[1] as radix scratch. + # Cluster handoff #2: leader's DSMEM gather of peer + # smem_keys/smem_vals. Skipped at do_cluster_sync=False. + if cutlass.const_expr(cluster_size > 1): + if do_cluster_sync: + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # Phase 4 runs on the leader only. const_expr (compile- + # time eliminated) split from runtime so cs=1 gets a flat + # code path with no leader/sync checks. + # Pre-init cand_count_p4 so CuTe DSL sees a stable Int32 type + # across the runtime ``if is_leader:`` branch in cs>1 mode + # (DSL forbids first-assigning a variable inside a dynamic if). + ck2 = cutlass.Int64(0) + if cutlass.const_expr(_P4_TAIL_DBG): + ck2 = cute.arch.clock64() + cand_count_p4 = cutlass.Int32(0) + if cutlass.const_expr(cluster_size == 1): + # cs=1: the single CTA per row IS the leader. + # Capture the P2 terminal BEFORE Phase 4: P4 reuses + # s_iscalars[1] as radix scratch. if tidx == cutlass.Int32(0): s_iscalars[6] = cutlass.Int32(0) if s_iscalars[1] == cutlass.Int32(3): @@ -5004,20 +7859,43 @@ def _run_phases( cute.arch.barrier() cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) if cutlass.const_expr(self.enable_p4_rank_scatter): - self.phase4_rank_scatter( - smem_keys, - smem_vals, - smem_hist, - smem_wcnt, - s_thr, - s_iscalars, - output_values_row, - output_indices_row, - cand_count_p4, - tidx, - warp_id, - lane, - ) + if cutlass.const_expr( + self.use_ext_cand and self.use_ext_counts and self.dtype == cutlass.Float32 + ): + # list rows carry a walk-staged range + pre-zeroed + # hist (flag = list_used; fallback rows take the + # stock minmax path inside) + self.phase4_rank_scatter( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ext_range_flag=list_used, + ext_min=cut_t, + ) + else: + self.phase4_rank_scatter( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) else: self.phase4_histogram_snap( smem_keys, @@ -5033,7 +7911,6 @@ def _run_phases( warp_id, lane, ) - # ---- plateau fill (done == 3): complete the row from the # bitwise-equal plateau class. The terminal is only set on an # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) @@ -5045,8 +7922,9 @@ def _run_phases( pv_lo = s_thr[1] pv_hi = s_thr[0] if tidx == cutlass.Int32(0): - # cand_count_p4 was captured BEFORE Phase 4; s_iscalars[0] - # is radix scratch by now (same hazard as the flag). + # cand_count_p4 was captured BEFORE Phase 4; + # s_iscalars[0] is radix scratch by now (same + # hazard as the flag). s_iscalars[7] = cand_count_p4 cute.arch.barrier() ifp = tidx @@ -5066,6 +7944,256 @@ def _run_phases( output_indices_row[pfill] = ifp ifp = ifp + cutlass.Int32(self.num_threads) cute.arch.barrier() + ck_sw0 = cutlass.Int64(0) + ck_sw1 = cutlass.Int64(0) + if cutlass.const_expr(_P4_SUB_DBG): + ck_sw0 = cute.arch.clock64() + if cutlass.const_expr( + self.self_scan and self.use_ext_counts and self.dtype == cutlass.Float32 + ): + # self_scan rows: the compact stored SEGMENT COORDS in + # the vals slots. Swap them for true positions with K + # fully-parallel gathers. Must precede the xstate + # publish, which reads output slot K-1 as a position. + # (ext_cand list rows translate EAGERLY in the take + # walk - the idx column rides the value ILP batch - + # so they never reach this loop.) + if list_used == cutlass.Int32(1): + io_r = tidx + while io_r < cutlass.Int32(self.top_k): + li_r = output_indices_row[io_r] + # slots are segmented offsets (may exceed the + # entry count); only sentinel -1 is invalid + if li_r >= cutlass.Int32(0): + ip_r = cute.make_ptr( + cutlass.Int32, + cand_idx_row.iterator.toint() + + cutlass.Int64(li_r) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + output_indices_row[io_r] = cute.make_tensor( + ip_r, cute.make_layout((1,)) + )[0] + io_r = io_r + cutlass.Int32(num_threads) + cute.arch.barrier() + if cutlass.const_expr(_P4_SUB_DBG): + ck_sw1 = cute.arch.clock64() + if cutlass.const_expr(self.emit_xstate): + # Closed-loop state (interface v2): [0] valid, [1] kth + # proxy (= accepted threshold; the tie-fill makes it a + # tight lower bound of the true kth), [2] accepted + # threshold, [3] cand_count. The next step derives its + # seed rung group from these. + if list_used == cutlass.Int32(1): + # the anchor below reads output_indices_row[K-1], + # written by peer threads in rank-scatter / tail + # repair; not every exit of that phase ends in a + # block barrier (the eager-position path dropped + # the swap loop's trailing one), so publish + # visibility explicitly. list_used is uniform + # (admission is decided from shared control + # words), so the barrier is block-safe. + cute.arch.barrier() + if tidx == cutlass.Int32(0): + xstate_row[0] = cutlass.Float32(1.0) + thr_pub = s_thr[0] + anch_pub = s_thr[0] + if list_used == cutlass.Int32(1): + # list rows: rank-scatter's output is rank- + # ordered, so slot K-1 holds the exact k-th + # boundary - a tighter, healthier closed-loop + # anchor than the loose collect line. + idx_k = output_indices_row[cutlass.Int32(self.top_k - 1)] + if idx_k >= cutlass.Int32(0) and idx_k < N: + thr_pub = cutlass.Float32(input_row[idx_k]) + xstate_row[1] = thr_pub + xstate_row[2] = anch_pub + if cutlass.const_expr(_P4_TAIL_DBG): + ck3 = cute.arch.clock64() + # [1] device total (entry->publish), [2] true + # in-kernel prologue (entry->walk start): wall + # minus [1] = host/launch, NOT kernel work + xstate_row[1] = cutlass.Float32(cutlass.Int32(ck3 - ckE)) + xstate_row[2] = cutlass.Float32(cutlass.Int32(ck0 - ckE)) + xstate_row[4] = cutlass.Float32(cutlass.Int32(ck1 - ck0)) # walk+flags + xstate_row[5] = cutlass.Float32(cutlass.Int32(ck2 - ck1)) # P2/P3 gap + xstate_row[6] = cutlass.Float32(cutlass.Int32(ck3 - ck2)) # Phase 4 + xstate_row[7] = s_thr[1] # cnt_strad + if cutlass.const_expr(_P4_SUB_DBG): + xstate_row[2] = cutlass.Float32(smem_wcnt_p1[7]) + if cutlass.const_expr(_P4_SUB_DBG): + # P4 sub-phase cycles staged by rank_scatter. + # Chain-safe layout: [2] (closed-loop anchor) + # untouched; [1] cnt_strad (tail class size), + # [4] fine, [5] scatter, [6] tail, [7] deferred- + # position swap. The small C-predictable phases + # (minmax/hist/coarse, wcnt[8..10]) are not + # published. + xstate_row[1] = s_thr[1] + if cutlass.const_expr(_P4_SUB_HEAD): + xstate_row[4] = cutlass.Float32(smem_wcnt[8]) + xstate_row[5] = cutlass.Float32(smem_wcnt[9]) + xstate_row[6] = cutlass.Float32(smem_wcnt[10]) + xstate_row[7] = cutlass.Float32(smem_wcnt[11]) + else: + xstate_row[4] = cutlass.Float32(smem_wcnt[11]) + xstate_row[5] = cutlass.Float32(smem_wcnt[12]) + xstate_row[6] = cutlass.Float32(smem_wcnt[13]) + xstate_row[7] = cutlass.Float32(cutlass.Int32(ck_sw1 - ck_sw0)) + # cand_count_p4 = pre-P4 snapshot (P4 repurposes + # the s_iscalars slots). + xstate_row[3] = cutlass.Float32(cand_count_p4) + if cutlass.const_expr( + self.ext_rungs and not _P4_SUB_DBG and not _P4_TAIL_DBG + ): + # closed-loop food: the three rung counts this + # step measured (exact, straight from the R0 + # multi-count) - the host derives next-step + # lines from these instead of re-counting. + xstate_row[4] = cutlass.Float32(s_mt_cnt[0]) + xstate_row[5] = cutlass.Float32(s_mt_cnt[1]) + xstate_row[6] = cutlass.Float32(s_mt_cnt[2]) + if cutlass.const_expr(_SKIP_DBG and self.enable_block_skip): + xstate_row[4] = cutlass.Float32(s_active_cnt[0]) + xstate_row[5] = cutlass.Float32(s_active_cnt[1]) + xstate_row[6] = cutlass.Float32(s_r0col[0]) + xstate_row[7] = cutlass.Float32(s_active_cnt[2]) + else: + # cs>1: only the leader (CTA 0 in cluster) runs Phase 4. + if is_leader: + if do_cluster_sync: + # DSMEM-gather peer candidates into the leader's + # smem_keys/smem_vals. Layout: leader's chunk goes + # to [0 .. leader_local_cnt); each peer r's chunk + # appends the next peer_r_local_cnt entries. + local_cnt_self = s_iscalars[5] + local_iscalars_ptr = s_iscalars.iterator + cutlass.Int32(5) + smem_keys_iter = smem_keys.iterator + smem_vals_iter = smem_vals.iterator + base_offset = local_cnt_self + for peer in cutlass.range_constexpr(1, cluster_size): + peer_iscalars_addr = mapa_shared_cluster( + local_iscalars_ptr, cutlass.Int32(peer) + ) + peer_cnt = ld_shared_cluster_i32(peer_iscalars_addr) + # Cap to kC (defense-in-depth vs. the + # done==2 bracket-exhaustion path). + peer_cnt = min(peer_cnt, cutlass.Int32(self.kC)) + i_gather = tidx + while i_gather < peer_cnt: + peer_key_addr = mapa_shared_cluster( + smem_keys_iter + i_gather, cutlass.Int32(peer) + ) + peer_val_addr = mapa_shared_cluster( + smem_vals_iter + i_gather, cutlass.Int32(peer) + ) + k_val = ld_shared_cluster_f32(peer_key_addr) + v_val = ld_shared_cluster_i32(peer_val_addr) + dst = base_offset + i_gather + if dst < cutlass.Int32(self.kC): + smem_keys[dst] = k_val + smem_vals[dst] = v_val + i_gather = i_gather + cutlass.Int32(num_threads) + base_offset = base_offset + peer_cnt + # Reset s_iscalars[0] to cluster-wide cand_count. + if tidx == cutlass.Int32(0): + s_iscalars[0] = base_offset + cute.arch.barrier() + # else: short-row degrade — leader (CTA 0) already + # holds the full row's candidates in its own + # smem_keys/smem_vals (no peers to gather from). + + # ---- Phase 4: histogram snap + writeback ---- + # Capture the P2 terminal BEFORE Phase 4: P4 + # reuses s_iscalars[1] as radix scratch. + if tidx == cutlass.Int32(0): + s_iscalars[6] = cutlass.Int32(0) + if s_iscalars[1] == cutlass.Int32(3): + s_iscalars[6] = cutlass.Int32(1) + cute.arch.barrier() + cand_count_p4 = min(s_iscalars[0], cutlass.Int32(self.kC)) + if cutlass.const_expr(self.enable_p4_rank_scatter): + self.phase4_rank_scatter( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) + else: + self.phase4_histogram_snap( + smem_keys, + smem_vals, + smem_hist, + smem_wcnt, + s_thr, + s_iscalars, + output_values_row, + output_indices_row, + cand_count_p4, + tidx, + warp_id, + lane, + ) + if cutlass.const_expr(self.emit_xstate): + # closed-loop state, leader-only at cs > 1 (same + # layout as the cs == 1 exit). + if tidx == cutlass.Int32(0): + xstate_row[0] = cutlass.Float32(1.0) + xstate_row[1] = s_thr[0] + xstate_row[2] = s_thr[0] + xstate_row[3] = cutlass.Float32(cand_count_p4) + if cutlass.const_expr( + self.ext_rungs and not _P4_SUB_DBG and not _P4_TAIL_DBG + ): + # cluster-merged rung counts (identical on + # every CTA after the multi-count DSMEM + # aggregation) + xstate_row[4] = cutlass.Float32(s_mt_cnt[0]) + xstate_row[5] = cutlass.Float32(s_mt_cnt[1]) + xstate_row[6] = cutlass.Float32(s_mt_cnt[2]) + + # ---- plateau fill (done == 3): complete the row from the + # bitwise-equal plateau class. The terminal is only set on an + # ADJACENT-FLOAT bracket, so every value in [s_thr[1], s_thr[0]) + # is bitwise-equal; Phase 4 has already emitted the + # cnt(>= s_thr[0]) sure winners, and ANY (K - count)-subset of + # the tie class is a valid tie-aware completion. Ticket counter + # lives in the DEDICATED s_iscalars[7]. + if s_iscalars[6] == cutlass.Int32(1): + pv_lo = s_thr[1] + pv_hi = s_thr[0] + if tidx == cutlass.Int32(0): + # cand_count_p4 was captured BEFORE Phase 4; + # s_iscalars[0] is radix scratch by now (same + # hazard as the flag). + s_iscalars[7] = cand_count_p4 + cute.arch.barrier() + ifp = tidx + while ifp < N: + vfp = cutlass.Float32(0.0) + if cutlass.const_expr(self.dtype == cutlass.Float32): + vfp = input_row[ifp] + else: + vfp = cutlass.Float32(input_row[ifp]) + if vfp >= pv_lo and vfp < pv_hi: + pfill = atomicAdd( + s_iscalars.iterator + cutlass.Int32(7), cutlass.Int32(1) + ) + if pfill < cutlass.Int32(self.top_k): + if cutlass.const_expr(self.return_output_values): + output_values_row[pfill] = self.dtype(vfp) + output_indices_row[pfill] = ifp + ifp = ifp + cutlass.Int32(self.num_threads) + cute.arch.barrier() # Final cluster barrier: keep peer CTAs (and their SMEM) alive # until the leader's gather + Phase 4 finish. Skipped at @@ -5089,6 +8217,13 @@ def __call__( output_indices: cute.Tensor, order_row: cute.Tensor, # or None when seqlen_sorted=False stream, + block_max: cute.Tensor = None, # block-skip bounds; None = disabled + seed_thr: cute.Tensor = None, # [num_rows, 3] fp32 (ext counts) + seed_counts: cute.Tensor = None, # [num_rows, 3] int32 (ext counts) + xstate: cute.Tensor = None, # [num_rows, 8] fp32 (emit_xstate) + cand_vals: cute.Tensor = None, # [num_rows, CAP] fp32 (ext cand) + cand_idx: cute.Tensor = None, # [num_rows, CAP] int32 (ext cand) + cand_ctl: cute.Tensor = None, # [num_rows, 2] int32 (ext cand) ): num_rows = input_data.shape[0] cluster_size = cutlass.const_expr(self.cluster_size) @@ -5110,6 +8245,13 @@ def __call__( output_values, output_indices, order_row, + block_max, + seed_thr, + seed_counts, + xstate, + cand_vals, + cand_idx, + cand_ctl, ).launch( grid=(total_ctas, 1, 1), block=(self.num_threads, 1, 1), @@ -5229,21 +8371,41 @@ def pick_config( num_candidates: int, max_seq_len: Optional[int] = None, num_sms: Optional[int] = None, + has_block_max: bool = False, ) -> dict: - """Launch-shape ctor kwargs for ``(dtype, BS, N)`` — the single - source of truth shared by the production runner - (``CuteDSLGvrTopKDecodeRunner``) and direct-drive users (tests, - benchmarks): composition of :meth:`pick_cluster_size` and - :meth:`pick_tuning`. + """Pick the launch-shape ctor kwargs for ``(dtype, BS, N)``. + + Mirrors the production runner policy (cluster_size auto-pick + + ``_pick_tuning``) so any caller instantiating the kernel directly + gets the same shapes the custom op would use. Rationale (B200, + nsys cold-L2, 2026-07-15 big-BS triage): a config frozen at the + BS=1 optimum (cs = N>=65536 ? 4 : 1, T=1024, mbpm=1) is geomean + 2.27x slower (max 6.0x) than the op-bench anchor at BS in + {64, 256, 1024}, while this policy is 0.95x (parity/better). + Multi-CTA splitting only pays while the grid is a single wave + (num_rows * cluster_size <= num_sms); past that, row parallelism + already saturates the SMs and per-row splitting is pure overhead. ``max_seq_len``: pass the peak runtime N under CUDA-graph capture so the variant is picked for the replay shape, not the capture - shape. + shape (same contract as the custom op's ``_pick_tuning``). + + Returns kwargs for ``GvrTopKKernel(...)``: ``cluster_size``, + ``num_threads``, ``use_256bit_load``, ``min_blocks_per_mp``, + ``enable_warp_parallel_reduce``. """ if num_sms is None: num_sms = GvrTopKKernel._device_num_sms() n_row = max_seq_len if max_seq_len is not None else num_candidates - cluster_size = GvrTopKKernel.pick_cluster_size(num_rows, n_row, num_sms) + if has_block_max and n_row >= 200_000: + # Block-skip requires cs == 1 with a large per-CTA slice + # (splitting shrinks each CTA's slice below the skip + # break-even and disables rung tightening). Below 200k the + # wrapper drops block_max anyway (skip_min_n gate) and the + # stock picks apply. + cluster_size = 1 + else: + cluster_size = GvrTopKKernel.pick_cluster_size(num_rows, n_row, num_sms) cfg = GvrTopKKernel.pick_tuning( torch_dtype, num_rows, diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 950b02c30f93..97e6cb6cc694 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -59,6 +59,13 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio + # emission-assisted GVR (opt-in via prepare_gvr_emission): the + # module owns the closed-loop emission state; the caller passes + # the returned kwargs to the scoring op, and the consume side is + # injected into the GVR Top-K call while the step stays armed + self._gvr_emission_state = None + self._gvr_emission_route = None + self._gvr_emission_armed = False def forward( self, @@ -334,6 +341,16 @@ def _forward_decode_gvr( ) else: assert max_seq_len is not None + emission_kwargs: dict = {} + if self._gvr_emission_armed: + state = self._gvr_emission_state + num_rows = scores.shape[0] + emission_kwargs = state.topk_ext_kwargs( + self._gvr_emission_route, + num_rows, + state.block_max[:num_rows] if state.block_max is not None else None, + ) + self._gvr_emission_armed = False torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, gvr_prior_indices, @@ -344,9 +361,71 @@ def _forward_decode_gvr( compress_ratio=self.compress_ratio, max_seq_len=max_seq_len, order_row=gvr_row_order, + **emission_kwargs, ) return output_indices + def prepare_gvr_emission( + self, + batch: int, + n_comp: int, + num_sms: int, + gvr_prior_indices: torch.Tensor, + ) -> dict: + """Plan the emission-assisted GVR tier for this decode step. + + Returns the emission kwargs for the paged-MQA scoring op (empty + when the planner declines this step); the matching consume-side + kwargs are injected into the next GVR Top-K call automatically. + Host arithmetic on engine-static shapes plus capturable device + ops, so a captured graph bakes the tier and replays refresh the + state buffers in place. + + Args: + batch: Number of decode requests this step. + n_comp: Engine-static compressed maximum sequence length. + num_sms: Device SM count. + gvr_prior_indices: Caller-owned previous-selection state; + defines the emission state's row capacity and device. + """ + # the emission/xstate writes are undeclared mutations (see the op's + # schema note), so the tier is eager / CUDA-graph only + if torch.compiler.is_dynamo_compiling(): + return {} + from ..cute_dsl_kernels.blackwell.top_k.gvr_emission import ( + LIST_EMIT_MIN_N, + GvrEmissionState, + ) + + if self._gvr_emission_state is None: + self._gvr_emission_state = GvrEmissionState( + max_rows=gvr_prior_indices.shape[0], + top_k=self.top_k, + device=gvr_prior_indices.device, + enable_list_tier=n_comp >= LIST_EMIT_MIN_N, + own_prior=False, + ) + state = self._gvr_emission_state + emit_tier, self._gvr_emission_route = state.plan( + batch, n_comp, num_sms, compress_ratio=max(self.compress_ratio, 1) + ) + self._gvr_emission_armed = self._gvr_emission_route.tier != "none" + if emit_tier in ("counts", "list", "rungs"): + state.update_seed_rows(batch, emit_tier) + kwargs: dict = {} + if emit_tier in ("counts", "list"): + kwargs = state.indexer_emit_kwargs(emit_tier, batch) + if self._gvr_emission_route.attach_block_max or emit_tier in ("counts", "list"): + kwargs["block_max_out"] = state.ensure_block_max(n_comp)[:batch] + return kwargs + + def reset_gvr_emission_rows(self, rows: slice) -> None: + """Cold-start the emission closed-loop state for reused request + slots (prefill-to-decode handoff): a zeroed xstate reads as + invalid and routes those rows to the stock path in-kernel.""" + if self._gvr_emission_state is not None: + self._gvr_emission_state.xstate[rows].zero_() + def update_gvr_prior_from_prefill( self, output_indices: torch.Tensor, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0a59fb6bafca..fc808d4c406b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -79,7 +79,8 @@ l0_b200: - test_e2e.py::test_ptp_quickstart_advanced_ngram[Llama-3.1-8B-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct] - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False] - test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] - - unittest/_torch/attention + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) - unittest/_torch/compilation - unittest/_torch/debugger - unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 79f37a2e49d1..234ece45eb9c 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -19,7 +19,7 @@ l0_b300: - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml index 2e9baee7978a..4b6822c6db5f 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml @@ -18,7 +18,7 @@ l0_dgx_b300: - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py - unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py - unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py - - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py + - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120) - unittest/_torch/executor # ------------- modules (multi-GPU) --------------- - unittest/_torch/modules/test_mla_helix.py diff --git a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py index 41cb23260cd4..73791781d666 100644 --- a/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py +++ b/tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py @@ -8,21 +8,22 @@ ``use_constant_hint``, ``compress_ratio``, ``max_seq_len`` hint) so bench scripts can override the heuristic. -**Not in CI** — this file imports the DSL kernel module directly +**Not in CI** - this file imports the DSL kernel module directly (no ``trtllm`` runtime dep), to enable knob-A/B development outside the production op. Two usage modes: -* `python -m pytest run_gvr_topk.py`` — exhaustive parameterized correctness sweep - (dtype × K × N × seed × next_n × T × V × warp-parallel-reduce). -* ``python run_gvr_topk.py --dtype bf16 --top_k 1024 --N 8192`` — +* `python -m pytest run_gvr_topk.py`` - exhaustive parameterized correctness sweep + (dtype x K x N x seed x next_n x T x V x warp-parallel-reduce). +* ``python run_gvr_topk.py --dtype bf16 --top_k 1024 --N 8192`` - single-case correctness verification on user-specified shape; knob overrides via ``--num_threads`` / ``--use_256bit_load`` / etc. """ import argparse import functools +import os import sys from pathlib import Path from typing import Optional @@ -64,6 +65,25 @@ def _compile( seqlen_sorted: bool = False, p4_warp_redundant: bool = True, p2_warp_redundant: bool = True, + enable_block_skip: bool = False, + pdl_wait_late: bool = True, + p4_tail_v3: "bool | None" = None, + p4_no_fine: "bool | None" = None, + p4_exact_tail: "bool | None" = None, + p4_tail_fast: "bool | None" = None, + p1r_rescue: bool = True, + num_bins: "int | None" = None, + p4_fine_rangetest: "bool | None" = None, + p4_scat_rangetest: bool = False, + use_ext_counts: bool = False, + emit_xstate: bool = False, + use_ext_cand: bool = False, + ext_rungs: bool = False, + cand_cap: int = 5120, + accept_cap: "int | None" = None, + kc_override: "int | None" = None, + self_scan: bool = False, + cap_c: "int | None" = None, ): """JIT-compile the GVR kernel for a specific knob combination. @@ -122,6 +142,57 @@ def _compile( if seqlen_sorted else None ) + block_max_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (n_rows, cute.sym_int()), + stride_order=(1, 0), + assumed_align=16, + ) + if enable_block_skip + else None + ) + # ext counts ride PACKED with the lines ([rows, 8] fp32: lines at + # [0..2], counts as floats at [3..5]) - one 32B sector per row + seed_thr_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (n_rows, 8 if use_ext_counts else 3), + stride_order=(1, 0), + assumed_align=4, + ) + if (use_ext_counts or ext_rungs) + else None + ) + seed_counts_fake = None + cand_vals_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4 + ) + if use_ext_cand + else None + ) + cand_idx_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (n_rows, cand_cap), stride_order=(1, 0), assumed_align=4 + ) + if (use_ext_cand or self_scan) + else None + ) + cand_ctl_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (n_rows, 4), stride_order=(1, 0), assumed_align=8 + ) + if use_ext_cand + else None + ) + xstate_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (n_rows, 8), stride_order=(1, 0), assumed_align=4 + ) + if emit_xstate + else None + ) fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = GvrTopKKernel( dtype=cute_dtype, @@ -140,6 +211,28 @@ def _compile( seqlen_sorted=seqlen_sorted, p4_warp_redundant=p4_warp_redundant, p2_warp_redundant=p2_warp_redundant, + enable_block_skip=enable_block_skip, + pdl_wait_late=pdl_wait_late, + p4_tail_v3=p4_tail_v3, + p4_no_fine=p4_no_fine, + p4_exact_tail=p4_exact_tail, + p4_tail_fast=p4_tail_fast, + p1r_rescue=p1r_rescue, + num_bins=num_bins, + p4_fine_rangetest=p4_fine_rangetest, + p4_scat_rangetest=p4_scat_rangetest, + use_ext_counts=use_ext_counts, + emit_xstate=emit_xstate, + use_ext_cand=use_ext_cand, + ext_rungs=ext_rungs, + cand_cap=cand_cap, + accept_cap=accept_cap, + kc_override=kc_override, + self_scan=self_scan, + cap_c=cap_c, + # ext counts need 3 rung slots (M_thr == 3): 2 qfracs + vseed; + # the qfrac values are unused here - only the slot count matters. + r0_qfracs=(0.85, 0.35) if (use_ext_counts or ext_rungs) else None, ) return cute.compile( kernel, @@ -150,10 +243,409 @@ def _compile( out_indices_fake, order_row_fake, stream=fake_stream, + block_max=block_max_fake, + seed_thr=seed_thr_fake, + seed_counts=seed_counts_fake, + xstate=xstate_fake, + cand_vals=cand_vals_fake, + cand_idx=cand_idx_fake, + cand_ctl=cand_ctl_fake, options="--enable-tvm-ffi", ) +_FLT_MAX = torch.finfo(torch.float32).max +_META_BLOCK = 128 +_META_RECS_PER_BLOCK = 4 # warp-partial records per block (indexer layout) + + +def _row_n_eff( + seq_lens: torch.Tensor, + num_rows: int, + next_n: int, + compress_ratio: int, +) -> torch.Tensor: + """Per-row effective scan length, mirroring the kernel formula.""" + dev = seq_lens.device + rows = torch.arange(num_rows, device=dev) + sl = seq_lens[rows // next_n].to(torch.int64) + actual = sl - next_n + (rows % next_n) + 1 + return actual // compress_ratio + + +def emu_block_max( + logits: torch.Tensor, + seq_lens: torch.Tensor, + next_n: int = 1, + compress_ratio: int = 1, + tail_mode: str = "pad_inf", + records: str = "positional", +) -> torch.Tensor: + """``[num_rows, nb_pad*4] fp32`` warp-partial upper-bound records. + + tail_mode: + "exact": tight bound - max over valid positions only. + "pad_inf": a partially-valid tail unit is forced to +FLT_MAX, the + worst legal inflation (the indexer masks by request-level + ctx >= N_eff, so tail positions can inflate the bound). + records: + "rotate": fold-correctness fixture - the 128-block max lands in + ONE slot rotated by blk % 4, the other 3 hold + -FLT_MAX. Valid ONLY for grain-128 consumers (slots + are NOT positional). + "positional": production semantics - record r is the exact max of + positions [r*32, r*32+32) (the indexer's TMEM T2R + partition gives warp w of a tile the contiguous + positions [tile*128 + w*32, +32)). Required for + skip_grain=32; also a legal grain-128 input (its + fold is the block max). + """ + assert tail_mode in ("exact", "pad_inf") + assert records in ("rotate", "positional") + R, C = logits.shape + nb = (C + _META_BLOCK - 1) // _META_BLOCK + dev = logits.device + n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio).unsqueeze(1) + lf = logits.to(torch.float32) + pad = nb * _META_BLOCK - C + if pad: + lf = torch.nn.functional.pad(lf, (0, pad), value=float("-inf")) + pos = torch.arange(nb * _META_BLOCK, device=dev).unsqueeze(0) + masked = torch.where(pos < n_eff, lf, torch.full_like(lf, float("-inf"))) + if records == "positional": + sub = _META_BLOCK // _META_RECS_PER_BLOCK # 32 positions/record + nrec = nb * _META_RECS_PER_BLOCK + rmax = masked.view(R, nrec, sub).amax(-1) + if tail_mode == "pad_inf": + rec_start = torch.arange(nrec, device=dev).unsqueeze(0) * sub + partial = (rec_start < n_eff) & (rec_start + sub > n_eff) + rmax = torch.where(partial, torch.full_like(rmax, _FLT_MAX), rmax) + return rmax.contiguous() + bmax = masked.view(R, nb, _META_BLOCK).amax(-1) + if tail_mode == "pad_inf": + blk_start = torch.arange(nb, device=dev).unsqueeze(0) * _META_BLOCK + partial = (blk_start < n_eff) & (blk_start + _META_BLOCK > n_eff) + bmax = torch.where(partial, torch.full_like(bmax, _FLT_MAX), bmax) + out = torch.full((R, nb, _META_RECS_PER_BLOCK), -_FLT_MAX, dtype=torch.float32, device=dev) + slot = torch.arange(nb, device=dev) % _META_RECS_PER_BLOCK + out[:, torch.arange(nb, device=dev), slot] = bmax + return out.reshape(R, nb * _META_RECS_PER_BLOCK).contiguous() + + +# Rung offsets below each 32-position record's max for the meta-seed +# metadata. +_META_DELTAS = (0.25, 0.5, 1.0, 2.0, 4.0, 8.0) + + +def emu_block_meta( + logits: torch.Tensor, + seq_lens: torch.Tensor, + compress_ratio: int = 1, + next_n: int = 1, +) -> torch.Tensor: + """Emulate the indexer-side per-32-block rung-count metadata. + + Record r packs, for each rung offset ``delta_j`` in + ``_META_DELTAS``, ``count(v >= record_max - delta_j)`` over + positions ``[32r, 32r+32)`` as a 5-bit saturating field at bits + ``[5j, 5j+5)`` (31 = "31 or 32"; the kernel decodes 31 as the safe + upper bound 32). Positions beyond the row's effective length are + excluded. Layout matches ``block_max``: ``[num_rows, nrec]`` int32 + with ``nrec = ceil(N/128)*4`` (one record per 32 positions). + """ + assert next_n == 1, "emu_block_meta: next_n == 1 only" + x = logits.float() + R, N = x.shape + nrec = ((N + _META_BLOCK - 1) // _META_BLOCK) * _META_RECS_PER_BLOCK + npad = nrec * 32 + if npad > N: + x = torch.nn.functional.pad(x, (0, npad - N), value=float("-inf")) + n_eff = seq_lens.long() // compress_ratio + ar = torch.arange(npad, device=x.device)[None, :] + x = x.masked_fill(ar >= n_eff[:, None], float("-inf")) + xb = x.view(R, nrec, 32) + m = xb.amax(2, keepdim=True) + meta = torch.zeros(R, nrec, dtype=torch.int32, device=x.device) + for j, d in enumerate(_META_DELTAS): + cj = (xb >= (m - d)).sum(2).clamp(max=31).to(torch.int32) + meta |= cj << (5 * j) + return meta.contiguous() + + +def derive_seed_rungs( + prev_thr: torch.Tensor, + prev_sthr: "torch.Tensor | None" = None, + prev_counts: "torch.Tensor | None" = None, + count_octaves: float = 2.0, + fallback_spread: float = 0.5, + top_k: "int | None" = None, +) -> torch.Tensor: + """Host-side slope-adaptive seed rung derivation (waterfall closed loop). + + Estimates the per-row local slope of log2(count) vs threshold from the + PREVIOUS step's 3 rung measurements and places the next step's guard + rungs ``count_octaves`` octaves away from the mid rung (= the previous + accepted threshold). + + Args: + prev_thr: [rows] previous accepted threshold (xstate[:, 2]). + prev_sthr: [rows, 3] previous step's rung thresholds (or None). + prev_counts: [rows, 3] previous step's rung counts (or None). + + Returns: + [rows, 3] fp32 seed thresholds (ascending). + """ + if prev_sthr is None or prev_counts is None: + d = torch.full_like(prev_thr, fallback_spread) + d_lo = d + else: + c_lo = prev_counts[:, 0].float().clamp(min=1.0) + c_hi = prev_counts[:, 2].float().clamp(min=1.0) + dthr = (prev_sthr[:, 2] - prev_sthr[:, 0]).clamp(min=1e-3) + slope = (torch.log2(c_lo) - torch.log2(c_hi)) / dthr + d = torch.where( + slope > 0.05, + count_octaves / slope.clamp(min=0.05), + torch.full_like(slope, fallback_spread), + ).clamp(0.1, 4.0) + # undershoot hysteresis: a row whose previous loose rung caught + # fewer than K widens only its next down-guard by +2 octaves. + oct_lo = torch.full_like(slope, count_octaves) + if top_k is not None: + oct_lo = torch.where(prev_counts[:, 0].float() < float(top_k), oct_lo + 2.0, oct_lo) + d_lo = torch.where( + slope > 0.05, + oct_lo / slope.clamp(min=0.05), + torch.full_like(slope, fallback_spread), + ).clamp(0.1, 6.0) + return torch.stack([prev_thr - d_lo, prev_thr, prev_thr + d], dim=1).contiguous() + + +def emu_cand_bucketed( + logits: torch.Tensor, + seq_lens: torch.Tensor, + seed_thr: torch.Tensor, + cap: int, + seg_cap: int = 8192, + next_n: int = 1, + compress_ratio: int = 1, + sentinel_pad: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Bucketed SoA candidate emission (v5 contract). + + Three fixed segments in one buffer: A = [0, seg_cap) holds >= t2, + B = [seg_cap, 2*seg_cap) holds [t1, t2), C = [2*seg_cap, 2*seg_cap + + cap) holds [t0, t1). A full segment spills to the next looser one + (never drops an entry), so the union always equals the full >= t0 + set, and a segment group is complete exactly when its line's count + fits the acceptance band (seg_cap = B*). Sentinel pads land in C. + ctl = {n0 (claimed incl pads), void, n1, n2}. + """ + R, _C = logits.shape + dev = logits.device + n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio) + lf = logits.to(torch.float32) + width = 2 * seg_cap + cap + cand_vals = torch.full((R, width), float("-inf"), dtype=torch.float32, device=dev) + cand_idx = torch.full((R, width), -1, dtype=torch.int32, device=dev) + ctl = torch.zeros((R, 4), dtype=torch.int32, device=dev) + for r in range(R): + ne = int(n_eff[r]) + row = lf[r, :ne] + hits = torch.nonzero(row >= seed_thr[r, 0], as_tuple=False).flatten() + cnt = hits.numel() + ctl[r, 2] = int((row >= seed_thr[r, 1]).sum()) + ctl[r, 3] = int((row >= seed_thr[r, 2]).sum()) + # emission order is value-blind: shuffle, then classify + perm = hits[torch.randperm(cnt, device=dev)] + v = row[perm] + seg = torch.where(v >= seed_thr[r, 2], 0, torch.where(v >= seed_thr[r, 1], 1, 2)) + # vectorized spill-to-looser: stream s = native entries + spill + # from s-1 (in emission order); ordinal beyond the cap spills on. + in_a = seg == 0 + ord_a = torch.cumsum(in_a.int(), 0) + stay_a = in_a & (ord_a <= seg_cap) + in_b = (seg == 1) | (in_a & ~stay_a) + ord_b = torch.cumsum(in_b.int(), 0) + stay_b = in_b & (ord_b <= seg_cap) + in_c = (seg == 2) | (in_b & ~stay_b) + ord_c = torch.cumsum(in_c.int(), 0) + stay_c = in_c & (ord_c <= cap) + voided = int(in_c.sum()) > cap + slot = torch.full_like(seg, -1) + slot[stay_a] = ord_a[stay_a] - 1 + slot[stay_b] = seg_cap + (ord_b[stay_b] - 1) + slot[stay_c] = 2 * seg_cap + (ord_c[stay_c] - 1) + live = slot >= 0 + cand_vals[r, slot[live].long()] = v[live] + cand_idx[r, slot[live].long()] = perm[live].int() + claimed = cnt + sentinel_pad + ctl[r, 0] = claimed + ctl[r, 1] = 1 if (voided or claimed > cap) else 0 + return cand_vals, cand_idx, ctl + + +def derive_seed_lines_v4( + prev_anchor: torch.Tensor, + prev_sthr: "torch.Tensor | None" = None, + prev_ctl: "torch.Tensor | None" = None, + targets: "tuple[float, float, float]" = (8192.0, 5120.0, 2048.0), + fallback_spread: float = 0.5, +) -> torch.Tensor: + """v4 host-side line placement: put [t0, t1, t2] at target COUNTS. + + Slope of log2(count) vs threshold is fit from the previous step's + (t0, n0) / (t2, n2) pairs (counts ride in the widened control words); + each new line lands where the fit predicts its target count. Targets + descend (t0 loosest / largest count, t2 tightest). + + Args: + prev_anchor: [rows] previous accepted cut value (xstate[:, 2]). + prev_sthr: [rows, 3] previous lines (or None -> fixed spread). + prev_ctl: [rows, 4] previous control words {n0, void, n1, n2}. + targets: (T0, T1, T2) target counts, T0 > T1 > T2. + + Returns: + [rows, 3] fp32 lines ascending [t0, t1, t2]. + """ + t0_t, t1_t, t2_t = targets + if prev_sthr is None or prev_ctl is None: + d = torch.full_like(prev_anchor, fallback_spread) + return torch.stack([prev_anchor - d, prev_anchor, prev_anchor + d], dim=1).contiguous() + c0 = prev_ctl[:, 0].float().clamp(min=1.0) + c2 = prev_ctl[:, 3].float().clamp(min=1.0) + dthr = (prev_sthr[:, 2] - prev_sthr[:, 0]).clamp(min=1e-3) + slope = ((torch.log2(c0) - torch.log2(c2)) / dthr).clamp(min=0.05, max=64.0) + # anchor count estimate: slide the anchor onto the prev line fit + anch_c = (c2 * torch.exp2(-(prev_anchor - prev_sthr[:, 2]) * slope)).clamp(min=1.0, max=1e6) + lines = [prev_anchor + torch.log2(anch_c / tgt) / slope for tgt in (t0_t, t1_t, t2_t)] + out = torch.stack(lines, dim=1) + # enforce strictly ascending (degenerate slope guards) + out[:, 1] = torch.maximum(out[:, 1], out[:, 0] + 1e-4) + out[:, 2] = torch.maximum(out[:, 2], out[:, 1] + 1e-4) + return out.contiguous() + + +def pack_seed( + seed_thr: torch.Tensor, + seed_counts: torch.Tensor, + block_max: torch.Tensor = None, +) -> torch.Tensor: + """Pack lines + exact counts into one [rows, 8] fp32 seed row. + + Lines land at [0..2], counts as floats at [3..5] (exact to 2^24); + col 6 optionally carries the adaptive-skip pass count (32-grain + block records clearing t_0; 0 = not provided). One 32B sector per + row. Build ONCE per step, outside any timed region. + """ + pack = torch.zeros((seed_thr.shape[0], 8), dtype=torch.float32, device=seed_thr.device) + pack[:, 0:3] = seed_thr + pack[:, 3:6] = seed_counts.float() + if block_max is not None: + pack[:, 6] = (block_max >= seed_thr[:, 0:1]).sum(dim=1).float() + return pack.contiguous() + + +def emu_seed_counts( + logits: torch.Tensor, + seq_lens: torch.Tensor, + seed_thr: torch.Tensor, + next_n: int = 1, + compress_ratio: int = 1, +) -> torch.Tensor: + """L1 emu: exact per-row threshold counts. + + counts[r][j] = |{i < N_eff(r) : logits[r, i] >= t_j}| on the + post-conversion values (contract: epilogue_topk_interface.md). + """ + R, C = logits.shape + n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio).unsqueeze(1) + pos = torch.arange(C, device=logits.device).unsqueeze(0) + lf = logits.to(torch.float32) + valid = pos < n_eff + counts = torch.empty((R, seed_thr.shape[1]), dtype=torch.int32, device=logits.device) + for j in range(seed_thr.shape[1]): + counts[:, j] = ((lf >= seed_thr[:, j : j + 1]) & valid).sum(-1, dtype=torch.int32) + return counts + + +def emu_cand( + logits: torch.Tensor, + seq_lens: torch.Tensor, + seed_thr: torch.Tensor, + cap: int, + next_n: int = 1, + compress_ratio: int = 1, + sentinel_pad: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """L2 emu: unordered candidate pre-collect (SoA). + + Unordered (value fp32-bits, index) pairs of all valid positions + >= t_0 = seed_thr[:, 0]; ctl = {claimed, void}. claimed may + over-approximate the true count (window sentinels, idx word = -1) - + ``sentinel_pad`` injects that legally. void=1 when claimed > cap; on + overflow only the first ``cap`` entries are materialized (contract v2: + consumers scan [0, min(claimed, cap)) skipping sentinels). + """ + R, C = logits.shape + dev = logits.device + n_eff = _row_n_eff(seq_lens, R, next_n, compress_ratio) + lf = logits.to(torch.float32) + cand_vals = torch.full((R, cap), float("-inf"), dtype=torch.float32, device=dev) + cand_idx = torch.full((R, cap), -1, dtype=torch.int32, device=dev) + ctl = torch.zeros((R, 4), dtype=torch.int32, device=dev) + for r in range(R): + ne = int(n_eff[r]) + hits = torch.nonzero(lf[r, :ne] >= seed_thr[r, 0], as_tuple=False).flatten() + cnt = hits.numel() + # emitter-side counts: two extra compares per EMITTED element + # (t1, t2 > t0 so counting over the list == counting over the row) + ctl[r, 2] = int((lf[r, :ne] >= seed_thr[r, 1]).sum()) + ctl[r, 3] = int((lf[r, :ne] >= seed_thr[r, 2]).sum()) + # unordered contract: shuffle, then interleave sentinels + perm = hits[torch.randperm(cnt, device=dev)] + ent = torch.full((cnt + sentinel_pad,), -1, dtype=torch.int64, device=dev) + if sentinel_pad: + slots = torch.randperm(cnt + sentinel_pad, device=dev)[:cnt] + slots = slots.sort().values + else: + slots = torch.arange(cnt, device=dev) + ent[slots] = perm + claimed = int(ent.numel()) + nwr = min(claimed, cap) + live = ent[:nwr] >= 0 + cand_idx[r, :nwr] = ent[:nwr].to(torch.int32) + cand_vals[r, :nwr][live] = lf[r, ent[:nwr][live]] + ctl[r, 0] = claimed + ctl[r, 1] = 1 if claimed > cap else 0 + return cand_vals, cand_idx, ctl + + +def enc_ordered_f32(t: torch.Tensor) -> torch.Tensor: + """Order-preserving int encoding of fp32 (an involution). + + Stored back in fp32 slots - matches the indexer's encoded-int atomic + min/max. + """ + bits = t.float().contiguous().view(torch.int32) + enc = torch.where(bits >= 0, bits, bits ^ 0x7FFFFFFF) + return enc.view(torch.float32) + + +def hit_agg_identities(num_rows: int, device) -> torch.Tensor: + """Identity-initialized per-row hit aggregate. + + {enc(+FLT_MAX), enc(-FLT_MAX), 0, 0} - the required initial state of + the buffer the indexer atomically merges into. + """ + ident = torch.tensor([_FLT_MAX, -_FLT_MAX], dtype=torch.float32, device=device) + enc = enc_ordered_f32(ident) + out = torch.zeros((num_rows, 4), dtype=torch.float32, device=device) + out[:, 0] = enc[0] + out[:, 1] = enc[1] + return out.contiguous() + + def gvr_topk_decode( logits: torch.Tensor, pre_idx: torch.Tensor, @@ -178,6 +670,25 @@ def gvr_topk_decode( order_row: Optional[torch.Tensor] = None, p4_warp_redundant: bool = True, p2_warp_redundant: bool = True, + pdl_wait_late: bool = True, + p4_tail_v3: "bool | None" = None, + p4_no_fine: "bool | None" = None, + p4_exact_tail: "bool | None" = None, + p4_tail_fast: "bool | None" = None, + p1r_rescue: bool = True, + num_bins: "int | None" = None, + p4_fine_rangetest: Optional[bool] = None, + p4_scat_rangetest: bool = False, + block_max: Optional[torch.Tensor] = None, + skip_min_n: Optional[int] = 200_000, + seed_thr: Optional[torch.Tensor] = None, + seed_counts: Optional[torch.Tensor] = None, + xstate: Optional[torch.Tensor] = None, + cand_vals: Optional[torch.Tensor] = None, + cand_idx: Optional[torch.Tensor] = None, + cand_ctl: Optional[torch.Tensor] = None, + self_scan: bool = False, + cap_c: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """CuTe DSL GVR Top-K wrapper with every tuning knob exposed. @@ -188,9 +699,9 @@ def gvr_topk_decode( Args: logits: ``[num_rows, max_S]`` float32 / bfloat16 / float16. pre_idx: ``[num_rows // next_n, pre_idx_count]`` int32. - ``pre_idx[..., 0]`` must be the argmax index — indexer invariant. + ``pre_idx[..., 0]`` must be the argmax index - indexer invariant. seq_lens: ``[num_rows // next_n]`` int32 (uncompressed-token space). - top_k: K ∈ {512, 1024, 2048} — compile-time specialized. + top_k: K in {512, 1024, 2048} - compile-time specialized. next_n: Temporal stride for V3.2 ``preIdxOffset = (row % next_n) + 1``. compress_ratio: KV-indexer compression factor (1 = DSv3.2, 4 = DSv4). When != 1, logits/preIdx live in compressed-token-index space: @@ -214,7 +725,7 @@ def gvr_topk_decode( round; every warp reduces the staged warp counts and replays the classify + secant update in registers. False restores the leader cadence. - order_row: Required iff ``seqlen_sorted=True``. Request-level — + order_row: Required iff ``seqlen_sorted=True``. Request-level - ``int32[batch_size = num_rows // next_n]`` on the same device as ``logits``; ``order_row[i]`` is the original request_id of the i-th-priority request. The kernel @@ -247,6 +758,113 @@ def gvr_topk_decode( cute_dtype = _DTYPE_TORCH_TO_CUTE[logits.dtype] num_rows = logits.shape[0] + # Host dispatch gate: below skip_min_n (compressed-index space, + # shape-based so no device sync) drop block_max; None disables the gate. + if block_max is not None and skip_min_n is not None and logits.shape[1] < skip_min_n: + block_max = None + # K > 512 at tiny batch: drop block_max (stock path) unless the admitted + # line is known up front (ext counts / packed seed / self_scan). + if ( + block_max is not None + and num_rows < 8 + and top_k > 512 + and not self_scan + and not (seed_thr is not None and seed_counts is not None) + and not (seed_thr is not None and seed_thr.shape[1] >= 6) + ): + block_max = None + if self_scan: + # fused self-contained mode: kernel scans/buckets the row itself. + # Inputs: seed_thr (three closed-loop lines) + a write-only gmem + # POSITION column passed through the cand_idx slot; seed_counts is + # a dummy (zeros never pass the [K, kC] admission). + assert seed_thr is not None, "self_scan requires seed_thr" + assert cand_vals is None and cand_ctl is None, ( + "self_scan excludes external candidate values/control" + ) + assert "GVR_BSTAR" in os.environ, ( + "self_scan requires GVR_BSTAR (accept_cap) to size the position column" + ) + if seed_counts is None: + seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=logits.device) + _bstar = int(os.environ["GVR_BSTAR"]) + _capc = cap_c if cap_c is not None else int(os.environ.get("GVR_CAPC", "16384")) + _segtot = 2 * _bstar + _capc + if cand_idx is None: + cand_idx = torch.empty((num_rows, _segtot), dtype=torch.int32, device=logits.device) + assert ( + cand_idx.dtype == torch.int32 + and cand_idx.is_cuda + and cand_idx.is_contiguous() + and cand_idx.shape == (num_rows, _segtot) + ), f"self_scan position column must be int32 [num_rows, {_segtot}]" + # packed seed row ([rows, >=6] fp32: lines + counts-as-floats) is the + # native ext-counts input; separate seed_counts is the compat path and + # pays a per-call pack build - pre-pack with pack_seed() instead. + pre_packed = seed_thr is not None and seed_thr.shape[1] >= 6 + use_ext_counts = seed_thr is not None and (seed_counts is not None or pre_packed) + # variant B (two-pass): thresholds without counts -> the kernel counts + # the rungs itself (stock R0 multi-count) and admits in-kernel + ext_rungs = seed_thr is not None and seed_counts is None and not pre_packed + if use_ext_counts: + assert ( + seed_thr.dtype == torch.float32 + and seed_thr.is_cuda + and seed_thr.is_contiguous() + and seed_thr.shape[0] == num_rows + and (pre_packed or seed_thr.shape[1] == 3) + ), "seed_thr must be contiguous CUDA fp32 [num_rows, 3|8]" + assert pre_packed or ( + seed_counts.dtype == torch.int32 + and seed_counts.is_cuda + and seed_counts.is_contiguous() + and seed_counts.shape == (num_rows, 3) + ), "seed_counts must be contiguous CUDA int32 [num_rows, 3]" + use_ext_cand = cand_vals is not None and cand_idx is not None and cand_ctl is not None + cand_cap = 5120 + if self_scan: + cand_cap = cand_idx.shape[1] + if use_ext_cand: + assert ( + cand_vals.dtype == torch.float32 + and cand_vals.is_cuda + and cand_vals.is_contiguous() + and cand_vals.dim() == 2 + and cand_vals.shape[0] == num_rows + ), "cand_vals must be contiguous CUDA fp32 [num_rows, CAP]" + assert ( + cand_idx.dtype == torch.int32 + and cand_idx.is_cuda + and cand_idx.is_contiguous() + and cand_idx.shape == cand_vals.shape + ), "cand_idx must be contiguous CUDA int32 [num_rows, CAP]" + assert ( + cand_ctl.dtype == torch.int32 + and cand_ctl.is_cuda + and cand_ctl.is_contiguous() + and cand_ctl.shape == (num_rows, 4) + ), "cand_ctl must be contiguous CUDA int32 [num_rows, 4]" + cand_cap = cand_vals.shape[1] + emit_xstate = xstate is not None + if emit_xstate: + assert ( + xstate.dtype == torch.float32 + and xstate.is_cuda + and xstate.is_contiguous() + and xstate.shape == (num_rows, 8) + ), "xstate must be contiguous CUDA fp32 [num_rows, 8]" + enable_block_skip = block_max is not None + if enable_block_skip: + assert ( + block_max.dtype == torch.float32 + and block_max.is_cuda + and block_max.is_contiguous() + and block_max.dim() == 2 + and block_max.shape[0] == num_rows + and block_max.shape[1] % 4 == 0 + and block_max.shape[1] >= (logits.shape[1] + 31) // 32 + ), "block_max must be contiguous CUDA fp32 [num_rows, nb_pad*4] covering the row" + if return_output_values: if out_values is None: out_values = torch.empty((num_rows, top_k), dtype=logits.dtype, device=logits.device) @@ -262,11 +880,20 @@ def gvr_topk_decode( N_cols = logits.shape[1] N_dec = max_seq_len if max_seq_len is not None else N_cols if num_threads_per_block is None: - if max_seq_len is not None and logits.dtype != torch.float32: - n_thresh_t = 131072 + if use_ext_cand and top_k <= 512: + # list-hit rows do O(list) work, not O(N): use 512 threads + # (K=1024 lists are big enough to keep the N-keyed pick). + num_threads_per_block = 512 + elif self_scan: + # self_scan scans the whole row in one CTA; the phase-0 + # cp.async pipeline scales with warp count at every N. + num_threads_per_block = 1024 else: - n_thresh_t = 65536 - num_threads_per_block = 1024 if (num_rows <= num_sms and N_dec >= n_thresh_t) else 512 + if max_seq_len is not None and logits.dtype != torch.float32: + n_thresh_t = 131072 + else: + n_thresh_t = 65536 + num_threads_per_block = 1024 if (num_rows <= num_sms and N_dec >= n_thresh_t) else 512 if use_256bit_load is None: use_256bit_load = logits.dtype == torch.float32 and N_dec >= 16384 if enable_warp_parallel_reduce is None: @@ -295,6 +922,13 @@ def gvr_topk_decode( else: min_blocks_per_mp = 1 + seed_pack = None + if use_ext_counts: + if pre_packed: + seed_pack = seed_thr + else: + seed_pack = pack_seed(seed_thr, seed_counts) + compiled = _compile( cute_dtype, top_k, @@ -312,6 +946,27 @@ def gvr_topk_decode( seqlen_sorted, p4_warp_redundant, p2_warp_redundant, + enable_block_skip, + pdl_wait_late, + p4_tail_v3, + p4_no_fine, + p4_exact_tail, + p4_tail_fast, + p1r_rescue, + num_bins, + p4_fine_rangetest, + p4_scat_rangetest, + use_ext_counts, + emit_xstate, + use_ext_cand, + ext_rungs, + cand_cap, + int(os.environ["GVR_BSTAR"]) if "GVR_BSTAR" in os.environ else None, + int(os.environ["GVR_KC"]) if "GVR_KC" in os.environ else None, + self_scan, + cap_c + if cap_c is not None + else (int(os.environ.get("GVR_CAPC", "16384")) if self_scan else None), ) # When return_output_values=False the kernel was compiled to skip # STG.value and accepts None for the value-output slot. @@ -324,6 +979,13 @@ def gvr_topk_decode( out_values if return_output_values else None, out_indices, order_row if seqlen_sorted else None, + block_max if enable_block_skip else None, + seed_pack if use_ext_counts else (seed_thr if ext_rungs else None), + None, + xstate if emit_xstate else None, + cand_vals if use_ext_cand else None, + cand_idx if (use_ext_cand or self_scan) else None, + cand_ctl if use_ext_cand else None, ) if return_output_values: return out_values, out_indices @@ -336,7 +998,7 @@ def gvr_topk_sort_prepare(seq_lens: torch.Tensor) -> torch.Tensor: Returns ``int32[batch_size]`` (= ``seq_lens.shape[0]`` = ``num_rows // next_n``) whose i-th entry is the original-batch index - of the i-th longest request — request-level, NOT row-level. The + of the i-th longest request - request-level, NOT row-level. The kernel expands to row level via ``order_row[req] * next_n + nn`` inside the const_expr ``seqlen_sorted`` branch. Run once per decode step; the same @@ -413,7 +1075,7 @@ def gvr_topk_lb_prepare( ) -> tuple[torch.Tensor, torch.Tensor]: """Run the LB prepare kernel. - ``seq_lens`` keeps its actual shape ``(batch_size,)`` — the kernel + ``seq_lens`` keeps its actual shape ``(batch_size,)`` - the kernel is compiled to match that exact shape; ``max_batch_size`` only determines the prepare kernel's block size and the ``order_row`` buffer length. ``long_threshold`` is in SCAN-LENGTH space @@ -428,8 +1090,8 @@ def gvr_topk_lb_prepare( """ assert seq_lens.is_cuda and seq_lens.dtype == torch.int32 # block_prefix_sum_kernel (used inside LB prepare) constraints: - # num_warps = max_batch_size / 32 must be > 1 and a power of 2 → - # max_batch_size ∈ {64, 128, 256, 512, 1024}. + # num_warps = max_batch_size / 32 must be > 1 and a power of 2 -> + # max_batch_size in {64, 128, 256, 512, 1024}. if not (64 <= max_batch_size <= 1024) or (max_batch_size & (max_batch_size - 1)) != 0: raise ValueError( f"max_batch_size must be a power of 2 in [64, 1024] " @@ -601,9 +1263,9 @@ def _make_inputs( """Build (logits, pre_idx, seq_lens) for a multi-row test. Shapes: - logits : [num_rows, N] — compressed-token-index space - pre_idx : [num_rows // next_n, top_k] — argmax in slot 0 (indexer invariant) - seq_lens: [num_rows // next_n] — UNCOMPRESSED-token space + logits : [num_rows, N] - compressed-token-index space + pre_idx : [num_rows // next_n, top_k] - argmax in slot 0 (indexer invariant) + seq_lens: [num_rows // next_n] - UNCOMPRESSED-token space Kernel divides ``seq_lens`` by ``compress_ratio`` internally. Setting ``seq_lens = N * cr`` makes the kernel's @@ -620,7 +1282,7 @@ def _make_inputs( logits_f32 = torch.randn(num_rows, N, dtype=torch.float32, device="cuda") * 2.0 logits = logits_f32.to(dtype) num_groups = num_rows // next_n - # argmax must come from the effective scan range, not full N — for + # argmax must come from the effective scan range, not full N - for # next_n>1 the kernel's row-0 N_eff is only (N - next_n + 1) cols. effective_len = N - next_n + 1 argmax_idx = logits[::next_n, :effective_len].argmax(dim=-1).int() @@ -668,7 +1330,7 @@ def _tie_aware_correct( actual_kv_len = int(seq_lens_host[row // next_n]) - next_n + ofs + 1 N_eff = actual_kv_len // compress_ratio if N_eff < top_k: - # Degenerate path — skip; caller's main() guards against this. + # Degenerate path - skip; caller's main() guards against this. continue row_logits = logits_f32[row, :N_eff] topk_vals, _ = torch.topk(row_logits, k=top_k, largest=True, sorted=True) @@ -726,7 +1388,7 @@ def test_gvr_topk_decode( ) -> None: # Kernel scans `N_eff = seq_lens[0] - next_n + (row_idx % next_n) + 1` # columns. Smallest row's N_eff = N - next_n + 1. Degenerate path - # (N_eff <= top_k) is a separate code branch — skip here. + # (N_eff <= top_k) is a separate code branch - skip here. if N - next_n + 1 < top_k: pytest.skip("N_eff < top_k is degenerate; the kernel requires N_eff >= top_k") seed = 42 diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py index 8dae8201ba89..7f0142af1fb8 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py @@ -547,6 +547,550 @@ def test_cute_dsl_fp4_paged_mqa_logits( ) +# --------------------------------------------------------------------------- +# Block-meta emission (emit_block_meta — fused-GVR support). +# --------------------------------------------------------------------------- + +_FLT_MAX_F32 = torch.finfo(torch.float32).max + + +def _enc_ordered_f32(t: torch.Tensor) -> torch.Tensor: + """Order-preserving int encoding of fp32 (involution; also decodes).""" + bits = t.float().contiguous().view(torch.int32) + enc = torch.where(bits >= 0, bits, bits ^ 0x7FFFFFFF) + return enc.view(torch.float32) + + +def _hit_agg_identities(num_rows: int, device) -> torch.Tensor: + ident = torch.tensor([_FLT_MAX_F32, -_FLT_MAX_F32], dtype=torch.float32, device=device) + enc = _enc_ordered_f32(ident) + out = torch.zeros((num_rows, 4), dtype=torch.float32, device=device) + out[:, 0] = enc[0] + out[:, 1] = enc[1] + return out.contiguous() + + +def _pack_hit_bitmap( + pre_idx: torch.Tensor, batch_size: int, num_words: int, device +) -> torch.Tensor: + """[B, num_words] int32; bit (pos % 32) of word (pos // 32) set per + valid pre_idx entry — the kernel's hit test layout.""" + bitmap = torch.zeros((batch_size, num_words), dtype=torch.int64, device=device) + for b in range(batch_size): + idx = pre_idx[b].to(torch.int64).unique() + idx = idx[(idx >= 0) & (idx < num_words * 32)] + bitmap[b].scatter_add_(0, idx >> 5, torch.ones_like(idx) << (idx & 31)) + # int64 -> int32 with bit-31 wraparound (torch refuses the overflow). + wrapped = bitmap & 0xFFFFFFFF + wrapped = torch.where(wrapped >= 2**31, wrapped - 2**32, wrapped) + return wrapped.to(torch.int32) + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 3]) +# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile. +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("fix_length", [True, False]) +@pytest.mark.parametrize("emit_hit_stats", [True, False]) +def test_cute_dsl_fp4_paged_mqa_logits_block_meta( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + fix_length, + emit_hit_stats, +): + """emit_block_meta correctness: block_max / hit_stats recomputed from + the KERNEL'S OWN logits output (fp4 numerics differ from the torch + reference logits, but the meta contract is defined on what the kernel + stores). NaN-prefilled buffers prove no writes land outside + [0, num_kv (+1 when odd)) per row.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner + + torch.manual_seed(7) + torch.cuda.manual_seed(7) + num_heads, head_dim, top_k = 64, 128, 512 + max_model_len = max(avg_ctx * 2, 2048) + device = "cuda" + + if fix_length: + context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device) + else: + lo = max(phys_block_kv, int(0.7 * avg_ctx)) + context_lens = torch.randint( + lo, int(1.3 * avg_ctx) + 1, (batch_size,), dtype=torch.int32, device=device + ).clamp(max=max_model_len) + + num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv) + num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2 + max_blocks_per_seq = int(num_blocks_per_seq.max().item()) + block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device) + pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32) + off = 0 + for i, n_blks in enumerate(num_blocks_per_seq.tolist()): + block_table[i, :n_blks] = pool[off : off + n_blks] + off += n_blks + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16) + kv_cache = torch.randn( + (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16 + ) + weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32) + + q_packed, sf_q_packed = per_token_cast_to_fp4( + q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2) + sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads) + remove_online_sf_transpose = phys_block_kv == 128 + kv_fused, _ = kv_cache_cast_to_fp4( + kv_cache, remove_online_sf_transpose=remove_online_sf_transpose + ) + + DG_METADATA_BLOCK_KV = 64 + num_sms = deep_gemm.get_num_sms() + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms + ) + + # pre_idx per request within [0, ctx) -> packed bitmap. + aligned_max_ctx = align(max_model_len, 256) + nb_pad = aligned_max_ctx // 128 + pre_idx = torch.zeros((batch_size, top_k), dtype=torch.int32, device=device) + for b in range(batch_size): + pre_idx[b] = torch.randint( + 0, int(context_lens[b].item()), (top_k,), dtype=torch.int32, device=device + ) + bitmap = _pack_hit_bitmap(pre_idx, batch_size, nb_pad * 4, device) + + # block_max: 4 warp-partial records per block; consumers fold. NaN + # prefill proves write coverage is exactly [0, written_hi*4) per row. + # hit_stats: per-row aggregate the kernel atomically merges into — + # MUST be identity-initialized by the caller. + nan = float("nan") + block_max = torch.full( + (batch_size * next_n, nb_pad * 4), nan, dtype=torch.float32, device=device + ) + hit_stats = _hit_agg_identities(batch_size * next_n, device) + + meta_kwargs = dict( + emit_block_meta=True, + emit_hit_stats=emit_hit_stats, + block_max_out=block_max, + ) + if emit_hit_stats: + meta_kwargs.update(hit_bitmap=bitmap, hit_stats_out=hit_stats) + logits, bm, hs = CuteDSLFP4PagedMQALogitsRunner.forward( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + num_epi_subtiles=1, + epi_dtype=torch.float32, + output_dtype=torch.bfloat16, + remove_online_sf_transpose=remove_online_sf_transpose, + **meta_kwargs, + ) + torch.cuda.synchronize() + + lf = logits.float() + for row in range(batch_size * next_n): + req = row // next_n + ctx = int(context_lens[req].item()) + num_kv = ceil_div(ctx, 128) + tag = f"row={row} req={req} ctx={ctx} next_n={next_n} pbk={phys_block_kv}" + + # Fold the kernel's 4 warp-partials per block (the consumer-side + # contract); per-warp partials themselves depend on the TMEM + # lane->row mapping and are not checked individually. + bm_fold = bm[row].view(nb_pad, 4).amax(-1) + + # block_max reference from the kernel's own stored logits. + padded = torch.full((nb_pad * 128,), -_FLT_MAX_F32, device=device) + padded[:ctx] = lf[row, :ctx] + ref_bmax = padded.view(nb_pad, 128).amax(-1) + torch.testing.assert_close( + bm_fold[:num_kv], + ref_bmax[:num_kv], + atol=0.0, + rtol=0.0, + msg=lambda m, tag=tag: f"block_max mismatch: {tag}\n{m}", + ) + + if emit_hit_stats: + # Per-row hit aggregate reference (bitmap semantics: dedup + + # pos < ctx). min/max slots are encoded (involution decodes). + idx = pre_idx[req].to(torch.int64).unique() + idx = idx[(idx >= 0) & (idx < ctx)] + got_min = _enc_ordered_f32(hs[row, 0:1])[0] + got_max = _enc_ordered_f32(hs[row, 1:2])[0] + got_sum = hs[row, 2] + got_cnt = hs[row, 3] + if idx.numel() > 0: + vals = lf[row, idx] + assert got_min.item() == vals.min().item(), f"hit_min: {tag}" + assert got_max.item() == vals.max().item(), f"hit_max: {tag}" + # Atomic-add merge order vs torch sum order: fp slack. + torch.testing.assert_close( + got_sum, + vals.sum(), + atol=1e-2, + rtol=1e-4, + msg=lambda m, tag=tag: f"hit_sum mismatch: {tag}\n{m}", + ) + assert got_cnt.item() == float(idx.numel()), f"hit_cnt: {tag}" + else: + assert got_min.item() == _FLT_MAX_F32, f"identity min: {tag}" + assert got_max.item() == -_FLT_MAX_F32, f"identity max: {tag}" + assert got_cnt.item() == 0.0, f"identity cnt: {tag}" + + # Odd num_kv: WG1's OOB tile writes pure identities into block + # slot num_kv (every lane invalid). + written_hi = num_kv + (num_kv % 2) + if written_hi > num_kv: + assert bm_fold[num_kv].item() == -_FLT_MAX_F32, tag + # No stray writes past the padding tile: NaN prefill intact. + assert bm[row, written_hi * 4 :].isnan().all(), f"stray block_max write: {tag}" + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 3]) +# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile. +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("fix_length", [True, False]) +@pytest.mark.parametrize("packed", [False, True]) +def test_cute_dsl_fp4_paged_mqa_logits_seed_counts( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + fix_length, + packed, +): + """emit_seed_counts exactness: per-row counts of logits >= threshold + recomputed from the KERNEL'S OWN logits output (the count contract is + defined on post-conversion values, same as block_max). Thresholds are + per-row quantiles of the row's own logits so each of the 3 counters + lands in a different regime (loose/mid/tight).""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner + + torch.manual_seed(11) + torch.cuda.manual_seed(11) + num_heads, head_dim = 64, 128 + max_model_len = max(avg_ctx * 2, 2048) + device = "cuda" + + if fix_length: + context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device) + else: + lo = max(phys_block_kv, int(0.7 * avg_ctx)) + context_lens = torch.randint( + lo, int(1.3 * avg_ctx) + 1, (batch_size,), dtype=torch.int32, device=device + ).clamp(max=max_model_len) + + num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv) + num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2 + max_blocks_per_seq = int(num_blocks_per_seq.max().item()) + block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device) + pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32) + off = 0 + for i, n_blks in enumerate(num_blocks_per_seq.tolist()): + block_table[i, :n_blks] = pool[off : off + n_blks] + off += n_blks + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16) + kv_cache = torch.randn( + (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16 + ) + weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32) + + q_packed, sf_q_packed = per_token_cast_to_fp4( + q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2) + sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads) + remove_online_sf_transpose = phys_block_kv == 128 + kv_fused, _ = kv_cache_cast_to_fp4( + kv_cache, remove_online_sf_transpose=remove_online_sf_transpose + ) + + DG_METADATA_BLOCK_KV = 64 + num_sms = deep_gemm.get_num_sms() + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms + ) + + aligned_max_ctx = align(max_model_len, 256) + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + + # First pass without seed counts to harvest per-row logits for + # threshold picking (post-conversion value domain). + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + common = dict( + num_epi_subtiles=1, + epi_dtype=torch.float32, + output_dtype=torch.bfloat16, + remove_online_sf_transpose=remove_online_sf_transpose, + ) + logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + **common, + ) + torch.cuda.synchronize() + lf0 = logits0.float() + + seed_thr = torch.empty((num_rows, 3), dtype=torch.float32, device=device) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + vals = lf0[row, :ctx] + # Loose / mid / tight rungs; ties on exact stored values are the + # point (>= must count them all). + seed_thr[row, 0] = torch.quantile(vals, 0.10) + seed_thr[row, 1] = torch.quantile(vals, 0.90) + seed_thr[row, 2] = torch.quantile(vals, 0.998) + + if packed: + # Packed contract: one [rows, 8] fp32 seed row, lines at cols + # 0..2, counts accumulate as fp32 at cols 3..5 (caller zeroes). + seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device) + seed_row[:, 0:3] = seed_thr + thr_arg, counts_arg = seed_row, None + else: + seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device) + thr_arg, counts_arg = seed_thr, seed_counts + block_max.fill_(nan) + logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + emit_seed_counts=True, + seed_thr=thr_arg, + seed_counts_out=counts_arg, + **common, + ) + torch.cuda.synchronize() + + lf = logits.float() + # compare valid prefixes only: past ctx the buffer is unwritten + # allocator garbage and differs run-to-run + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + torch.testing.assert_close(lf[row, :ctx], lf0[row, :ctx], atol=0.0, rtol=0.0) + if packed: + assert torch.equal(seed_row[:, 0:3], seed_thr), "lines clobbered" + # col 6 carries the adaptive-skip pass count (lane0-accumulated + # diagnostic; the top-k consumer reads it when block_max rides + # along), so it is a legitimate output here - bounded by the + # block-max record count. Only col 7 must stay untouched. + assert (seed_row[:, 7] == 0).all(), "stray write past counts" + nrec = block_max.shape[1] + assert ((seed_row[:, 6] >= 0) & (seed_row[:, 6] <= nrec)).all(), ( + "adaptive-skip pass count out of range" + ) + seed_counts = seed_row[:, 3:6].to(torch.int32) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + tag = f"row={row} ctx={ctx} next_n={next_n} pbk={phys_block_kv}" + ref = (lf[row, :ctx].unsqueeze(0) >= seed_thr[row].unsqueeze(1)).sum(-1) + got = seed_counts[row].to(torch.int64) + assert torch.equal(got.cpu(), ref.cpu().to(torch.int64)), ( + f"seed_counts mismatch: {tag} got={got.tolist()} ref={ref.tolist()}" + ) + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 3]) +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("cap_mode", ["roomy", "tight"]) +def test_cute_dsl_fp4_paged_mqa_logits_cand( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + cap_mode, +): + """emit_cand correctness: the unordered (value, index) pre-collect at + t_0 must contain EXACTLY the set {i < ctx : logits[r, i] >= t_0} when + it fits (void == 0), and degrade safely on overflow (void == 1, all + written slots valid + unique, claimed and counts[0] still exact).""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner + + torch.manual_seed(13) + torch.cuda.manual_seed(13) + num_heads, head_dim = 64, 128 + max_model_len = max(avg_ctx * 2, 2048) + device = "cuda" + + context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device) + num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv) + num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2 + max_blocks_per_seq = int(num_blocks_per_seq.max().item()) + block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device) + pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32) + off = 0 + for i, n_blks in enumerate(num_blocks_per_seq.tolist()): + block_table[i, :n_blks] = pool[off : off + n_blks] + off += n_blks + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16) + kv_cache = torch.randn( + (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16 + ) + weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32) + q_packed, sf_q_packed = per_token_cast_to_fp4( + q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2) + sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads) + remove_online_sf_transpose = phys_block_kv == 128 + kv_fused, _ = kv_cache_cast_to_fp4( + kv_cache, remove_online_sf_transpose=remove_online_sf_transpose + ) + DG_METADATA_BLOCK_KV = 64 + num_sms = deep_gemm.get_num_sms() + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms + ) + aligned_max_ctx = align(max_model_len, 256) + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + common = dict( + num_epi_subtiles=1, + epi_dtype=torch.float32, + output_dtype=torch.bfloat16, + remove_online_sf_transpose=remove_online_sf_transpose, + ) + base_args = ( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) + + # Pass 1: harvest logits for threshold picking. + logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + **common, + ) + torch.cuda.synchronize() + lf0 = logits0.float() + + seed_thr = torch.empty((num_rows, 3), dtype=torch.float32, device=device) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + vals = lf0[row, :ctx] + seed_thr[row, 0] = torch.quantile(vals, 0.90) # t_0: ~10% of ctx + seed_thr[row, 1] = torch.quantile(vals, 0.97) + seed_thr[row, 2] = torch.quantile(vals, 0.998) + + # Window claiming over-claims by up to ~CAND_WIN per epilogue warp + # touching the row (sentinel-filled tails); B=1 rows spread over many + # CTAs, so roomy needs slack well beyond the ~410 true hits. + cap = 4096 if cap_mode == "roomy" else 128 + seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device) + cand = torch.full((num_rows, cap * 2), -1, dtype=torch.int32, device=device) + ctl = torch.zeros((num_rows, 2), dtype=torch.int32, device=device) + block_max.fill_(nan) + logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + emit_seed_counts=True, + seed_thr=seed_thr, + seed_counts_out=seed_counts, + emit_cand=True, + cand_out=cand, + cand_ctl_out=ctl, + **common, + ) + torch.cuda.synchronize() + lf = logits.float() + torch.testing.assert_close(lf, lf0, atol=0.0, rtol=0.0) + + pairs = cand.view(num_rows, cap, 2) + vals_bits = pairs[..., 0] + idxs = pairs[..., 1] + vals = vals_bits.view(torch.float32) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + t0 = seed_thr[row, 0] + ref_mask = lf[row, :ctx] >= t0 + ref_count = int(ref_mask.sum()) + ref_idx = set(torch.nonzero(ref_mask, as_tuple=False).flatten().tolist()) + tag = f"row={row} ctx={ctx} cap={cap} ref={ref_count} mode={cap_mode}" + claimed = int(ctl[row, 0]) + void = int(ctl[row, 1]) + # counts[0] is the exact count regardless of windows/overflow; + # claimed >= true count (sentinel-padded window tails). + assert int(seed_counts[row, 0]) == ref_count, f"counts0: {tag}" + assert claimed >= ref_count, f"claimed < true count: {tag} got={claimed}" + n_written = min(claimed, cap) + got_idx = idxs[row, :n_written].long() + live = got_idx >= 0 + got_list = got_idx[live].tolist() + assert len(set(got_list)) == len(got_list), f"duplicate idx: {tag}" + assert set(got_list).issubset(ref_idx), f"non-member idx: {tag}" + # pair integrity on live entries: value word == stored logit bits. + live_idx = got_idx[live] + torch.testing.assert_close( + vals[row, :n_written][live], + lf[row, live_idx], + atol=0.0, + rtol=0.0, + msg=lambda m, tag=tag: f"pair value mismatch: {tag}\n{m}", + ) + if cap_mode == "roomy": + assert void == 0, f"void set without overflow: {tag} claimed={claimed}" + assert claimed <= cap, f"claimed past cap without void: {tag}" + assert set(got_list) == ref_idx, f"set mismatch: {tag}" + # every claimed slot is live or sentinel; unclaimed tail untouched + assert bool((idxs[row, claimed:] == -1).all()), f"stray write: {tag}" + assert int(live.sum()) == ref_count, f"live count: {tag}" + else: + assert ref_count > cap, f"test setup wants overflow: {tag}" + assert void == 1, f"void not set on overflow: {tag}" + + # --------------------------------------------------------------------------- # Benchmarking entry point (run module directly). # --------------------------------------------------------------------------- @@ -1043,3 +1587,174 @@ def dg_fn(data=data, dg_ctx_2d=dg_ctx_2d, q_fp4_dg=q_fp4_dg): varlen=args.varlen, block_kv=args.block_kv, ) + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("cap_mode", ["roomy", "tight"]) +def test_cute_dsl_fp4_paged_mqa_logits_cand_bucketed( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + cap_mode, +): + """emit_cand_bucketed (v5 SoA contract): three fixed segments with + pad-free A/B prefixes and spill-to-looser, ctl {n0, void, n1, n2} + with n1/n2 mirrored from the seed counters, C-window pads carrying + score -inf / idx -1. All invariants recomputed from the kernel's + own logits.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLFP4PagedMQALogitsRunner + + torch.manual_seed(23) + torch.cuda.manual_seed(23) + num_heads, head_dim = 64, 128 + max_model_len = max(avg_ctx * 2, 2048) + device = "cuda" + context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device) + num_blocks_per_seq = ceil_div_tensor(context_lens, phys_block_kv) + num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2 + max_blocks_per_seq = int(num_blocks_per_seq.max().item()) + block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device) + pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32) + off = 0 + for i, n_blks in enumerate(num_blocks_per_seq.tolist()): + block_table[i, :n_blks] = pool[off : off + n_blks] + off += n_blks + q = torch.randn((batch_size, next_n, num_heads, head_dim), device=device, dtype=torch.bfloat16) + kv_cache = torch.randn( + (num_total_blocks, phys_block_kv, 1, head_dim), device=device, dtype=torch.bfloat16 + ) + weights = torch.randn((batch_size * next_n, num_heads), device=device, dtype=torch.float32) + q_packed, sf_q_packed = per_token_cast_to_fp4( + q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + q_fp4 = q_packed.view(torch.uint8).view(batch_size, next_n, num_heads, head_dim // 2) + sf_q = sf_q_packed.view(torch.int32).view(batch_size, next_n, num_heads) + remove_online_sf_transpose = phys_block_kv == 128 + kv_fused, _ = kv_cache_cast_to_fp4( + kv_cache, remove_online_sf_transpose=remove_online_sf_transpose + ) + DG_METADATA_BLOCK_KV = 64 + num_sms = deep_gemm.get_num_sms() + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms + ) + aligned_max_ctx = align(max_model_len, 256) + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + common = dict( + num_epi_subtiles=1, + epi_dtype=torch.float32, + output_dtype=torch.bfloat16, + remove_online_sf_transpose=remove_online_sf_transpose, + ) + logits0, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + **common, + ) + torch.cuda.synchronize() + lf0 = logits0.float() + seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + vals = lf0[row, :ctx] + seed_row[row, 0] = torch.quantile(vals, 0.60) + seed_row[row, 1] = torch.quantile(vals, 0.90) + seed_row[row, 2] = torch.quantile(vals, 0.99) + # segment caps: roomy fits everything; tight forces A/B spill and a + # C-window void + if cap_mode == "roomy": + segA, capC = 2048, 4096 + else: + segA, capC = 32, 128 + W = 2 * segA + capC + cand_vals = torch.full((num_rows, W), nan, dtype=torch.float32, device=device) + cand_idx = torch.full((num_rows, W), -7, dtype=torch.int32, device=device) + cand_ctl = torch.zeros((num_rows, 4), dtype=torch.int32, device=device) + cand_cur = torch.zeros((num_rows, 4), dtype=torch.int32, device=device) + block_max.fill_(nan) + logits, _, _ = CuteDSLFP4PagedMQALogitsRunner.forward( + q_fp4, + sf_q, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + emit_hit_stats=False, + block_max_out=block_max, + emit_seed_counts=True, + seed_thr=seed_row, + emit_cand_bucketed=True, + accept_cap=segA, + cand_out=cand_vals, + cand_idx_out=cand_idx, + cand_ctl_out=cand_ctl, + cand_cur_out=cand_cur, + **common, + ) + torch.cuda.synchronize() + lf = logits.float() + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + t0, t1, t2 = (float(seed_row[row, j]) for j in range(3)) + v = lf[row, :ctx] + n0_ref = int((v >= t0).sum()) + n1_ref = int((v >= t1).sum()) + n2_ref = int((v >= t2).sum()) + n0c, voidc, n1c, n2c = (int(cand_ctl[row, j]) for j in range(4)) + tag = f"row={row} caps=({segA},{capC}) refs=({n0_ref},{n1_ref},{n2_ref})" + assert n1c == n1_ref and n2c == n2_ref, f"n1/n2 mismatch {tag} got {n1c},{n2c}" + curA, curB, curC = (int(cand_cur[row, j]) for j in range(3)) + lenA = min(n2_ref, segA) + lenB = min(n1_ref - n2_ref + max(n2_ref - segA, 0), segA) + assert min(curA, segA) >= lenA or curA == n2_ref, f"curA {curA} {tag}" + # A prefix: pad-free, every entry >= t2, positions valid + unique + pa = cand_idx[row, :lenA] + va = cand_vals[row, :lenA] + assert (pa >= 0).all() and (pa < ctx).all(), f"A idx {tag}" + assert (va >= t2).all(), f"A vals {tag}" + got_a = lf[row, pa.long()] + torch.testing.assert_close(got_a, va, atol=0.0, rtol=0.0) + # B prefix: pad-free, [t1, t2) or A-spill (>= t2) + pb = cand_idx[row, segA : segA + lenB] + vb = cand_vals[row, segA : segA + lenB] + assert (pb >= 0).all() and (pb < ctx).all(), f"B idx {tag}" + assert (vb >= t1).all(), f"B vals {tag}" + torch.testing.assert_close(lf[row, pb.long()], vb, atol=0.0, rtol=0.0) + if voidc == 0: + # full coverage: union of live entries == the >= t0 set + lenC = n0c - lenA - lenB + pc = cand_idx[row, 2 * segA : 2 * segA + lenC] + vc = cand_vals[row, 2 * segA : 2 * segA + lenC] + live = pc >= 0 + assert (vc[live] >= t0).all(), f"C vals {tag}" + # pads carry -FLT_MAX (never ranks; the emu uses -inf, the + # kernel the finite sentinel - both satisfy the contract) + assert (vc[~live] <= -3e38).all(), f"C pads {tag}" + allp = torch.cat([pa, pb, pc[live]]) + assert allp.unique().numel() == allp.numel() == n0_ref, ( + f"coverage {tag}: {allp.unique().numel()} vs {n0_ref}" + ) + else: + assert cap_mode == "tight", f"unexpected void {tag}" + if cap_mode == "tight": + assert int(cand_ctl[:, 1].sum()) > 0, "tight caps never voided" diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py index 72afad0e6e49..608dcfc0addb 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py @@ -1272,6 +1272,466 @@ def test_cute_dsl_gvr_topk_decode_launch_autoconfig(dtype, top_k, N, batch_size, _assert_index_sets_equal_tie_aware(out, out_sec, logits) +# =========================================================================== +# Degenerate preIdx states (P1r data reseed). +# +# The stock seed path derives the P2 refine bracket from the preIdx gather; +# its exactness invariant (count(>= gather min) >= K) holds only when the +# preIdx row carries K DISTINCT in-range positions. Production-reachable +# violations: the first decode step of a request feeds the zero-init +# prev_topk feedback buffer (all-duplicate index 0); a reused batch slot can +# carry stale indices past the new row's N_eff (all-invalid); FP4-quantized +# logits can tie the whole gather (zero-width bracket). The old degenerate +# shortcut emitted identity indices [0, K) — NOT the top-K on real data. +# P1r rebuilds the bracket from the row itself, so all cases must pass the +# strict multiset check. +# =========================================================================== + + +@skip_not_sm100 +@pytest.mark.parametrize( + "dtype,top_k", + [ + (torch.bfloat16, 512), + (torch.float32, 2048), + ], +) +@pytest.mark.parametrize("compress_ratio", [1, 4]) +@pytest.mark.parametrize("pre_mode", ["zero", "dup", "oob"]) +@pytest.mark.parametrize("data_mode", ["random", "all_tied", "tie_flood"]) +def test_cute_dsl_gvr_topk_decode_degenerate_preidx( + dtype, top_k, compress_ratio, pre_mode, data_mode, tie_aware_check +): + """Degenerate preIdx rows must still produce an exact top-K.""" + N = 4096 + num_rows = 4 + torch.manual_seed(3) + torch.cuda.manual_seed(3) + if data_mode == "random": + logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(dtype) + elif data_mode == "all_tied": + # rescue re-degenerates (row min == max) -> identity output, which + # is exact here because every value is identical + logits = torch.full((num_rows, N), 5.0, device="cuda", dtype=dtype) + else: # tie_flood: kth sits inside a large tie class (within kC — + # count(>= tie value) must stay under the candidate capacity; the + # beyond-kC flood is a known pre-existing limitation, see the + # xfail test below) + logits = (torch.rand(num_rows, N, device="cuda") * 0.5).to(dtype) + # tie class sized to 1.5*K so the case stays a genuine flood for + # every top_k parametrization (2*K would cover the whole row at + # top_k=2048/N=4096 and collapse into all_tied) while + # count(>= 1.0) = 1.75*K stays inside the candidate capacity + tie_n = top_k + top_k // 2 + for r in range(num_rows): + perm = torch.randperm(N, device="cuda") + logits[r, perm[:tie_n]] = 1.0 + logits[r, perm[tie_n : tie_n + top_k // 4]] = 2.0 + + if pre_mode == "zero": + pre_idx = torch.zeros(num_rows, top_k, dtype=torch.int32, device="cuda") + elif pre_mode == "dup": + pre_idx = torch.full((num_rows, top_k), 37, dtype=torch.int32, device="cuda") + else: # oob: every slot past N_eff (stale-slot state, pcnt == 0) + pre_idx = torch.full((num_rows, top_k), N + 7, dtype=torch.int32, device="cuda") + + seq_lens = torch.full((num_rows,), N * compress_ratio, dtype=torch.int32, device="cuda") + out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + next_n=1, + compress_ratio=compress_ratio, + ) + torch.cuda.synchronize() + _gvr_check( + tie_aware_check, out_indices, logits, seq_lens, top_k, 1, compress_ratio=compress_ratio + ) + + +@skip_not_sm100 +def test_cute_dsl_gvr_topk_decode_degenerate_preidx_cs4(tie_aware_check): + """cs>1 rows run the rescue per-CTA (redundant full-row scan) — the + cluster path must stay exact for the zero-init cold-start state too.""" + N = 65536 + top_k = 512 + num_rows = 2 + torch.manual_seed(5) + logits = (torch.randn(num_rows, N, device="cuda") * 2.0).to(torch.bfloat16) + pre_idx = torch.zeros(num_rows, top_k, dtype=torch.int32, device="cuda") + seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda") + out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + next_n=1, + compress_ratio=1, + cluster_size=4, + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1, compress_ratio=1) + + +@skip_not_sm100 +@pytest.mark.xfail( + reason="known pre-existing limitation (upstream lineage, reproduces on " + "the PR-tip kernel unmodified): when the kth tie class alone exceeds " + "the candidate capacity kC, no threshold lands in [K, kC]; the selected " + "VALUE multiset is still exact but the index list can contain " + "duplicate / unwritten (-1) slots. Requires >kC exactly-equal scores " + "at the boundary — unreachable for real FP4 indexer logits observed " + "so far. Tracked as a follow-up; independent of the P1r rescue.", + strict=False, +) +def test_cute_dsl_gvr_topk_decode_tie_flood_beyond_capacity(tie_aware_check): + N = 4096 + top_k = 512 + num_rows = 4 + torch.manual_seed(3) + torch.cuda.manual_seed(3) + logits = torch.ones(num_rows, N, device="cuda", dtype=torch.bfloat16) + for r in range(num_rows): + hot = torch.randperm(N, device="cuda")[: top_k // 4] + logits[r, hot] = 2.0 + # healthiest possible pre (true previous top-k) — the flood defect is + # independent of preIdx quality + pre_idx = torch.topk(logits.float(), top_k, dim=-1).indices.int().contiguous() + seq_lens = torch.full((num_rows,), N, dtype=torch.int32, device="cuda") + out_indices = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + next_n=1, + compress_ratio=1, + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1, compress_ratio=1) + + +# --------------------------------------------------------------------------- +# Emission-assisted (ext) tiers: packed seed row / candidate list / block max. +# Inputs emulate the indexer epilogue host-side against the layout contracts +# in ``gvr_emission`` (segments at bases 0 / LIST_SEG_A / 2*LIST_SEG_A, packed row +# = lines at [0..2] + exact counts at [3..5] + skip pass count at [6]). +# --------------------------------------------------------------------------- + +_FLT_MAX = 3.4028234663852886e38 + + +def _lines_at_counts(logits_f32, n_eff, targets): + """Per-row threshold lines placed at exact counts (descending targets + -> ascending line values). count(logits >= line[j]) == targets[j].""" + num_rows = logits_f32.shape[0] + lines = torch.empty((num_rows, len(targets)), dtype=torch.float32, device=logits_f32.device) + for r in range(num_rows): + ne = int(n_eff[r]) + row = logits_f32[r, :ne] + for j, c in enumerate(targets): + c = min(int(c), ne) + lines[r, j] = torch.kthvalue(row, ne - c + 1).values + return lines + + +def _pack_seed_row(logits_f32, n_eff, lines, block_max=None): + """[rows, 8] fp32 packed seed row: lines + exact counts (+ skip count).""" + num_rows, N = logits_f32.shape + pos = torch.arange(N, device=logits_f32.device)[None, :] + valid = pos < n_eff[:, None] + counts = torch.stack( + [((logits_f32 >= lines[:, j : j + 1]) & valid).sum(-1) for j in range(3)], 1 + ).int() + pack = torch.zeros((num_rows, 8), dtype=torch.float32, device=logits_f32.device) + pack[:, 0:3] = lines + pack[:, 3:6] = counts.float() + if block_max is not None: + pack[:, 6] = (block_max >= lines[:, 0:1]).sum(dim=1).float() + return pack.contiguous() + + +@skip_not_sm100 +# One K only: the four modes exercise count-based admission, which is +# K-independent, while K is a compile-time parameter - sweeping it cost +# two extra kernel compiles (~43s) for no extra coverage. +@pytest.mark.parametrize("top_k", [2048]) +@pytest.mark.parametrize("mode", ["band", "fat", "miss", "inf"]) +def test_cute_dsl_gvr_topk_decode_ext_counts(top_k, mode, tie_aware_check): + """Packed seed row ([rows, 8]: lines + exact counts) consumption. + + band: one line's count sits inside the admission band (direct path). + fat: every count overshoots the candidate capacity -> full fallback. + miss: lines above the row max (count 0) -> seed rejected. + inf: non-finite lines (production cold start) -> validity guard. + """ + N, batch = 131072, 4 + logits, pre_idx, seq_lens = _make_inputs(batch, N, top_k, torch.float32, 1, seed=7, varlen=True) + n_eff = seq_lens.to(device=logits.device, dtype=torch.long) + if mode == "band": + lines = _lines_at_counts(logits, n_eff, (4 * top_k, 2 * top_k, top_k + top_k // 4)) + elif mode == "fat": + lines = _lines_at_counts(logits, n_eff, (32768, 24576, 16384)) + elif mode == "miss": + pos = torch.arange(N, device=logits.device)[None, :] + rowmax = torch.where(pos < n_eff[:, None], logits, float("-inf")).amax(-1) + lines = rowmax[:, None] + torch.tensor([1.0, 2.0, 3.0], device=logits.device) + else: + lines = torch.full((batch, 3), float("inf"), device=logits.device) + seed_row = _pack_seed_row(logits, n_eff, lines) + xstate = torch.zeros((batch, 8), dtype=torch.float32, device=logits.device) + out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + cluster_size=1, + seed_thr=seed_row, + xstate=xstate, + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1) + if mode == "band": + # the closed loop must republish valid state for the next step + assert bool((xstate[:, 0] > 0).all().item()) + + +@skip_not_sm100 +@pytest.mark.parametrize("mode", ["hit", "pads", "hist", "void", "bucketed", "starved"]) +def test_cute_dsl_gvr_topk_decode_ext_list(mode, tie_aware_check): + """Candidate-list tier at the production geometry (accept_cap = + LIST_SEG_A, width = LIST_WIDTH). + + hit: parked lines (production shape), claimed inside [K+64, B*] + -> line cut, single mapped load. + pads: same + interleaved idx=-1 window sentinels (claimed counts + them; the K+64 slack absorbs them). + hist: claimed past B* but list complete -> clamped-histogram + fallback over segment C. + void: collection overflows LIST_CAP_C -> void=1 -> full scan. + bucketed: three live lines spread across segments A/B/C, cut at the + tightest line inside the band. + starved: fewer than K real candidates, but sentinel pads lift the + claim into the admission band -> the line-cut copy must + re-measure and demote (exactness regression). + """ + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import ( + LIST_CAP_C, + LIST_PARK_LINE, + LIST_SEG_A, + LIST_WIDTH, + ) + + top_k, N, batch = 512, 131072, 2 + logits, pre_idx, seq_lens = _make_inputs( + batch, N, top_k, torch.float32, 1, seed=11, varlen=True + ) + dev = logits.device + n_eff = seq_lens.to(device=dev, dtype=torch.long) + if mode == "bucketed": + lines = _lines_at_counts(logits, n_eff, (20000, 4000, 600)) + else: + n0 = {"hit": 4096, "pads": 4096, "hist": 12000, "void": 30000, "starved": 400}[mode] + l0 = _lines_at_counts(logits, n_eff, (n0,)) + lines = torch.cat( + [l0, torch.full_like(l0, LIST_PARK_LINE), torch.full_like(l0, 2 * LIST_PARK_LINE)], 1 + ) + seed_row = _pack_seed_row(logits, n_eff, lines) + cand_vals = torch.full((batch, LIST_WIDTH), float("-inf"), dtype=torch.float32, device=dev) + cand_idx = torch.full((batch, LIST_WIDTH), -1, dtype=torch.int32, device=dev) + cand_ctl = torch.zeros((batch, 4), dtype=torch.int32, device=dev) + pads = {"pads": 64, "starved": 200}.get(mode, 0) + for r in range(batch): + ne = int(n_eff[r]) + row = logits[r, :ne] + hits = torch.nonzero(row >= lines[r, 0], as_tuple=False).flatten() + cand_ctl[r, 2] = int((row >= lines[r, 1]).sum()) + cand_ctl[r, 3] = int((row >= lines[r, 2]).sum()) + # emission order is value-blind: shuffle, then classify by the + # tightest line passed; a full segment spills to the looser one + perm = hits[torch.randperm(hits.numel(), device=dev)] + v = row[perm] + seg = torch.where(v >= lines[r, 2], 0, torch.where(v >= lines[r, 1], 1, 2)) + in_a = seg == 0 + ord_a = torch.cumsum(in_a.int(), 0) + stay_a = in_a & (ord_a <= LIST_SEG_A) + in_b = (seg == 1) | (in_a & ~stay_a) + ord_b = torch.cumsum(in_b.int(), 0) + stay_b = in_b & (ord_b <= LIST_SEG_A) + in_c = (seg == 2) | (in_b & ~stay_b) + ord_c = torch.cumsum(in_c.int(), 0) + stay_c = in_c & (ord_c <= LIST_CAP_C - pads) + slot = torch.full_like(seg, -1) + slot[stay_a] = ord_a[stay_a] - 1 + slot[stay_b] = LIST_SEG_A + (ord_b[stay_b] - 1) + slot[stay_c] = 2 * LIST_SEG_A + (ord_c[stay_c] - 1) + if pads: + # sentinels displace C entries later in emission order: the + # kept ordinals shift by how many sentinels landed before them + pad_slots = torch.randperm(int(stay_c.sum()) + pads, device=dev)[:pads] + keep = slot[stay_c] - 2 * LIST_SEG_A + shift = (pad_slots[None, :] <= keep[:, None]).sum(-1) + slot[stay_c] = 2 * LIST_SEG_A + keep + shift + live = slot >= 0 + cand_vals[r, slot[live].long()] = v[live] + cand_idx[r, slot[live].long()] = perm[live].int() + cand_ctl[r, 0] = int(hits.numel()) + pads + cand_ctl[r, 1] = 1 if int(in_c.sum()) > LIST_CAP_C - pads else 0 + out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda") + # xstate is passed even though this test does not chain: it is part of + # the kernel's compile key, so sharing it with ext_closed_loop's list + # shape lets both reuse one compiled kernel instead of paying two. + xstate = torch.zeros((batch, 8), dtype=torch.float32, device=dev) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + cluster_size=1, + seed_thr=seed_row, + cand_vals=cand_vals, + cand_idx=cand_idx, + cand_ctl=cand_ctl, + accept_cap=LIST_SEG_A, + num_threads=512, + xstate=xstate, + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1) + + +@skip_not_sm100 +@pytest.mark.parametrize("tail_mode", ["exact", "pad_inf"]) +def test_cute_dsl_gvr_topk_decode_ext_block_max(tail_mode, tie_aware_check): + """32-grain positional upper-bound records + packed seed row. + + The skip walk may only skip units whose bound clears no line, so it + must be exact under both legal tail bounds: the tight max over valid + positions and the worst legal inflation (+FLT_MAX on a partially + valid record). + """ + top_k, N, batch = 1024, 262144, 2 + seq_lens = torch.tensor([N, N - 37], dtype=torch.int32, device="cuda") + logits, pre_idx, seq_lens = _make_inputs( + batch, N, top_k, torch.float32, 1, seed=13, seq_lens=seq_lens + ) + dev = logits.device + n_eff = seq_lens.to(device=dev, dtype=torch.long) + pos = torch.arange(N, device=dev)[None, :] + masked = torch.where(pos < n_eff[:, None], logits, float("-inf")) + records = masked.view(batch, N // 32, 32).amax(-1) + if tail_mode == "pad_inf": + rec_start = torch.arange(N // 32, device=dev)[None, :] * 32 + partial = (rec_start < n_eff[:, None]) & (rec_start + 32 > n_eff[:, None]) + records = torch.where(partial, torch.full_like(records, _FLT_MAX), records) + records = records.contiguous() + lines = _lines_at_counts(logits, n_eff, (4 * top_k, 2 * top_k, top_k + top_k // 4)) + seed_row = _pack_seed_row(logits, n_eff, lines, block_max=records) + out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, + pre_idx, + seq_lens, + out_indices, + top_k=top_k, + cluster_size=1, + seed_thr=seed_row, + block_max=records, + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1) + + +def _emulate_emission(logits, n_eff, st, tier, top_k): + """Host-side stand-in for the indexer epilogue: fill the packed-row + counts (and the candidate list on the list tier) against the CURRENT + seed lines, exactly as the production emitter would.""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import ( + LIST_CAP_C, + LIST_SEG_A, + ) + + batch, N = logits.shape + dev = logits.device + lines = st.seed_row[:batch, 0:3] + pos = torch.arange(N, device=dev)[None, :] + valid = pos < n_eff[:, None] + finite = torch.isfinite(lines[:, 0]) + counts = torch.stack( + [((logits >= lines[:, j : j + 1]) & valid).sum(-1) for j in range(3)], 1 + ).float() + st.seed_row[:batch, 3:6] = torch.where(finite[:, None], counts, torch.zeros_like(counts)) + if tier == "list" and st.cand_vals is not None: + st.cand_vals[:batch].fill_(float("-inf")) + st.cand_idx[:batch].fill_(-1) + st.cand_ctl[:batch].zero_() + for r in range(batch): + if not bool(finite[r]): + continue + ne = int(n_eff[r]) + row = logits[r, :ne] + hits = torch.nonzero(row >= lines[r, 0], as_tuple=False).flatten() + cnt = int(hits.numel()) + nwr = min(cnt, LIST_CAP_C) + base = 2 * LIST_SEG_A + st.cand_idx[r, base : base + nwr] = hits[:nwr].int() + st.cand_vals[r, base : base + nwr] = row[hits[:nwr]] + st.cand_ctl[r, 0] = cnt + st.cand_ctl[r, 1] = 1 if cnt > LIST_CAP_C else 0 + + +@skip_not_sm100 +@pytest.mark.parametrize( + "tier_shape", [("list", 2, 131072), ("counts", 8, 131072), ("rungs", 1, 32768)] +) +def test_cute_dsl_gvr_topk_decode_ext_closed_loop(tier_shape, tie_aware_check): + """Chained multi-step closed loop: kernel xstate publish -> + update_seed_rows -> next step's emission and admission. + + Step 2's logits shift the k-th value far beyond any fixed guard + width, so line placement must come from the fitted slope; every step + must stay exact regardless of which internal path admission picks. + """ + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import GvrEmissionState + + want_tier, batch, N = tier_shape + top_k = 512 + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + logits0, pre_idx, seq_lens = _make_inputs( + batch, N, top_k, torch.float32, 1, seed=17, varlen=False + ) + dev = logits0.device + n_eff = seq_lens.to(device=dev, dtype=torch.long) + st = GvrEmissionState(max_rows=batch, top_k=top_k, device=dev, enable_list_tier=True) + # k-th drift per step: step1 -> step2 rises by ~0.05 (>> any fixed + # guard), step3 falls back near step1's level + shifts = (0.0, 0.05, -0.03) + for step, shift in enumerate(shifts): + logits = (logits0 + shift).contiguous() + tier, route = st.plan(batch, N, num_sms, compress_ratio=1) + assert tier == want_tier + st.update_seed_rows(batch, tier) + _emulate_emission(logits, n_eff, st, tier, top_k) + pre = torch.topk(logits.float(), top_k, dim=-1).indices.int().contiguous() + out_indices = torch.empty(batch, top_k, dtype=torch.int32, device="cuda") + kw = st.topk_ext_kwargs(route, batch, None) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits, pre, seq_lens, out_indices, top_k=top_k, **kw + ) + torch.cuda.synchronize() + _gvr_check(tie_aware_check, out_indices, logits, seq_lens, top_k, 1) + assert bool((st.xstate[:batch, 0] > 0).all().item()), f"step {step}: publish missing" + + @skip_not_sm100 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_cute_dsl_gvr_topk_decode_p4_exact_tail_16bit(dtype): diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 9682d7981ab5..1ef992746f9e 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -328,3 +328,33 @@ def test_unsupported_prefill_implementation_raises() -> None: row_starts=torch.zeros(1, dtype=torch.int32), row_ends=torch.ones(1, dtype=torch.int32), ) + + +def test_gvr_emission_reset_parks_reused_slots(monkeypatch) -> None: + """Cold-started rows must carry non-finite lines. + + Slot turnover under continuous batching can hand a request another + request's emission state. Exactness never rides on the lines - the + consumer admits on the counts the emitter measures for the current + query - but a reset row must park onto the stock path rather than + inherit finite thresholds, so the closed loop restarts cleanly. + """ + gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) + monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + prior_indices = torch.zeros(4, 2, dtype=torch.int32) + + top_k.prepare_gvr_emission(4, 1 << 17, 148, prior_indices) + state = top_k._gvr_emission_state + # emulate a warmed closed loop: every slot carries finite lines + state.xstate[:, 0] = 1.0 + state.seed_row[:, :3] = torch.tensor([1.0, 2.0, 3.0]) + + top_k.reset_gvr_emission_rows(slice(1, 3)) + top_k.prepare_gvr_emission(4, 1 << 17, 148, prior_indices) + + lines = state.seed_row[:, 0] + assert torch.isinf(lines[1:3]).all(), "reset rows must park on non-finite lines" + assert torch.isfinite(lines[0]) and torch.isfinite(lines[3]), ( + "untouched slots must keep their closed-loop state" + )