perf: using multi-cta optimization for top-k/top-p

csrc/flashinfer_sampling_binding.cu

@@ -55,7 +55,8 @@ void top_k_renorm_probs(TensorView probs, TensorView renorm_probs,
                         Optional<TensorView> maybe_top_k_arr, int64_t top_k_val);
 
 void top_k_mask_logits(TensorView logits, TensorView mask_logits,
-                       Optional<TensorView> maybe_top_k_arr, int64_t top_k_val);
+                       Optional<TensorView> maybe_top_k_arr, int64_t top_k_val,
+                       TensorView row_states_buffer);
 
 void chain_speculative_sampling(TensorView draft_probs, TensorView draft_token_ids,
                                 TensorView target_probs, TensorView output_token_ids,

csrc/renorm.cu

@@ -59,19 +59,25 @@ void top_k_renorm_probs(TensorView probs, TensorView renorm_probs,
 }
 
 void top_k_mask_logits(TensorView logits, TensorView mask_logits,
-                       Optional<TensorView> maybe_top_k_arr, int64_t top_k_val) {
+                       Optional<TensorView> maybe_top_k_arr, int64_t top_k_val,
+                       TensorView row_states_buffer) {
   CHECK_INPUT(logits);
+  CHECK_INPUT(row_states_buffer);
   CHECK_DIM(2, logits);  // logits: (batch_size, vocab_size)
   unsigned int batch_size = logits.size(0);
   unsigned int vocab_size = logits.size(1);
   bool has_top_k_arr = maybe_top_k_arr.has_value();
 
   cudaSetDevice(logits.device().device_id);
   auto stream = get_stream(logits.device());
-  cudaError_t status = sampling::TopKMaskLogits<float>(
+
+  cudaError_t status;
+  // Use multi-CTA kernel
+  status = sampling::TopKMaskLogitsMultiCTA<float, int>(
       static_cast<float*>(logits.data_ptr()), static_cast<float*>(mask_logits.data_ptr()),
       has_top_k_arr ? static_cast<int*>(maybe_top_k_arr.value().data_ptr()) : nullptr, batch_size,
-      top_k_val, vocab_size, stream);
+      top_k_val, vocab_size,
+      static_cast<sampling::RowReductionState<float>*>(row_states_buffer.data_ptr()), stream);
 
   TVM_FFI_ICHECK(status == cudaSuccess)
       << "TopKMaskLogits failed with error code " << cudaGetErrorString(status);

flashinfer/sampling.py

@@ -377,20 +377,25 @@ def _fake_top_k_renorm_probs(
 
     # torch library for top_k_mask_logits
 
-    @register_custom_op("flashinfer::top_k_mask_logits", mutates_args=())
+    @register_custom_op(
+        "flashinfer::top_k_mask_logits", mutates_args=("row_states_buffer",)
+    )
     def top_k_mask_logits(
         logits: torch.Tensor,
         maybe_top_k_arr: Optional[torch.Tensor],
         top_k_val: int,
+        row_states_buffer: torch.Tensor,
     ) -> torch.Tensor:
         logits = logits.float()
         maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
         mask_logits = torch.empty_like(logits)
+
         module.top_k_mask_logits(
             logits,
             mask_logits,
             maybe_top_k_arr,
             top_k_val,
+            row_states_buffer,
         )
         return mask_logits
 
@@ -399,8 +404,9 @@ def _fake_top_k_mask_logits(
         logits: torch.Tensor,
         maybe_top_k_arr: Optional[torch.Tensor],
         top_k_val: int,
+        row_states_buffer: torch.Tensor,
     ) -> torch.Tensor:
-        return torch.empty_like(logits)
+        return torch.empty_like(logits, dtype=torch.float32)
 
     # torch library for chain_speculative_sampling
 
@@ -1346,8 +1352,21 @@ def top_k_mask_logits(
     top_k_renorm_probs
     """
     _check_tensor_param(top_k, logits)
+
+    # Allocate row_states buffer for multi-CTA kernel (1MB is enough for any GPU)
+    buffer_bytes = 1024 * 1024  # 1MB
+    row_states_buffer = _get_cache_buf(
+        f"top_k_mask_logits_row_states_{logits.device}",
+        buffer_bytes,
+        logits.device,
+        zero_init=True,
+    )
+
+    # Note: row_states_buffer is zero-initialized on first allocation by _get_cache_buf
+    # Kernel will reset arrival_counter to 0 at the end of each launch
+
     return get_sampling_module().top_k_mask_logits(
-        logits, *_to_tensor_scalar_tuple(top_k)
+        logits, *_to_tensor_scalar_tuple(top_k), row_states_buffer
     )
 
 

flashinfer/utils.py

@@ -203,11 +203,16 @@ def get_alibi_slopes(n_heads: int) -> torch.Tensor:
 _cache_buf: Dict[Tuple[str, torch.device], torch.Tensor] = {}
 
 
-def _get_cache_buf(name: str, bytes: int, device: torch.device) -> torch.Tensor:
+def _get_cache_buf(
+    name: str, bytes: int, device: torch.device, zero_init: bool = False
+) -> torch.Tensor:
     key = (name, device)
     buf = _cache_buf.get(key)
     if buf is None or buf.size(0) < bytes:
-        buf = torch.empty(bytes, dtype=torch.uint8, device=device)
+        if zero_init:
+            buf = torch.zeros(bytes, dtype=torch.uint8, device=device)
+        else:
+            buf = torch.empty(bytes, dtype=torch.uint8, device=device)
         _cache_buf[key] = buf
     return buf