"tests/vscode:/vscode.git/clone" did not exist on "c375903db58826494d858e02b44d21b42669ff5e"
Unverified Commit 971a0dfa authored by Baizhou Zhang's avatar Baizhou Zhang Committed by GitHub
Browse files

Extend cuda graph capture bs for B200 (#6937)

parent 2fc12995
......@@ -139,6 +139,8 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
gpu_mem = get_device_memory_capacity()
if gpu_mem is not None and gpu_mem > 96 * 1024:
capture_bs += list(range(160, 257, 8))
if gpu_mem is not None and gpu_mem > 180 * 1000:
capture_bs += list(range(256, 513, 16))
if max(capture_bs) > model_runner.req_to_token_pool.size:
# In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment