The performance is very unefficient

[ynicle]

By running the sample code , 4096 context + 768 steps, it costs 8min for one question on H20 gpu and about 50G GRAM is occupied.

• 4096 context + 768 steps: 8min + 50G GRAM
• 2048 context + 768 steps: 4min + 31G GRAM
• 768 context + 768 steps: 2min + 20G GRAM

[BrianPulfer]

Hi @ynicle! I am experiencing the same issue. Did you find the cause / any workaround?

On an RTX 4090, 768 context + 768 steps takes:

• ~90 seconds with the entropy sampling algorithm
• ~70 seconds with the origin sampling algorithm

Here's a quick trace when using alg=entropy (similar results when using torch.compile):

-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                           aten::linear         0.30%     257.858ms         3.52%        3.036s      20.069us       0.000us         0.00%       63.111s     417.134us           0 B           0 B    1827.78 GB     -11.32 MB        151296  
                                           aten::matmul         0.22%     188.263ms         1.96%        1.689s      19.297us       0.000us         0.00%       58.874s     672.441us           0 B           0 B    1675.04 GB     -35.58 MB         87552  
                                               aten::mm         0.77%     659.628ms         1.55%        1.338s      15.412us       54.576s        61.17%       58.873s     678.380us           0 B           0 B    1674.89 GB    1674.87 GB         86784  
void cutlass::Kernel2<cutlass_80_tensorop_bf16_s1681...         0.00%       0.000us         0.00%       0.000us       0.000us       31.701s        35.53%       31.701s     737.107us           0 B           0 B           0 B           0 B         43008  
ampere_bf16_s1688gemm_bf16_128x64_sliced1x2_ldg8_f2f...         0.00%       0.000us         0.00%       0.000us       0.000us       18.147s        20.34%       18.147s     421.953us           0 B           0 B           0 B           0 B         43008  
                                             aten::sort         0.01%      12.601ms         0.04%      31.956ms      41.609us        9.597s        10.76%       10.027s      13.056ms           0 B           0 B     527.70 GB     526.44 GB           768  
void at_cuda_detail::cub::DeviceSegmentedRadixSortKe...         0.00%       0.000us         0.00%       0.000us       0.000us        6.002s         6.73%        6.002s       2.270ms           0 B           0 B           0 B           0 B          2644  
         ampere_bf16_s16816gemm_bf16_128x64_ldg8_f2f_tn         0.00%       0.000us         0.00%       0.000us       0.000us        4.708s         5.28%        4.708s       6.130ms           0 B           0 B           0 B           0 B           768  
                                       cudaLaunchKernel         2.03%        1.748s         3.17%        2.730s       2.737us       0.000us         0.00%        4.299s       4.310us           0 B           0 B     -73.00 KB     -73.00 KB        997347  
                                    Command Buffer Full         1.09%     937.461ms         1.09%     937.461ms     368.065us        4.289s         4.81%        4.289s       1.684ms           0 B           0 B           0 B           0 B          2547  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 86.220s
Self CUDA time total: 89.214s

So a lot of time (albeit not most of the time) goes into the aten::sort for some sampling algorithms.

[ynicle]

Hi @ynicle! I am experiencing the same issue. Did you find the cause / any workaround?

On an RTX 4090, 768 context + 768 steps takes:

• ~90 seconds with the entropy sampling algorithm
• ~70 seconds with the origin sampling algorithm

Here's a quick trace when using alg=entropy (similar results when using torch.compile):

-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                           aten::linear         0.30%     257.858ms         3.52%        3.036s      20.069us       0.000us         0.00%       63.111s     417.134us           0 B           0 B    1827.78 GB     -11.32 MB        151296  
                                           aten::matmul         0.22%     188.263ms         1.96%        1.689s      19.297us       0.000us         0.00%       58.874s     672.441us           0 B           0 B    1675.04 GB     -35.58 MB         87552  
                                               aten::mm         0.77%     659.628ms         1.55%        1.338s      15.412us       54.576s        61.17%       58.873s     678.380us           0 B           0 B    1674.89 GB    1674.87 GB         86784  
void cutlass::Kernel2<cutlass_80_tensorop_bf16_s1681...         0.00%       0.000us         0.00%       0.000us       0.000us       31.701s        35.53%       31.701s     737.107us           0 B           0 B           0 B           0 B         43008  
ampere_bf16_s1688gemm_bf16_128x64_sliced1x2_ldg8_f2f...         0.00%       0.000us         0.00%       0.000us       0.000us       18.147s        20.34%       18.147s     421.953us           0 B           0 B           0 B           0 B         43008  
                                             aten::sort         0.01%      12.601ms         0.04%      31.956ms      41.609us        9.597s        10.76%       10.027s      13.056ms           0 B           0 B     527.70 GB     526.44 GB           768  
void at_cuda_detail::cub::DeviceSegmentedRadixSortKe...         0.00%       0.000us         0.00%       0.000us       0.000us        6.002s         6.73%        6.002s       2.270ms           0 B           0 B           0 B           0 B          2644  
         ampere_bf16_s16816gemm_bf16_128x64_ldg8_f2f_tn         0.00%       0.000us         0.00%       0.000us       0.000us        4.708s         5.28%        4.708s       6.130ms           0 B           0 B           0 B           0 B           768  
                                       cudaLaunchKernel         2.03%        1.748s         3.17%        2.730s       2.737us       0.000us         0.00%        4.299s       4.310us           0 B           0 B     -73.00 KB     -73.00 KB        997347  
                                    Command Buffer Full         1.09%     937.461ms         1.09%     937.461ms     368.065us        4.289s         4.81%        4.289s       1.684ms           0 B           0 B           0 B           0 B          2547  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 86.220s
Self CUDA time total: 89.214s

So a lot of time (albeit not most of the time) goes into the aten::sort for some sampling algorithms.

try this project : https://github.com/NVlabs/Fast-dLLM. this can help to speed up the infer process.