diff --git a/docs_new/docs/hardware-platforms/apple_metal.mdx b/docs_new/docs/hardware-platforms/apple_metal.mdx index bac5f4fc2..f1eb7de59 100644 --- a/docs_new/docs/hardware-platforms/apple_metal.mdx +++ b/docs_new/docs/hardware-platforms/apple_metal.mdx @@ -60,6 +60,8 @@ SGLANG_USE_MLX=1 python -m sglang.launch_server \ 2. `--disable-cuda-graph` - Disables usage of CUDA graph, which is not relevant for Apple Metal. 3. `--disable-overlap-schedule` - Disables overlap scheduling (enabled/not present by default) achieved using MLX's `async_eval()` 4. `SGLANG_MLX_USE_CUSTOM_ROPE=1` - Enables the optional custom Metal RoPE kernel. It is disabled by default, so the MLX backend uses the standard RoPE path unless you opt in for A/B testing. +5. `SGLANG_MLX_FUSE_SWIGLU=1` - Enables the use of fused Swish-Gated Linear Unit kernel (disabled by default) +6. `SGLANG_MLX_CLEAR_CACHE_STEPS=256` - Sets the number of decode steps before clearing the MLX cache (256 by default) ## Quantization diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a34b1b11a..2f3a15aae 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -472,6 +472,9 @@ class Envs: SGLANG_USE_MLX = EnvBool(False) SGLANG_MLX_USE_CUSTOM_ROPE = EnvBool(False) SGLANG_MLX_FUSE_SWIGLU = EnvBool(False) + # Number of decode steps between periodic mx.clear_cache() calls. + # Set to 0 to disable cache clearing entirely. + SGLANG_MLX_CLEAR_CACHE_STEPS = EnvInt(256) # NPU SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index b0458ab09..07cd409e8 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -131,6 +131,7 @@ class MlxModelRunner: self._mem_fraction_static = mem_fraction_static # Counter used to trigger periodic mx.clear_cache() calls. self._decode_step_ct: int = 0 + self._clear_steps = envs.SGLANG_MLX_CLEAR_CACHE_STEPS.get() # On-the-fly quantization preset (e.g. "mlx_q4"). None = no on-load quantization. # Pre-quantized HF repos load correctly regardless of this setting: # mlx_lm.load() detects the config and instantiates QuantizedLinear @@ -1244,8 +1245,7 @@ class MlxModelRunner: self._req_token_ids[rid].append(next_tokens[i]) self._decode_step_ct += 1 - # TODO (changminbark): allow for flag configuration for clearing mx cache - if self._decode_step_ct % 256 == 0: + if self._clear_steps > 0 and self._decode_step_ct % self._clear_steps == 0: mx.clear_cache() return next_tokens