Tune memory arguments on B200 (#6718)

2025-05-29 00:03:22 -07:00
parent c459536b0f
commit f2bd3515fb
2 changed files with 8 additions and 2 deletions
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -149,6 +149,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, 528, 16))

    if max(capture_bs) > model_runner.req_to_token_pool.size:
        # In some case (e.g., with a small GPU or --max-running-requests), the #max-running-requests