Update symm mem constants and funcs

Signed-off-by: ilmarkov <imarkov@redhat.com>

Author: ilmarkov

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -49,13 +49,13 @@ def is_weak_contiguous(inp: torch.Tensor):
 class CustomAllreduce:
 
     _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
-    MB = 1024 * 1024
+    MiB = 1024 * 1024
     # Max sizes for each world size in case symmetric memory is available
     _MAX_SIZES = {
-        2: MB,  # 1 MB
-        4: MB,  # 1 MB
-        6: MB // 2,  # 512 KB
-        8: MB // 2,  # 512 KB
+        2: 2 * MiB,  # 1 MB
+        4: 2 * MiB,  # 1 MB
+        6: MiB,  # 512 KB
+        8: MiB // 2,  # 512 KB
     }
 
     # max_size: max supported allreduce size

vllm/distributed/device_communicators/symm_mem.py

@@ -19,13 +19,13 @@
 
 
 class SymmMemCommunicator:
-    MB = 1024 * 1024
+    MiB = 1024 * 1024
     # Max sizes for each world size
     _MAX_SIZES = {
-        2: 8 * MB,
-        4: 32 * MB,
-        6: 64 * MB,
-        8: 256 * MB,
+        2: 8 * MiB,
+        4: 32 * MiB,
+        6: 128 * MiB,
+        8: 128 * MiB,
     }
 
     def __init__(self, group: ProcessGroup, device: Union[int, str,