Turn on shm_allreduce and shm_allgather for fp16 (#10725)

2025-10-18 03:35:20 +08:00
parent f440baa136
commit 8fcc69e7c4
2 changed files with 119 additions and 1 deletions
--- a/python/sglang/srt/utils/common.py
+++ b/python/sglang/srt/utils/common.py
@@ -2921,7 +2921,7 @@ def get_cpu_ids_by_node():
 def is_shm_available(dtype, world_size, local_size):
    return (
        cpu_has_amx_support()
-        and dtype in [torch.bfloat16, torch.float]
+        and dtype in [torch.bfloat16, torch.float16, torch.float]
        and world_size >= 1
        and world_size == local_size
    )