[TRTLLM-35882][feat] cute dsl gvr-topk load-balance optimization

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -5734,11 +5734,13 @@ def _compile(
             compress_ratio: int,
             return_output_values: bool,
             cluster_size: int,
+            seqlen_sorted: bool,
         ) -> None:
             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)
+                   compress_ratio, return_output_values, cluster_size,
+                   seqlen_sorted)
             if key in cls.kernel_cache:
                 return
             n_rows = cute.sym_int()
@@ -5769,6 +5771,14 @@ def _compile(
                 cutlass.Int32, (n_rows, top_k),
                 stride_order=(1, 0),
                 assumed_align=16)
+            # seqlen_sorted=False compiles the LJF indirection out via
+            # const_expr → no order_row read → fake = None (mirrors the
+            # out_values optional pattern). seqlen_sorted=True compiles
+            # the indirection in; fake shape = batch (n_batch) since
+            # order_row is request-level, not row-level.
+            order_row_fake = (cute.runtime.make_fake_compact_tensor(
+                cutlass.Int32,
+                (n_batch, ), stride_order=(0, )) if seqlen_sorted else None)
             fake_stream = cute.runtime.make_fake_stream(
                 use_tvm_ffi_env_stream=True)
 
@@ -5786,6 +5796,7 @@ def _compile(
                 compress_ratio=compress_ratio,
                 return_output_values=return_output_values,
                 cluster_size=cluster_size,
+                seqlen_sorted=seqlen_sorted,
             )
             cls.kernel_cache[key] = cute.compile(
                 kernel,
@@ -5794,6 +5805,7 @@ def _compile(
                 seq_lens_fake,
                 out_values_fake,
                 out_indices_fake,
+                order_row_fake,
                 stream=fake_stream,
                 options="--enable-tvm-ffi",
             )
@@ -5811,6 +5823,7 @@ def forward(
             compress_ratio: int = 1,
             max_seq_len: Optional[int] = None,
             cluster_size: Optional[int] = None,
+            order_row: Optional[torch.Tensor] = None,
         ) -> None:
             cute_dtype = _TORCH_TO_CUTLASS_DTYPE[logits.dtype]
             num_rows = logits.shape[0]
@@ -5947,6 +5960,19 @@ def forward(
                 else:
                     min_blocks_per_mp = 1
 
+            # seqlen_sorted = (order_row is not None). When True the
+            # kernel resolves the owning row via order_row[req] * next_n
+            # + nn so longer rows hit earlier waves (LJF). order_row is
+            # request-level (shape == seq_lens.shape).
+            seqlen_sorted = order_row is not None
+            if seqlen_sorted:
+                assert (
+                    order_row.dtype == torch.int32 and order_row.is_cuda
+                    and order_row.shape == seq_lens.shape
+                ), ("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)}")
+
             cls._compile(
                 cute_dtype,
                 top_k,
@@ -5961,18 +5987,23 @@ def forward(
                 compress_ratio,
                 return_output_values,
                 cluster_size,
+                seqlen_sorted,
             )
             key = (cute_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)
+                   return_output_values, cluster_size, seqlen_sorted)
             # TVM FFI: pass raw torch tensors directly, env stream picked
             # up automatically (no from_dlpack, no stream argument).
-            # ``output_values=None`` matches the kernel's compile-time
-            # ``return_output_values=False`` constexpr — no writes for top-k values.
+            # ``out_values = None``  — return_output_values=False, kernel
+            #                          skips STG.value.
+            # ``order_row``           — None when seqlen_sorted=False (no
+            #                          LJF indirection), else the
+            #                          caller-provided request-level
+            #                          dispatch order.
             cls.kernel_cache[key](logits, pre_idx, seq_lens, None,
-                                  output_indices)
+                                  output_indices, order_row)
 
     @torch.library.custom_op("trtllm::cute_dsl_gvr_topk_decode",
                              mutates_args=("output_indices", ),
@@ -5987,6 +6018,7 @@ def cute_dsl_gvr_topk_decode(
         compress_ratio: int = 1,
         max_seq_len: Optional[int] = None,
         cluster_size: Optional[int] = None,
+        order_row: Optional[torch.Tensor] = None,
     ) -> None:
         """CuTe DSL GVR (Guess-Verify-Refine) Top-K decode for Blackwell.
 
@@ -6013,6 +6045,20 @@ def cute_dsl_gvr_topk_decode(
                 processing one row via DSMEM + cluster sync). ``None``
                 (default) delegates to the wrapper's auto-dispatch heuristic
                 based on (N, BS); pass an explicit int to pin the choice.
+            order_row: Optional request-level LJF (longest-job-first)
+                dispatch order. When provided, must be ``int32`` on CUDA
+                with ``shape == seq_lens.shape``; ``order_row[i]`` is the
+                original request_id of the i-th-priority request. The
+                kernel then resolves the owning row per CTA as
+                ``order_row[cluster_id // next_n] * next_n +
+                (cluster_id % next_n)`` so longer rows hit earlier waves
+                — mitigates wave-tail imbalance when ``num_rows >
+                num_sms / cluster_size``. ``None`` (default) keeps the
+                legacy static mapping and the kernel const_expr elides
+                the indirection (legacy SASS unchanged). Typically built
+                once per decode step via :func:`gvr_topk_sort_prepare`
+                and reused across layers (``seq_lens`` is layer-
+                invariant).
         """
         if not is_sm_100f():
             raise ValueError(
@@ -6054,6 +6100,7 @@ def cute_dsl_gvr_topk_decode(
             compress_ratio=compress_ratio,
             max_seq_len=max_seq_len,
             cluster_size=cluster_size,
+            order_row=order_row,
         )
 
     @torch.library.register_fake("trtllm::cute_dsl_gvr_topk_decode")
@@ -6067,6 +6114,7 @@ def _(
         compress_ratio: int = 1,
         max_seq_len: Optional[int] = None,
         cluster_size: Optional[int] = None,
+        order_row: Optional[torch.Tensor] = None,
     ) -> None:
         return None
 

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

@@ -230,6 +230,7 @@ def __init__(
         cluster_size: int = 1,
         enable_smem_cache: bool = False,
         smem_cache_elems: int = 32768,
+        seqlen_sorted: bool = False,
     ):
         # cluster_size: number of CTAs cooperating per row. 1 = single-CTA
         # path; 2/4 = thread-block cluster with DSMEM aggregation. Capped at
@@ -239,6 +240,16 @@ def __init__(
                 f"cluster_size must be in [1, 16] (B200 GPC limit); got {cluster_size}"
             )
         self.cluster_size = cluster_size
+        # When True, the kernel resolves the owning row per CTA via a
+        # caller-provided ``order_row`` indirection — an LJF host-side
+        # dispatch order so longer rows hit earlier waves. ``order_row``
+        # is REQUEST-level: int32[batch_size = num_rows / next_n],
+        # typically a descending argsort of seq_lens. The kernel
+        # expands to row level as ``order_row[req] * next_n + nn`` so a
+        # request's ``next_n`` rows stay contiguous. Compatible with
+        # cluster_size > 1: all cs CTAs in a cluster see the same
+        # cluster_id (= bidx // cluster_size), hence the same row.
+        self.seqlen_sorted = seqlen_sorted
         # SMEM slice cache (optional): pre-stage each CTA's slice into SMEM
         # once between Phase 1 and Phase 2, so Phase 2/3's GE-count scans
         # read LDS instead of re-streaming GMEM. Requires
@@ -1722,25 +1733,51 @@ def gvr_topk_kernel(
         seq_lens: cute.Tensor,  # [numRows / next_n] int32
         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)
     ):
-        """Thin entry: bidx → (row_idx, cta_in_cluster), then run_one_row.
-
-        grid = (num_rows * cluster_size,); 1 row per thread-block cluster.
-        cluster_id == row index, cta_in_cluster ∈ [0, cluster_size). CTA r
-        scans row[r * N / cs : (r+1) * N / cs] in Phase 2, so the per-row
-        GE-count scales as 1 / cs. At cluster_size == 1 this collapses to
-        one CTA per row scanning the whole row.
+        """Thin entry: bidx → row_idx → run_one_row.
+
+        grid = (num_rows * cluster_size,) where num_rows = batch_size *
+        next_n. cluster_id = bidx // cluster_size, cta_in_cluster ∈
+        [0, cluster_size). CTA r scans row[r * N / cs : (r+1) * N / cs]
+        in Phase 2, so the per-row GE-count scales as 1 / cs. At
+        cluster_size == 1 this collapses to one CTA per row scanning
+        the whole row.
+
+        When ``self.seqlen_sorted`` is True, the LJF dispatch order
+        operates at REQUEST granularity (``order_row`` has length
+        batch_size = num_rows / next_n). The owning row is resolved as
+        ``order_row[cluster_id // next_n] * next_n + cluster_id % next_n``
+        so the ``next_n`` rows of one request stay contiguous in
+        dispatch order. All ``cluster_size`` CTAs within a cluster see
+        the same ``cluster_id`` and therefore the same row, preserving
+        cluster-sync semantics.
 
         Body is extracted into :meth:`run_one_row` so other entries (e.g.
         the LB load-balance variant) can resolve ``row_idx`` differently
-        from the static ``bidx // cluster_size`` mapping used here.
+        from the mappings used here.
         """
         bidx, _, _ = cute.arch.block_idx()
         cluster_size = cutlass.const_expr(self.cluster_size)
+        seqlen_sorted = cutlass.const_expr(self.seqlen_sorted)
+        next_n = cutlass.const_expr(self.next_n)
         if cutlass.const_expr(cluster_size > 1):
-            row_idx = bidx // cluster_size
+            cluster_id = bidx // cluster_size
+        else:
+            cluster_id = bidx
+        if cutlass.const_expr(seqlen_sorted):
+            # order_row is request-level (batch_size); expand to row-level
+            # via req_id * next_n + nn so a request's next_n rows stay
+            # contiguous in dispatch order (mirrors the LB main entry).
+            if cutlass.const_expr(next_n == 1):
+                row_idx = order_row[cluster_id]
+            else:
+                req_offset = cluster_id // cutlass.Int32(next_n)
+                nn = cluster_id % cutlass.Int32(next_n)
+                req_id = order_row[req_offset]
+                row_idx = req_id * cutlass.Int32(next_n) + nn
         else:
-            row_idx = bidx
+            row_idx = cluster_id
         self.run_one_row(
             row_idx,
             input_data,
@@ -2168,6 +2205,7 @@ def __call__(
         seq_lens: cute.Tensor,
         output_values: cute.Tensor,  # or None.
         output_indices: cute.Tensor,
+        order_row: cute.Tensor,  # or None when seqlen_sorted=False
         stream,
     ):
         num_rows = input_data.shape[0]
@@ -2183,6 +2221,7 @@ def __call__(
             seq_lens,
             output_values,
             output_indices,
+            order_row,
         ).launch(
             grid=(total_ctas, 1, 1),
             block=(self.num_threads, 1, 1),