diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 1722b524ee..cdfa83002a 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -119,6 +119,11 @@ FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = { 8: 2, # 2MB 16: 64, # 64MB (mnnvl multi-node) }, + 107: { + 2: 64, # 64MB + 4: 64, # 64MB + 8: 2, # 2MB + }, } # Max size of the input tensor per world size per device capability @@ -140,6 +145,11 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = { 4: 4, # 4MB 8: 2, # 2MB }, + 107: { + 2: 32, # 32MB + 4: 4, # 4MB + 8: 2, # 2MB + }, } MiB = 1024 * 1024 diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index 423cc53734..86b709ed41 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -47,6 +47,12 @@ CUSTOM_ALL_REDUCE_MAX_SIZES = { 6: 8 * MiB, # 8 MB 8: 4 * MiB, # 4 MB }, + "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds + 2: 4 * MiB, # 4 MB + 4: 4 * MiB, # 4 MB + 6: 8 * MiB, # 8 MB + 8: 4 * MiB, # 4 MB + }, } SYMM_MEM_ALL_REDUCE_MAX_SIZES = { @@ -68,6 +74,12 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = { 6: 32 * MiB, # 32 MB 8: 64 * MiB, # 64 MB }, + "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds + 2: 4 * MiB, # 4 MB + 4: 32 * MiB, # 32 MB + 6: 32 * MiB, # 32 MB + 8: 64 * MiB, # 64 MB + }, } # NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks. diff --git a/vllm/distributed/device_communicators/symm_mem.py b/vllm/distributed/device_communicators/symm_mem.py index 8c174602c3..26c6ec1ff4 100644 --- a/vllm/distributed/device_communicators/symm_mem.py +++ b/vllm/distributed/device_communicators/symm_mem.py @@ -27,6 +27,7 @@ class SymmMemCommunicator: "9.0": [4, 6, 8], "10.0": [6, 8], "10.3": [6, 8], + "10.7": [6, 8], # sm_107 (Rubin): reuse 10.3 thresholds } def __init__(