[Perf][Kernel] Add sampled filtering for persistent top-k

benchmarks/kernels/benchmark_persistent_topk.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows.
+
+Run on a reserved GPU, for example:
+    chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py
+
+Use --capacity to emulate graphs captured at a larger maximum context length.
+Run the same command on the base and candidate builds to compare JSON results.
+"""
+
+import argparse
+import itertools
+import json
+import statistics
+from pathlib import Path
+
+import torch
+from flashinfer.testing import bench_gpu_time_with_cupti
+
+import vllm._custom_ops  # noqa: F401
+from vllm.platforms import current_platform
+
+
+def benchmark(rows: int, width: int, k: int, args) -> dict:
+    capacity = args.capacity or width
+    if capacity < width:
+        raise ValueError("Capacity must be at least the valid-length upper bound")
+    torch.manual_seed(args.seed)
+    logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
+    if args.full_length:
+        lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32)
+    else:
+        lengths = torch.randint(
+            int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32
+        )
+    invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None]
+    logits.masked_fill_(invalid, float("nan"))
+    out = torch.empty((rows, k), device="cuda", dtype=torch.int32)
+    workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8)
+    cooperative = (
+        args.backend == "auto"
+        and rows <= 64
+        and capacity % 4 == 0
+        and current_platform.has_device_capability(90)
+        and not current_platform.is_device_capability_family(120)
+    )
+    op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk
+
+    def run():
+        op(logits, lengths, out, workspace, k, capacity)
+
+    run()
+    valid = torch.arange(k, device="cuda")[None] < lengths[:, None]
+    assert torch.all(out[~valid] == -1)
+    assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid)
+    sorted_indices = out.sort(dim=1).values
+    assert torch.all(
+        (sorted_indices[:, 1:] != sorted_indices[:, :-1])
+        | (sorted_indices[:, 1:] == -1)
+    )
+    selected = logits.gather(1, out.clamp_min(0).long()).masked_fill(
+        ~valid, -float("inf")
+    )
+    reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values
+    torch.testing.assert_close(
+        selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0
+    )
+    trials = [
+        statistics.median(
+            bench_gpu_time_with_cupti(
+                run,
+                use_cuda_graph=True,
+                cold_l2_cache=not args.warm,
+                dry_run_iters=5,
+                repeat_iters=args.repeat,
+            )
+        )
+        * 1000
+        for _ in range(args.trials)
+    ]
+    us = statistics.median(trials)
+    return dict(
+        rows=rows,
+        width=width,
+        capacity=capacity,
+        k=k,
+        backend="cooperative" if cooperative else "persistent",
+        us=us,
+        trials_us=trials,
+        # Minimum useful traffic: one valid-score read and one index write.
+        effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024])
+    parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840])
+    parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048])
+    parser.add_argument("--capacity", type=int)
+    parser.add_argument(
+        "--full-length", action="store_true", help="Use exactly --widths valid scores"
+    )
+    parser.add_argument("--backend", choices=["auto", "persistent"], default="auto")
+    parser.add_argument(
+        "--warm", action="store_true", help="Keep L2 warm between replays"
+    )
+    parser.add_argument("--repeat", type=int, default=60)
+    parser.add_argument("--trials", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--output", type=Path)
+    args = parser.parse_args()
+    metadata = dict(
+        gpu=str(torch.cuda.get_device_properties(0)),
+        torch=torch.__version__,
+        cuda=torch.version.cuda,
+        args=vars(args),
+    )
+    print(json.dumps(metadata, default=str))
+    results = []
+    for rows, width, k in itertools.product(args.rows, args.widths, args.top_k):
+        result = benchmark(rows, width, k, args)
+        results.append(result)
+        print(json.dumps(result), flush=True)
+    if args.output:
+        args.output.write_text(
+            json.dumps(dict(metadata=metadata, results=results), indent=2, default=str)
+            + "\n"
+        )
+
+
+if __name__ == "__main__":
+    main()

csrc/libtorch_stable/persistent_topk.cuh

@@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
 constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
     sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE;  // 128KB
 
-/*!
- * \brief Filtered Top-K kernel for ragged sequences.
- *
- * \tparam DType Data type (float, half, nv_bfloat16)
- * \tparam IdType Index type (int32_t)
- * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
- */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
-          bool UsePredicatedShortLoads = false>
-__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
-    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
-                              IdType* __restrict__ output,
-                              const IdType* __restrict__ lengths,
-                              uint32_t num_rows, uint32_t top_k,
-                              uint32_t max_len) {
+template <uint32_t MAX_K>
+struct FilteredTopKStorage {
+  alignas(128) int histogram[2][256 + 128];
+  alignas(128) int counter;
+  alignas(128) int threshold_bin;
+  alignas(128) int num_input[2];
+  alignas(128) int indices[MAX_K];
+  int last_remain;
+};
+
+// With CheckOverflow, return false before using a truncated stash so the caller
+// can retry with exact full-row selection.
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K,
+          bool UsePredicatedShortLoads, bool CheckOverflow = false>
+__device__ __forceinline__ bool filtered_topk_row(
+    const DType* score, IdType* dst, int length, uint32_t top_k,
+    FilteredTopKStorage<MAX_K>& storage) {
   constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
   constexpr int RADIX = 256;
   constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
-
-  const uint32_t bid = blockIdx.x;
   const int tx = threadIdx.x;
 
-  if (bid >= num_rows) return;
-
-  const int length =
-      (lengths != nullptr) ? lengths[bid] : static_cast<int>(max_len);
-  const DType* score = input + bid * max_len;
-  IdType* dst = output + bid * top_k;
-
   // Trivial case: length <= top_k
   if (length <= static_cast<int>(top_k)) {
     for (int i = tx; i < static_cast<int>(top_k); i += BLOCK_SIZE) {
       dst[i] = (i < length) ? static_cast<IdType>(i) : static_cast<IdType>(-1);
     }
-    return;
+    return true;
   }
 
   // Short path
@@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
       hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
                                                   _smem_reg);
     }
-    return;
+    return true;
   }
 
-  // Static shared memory
-  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
-  alignas(128) __shared__ int s_counter;
-  alignas(128) __shared__ int s_threshold_bin_id;
-  alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[MAX_K];
-
+  auto& s_histogram_buf = storage.histogram;
+  auto& s_counter = storage.counter;
+  auto& s_threshold_bin_id = storage.threshold_bin;
+  auto& s_num_input = storage.num_input;
+  auto& s_indices = storage.indices;
   auto& s_histogram = s_histogram_buf[0];
 
   // Dynamic shared memory for input double buffer
@@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     // Stage 2: refine with 8bit radix passes
 #pragma unroll
     for (int round = 0; round < NUM_ROUNDS; ++round) {
-      __shared__ int s_last_remain;
+      auto& s_last_remain = storage.last_remain;
       const auto r_idx = round % 2;
 
       const auto _raw_num_input = s_num_input[r_idx];
+      if constexpr (CheckOverflow) {
+        if (_raw_num_input > SMEM_INPUT_SIZE) return false;
+      }
       const auto num_input =
           (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE;
 
@@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     const int idx = s_indices[base];
     dst[base] = static_cast<IdType>(idx);
   }
+  return true;
+}
+
+/*!
+ * \brief Filtered Top-K kernel for ragged sequences.
+ *
+ * \tparam DType Data type (float, half, nv_bfloat16)
+ * \tparam IdType Index type (int32_t)
+ * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
+ */
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
+          bool UsePredicatedShortLoads = false>
+__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
+    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
+                              IdType* __restrict__ output,
+                              const IdType* __restrict__ lengths,
+                              uint32_t num_rows, uint32_t top_k,
+                              uint32_t max_len) {
+  const uint32_t bid = blockIdx.x;
+  if (bid >= num_rows) return;
+  const int length = lengths ? lengths[bid] : static_cast<int>(max_len);
+  __shared__ FilteredTopKStorage<MAX_K> storage;
+  filtered_topk_row<DType, IdType, VEC_SIZE, MAX_K, UsePredicatedShortLoads>(
+      input + bid * max_len, output + bid * top_k, length, top_k, storage);
 }
 
 // Helper to compute GCD for VEC_SIZE selection