mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-09 21:35:01 +08:00
[Hardware][Apple-CPU] Enable native bfloat16 on Apple Silicon (M2 and later) (#24129)
Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
This commit is contained in:
parent
7e7db04310
commit
3c2156b3af
@ -75,12 +75,12 @@ class CpuPlatform(Platform):
|
||||
def supported_dtypes(self) -> list[torch.dtype]:
|
||||
if self.get_cpu_architecture() == CpuArchEnum.POWERPC:
|
||||
return [torch.bfloat16, torch.float32]
|
||||
elif sys.platform.startswith(
|
||||
"darwin") and self.get_cpu_architecture() == CpuArchEnum.ARM:
|
||||
# TODO: change this condition to check if the platform support bf16
|
||||
# instead of checking the OS. For instance M2 shall supports bf16
|
||||
# already. But we need to modify `cpu_extension.cmake` to activate
|
||||
# the feature in the build.
|
||||
elif (self.get_cpu_architecture() == CpuArchEnum.ARM
|
||||
and sys.platform.startswith("darwin")):
|
||||
if (subprocess.check_output(
|
||||
["sysctl -n hw.optional.arm.FEAT_BF16"],
|
||||
shell=True).strip() == b"1"):
|
||||
return [torch.bfloat16, torch.float16, torch.float32]
|
||||
return [torch.float16, torch.float32]
|
||||
# x86/aarch64 CPU has supported both bf16 and fp16 natively.
|
||||
return [torch.bfloat16, torch.float16, torch.float32]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user