File size: 1,330 Bytes
f36843f 728caeb f36843f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """
Unified Engine Router for Parallel Constrained Decoding.
Automatically selects MLX backend on Apple Silicon macOS,
or the PyTorch / MPS backend on Apple Silicon.
"""
import os
import platform
USE_MLX = False
if platform.system() == "Darwin" and os.environ.get("BACKEND", "").lower() != "torch":
try:
import mlx.core as mx
import mlx_lm
USE_MLX = True
except Exception:
USE_MLX = False
if USE_MLX:
from core.engine_mlx import (
get_engine,
run_parallel_generation,
run_naive_generation,
stream_naive_generation,
run_rlcd_generation,
)
else:
from core.engine_torch import (
get_torch_engine as get_engine,
run_parallel_generation_torch as run_parallel_generation,
run_naive_generation_torch as run_naive_generation,
stream_naive_generation_torch as stream_naive_generation,
)
if os.environ.get("RLCD_ATTENTION", "tree").lower() != "batch":
from core.engine_tree import run_parallel_generation_tree as run_parallel_generation
run_rlcd_generation = run_parallel_generation
__all__ = [
"get_engine",
"run_parallel_generation",
"run_naive_generation",
"stream_naive_generation",
"run_rlcd_generation",
"USE_MLX",
]
|