mirror/exo
1
0
Fork 0
mirror of https://github.com/exo-explore/exo.git synced 2026-02-18 23:06:23 -05:00
Code Issues Packages Projects Releases Wiki Activity

format and lint

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-02-18 11:51:30 +00:00
parent 82ac16cc0f
commit 0287ad71a1
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/exo/worker/engines/mlx/auto_parallel.py
View File

@@ -165,12 +165,12 @@ class PipelineLastLayer(CustomMlxLayer):
if cache is not None:
# CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
# doesn't have .keys directly; access via first sub-cache.
dep_cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore
dep_cache.keys = mx.depends(dep_cache.keys, output) # type: ignore[reportUnknownMemberType]
_cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore
_cache.keys = mx.depends(_cache.keys, output) # type: ignore
if self.is_prefill:
mx.eval(output)
if cache is not None:
mx.eval(dep_cache.keys) # type: ignore
mx.eval(_cache.keys) # type: ignore
if not self.is_prefill:
output = mx.distributed.all_gather(output, group=self.group)[