github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 10:06:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat] add use_cuda_kernel option

Browse Source
This commit is contained in:
Runyu Lu
2024-03-19 13:24:25 +08:00
parent 6e30248683
commit aabc9fb6aa
3 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_infer/test_cuda_graph.py
View File

@@ -49,6 +49,7 @@ def check_inference_engine(use_cuda_graph=False, batch_size=32):
max_batch_size=batch_size,
max_input_len=input_len,
max_output_len=output_len,
use_cuda_kernel=False,
use_cuda_graph=True,
block_size=16,
)
@@ -57,6 +58,7 @@ def check_inference_engine(use_cuda_graph=False, batch_size=32):
max_batch_size=batch_size,
max_input_len=input_len,
max_output_len=output_len,
use_cuda_kernel=False,
use_cuda_graph=False,
block_size=16,
)