[None][feat] Raise the CuTE-DSL top-k decode limit to 16384 and support odd top_k

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -5252,7 +5252,7 @@ def cute_dsl_topk_decode_blackwell(
         Args:
             input_values: Input logits tensor [batch_size * next_n, vocab_size]
             seq_lens: Sequence lengths for each batch [batch_size]
-            top_k: Number of top elements to select (max 2048)
+            top_k: Number of top elements to select (max 16384)
             next_n: Number of candidates per sequence (for speculative decoding)
             num_copy_bits: Number of bits for vectorized memory copy (128 or 256)
             load_balance: Enable persistent dynamic scheduling for load balancing
@@ -5262,7 +5262,7 @@ def cute_dsl_topk_decode_blackwell(
 
         Note:
             This function requires Blackwell architecture (SM100+) and CuTE DSL support.
-            Maximum supported top_k is 2048.
+            Maximum supported top_k is 16384.
         """
         # Validate SM version
         sm_version = get_sm_version()
@@ -5272,10 +5272,10 @@ def cute_dsl_topk_decode_blackwell(
                 "Use standard top-k implementation for older architectures.")
 
         # Validate inputs
-        if top_k <= 0 or top_k > 2048:
+        if top_k <= 0 or top_k > 16384:
             raise ValueError(
-                f"top_k must be in range [1, 2048], got {top_k}. "
-                "Maximum supported top_k is 2048 for Blackwell architecture.")
+                f"top_k must be in range [1, 16384], got {top_k}. "
+                "Maximum supported top_k is 16384 for Blackwell architecture.")
 
         if next_n <= 0:
             raise ValueError(f"next_n must be positive, got {next_n}")
@@ -5987,7 +5987,7 @@ def cute_dsl_topk_decode_multi_cta_blackwell(
         Args:
             input_values: Input logits tensor [batch_size * next_n, vocab_size]
             seq_lens: Sequence lengths for each batch [batch_size]
-            top_k: Number of top elements to select (max 2048)
+            top_k: Number of top elements to select (max 16384)
             next_n: Number of candidates per sequence (for speculative decoding)
             num_copy_bits: Number of bits for vectorized memory copy (128 or 256)
             chunk_size_per_cta: Number of columns each CTA processes
@@ -6007,10 +6007,10 @@ def cute_dsl_topk_decode_multi_cta_blackwell(
                 "Use standard top-k implementation for older architectures.")
 
         # Validate inputs
-        if top_k <= 0 or top_k > 2048:
+        if top_k <= 0 or top_k > 16384:
             raise ValueError(
-                f"top_k must be in range [1, 2048], got {top_k}. "
-                "Maximum supported top_k is 2048 for Blackwell architecture.")
+                f"top_k must be in range [1, 16384], got {top_k}. "
+                "Maximum supported top_k is 16384 for Blackwell architecture.")
 
         if next_n <= 0:
             raise ValueError(f"next_n must be positive, got {next_n}")
@@ -6118,14 +6118,20 @@ def cute_dsl_indexer_topk_decode(
             input_values: Input logits tensor [batch_size * next_n, vocab_size]
             seq_lens: Sequence lengths for each batch [batch_size]
             output_indices: Pre-allocated output buffer [batch_size * next_n, top_k]
-            top_k: Number of top elements to select (max 2048)
+            top_k: Number of top elements to select (max 16384)
             next_n: Number of candidates per sequence (for speculative decoding)
             num_copy_bits: Number of bits for vectorized memory copy (128 or 256)
             dynamic: Use dynamic multi-CTA scheduling (for 2-pass multi-CTA)
             single_pass_multi_cta: Use single-pass multi-CTA radix top-k
             single_pass_multi_cta_cluster: Force cluster-accelerated variant
                 (only effective when single_pass_multi_cta=True)
         """
+        # Validate inputs
+        if top_k <= 0 or top_k > 16384:
+            raise ValueError(
+                f"top_k must be in range [1, 16384], got {top_k}. "
+                "Maximum supported top_k is 16384 for Blackwell architecture.")
+
         num_rows = input_values.shape[0]
         num_tokens = input_values.shape[1]
 

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

@@ -56,7 +56,6 @@
       --top_k 2048 --do_ref_check --return_val --do_benchmark
 
 Constraints for this example:
-* The problem size of top_k <= 2048.
 * The input tensor has data contiguous on the n dimension (row-major).
 * The supported input data types are Float32, Float16, or BFloat16.
 """
@@ -346,7 +345,7 @@ def filtered_topk_kernel(
             g_num_input = None
         s_indices = smem.allocate_tensor(
             element_type=self.index_type,
-            layout=cute.make_ordered_layout((self.filtered_topk_max_k,), order=(0)),
+            layout=cute.make_ordered_layout((self.top_k,), order=(0)),
             byte_alignment=128,
         )
         s_input_idx = smem.allocate_tensor(
@@ -1393,9 +1392,6 @@ def run_topk_decode(
     parser.add_argument("--use_cold_l2", action="store_true", default=True, help="Use cold L2")
 
     args = parser.parse_args()
-    if args.top_k % 2 != 0:
-        parser.error("top_k must be a multiple of 2 (got top_k={})".format(args.top_k))
-
     run_topk_decode(
         dtype=args.dtype,
         batch_size=args.batch_size,

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

@@ -58,8 +58,17 @@ def __init__(
         self.num_ctas_per_row = num_ctas_per_row
         self.merge_blocks = merge_blocks
 
-        # Note: now we only support top_k <= 2048, we could change the code here to support larger top_k.
-        self.filtered_topk_max_k = 2048
+        # top_k sizes the shared-memory index staging (s_indices), so bound it
+        # here. Direct callers and the run_topk_decode CLI bypass the decode
+        # wrappers, and an oversized top_k would otherwise surface as an opaque
+        # smem launch failure. 16384 matches the wrapper guards in
+        # cute_dsl_custom_ops.py.
+        if top_k <= 0 or top_k > 16384:
+            raise ValueError(
+                f"top_k must be in range [1, 16384], got {top_k}. "
+                "Maximum supported top_k is 16384 for Blackwell architecture."
+            )
+
         # 8 bits for radix-based filter.
         self.radix = 256
 
@@ -963,13 +972,14 @@ def filtered_topk_kernel_per_row(
                             cute.arch.barrier()
 
             # Phase 3: Output phase
+            output_vector_width = 2 if self.top_k % 2 == 0 else 1
             vecsize_out = cutlass.const_expr(
                 min(
                     self.top_k,
                     cute.ceil_div(self.top_k, self.num_threads_per_cta),
                     self.num_copy_bits // self.dtype.width,
                     # TODO: only tested for float32. need to check for other dtypes.
-                    2,
+                    output_vector_width,
                 )
             )
             assert self.top_k % vecsize_out == 0

tests/unittest/_torch/thop/parallel/test_indexer_topk.py

@@ -811,6 +811,111 @@ def run_fn(logits, seq_lens):
     )
 
 
+@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available")
+@skip_pre_blackwell
+@pytest.mark.parametrize("batch_size", [1, 8])
+@pytest.mark.parametrize("index_topk", [1023, 2047, 4095, 4096, 8192, 16384])
+@pytest.mark.parametrize("num_tokens", [32768, 131072])
+def test_cute_dsl_topk_decode_high_and_odd_k(batch_size, index_topk, num_tokens):
+    """Large (>4096) and odd top_k stay bit-exact on the decode entry points.
+
+    The indexer entry point is the path KV-cache eviction uses with large or
+    odd keep budgets, so it is checked for every top_k across all three dtypes
+    (the scalar output write at vecsize_out=1 is dtype dependent). The
+    single-CTA and multi-CTA wrappers, whose raised guard this test backs, only
+    support even top_k, so they are checked on the even values in fp32.
+    """
+
+    def run_single_cta(logits, seq_lens):
+        return torch.ops.trtllm.cute_dsl_topk_decode_blackwell(
+            input_values=logits,
+            seq_lens=seq_lens,
+            top_k=index_topk,
+            next_n=1,
+            num_copy_bits=256,
+            load_balance=False,
+        )
+
+    def run_multi_cta(logits, seq_lens):
+        return torch.ops.trtllm.cute_dsl_topk_decode_multi_cta_blackwell(
+            input_values=logits,
+            seq_lens=seq_lens,
+            top_k=index_topk,
+            next_n=1,
+            num_copy_bits=256,
+            chunk_size_per_cta=16384,
+            dynamic=False,
+        )
+
+    def run_indexer(logits, seq_lens):
+        output_indices = torch.empty(batch_size, index_topk, dtype=torch.int32, device="cuda")
+        torch.ops.trtllm.cute_dsl_indexer_topk_decode(
+            input_values=logits,
+            seq_lens=seq_lens,
+            output_indices=output_indices,
+            top_k=index_topk,
+            next_n=1,
+            num_copy_bits=256,
+        )
+        return output_indices
+
+    if index_topk % 2 == 0:
+        # Even top_k: all three entry points are defined; the raised guard this
+        # test backs is exercised on the wrappers in fp32.
+        dtype_runs = [(torch.float32, (run_single_cta, run_multi_cta, run_indexer))]
+    else:
+        # Odd top_k: only the indexer path supports it, checked across dtypes.
+        dtype_runs = [
+            (torch.float32, (run_indexer,)),
+            (torch.float16, (run_indexer,)),
+            (torch.bfloat16, (run_indexer,)),
+        ]
+
+    for dtype, run_fns in dtype_runs:
+        for run_fn in run_fns:
+            _run_cute_dsl_topk_test(
+                batch_size,
+                1,
+                index_topk,
+                num_tokens,
+                dtype,
+                run_fn,
+            )
+
+
+@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available")
+@skip_pre_blackwell
+@pytest.mark.parametrize("top_k", [1023, 2047, 2048])
+@pytest.mark.parametrize("dtype_name", ["float32", "float16", "bfloat16"])
+def test_filtered_topk_varlen_odd_k(top_k, dtype_name):
+    """The filtered varlen kernel supports odd top_k (scalar output tail).
+
+    The reference check inside the runner asserts bitwise agreement with a
+    torch reference; 2048 is the even control. The scalar tail write is dtype
+    dependent, so all three supported dtypes are exercised.
+    """
+    import cutlass
+
+    from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.filtered_top_k_decode_varlen import (
+        run_topk_decode,
+    )
+
+    dtype = {
+        "float32": cutlass.Float32,
+        "float16": cutlass.Float16,
+        "bfloat16": cutlass.BFloat16,
+    }[dtype_name]
+
+    run_topk_decode(
+        dtype,
+        batch_size=16,
+        max_num_cols=4096,
+        top_k=top_k,
+        next_n=3,
+        do_benchmark=False,
+    )
+
+
 @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available")
 @skip_pre_blackwell
 @pytest.mark.parametrize("batch_size", [1, 4, 8, 16, 256])
@@ -1452,6 +1557,52 @@ def run_fn(logits, seq_lens):
     )
 
 
+@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available")
+@skip_pre_blackwell
+@pytest.mark.parametrize("batch_size", [1, 8])
+@pytest.mark.parametrize("index_topk", [2047, 8192])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_cluster", [False, True])
+def test_cute_dsl_single_pass_multi_cta_high_and_odd_k(batch_size, index_topk, dtype, use_cluster):
+    """Large (>2048) and odd top_k stay bit-exact on the single-pass multi-CTA
+    and cluster decode paths.
+
+    The raised wrapper guard forwards these values to both dispatch paths, so
+    this covers the odd-K scalar output write on routes the other tests only
+    exercise at top_k=2048.
+    """
+    num_tokens = 131072
+    next_n = 1
+
+    if use_cluster:
+        runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSinglePassMultiCTAClusterRunner
+    else:
+        runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSinglePassMultiCTARunner
+
+    def run_fn(logits, seq_lens):
+        result = runner.forward(
+            input_values=logits,
+            seq_lens=seq_lens,
+            top_k=index_topk,
+            next_n=next_n,
+            return_val=False,
+            num_copy_bits=256,
+        )
+        # The cluster runner returns None when the problem exceeds its capacity.
+        if result[0] is None:
+            pytest.skip("Problem size exceeds cluster kernel capacity")
+        return result[0]
+
+    _run_cute_dsl_topk_test(
+        batch_size,
+        next_n,
+        index_topk,
+        num_tokens,
+        dtype,
+        run_fn,
+    )
+
+
 # ============================================================================
 # Heuristic Decode Distribution-Parameterised Tests
 # ============================================================================