Implement custom kernels for top-k and top-p sampling

[WoosukKwon]

As mentioned in #81 (comment), the current PyTorch-based top-k and top-p implementation is memory-inefficient. This can be improved by introducing custom kernels.

[SuperCB]

Are you referring to implementing a more efficient CUDA kernel for this function?

[SuperCB]

Or try to implement a sampling operator that can simultaneously support multiple sampling parameters?

[WoosukKwon]

Hi @SuperCB, more specifically, this includes

1. Replacing torch.sort and torch.topk with more (memory-)efficient CUDA kernels.
2. A sampling kernel that can simultaneously support multiple sampling parameters, as you said.

[SuperCB]

Is using nvidia's cub 's sort is a good choice?

[SuperCB]

In addition, I am currently working on writing a sort function that does not use cub

[SuperCB]

I compared the sort function from NVIDIA's cub library with PyTorch's sort function in a benchmark test. The implementation results are as follows: it's evident that cub's sort function is slower than PyTorch's sort function in terms of speed. However, cub's sort function consumes significantly less GPU memory compared to PyTorch's sort. Given that the sort function contributes relatively less to the overall inference latency of the large language model, have you considered replacing PyTorch's sort function with cub's sort for performance optimization?

[SuperCB]

import torch
import vllm_sort
from loguru import logger


class CudaTimer:

    def __init__(self,name ):
        self.name = name
        self.start_event = torch.cuda.Event(enable_timing=True)
        self.end_event = torch.cuda.Event(enable_timing=True)

    def __enter__(self):
        self.start_event.record()

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        self.end_event.record()
        self.end_event.synchronize()
        self.elapsed_time_ms = self.start_event.elapsed_time(self.end_event)  
        elapsed_time_seconds = self.elapsed_time_ms / 1000  
        formatted_float = "{:.8f}".format(elapsed_time_seconds)
        logger.info(f"{self.name} cost is : {formatted_float}s")
        return True
def main(size):
    logger.error(f'size is {size}')
    out= torch.rand((1,size), device='cuda:0')
    tensor = torch.rand((1,size), device='cuda:0')
    with CudaTimer('pytorch sort'):
        torch.sort(tensor)
    with CudaTimer('cub sort'):
        vllm_sort.vllm_sort(tensor,out)

if __name__=="__main__":
    for i in range(1,9):
        main(i*10**6)

This is the code I used for testing, where vllm_sort internally calls the DeviceSegmentedSort function provided by cub, and I implemented this logic.

[SuperCB]

@WoosukKwon

[WoosukKwon]

@SuperCB Thanks for your effort. Could you also measure the memory usage of the two implementations?

[calvin1978]

TRT-LLM has these kernels, why not just ........

[SuperCB]

has these kernels, why not just ........

I couldn't find the open-source code for the project 'TRT-LLM'.

[SuperCB]

@SuperCB Thanks for your effort. Could you also measure the memory usage of the two implementations?
The following are the memory usage percentages for two algorithms when sorting 10^8 numbers.