From 723fe2a32dd7cabd628e6471d65f71505ab24180 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 05:25:57 -0700 Subject: [PATCH 1/4] [None][feat] Raise the CuTE-DSL top-k decode limit to 16384 Signed-off-by: tianruih --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 20 ++++--- .../_torch/thop/parallel/test_indexer_topk.py | 57 +++++++++++++++++++ 2 files changed, 68 insertions(+), 9 deletions(-) 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 95ac81158e8d..d727d5bc7ab4 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/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,11 @@ 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}. " + "16384 is the largest top_k verified bit-exact against torch.topk " + "on Blackwell (unit-tested at 8192 and 16384).") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -5987,7 +5988,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 +6008,11 @@ 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}. " + "16384 is the largest top_k verified bit-exact against torch.topk " + "on Blackwell (unit-tested at 8192 and 16384).") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index ee0e2bc2129c..7234f132f4a8 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -811,6 +811,63 @@ 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", [8192, 16384]) +@pytest.mark.parametrize("num_tokens", [32768, 131072]) +def test_cute_dsl_topk_decode_high_k(batch_size, index_topk, num_tokens): + """top_k above 4096 stays bit-exact (wrapper guard raised to 16384). + + Covers the three decode entry points: the single-CTA and multi-CTA + Blackwell wrappers (whose top_k guard this test backs) and the indexer + variant used by KV-cache eviction with large keep budgets. + """ + + 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 + + for run_fn in (run_single_cta, run_multi_cta, run_indexer): + _run_cute_dsl_topk_test( + batch_size, + 1, + index_topk, + num_tokens, + torch.float32, + run_fn, + ) + + @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]) From 13c083c30e7edc005df34e7d788292fdf1890c3f Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 07:17:40 -0700 Subject: [PATCH 2/4] [None][fix] Support odd top-k in the filtered varlen decode kernel The output phase copied results with a hard two-wide vector, so odd top-k overran the row tail and the CLI rejected odd values. Odd top-k now degrades the output copy to scalar width (even top-k is unchanged), and the shared-memory staging is sized by the actual top-k instead of a fixed 2048 constant, which also removes the 2048 cap. Verified bitwise against the torch reference at top_k 2047 and 1023 with a 2048 control. Signed-off-by: tianruih --- .../blackwell/top_k/filtered_top_k_decode_varlen.py | 6 +----- .../blackwell/top_k/filtered_top_k_varlen_util.py | 5 ++--- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py index 66b79c6986f8..ab16d033dba2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py +++ b/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, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py index be796763e8e2..71d8bcbe79f2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py @@ -58,8 +58,6 @@ 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 # 8 bits for radix-based filter. self.radix = 256 @@ -963,13 +961,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 From 1e3e5257e894ea183b7b2bf57ea6372436b5b25b Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 07:20:37 -0700 Subject: [PATCH 3/4] [None][test] Cover 4096 and odd top_k on the CuTE-DSL decode paths Extends the high-k battery with 4096 and adds odd top_k coverage: the indexer decode entry at 1023/2047/4095 against torch.topk, and the filtered varlen kernel at 1023/2047 with a 2048 control through its own bitwise reference check. The thop/parallel directory already runs in the B200 test list, so these register with it. Signed-off-by: tianruih --- .../_torch/thop/parallel/test_indexer_topk.py | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 7234f132f4a8..970b79257d4e 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -814,7 +814,7 @@ 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", [8192, 16384]) +@pytest.mark.parametrize("index_topk", [4096, 8192, 16384]) @pytest.mark.parametrize("num_tokens", [32768, 131072]) def test_cute_dsl_topk_decode_high_k(batch_size, index_topk, num_tokens): """top_k above 4096 stays bit-exact (wrapper guard raised to 16384). @@ -868,6 +868,61 @@ def run_indexer(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]) +def test_cute_dsl_indexer_topk_decode_odd_k(batch_size, index_topk): + """Odd top_k stays bit-exact on the indexer decode entry point.""" + num_tokens = 32768 + + 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 + + _run_cute_dsl_topk_test( + batch_size, + 1, + index_topk, + num_tokens, + torch.float32, + run_indexer, + ) + + +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") +@skip_pre_blackwell +@pytest.mark.parametrize("top_k", [1023, 2047, 2048]) +def test_filtered_topk_varlen_odd_k(top_k): + """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. + """ + import cutlass + + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.filtered_top_k_decode_varlen import ( + run_topk_decode, + ) + + run_topk_decode( + cutlass.Float32, + 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]) From 843342caf38dafd13a688d39daee886955686f0a Mon Sep 17 00:00:00 2001 From: Tianrui Hu Date: Mon, 20 Jul 2026 09:13:46 -0700 Subject: [PATCH 4/4] [None][fix] Address review comments on CuTE-DSL top-k decode limit - Guard cute_dsl_indexer_topk_decode with [1, 16384] and fix its docstring - Bound top_k in FilteredTopKKernelVarlen for direct and CLI callers - Simplify the out-of-range error messages - Cover large and odd top_k on the single-pass multi-CTA and cluster paths - Parameterize the odd-K tests over fp32/fp16/bf16 - Merge the high-K and odd-K decode tests into one parametrized test Signed-off-by: Tianrui Hu --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 14 +- .../top_k/filtered_top_k_varlen_util.py | 11 ++ .../_torch/thop/parallel/test_indexer_topk.py | 137 +++++++++++------- 3 files changed, 108 insertions(+), 54 deletions(-) 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 d727d5bc7ab4..b373a8a04c2d 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5275,8 +5275,7 @@ def cute_dsl_topk_decode_blackwell( if top_k <= 0 or top_k > 16384: raise ValueError( f"top_k must be in range [1, 16384], got {top_k}. " - "16384 is the largest top_k verified bit-exact against torch.topk " - "on Blackwell (unit-tested at 8192 and 16384).") + "Maximum supported top_k is 16384 for Blackwell architecture.") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -6011,8 +6010,7 @@ def cute_dsl_topk_decode_multi_cta_blackwell( if top_k <= 0 or top_k > 16384: raise ValueError( f"top_k must be in range [1, 16384], got {top_k}. " - "16384 is the largest top_k verified bit-exact against torch.topk " - "on Blackwell (unit-tested at 8192 and 16384).") + "Maximum supported top_k is 16384 for Blackwell architecture.") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -6120,7 +6118,7 @@ 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) @@ -6128,6 +6126,12 @@ def cute_dsl_indexer_topk_decode( 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] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py index 71d8bcbe79f2..6725526407c2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py @@ -58,6 +58,17 @@ def __init__( self.num_ctas_per_row = num_ctas_per_row self.merge_blocks = merge_blocks + # 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 diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 970b79257d4e..3634336aa961 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -814,14 +814,16 @@ 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", [4096, 8192, 16384]) +@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_k(batch_size, index_topk, num_tokens): - """top_k above 4096 stays bit-exact (wrapper guard raised to 16384). - - Covers the three decode entry points: the single-CTA and multi-CTA - Blackwell wrappers (whose top_k guard this test backs) and the indexer - variant used by KV-cache eviction with large keep budgets. +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): @@ -857,55 +859,40 @@ def run_indexer(logits, seq_lens): ) return output_indices - for run_fn in (run_single_cta, run_multi_cta, run_indexer): - _run_cute_dsl_topk_test( - batch_size, - 1, - index_topk, - num_tokens, - torch.float32, - run_fn, - ) - - -@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]) -def test_cute_dsl_indexer_topk_decode_odd_k(batch_size, index_topk): - """Odd top_k stays bit-exact on the indexer decode entry point.""" - num_tokens = 32768 - - 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 - - _run_cute_dsl_topk_test( - batch_size, - 1, - index_topk, - num_tokens, - torch.float32, - run_indexer, - ) + 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]) -def test_filtered_topk_varlen_odd_k(top_k): +@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. + torch reference; 2048 is the even control. The scalar tail write is dtype + dependent, so all three supported dtypes are exercised. """ import cutlass @@ -913,8 +900,14 @@ def test_filtered_topk_varlen_odd_k(top_k): run_topk_decode, ) + dtype = { + "float32": cutlass.Float32, + "float16": cutlass.Float16, + "bfloat16": cutlass.BFloat16, + }[dtype_name] + run_topk_decode( - cutlass.Float32, + dtype, batch_size=16, max_num_cols=4096, top_k=top_k, @@ -1564,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 # ============================================================================